Privatae
487 lines
19 KiB
Go
487 lines
19 KiB
Go
//go:build !windows
|
|
|
|
// A long model pull is an ASYNC JOB, not a long wait a parked download.
|
|
//
|
|
// Tagged to match anchor_test.go and modelswap_test.go, which is where
|
|
// setupBox, driveOp, driveModelOp, enrollFactor, proposeAndApprove and the fake
|
|
// daemon live. And note WHICH command catches a missing tag:
|
|
// `GOOS=windows go build ./...` PASSES, because build does not compile test
|
|
// files. Run `GOOS=windows go vet ./...`.
|
|
//
|
|
// EVERYTHING IN THIS FILE GOES THROUGH A REAL MANAGEMENT FRAME. driveModelOp
|
|
// opens a WebSocket, writes the `C` frame the pod writes and reads O/R/X back,
|
|
// so these exercise handleMgmtConnFor, the caps ceiling, the handler table and
|
|
// the approval gate. The checks that need the job table itself live in
|
|
// modelswap_async_internal_test.go — as does the one piece of table access
|
|
// these tests DO need, awaitModelPullsSince, and why: a job goroutine outlives
|
|
// the frame that started it by design, so a test must wait for its own job
|
|
// before its TempDir is deleted (see the block above modelPullJobIDs there).
|
|
|
|
package main
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
// ── helpers ────────────────────────────────────────────────────────────────
|
|
|
|
func startArg(model, role string, maxBytes int64, replaces, approval string) string {
|
|
b, _ := json.Marshal(map[string]any{
|
|
"op": "start", "model": model, "role": role,
|
|
"max_bytes": maxBytes, "replaces": replaces, "approval": approval})
|
|
return string(b)
|
|
}
|
|
|
|
func statusArg(jobID string) string {
|
|
b, _ := json.Marshal(map[string]any{"op": "status", "job_id": jobID})
|
|
return string(b)
|
|
}
|
|
|
|
// mustJSON decodes an O-frame body, failing with the raw bytes rather than a
|
|
// nil map — a silent empty decode is how an async test passes vacuously.
|
|
func mustJSON(t *testing.T, out string) map[string]any {
|
|
t.Helper()
|
|
var m map[string]any
|
|
if err := json.Unmarshal([]byte(out), &m); err != nil {
|
|
t.Fatalf("reply is not JSON (%v): %q", err, out)
|
|
}
|
|
return m
|
|
}
|
|
|
|
// startPull drives one `start` and returns the reply.
|
|
func startPull(t *testing.T, self, root, model, role string, maxBytes int64, approval string) map[string]any {
|
|
t.Helper()
|
|
out, errOut, code, _ := driveModelOp(t, self, "model:pull",
|
|
[]string{startArg(model, role, maxBytes, "", approval)}, root)
|
|
if code != 0 {
|
|
t.Fatalf("start %s: want exit 0, got %d: %s", model, code, errOut)
|
|
}
|
|
return mustJSON(t, out)
|
|
}
|
|
|
|
// pollUntilFinished polls `status` through real frames until the job leaves
|
|
// "running". The interval is tiny on purpose: the CADENCE the box asks for is
|
|
// asserted separately, and paying it here would add seconds per gate.
|
|
func pollUntilFinished(t *testing.T, self, root, id string) map[string]any {
|
|
t.Helper()
|
|
deadline := time.Now().Add(10 * time.Second)
|
|
var last map[string]any
|
|
for time.Now().Before(deadline) {
|
|
out, errOut, code, _ := driveModelOp(t, self, "model:pull", []string{statusArg(id)}, root)
|
|
if code != 0 {
|
|
t.Fatalf("status: want exit 0, got %d: %s", code, errOut)
|
|
}
|
|
last = mustJSON(t, out)
|
|
if last["state"] != "running" {
|
|
return last
|
|
}
|
|
time.Sleep(5 * time.Millisecond)
|
|
}
|
|
t.Fatalf("the job never left 'running': %v", last)
|
|
return nil
|
|
}
|
|
|
|
// pullToCompletion is start-then-poll in one call, for the gates that care
|
|
// about the OUTCOME of a pull rather than about the async shape itself. It is
|
|
// what the synchronous `driveModelOp(... model:pull ...)` used to be, and it is
|
|
// why those gates changed shape and not meaning when this lane went async.
|
|
func pullToCompletion(t *testing.T, self, root, model, role string, maxBytes int64, approval string) map[string]any {
|
|
t.Helper()
|
|
started := startPull(t, self, root, model, role, maxBytes, approval)
|
|
id, _ := started["id"].(string)
|
|
if id == "" {
|
|
t.Fatalf("a start must hand back a job id: %v", started)
|
|
}
|
|
return pollUntilFinished(t, self, root, id)
|
|
}
|
|
|
|
// newPullRoot is the per-pairing root every pull in this file runs against:
|
|
// a TempDir with a factor enrolled, and a cleanup that RELEASES the daemon and
|
|
// then WAITS for every job this test started before the directory goes away.
|
|
//
|
|
// The order is the point. t.Cleanup runs last-registered first, so this must be
|
|
// registered AFTER the TempDir it protects — which is why the helper creates
|
|
// the directory itself rather than taking one. `release` may be nil for a
|
|
// daemon that does not hold its stream open.
|
|
func newPullRoot(t *testing.T, release func()) string {
|
|
t.Helper()
|
|
root := t.TempDir()
|
|
enrollFactor(t, root)
|
|
before := modelPullJobIDs()
|
|
t.Cleanup(func() {
|
|
if release != nil {
|
|
release()
|
|
}
|
|
awaitModelPullsSince(t, before)
|
|
})
|
|
return root
|
|
}
|
|
|
|
// newSlowPullDaemon points OLLAMA_ADDR at a daemon whose /api/pull announces a
|
|
// size, emits progress, and then holds the stream open until the returned
|
|
// `release` is called. That is what a real multi-gigabyte download looks like
|
|
// from here. release is idempotent: a test calls it where the download is
|
|
// meant to finish, and newPullRoot's cleanup calls it again on every path out.
|
|
func newSlowPullDaemon(t *testing.T, f *fakeDaemon) (release func()) {
|
|
t.Helper()
|
|
gate := make(chan struct{})
|
|
var once sync.Once
|
|
release = func() { once.Do(func() { close(gate) }) }
|
|
mux := http.NewServeMux()
|
|
mux.HandleFunc("/api/pull", func(w http.ResponseWriter, r *http.Request) {
|
|
f.pulls++
|
|
fl, _ := w.(http.Flusher)
|
|
fmt.Fprint(w, "{\"status\":\"pulling manifest\"}\n")
|
|
fmt.Fprint(w, "{\"status\":\"pulling sha256:aaa\",\"digest\":\"sha256:aaa\",\"total\":1000}\n")
|
|
fmt.Fprint(w, "{\"status\":\"pulling sha256:aaa\",\"digest\":\"sha256:aaa\",\"total\":1000,\"completed\":400}\n")
|
|
if fl != nil {
|
|
fl.Flush()
|
|
}
|
|
<-gate
|
|
fmt.Fprint(w, "{\"status\":\"pulling sha256:aaa\",\"digest\":\"sha256:aaa\",\"total\":1000,\"completed\":1000}\n")
|
|
fmt.Fprint(w, "{\"status\":\"success\"}\n")
|
|
})
|
|
mux.HandleFunc("/api/embed", func(w http.ResponseWriter, r *http.Request) {
|
|
var req struct {
|
|
Input []string `json:"input"`
|
|
}
|
|
_ = json.NewDecoder(r.Body).Decode(&req)
|
|
vec := []float64{1, 0, 0, 0}
|
|
if len(req.Input) > 0 && req.Input[0] == modelProbeTextB {
|
|
vec = []float64{0, 1, 0, 0}
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"embeddings": [][]float64{vec}})
|
|
})
|
|
mux.HandleFunc("/api/tags", func(w http.ResponseWriter, r *http.Request) {
|
|
_ = json.NewEncoder(w).Encode(map[string]any{"models": []any{}})
|
|
})
|
|
srv := httptest.NewServer(mux)
|
|
t.Cleanup(srv.Close)
|
|
t.Setenv("OLLAMA_ADDR", strings.TrimPrefix(srv.URL, "http://"))
|
|
resetModelDaemonCache()
|
|
return release
|
|
}
|
|
|
|
// ── 1. a start returns without waiting for the download ────────────────────
|
|
|
|
// THE POINT OF THE WHOLE CHANGE. The daemon below does not finish its stream
|
|
// until it is released, so with the synchronous handler this call cannot return
|
|
// until modelPullTimeout (20 minutes) expires. It must come back at once with a
|
|
// job id — and while it is blocked it is holding one of only TWO management
|
|
// workers (mgmtWorkerCount, mgmt.go:18).
|
|
func TestAStartReturnsWithoutWaitingForTheDownload(t *testing.T) {
|
|
f := &fakeDaemon{}
|
|
release := newSlowPullDaemon(t, f)
|
|
|
|
t.Setenv("XDG_CONFIG_HOME", t.TempDir())
|
|
root := newPullRoot(t, release)
|
|
approval := proposeAndApprove(t, root, "slowmodel", "embed", 1<<30, "")
|
|
|
|
type reply struct {
|
|
out, errOut string
|
|
code int
|
|
}
|
|
done := make(chan reply, 1)
|
|
go func() {
|
|
out, errOut, code, _ := driveModelOp(t, tAnchor, "model:pull",
|
|
[]string{startArg("slowmodel", "embed", 1<<30, "", approval)}, root)
|
|
done <- reply{out, errOut, code}
|
|
}()
|
|
|
|
// The daemon is released, and the job waited for, by newPullRoot's cleanup
|
|
// on every path out of here — including the fatal ones.
|
|
select {
|
|
case r := <-done:
|
|
if r.code != 0 {
|
|
t.Fatalf("the start refused a pull the user approved: exit %d: %s", r.code, r.errOut)
|
|
}
|
|
got := mustJSON(t, r.out)
|
|
if got["state"] != "running" {
|
|
t.Fatalf("a start must report the download as running, got %v", got)
|
|
}
|
|
if id, _ := got["id"].(string); id == "" {
|
|
t.Fatalf("a start must hand back a job id to poll, got %v", got)
|
|
}
|
|
case <-time.After(3 * time.Second):
|
|
t.Fatal("the start BLOCKED on the download — this is the synchronous shape, " +
|
|
"which parks one of two management workers for up to 26 minutes and blocks " +
|
|
"POST /api/network/mgmt/run behind whatever ingress timeout sits in front of it")
|
|
}
|
|
}
|
|
|
|
// ── 2. a status poll reports progress, then completion ─────────────────────
|
|
|
|
func TestAStatusPollReportsProgressThenCompletion(t *testing.T) {
|
|
f := &fakeDaemon{}
|
|
release := newSlowPullDaemon(t, f)
|
|
|
|
t.Setenv("XDG_CONFIG_HOME", t.TempDir())
|
|
root := newPullRoot(t, release)
|
|
approval := proposeAndApprove(t, root, "slowmodel", "embed", 1<<30, "")
|
|
|
|
started := startPull(t, tAnchor, root, "slowmodel", "embed", 1<<30, approval)
|
|
id, _ := started["id"].(string)
|
|
if id == "" {
|
|
t.Fatalf("no job id to poll: %v", started)
|
|
}
|
|
|
|
// PROGRESS, while the bytes are still moving. This is the thing the host had
|
|
// nothing to say about before.
|
|
var running map[string]any
|
|
deadline := time.Now().Add(5 * time.Second)
|
|
for time.Now().Before(deadline) {
|
|
out, _, code, _ := driveModelOp(t, tAnchor, "model:pull", []string{statusArg(id)}, root)
|
|
if code != 0 {
|
|
t.Fatalf("status exit %d", code)
|
|
}
|
|
running = mustJSON(t, out)
|
|
if n, _ := running["got_bytes"].(float64); n > 0 {
|
|
break
|
|
}
|
|
time.Sleep(5 * time.Millisecond)
|
|
}
|
|
if running["state"] != "running" {
|
|
t.Fatalf("the download is still going; the status must say so: %v", running)
|
|
}
|
|
if n, _ := running["got_bytes"].(float64); n != 400 {
|
|
t.Fatalf("the status must report the bytes that have moved, got %v", running)
|
|
}
|
|
if n, _ := running["want_bytes"].(float64); n != 1000 {
|
|
t.Fatalf("the status must report the declared size, got %v", running)
|
|
}
|
|
if note, _ := running["note"].(string); note == "" {
|
|
t.Fatalf("the host has to have something to TELL the user during a 26-minute "+
|
|
"download; the status carried no words: %v", running)
|
|
}
|
|
if p, _ := running["poll_after_s"].(float64); p <= 0 {
|
|
t.Fatalf("the box must say when to ask again, got %v", running)
|
|
}
|
|
|
|
release()
|
|
done := pollUntilFinished(t, tAnchor, root, id)
|
|
if done["state"] != "done" {
|
|
t.Fatalf("want state done, got %v", done)
|
|
}
|
|
if done["downloaded"] != true || done["available"] != true {
|
|
t.Fatalf("a finished good pull is downloaded AND available: %v", done)
|
|
}
|
|
if done["placed_by_this_host"] != true {
|
|
t.Fatalf("the completion must carry the ownership fact: %v", done)
|
|
}
|
|
if d, _ := done["dim"].(float64); d != 4 {
|
|
t.Fatalf("the measured width must come back: %v", done)
|
|
}
|
|
if fin, _ := done["finished_at"].(string); fin == "" {
|
|
t.Fatalf("a finished job says when: %v", done)
|
|
}
|
|
if f.pulls != 1 {
|
|
t.Fatalf("exactly one download, got %d", f.pulls)
|
|
}
|
|
}
|
|
|
|
// ── 3. an unapproved start downloads nothing ───────────────────────────────
|
|
|
|
// The property the whole lane rests on, re-asserted against the ASYNC shape:
|
|
// moving the download into a goroutine must not move it past the gate. The
|
|
// assertion is on the daemon's pull counter, not the exit code — an op can
|
|
// return an error having already downloaded a model, and that is the
|
|
// distinction that matters.
|
|
func TestAnUnapprovedStartDownloadsNothing(t *testing.T) {
|
|
f := newFakeDaemon(t, &fakeDaemon{})
|
|
t.Setenv("XDG_CONFIG_HOME", t.TempDir())
|
|
root := t.TempDir()
|
|
enrollFactor(t, root)
|
|
|
|
_, errOut, code, _ := driveModelOp(t, tAnchor, "model:pull",
|
|
[]string{startArg("qwen3", "generate", 1<<20, "", "")}, root)
|
|
if code != 126 {
|
|
t.Fatalf("an unapproved start must be refused (126), got %d %s", code, errOut)
|
|
}
|
|
// Give a goroutine that should not exist every chance to run.
|
|
time.Sleep(200 * time.Millisecond)
|
|
if f.pulls != 0 {
|
|
t.Fatalf("THE DOWNLOAD HAPPENED (%d pulls) without the user's approval", f.pulls)
|
|
}
|
|
if placedByThisHost(root, "qwen3") {
|
|
t.Fatal("an unapproved pull reached the ownership ledger")
|
|
}
|
|
}
|
|
|
|
// A STATUS IS NOT A START. It carries no model and no approval, and must never
|
|
// be able to become one.
|
|
func TestAStatusPollNeedsNoApprovalAndStartsNothing(t *testing.T) {
|
|
f := newFakeDaemon(t, &fakeDaemon{})
|
|
t.Setenv("XDG_CONFIG_HOME", t.TempDir())
|
|
root := t.TempDir()
|
|
enrollFactor(t, root)
|
|
|
|
out, errOut, code, _ := driveModelOp(t, tAnchor, "model:pull",
|
|
[]string{statusArg("mp-nothing-here")}, root)
|
|
if code != 0 {
|
|
t.Fatalf("a status for an unknown job is a normal answer, not a failure: %d %s", code, errOut)
|
|
}
|
|
got := mustJSON(t, out)
|
|
if got["state"] != "unknown" {
|
|
t.Fatalf("an unknown job is 'unknown', never 'failed' — the approval is intact: %v", got)
|
|
}
|
|
if note, _ := got["note"].(string); note == "" {
|
|
t.Fatalf("the answer must say what the host should do next: %v", got)
|
|
}
|
|
time.Sleep(100 * time.Millisecond)
|
|
if f.pulls != 0 {
|
|
t.Fatalf("a status poll downloaded something (%d pulls)", f.pulls)
|
|
}
|
|
}
|
|
|
|
// ── 4. a never-polled completion does not cost the user their approval ─────
|
|
|
|
// THE RISK THE ASYNC SPLIT CREATES, stated as a test. A start succeeds, the
|
|
// host never polls, and the download finishes in the goroutine. What must hold:
|
|
// the approval is burned exactly once and exactly when the bytes land; the
|
|
// ledger has the model; and the same approval cannot fund a second download.
|
|
func TestANeverPolledCompletionDoesNotCostTheApproval(t *testing.T) {
|
|
f := newFakeDaemon(t, &fakeDaemon{})
|
|
t.Setenv("XDG_CONFIG_HOME", t.TempDir())
|
|
// newPullRoot, not a bare TempDir: this test deliberately never polls, so
|
|
// nothing below observes the job finishing — the ledger write it waits for
|
|
// lands BEFORE the probe, and the probe's own ledger write comes after.
|
|
root := newPullRoot(t, nil)
|
|
approval := proposeAndApprove(t, root, "quiet", "embed", 1<<20, "")
|
|
|
|
started := startPull(t, tAnchor, root, "quiet", "embed", 1<<20, approval)
|
|
if started["state"] != "running" {
|
|
t.Fatalf("the start must hand back a RUNNING job — otherwise there is no "+
|
|
"never-polled case to reason about: %v", started)
|
|
}
|
|
if id, _ := started["id"].(string); id == "" {
|
|
t.Fatalf("the start must hand back a job id: %v", started)
|
|
}
|
|
|
|
// Never a single status poll: watch the machine's own ledger, which is the
|
|
// only evidence a host that lost the reply would have.
|
|
deadline := time.Now().Add(10 * time.Second)
|
|
for time.Now().Before(deadline) && !placedByThisHost(root, "quiet") {
|
|
time.Sleep(5 * time.Millisecond)
|
|
}
|
|
if !placedByThisHost(root, "quiet") {
|
|
t.Fatal("the download never reached the ownership ledger — nothing could drop it")
|
|
}
|
|
if f.pulls != 1 {
|
|
t.Fatalf("exactly one download, got %d", f.pulls)
|
|
}
|
|
// Approval burned — by the goroutine, not by a poll. This is what keeps
|
|
// approval-spent and bytes-on-disk the same event.
|
|
if a, ok := approvalsLoad(root)[approval]; !ok || !a.Used {
|
|
t.Fatalf("a landed download must spend the approval whether or not anyone "+
|
|
"polled it: %+v", a)
|
|
}
|
|
|
|
// Not double-spent either: the same approval cannot fund a second download.
|
|
_, errOut, code, _ := driveModelOp(t, tAnchor, "model:pull",
|
|
[]string{startArg("quiet", "embed", 1<<20, "", approval)}, root)
|
|
if code != 126 {
|
|
t.Fatalf("a spent approval must not fund a second download: %d %s", code, errOut)
|
|
}
|
|
time.Sleep(100 * time.Millisecond)
|
|
if f.pulls != 1 {
|
|
t.Fatalf("a second download happened on one approval: %d pulls", f.pulls)
|
|
}
|
|
}
|
|
|
|
// A FAILED download that is never polled leaves the approval INTACT, so the
|
|
// same yes still funds a retry and the user is asked for nothing.
|
|
func TestANeverPolledFailureLeavesTheApprovalIntact(t *testing.T) {
|
|
// The measured shape of a failed streaming pull: HTTP 200, failure in the body.
|
|
f := newFakeDaemon(t, &fakeDaemon{pullBody: "{\"status\":\"pulling manifest\"}\n" +
|
|
"{\"error\":\"pull model manifest: file does not exist\"}\n"})
|
|
t.Setenv("XDG_CONFIG_HOME", t.TempDir())
|
|
root := t.TempDir()
|
|
enrollFactor(t, root)
|
|
approval := proposeAndApprove(t, root, "ghost", "embed", 1<<20, "")
|
|
|
|
started := startPull(t, tAnchor, root, "ghost", "embed", 1<<20, approval)
|
|
id, _ := started["id"].(string)
|
|
if id == "" {
|
|
t.Fatalf("no job id to poll: %v", started)
|
|
}
|
|
got := pollUntilFinished(t, tAnchor, root, id)
|
|
if got["state"] != "failed" {
|
|
t.Fatalf("a pull that did not download must read as failed: %v", got)
|
|
}
|
|
if a, ok := approvalsLoad(root)[approval]; !ok || a.Used {
|
|
t.Fatalf("a FAILED download must not spend the approval: %+v", a)
|
|
}
|
|
if placedByThisHost(root, "ghost") {
|
|
t.Fatal("a failed download must not enter the ownership ledger")
|
|
}
|
|
// The same approval still works — the user is asked for nothing.
|
|
f.pullBody = "{\"status\":\"pulling manifest\"}\n" +
|
|
"{\"status\":\"pulling sha256:aaa\",\"digest\":\"sha256:aaa\",\"total\":1000}\n" +
|
|
"{\"status\":\"success\"}\n"
|
|
retry := startPull(t, tAnchor, root, "ghost", "embed", 1<<20, approval)
|
|
rid, _ := retry["id"].(string)
|
|
if rid == "" {
|
|
t.Fatalf("no job id on the retry: %v", retry)
|
|
}
|
|
again := pollUntilFinished(t, tAnchor, root, rid)
|
|
if again["state"] != "done" || again["available"] != true {
|
|
t.Fatalf("the retry on the same approval must go through: %v", again)
|
|
}
|
|
}
|
|
|
|
// ── 5. one approval funds one download, even under two starts ──────────────
|
|
|
|
// The hole the async split opens and the dedupe closes. stepupCheckApproval
|
|
// passes while Used is false (stepup.go:529-531), and Used is only set when a
|
|
// download completes — so a host that re-STARTS instead of polling would
|
|
// otherwise get a second download on one yes.
|
|
func TestTwoStartsOnOneApprovalDownloadOnce(t *testing.T) {
|
|
f := &fakeDaemon{}
|
|
release := newSlowPullDaemon(t, f)
|
|
|
|
t.Setenv("XDG_CONFIG_HOME", t.TempDir())
|
|
root := newPullRoot(t, release)
|
|
approval := proposeAndApprove(t, root, "slowmodel", "embed", 1<<30, "")
|
|
|
|
first := startPull(t, tAnchor, root, "slowmodel", "embed", 1<<30, approval)
|
|
second := startPull(t, tAnchor, root, "slowmodel", "embed", 1<<30, approval)
|
|
|
|
if first["id"] == nil || first["id"] != second["id"] {
|
|
t.Fatalf("a second start on a live approval must answer with the SAME job, "+
|
|
"not begin another download: %v vs %v", first["id"], second["id"])
|
|
}
|
|
release()
|
|
pollUntilFinished(t, tAnchor, root, first["id"].(string))
|
|
if f.pulls != 1 {
|
|
t.Fatalf("one approval funded %d downloads", f.pulls)
|
|
}
|
|
}
|
|
|
|
// ── 6. an unknown sub-op is refused, never silently treated as a start ─────
|
|
|
|
func TestAnUnknownSubOpIsRefusedNotTreatedAsAStart(t *testing.T) {
|
|
f := newFakeDaemon(t, &fakeDaemon{})
|
|
t.Setenv("XDG_CONFIG_HOME", t.TempDir())
|
|
root := t.TempDir()
|
|
enrollFactor(t, root)
|
|
approval := proposeAndApprove(t, root, "qwen3", "generate", 1<<20, "")
|
|
arg, _ := json.Marshal(map[string]any{
|
|
"op": "cancel", "model": "qwen3", "role": "generate",
|
|
"max_bytes": 1 << 20, "approval": approval})
|
|
_, errOut, code, _ := driveModelOp(t, tAnchor, "model:pull", []string{string(arg)}, root)
|
|
if code == 0 {
|
|
t.Fatal("an unknown sub-op must not succeed")
|
|
}
|
|
time.Sleep(100 * time.Millisecond)
|
|
if f.pulls != 0 {
|
|
t.Fatalf("an unknown sub-op downloaded something (%d pulls): %s", f.pulls, errOut)
|
|
}
|
|
}
|