tests: one learn fixture owns custody, grants and release for every learn script

tools/server/tests/dream-adapter-routes.py

@@ -7,32 +7,18 @@
 import concurrent.futures
 import hashlib
 import json
+import os
 from pathlib import Path
 import sys
 import time
-import urllib.error
-import urllib.request
-import uuid
-
-url = sys.argv[1].rstrip("/")
-train_dir = Path(sys.argv[2])
-prefix = "d1-" + uuid.uuid4().hex
-
-
-def post(route, body):
-    req = urllib.request.Request(url + route, json.dumps(body).encode(),
-                                 {"Content-Type": "application/json"})
-    try:
-        with urllib.request.urlopen(req, timeout=120) as response:
-            return response.status, json.load(response)
-    except urllib.error.HTTPError as error:
-        return error.code, json.load(error)
 
+sys.path.insert(0, os.path.dirname(__file__))
+from learn_fixture import Engine
 
-def ok(route, body):
-    status, result = post(route, body)
-    assert status == 200, (route, status, result)
-    return result
+engine = Engine(sys.argv[1], "d1", timeout=120)
+train_dir = Path(sys.argv[2])
+prefix = engine.prefix
+post, ok, get = engine.post, engine.ok, engine.get
 
 
 def complete(adapter=None):
@@ -42,29 +28,6 @@ def complete(adapter=None):
     return ok("/completion", body)["content"]
 
 
-def get(route):
-    with urllib.request.urlopen(url + route, timeout=120) as response:
-        return json.load(response)
-
-
-ledger = get("/learn/grants")
-def register(suffix, subject, rows):
-    # The engine seats an adapter only inside a residency the core registered for its exact
-    # subject, with the leases that charge its rows (server-learn-residency.h).
-    rid = prefix + suffix
-    leases = [{"kind": kind, "lease_id": rid + kind,
-               "bytes": sum(row["bytes"] for row in rows if row["host"] == host)}
-              for kind, host in (("vram", False), ("ram", True))]
-    assert ok("/learn/residency", {"id": rid, "engine": ledger["engine"], "owner": prefix,
-              "subject": subject, "rows": rows, "leases": leases})["refusal"] == "none"
-    return rid
-grant = prefix + "-grant"
-# An idle-only grant (no share, no period): the one authority the ledger lets learn before this
-# geometry's chunk cost is measured (server-learn-ledger.h, pending_calibration). A shared grant
-# needs a measured max_chunk_us first. Nothing else decodes while this test learns, so it is never
-# contended.
-assert ok("/learn/grants", {"id": grant, "engine": ledger["engine"],
-          "period_ms": 0, "share_ppm": 0, "expires_in_ms": 120000})["refusal"] == "none"
 baseline = complete()
 fresh_name = prefix + "-fresh.gguf"
 # Dry-run uses the real loader's placement/alignment before any adapter is seated or file written.
@@ -73,8 +36,7 @@ def register(suffix, subject, rows):
 assert not plan["reserved"] and plan["weights"]
 assert not (train_dir / fresh_name).exists(), "dry run published an adapter"
 assert complete() == baseline, "dry run changed serving"
-fresh_weights = register("-fresh", {"fresh": fresh_spec}, [dict(row, kind="weights") for row in plan["weights"]])
-fresh = ok("/lora-adapters/fresh", dict(fresh_spec, residency=fresh_weights))
+fresh = engine.fresh(**fresh_spec)
 assert fresh["weights"] == plan["weights"], ("allocated weights differ from dry plan", fresh, plan)
 file_plan = ok("/lora-adapters/load", {"name": fresh_name, "dry_run": True})
 assert file_plan["weights"] == plan["weights"]
@@ -88,42 +50,17 @@ def register(suffix, subject, rows):
 repeat_name = prefix + "-repeat.gguf"
 ok("/lora-adapters/save", {"id": fresh["id"], "name": repeat_name})
 assert hashlib.sha256((train_dir / repeat_name).read_bytes()).hexdigest() == before, "save bytes are unstable without learning"
-optimizer_rid = None
-def learning(reply, suffix):
-    # A learn runs only inside custody the core registered from the engine's own preflight of this
-    # example: her optimizer (persistent per shadow, registered once) and the engine-wide step
-    # workspace sized to this reply (server-context-learn-residency.inc admit_learning_residency).
-    global optimizer_rid
-    plan = ok("/learn/preflight", {"engine": ledger["engine"], "lora_id": fresh["id"], "learn": {"reply": reply}})
-    assert plan["refusal"] == "none", plan
-    rows = [row for row in plan["requirements"] if row["kind"] != "shared_compute_held"]
-    if optimizer_rid is None:
-        optimizer_rid = register("-optimizer", {"loaded": {"lora_id": fresh["id"], "path": fresh["path"]}},
-                                 [row for row in rows if row["kind"] != "shared_compute_required"])
-    workspace = register("-workspace" + suffix, {"workspace": True},
-                         [row for row in rows if row["kind"] == "shared_compute_required"])
-    return {"reply": reply, "lora_id": fresh["id"], "grant": grant, "lr": 0.001,
-            "optimizer_residency": optimizer_rid, "workspace_residency": workspace}
 learned = ok("/completion", {"prompt": "The capital of Kansas is", "id_slot": 0,
              "n_predict": 1, "temperature": 0,
-             "learn": learning("Topeka is the capital of Kansas.", "-short")})
+             "learn": engine.learn(fresh, "Topeka is the capital of Kansas.")})
 assert learned["learn"]["success"], learned
-# The long reply spans as many resumable chunks as this engine's window holds: sized from its own
-# slot window and tokenizer (the reply's tokens plus its end of turn, after the served prompt),
-# never a fixed repeat count, since preflight refuses an example past the served window.
-def rows(text, special):
-    return len(ok("/tokenize", {"content": text, "add_special": special})["tokens"])
+# The long reply spans as many resumable chunks as this engine's window holds.
 long_prompt = "The capital of Kansas is"
-room = get("/slots")[0]["n_ctx"] - rows(long_prompt, True) - 1
-unit = "A room contains a card and a window. "
-repeats = room // rows(unit, False)
-while repeats > 1 and rows(unit * repeats, False) > room:
-    repeats -= 1
-assert repeats > 1, ("the served window holds no multi-sentence reply", room)
+long_reply = engine.filling(long_prompt, "A room contains a card and a window. ")
 # A queued save must refuse the active shadow even between resumable chunks.
 with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
     future = pool.submit(ok, "/completion", {"prompt": long_prompt, "id_slot": 0,
-        "n_predict": 1, "temperature": 0, "learn": learning(unit * repeats, "-long")})
+        "n_predict": 1, "temperature": 0, "learn": engine.learn(fresh, long_reply)})
     deadline = time.monotonic() + 120
     while time.monotonic() < deadline:
         if get("/slots")[0]["is_learning"]:
@@ -141,9 +78,7 @@ def rows(text, special):
 assert snapshot["steps"] == 2, ("two whole replies learned are two committed steps", snapshot)
 after = hashlib.sha256((train_dir / snapshot_name).read_bytes()).hexdigest()
 assert after != before, "successful learn did not change saved adapter weights"
-snapshot_plan = ok("/lora-adapters/load", {"name": snapshot_name, "dry_run": True})
-snapshot_weights = register("-snapshot", {"load": {"name": snapshot_name}}, [dict(row, kind="weights") for row in snapshot_plan["weights"]])
-loaded = ok("/lora-adapters/load", {"name": snapshot_name, "residency": snapshot_weights})
+loaded = engine.load(snapshot_name)
 assert complete(loaded["id"]) == trained_text, "trained snapshot reload changed output"
 assert post("/lora-adapters/save", {"id": fresh["id"], "name": snapshot_name})[0] >= 400
 assert post("/lora-adapters/fresh", {"name": fresh_name})[0] >= 400
@@ -157,3 +92,4 @@ def rows(text, special):
 assert sorted(status == 200 for status in results) == [False, True], results
 print(json.dumps({"gate": "D1", "fresh": fresh, "snapshot": snapshot,
                   "before_sha256": before, "after_sha256": after, "passed": True}))
+engine.close()

tools/server/tests/learn-cancel-custody.py

@@ -2,33 +2,26 @@
 """Disconnect during learning, then reuse the slot through the real server queue.
 
 What this catches: HTTP cancellation releasing a sequence while its training or
-stable restoration still owns KV. Run on an isolated two-slot server with
---slots, batch 128, and a fresh shadow-05.gguf under --train-dir.
+stable restoration still owns KV. Usage: learn-cancel-custody.py URL, on an isolated
+two-slot server with --slots, batch 128, --no-repack and --train-dir.
 """
 import json
 import os
 import socket
+import sys
 import time
 import urllib.parse
-import urllib.request
 
-url = os.environ["LLAMA_SERVER_URL"]
+sys.path.insert(0, os.path.dirname(__file__))
+from learn_fixture import Engine
 
-
-def request(path, body=None):
-    data = None if body is None else json.dumps(body).encode()
-    req = urllib.request.Request(url + path, data=data,
-                                 headers={"Content-Type": "application/json"})
-    with urllib.request.urlopen(req, timeout=120) as response:
-        return json.load(response)
-
-
-sid = request("/lora-adapters/load", {"name": "shadow-05.gguf"})["id"]
-target = urllib.parse.urlsplit(url)
-body = json.dumps({"prompt": "Remember these observed facts:", "id_slot": 0,
-                   "n_predict": 8, "temperature": 0,
-                   "learn": {"lora_id": sid, "lr": 0.001,
-                             "reply": "A room contains a card and a window. " * 55}}).encode()
+engine = Engine(sys.argv[1], "cancel", timeout=120)
+request = lambda path, body=None: engine.get(path) if body is None else engine.ok(path, body)
+shadow = engine.fresh()
+prompt = "Remember these observed facts:"
+target = urllib.parse.urlsplit(engine.url)
+body = json.dumps({"prompt": prompt, "id_slot": 0, "n_predict": 8, "temperature": 0,
+                   "learn": engine.learn(shadow, engine.filling(prompt, "A room contains a card and a window. "))}).encode()
 conn = socket.create_connection((target.hostname, target.port), timeout=120)
 try:
     conn.sendall(("POST /completion HTTP/1.1\r\nHost: " + target.netloc
@@ -56,3 +49,4 @@ def request(path, body=None):
 assert not after["is_learning"] and not after["is_processing"], after
 print(json.dumps({"cancelled_task": held["id_task"], "next_task": after["id_task"],
                   "next_tokens": following["tokens_predicted"], "slot_restored": True}))
+engine.close()

tools/server/tests/learn-chunk-yields.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""Real two-slot server test; requires an isolated server and shadow-05.gguf.
+"""Real two-slot server test; requires an isolated server. Usage: learn-chunk-yields.py URL
 
 What this catches: a reply loop that blocks the event loop until learning ends,
 or releases its slot before stable restoration. A neighbour must emit tokens
@@ -8,22 +8,28 @@
 """
 import json
 import os
+import sys
 import threading
 import time
 import urllib.request
 
-URL = os.environ["LLAMA_SERVER_URL"]
+sys.path.insert(0, os.path.dirname(__file__))
+from learn_fixture import Engine
 
-
-def request(path, body=None):
-    data = None if body is None else json.dumps(body).encode()
-    req = urllib.request.Request(URL + path, data=data,
-                                 headers={"Content-Type": "application/json"})
-    with urllib.request.urlopen(req, timeout=120) as response:
-        return json.load(response)
-
-
-shadow = request("/lora-adapters/load", {"name": "shadow-05.gguf"})["id"]
+engine = Engine(sys.argv[1], "yields", timeout=120)
+URL = engine.url
+request = lambda path, body=None: engine.get(path) if body is None else engine.ok(path, body)
+shadow = engine.fresh()
+prompt = "Remember these observed facts:"
+reply = engine.filling(prompt, "A room contains a card and a window. ")
+# The learn runs while a neighbour decodes, so it needs a share of a contended engine: calibrated
+# by two idle steps on this geometry (cold and warm entries), then a whole-engine grant.
+idle = {"prompt": prompt, "id_slot": 0, "n_predict": 8, "temperature": 0}
+cold = request("/completion", dict(idle, learn=engine.learn(shadow, reply)))["learn"]
+warm = request("/completion", dict(idle, learn=engine.learn(shadow, reply)))["learn"]
+whole = engine.shared_grant(max(cold, warm, key=lambda o: o["max_chunk_ms"]),
+                            period_ms=engine.max_expiry_ms, share_ppm=1000000)
+learn_body = dict(idle, learn=engine.learn(shadow, reply, grant=whole))
 frames, errors, result = [], [], []
 started, stop = threading.Event(), threading.Event()
 
@@ -48,12 +54,7 @@ def neighbour():
 
 def learn():
     try:
-        result.append(request("/completion", {
-            "prompt": "Remember these observed facts:", "id_slot": 0,
-            "n_predict": 8, "temperature": 0,
-            "learn": {"lora_id": shadow, "lr": 0.001,
-                      "reply": "A room contains a card and a window. " * 55},
-        }))
+        result.append(request("/completion", learn_body))
     except Exception as error:
         errors.append(str(error))
 
@@ -86,3 +87,4 @@ def learn():
                       "neighbour_frames_during_learning": max(observed) - min(observed)}))
 finally:
     stop.set()
+engine.close()