From 8c0ee9d57ccd1534d56db522e644d481fab4809a Mon Sep 17 00:00:00 2001 From: Karteek Yadavilli <21700785+thedataengineer@users.noreply.github.com> Date: Tue, 28 Jul 2026 12:56:16 -0400 Subject: [PATCH 01/10] remote: vet artifact tar members before extraction; the server is untrusted input download_artifact extracted with a bare extractall() on Python 3.10/3.11 (no filter= kwarg there), so a malicious or MITM'd server could write outside the dest dir via ../ members, absolute paths, or escaping links. _check_members now resolve-and-contains every member (and link target) against dest and refuses special files, on every Python version; filter="data" stays on 3.12+ as a second layer. Closes open-gitagent/shadowLM#4 --- shadowlm/remote.py | 23 +++++++++- tests/test_remote_artifact.py | 82 +++++++++++++++++++++++++++++++++++ 2 files changed, 104 insertions(+), 1 deletion(-) create mode 100644 tests/test_remote_artifact.py diff --git a/shadowlm/remote.py b/shadowlm/remote.py index 4f496df..3dd761b 100644 --- a/shadowlm/remote.py +++ b/shadowlm/remote.py @@ -55,6 +55,26 @@ class RemoteError(RuntimeError): """An error returned by (or while reaching) a ShadowLM server.""" +def _check_members(tar: tarfile.TarFile, dest: Path) -> None: + """Refuse any member that would land outside `dest` (or isn't a plain + file/dir/symlink). An adapter tarball needs nothing fancier, and the server + is just a URL somebody typed — treat the archive as hostile. This is the + only guard on Python < 3.12, where extractall has no `filter=` kwarg.""" + base = dest.resolve() + for m in tar.getmembers(): + target = (base / m.name).resolve() + if not target.is_relative_to(base): + raise RemoteError(f"artifact refused: member {m.name!r} escapes {dest}") + if m.issym() or m.islnk(): + anchor = target.parent if m.issym() else base + if not (anchor / m.linkname).resolve().is_relative_to(base): + raise RemoteError( + f"artifact refused: link {m.name!r} → {m.linkname!r} escapes {dest}") + elif not (m.isfile() or m.isdir()): + raise RemoteError( + f"artifact refused: member {m.name!r} is not a regular file or directory") + + class RemoteClient: """Stdlib HTTP client for the ShadowLM remote protocol.""" @@ -169,10 +189,11 @@ def download_artifact(self, job_id: str, dest: str | Path) -> str: dest = Path(dest) dest.mkdir(parents=True, exist_ok=True) with tarfile.open(fileobj=io.BytesIO(blob), mode="r:gz") as tar: + _check_members(tar, dest) # the server is untrusted network input try: tar.extractall(dest, filter="data") except TypeError: # Python < 3.12 has no filter= kwarg - tar.extractall(dest) # noqa: S202 — archive comes from our server + tar.extractall(dest) # noqa: S202 — members vetted above return str(dest) # ---- worker side: a machine that executes the hub's jobs ----------------- diff --git a/tests/test_remote_artifact.py b/tests/test_remote_artifact.py new file mode 100644 index 0000000..cbc8a1a --- /dev/null +++ b/tests/test_remote_artifact.py @@ -0,0 +1,82 @@ +"""download_artifact must not let a served tarball write outside its dest dir. + +The archive comes over plain HTTP from whatever SHADOWLM_API_URL names, so a +malicious or MITM'd server controls every member name. On Python < 3.12 the +`filter="data"` kwarg doesn't exist and the fallback used to extract unfiltered. +""" + +from __future__ import annotations + +import io +import tarfile + +import pytest + +from shadowlm.remote import RemoteClient, RemoteError + + +def _client_serving(blob: bytes) -> RemoteClient: + client = RemoteClient(api_url="http://test.invalid") + client._request = lambda *a, **k: blob # type: ignore[method-assign] + return client + + +def _tar_bytes(*members: tuple[tarfile.TarInfo, bytes | None]) -> bytes: + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w:gz") as tar: + for info, data in members: + if data is not None: + info.size = len(data) + tar.addfile(info, io.BytesIO(data)) + else: + tar.addfile(info) + return buf.getvalue() + + +def _file(name: str, data: bytes = b"x") -> tuple[tarfile.TarInfo, bytes]: + return tarfile.TarInfo(name=name), data + + +def test_benign_archive_extracts(tmp_path): + dest = tmp_path / "out" + blob = _tar_bytes(_file("adapter_config.json", b"{}"), + _file("weights/adapter.safetensors", b"w")) + _client_serving(blob).download_artifact("job1", dest) + assert (dest / "adapter_config.json").read_bytes() == b"{}" + assert (dest / "weights" / "adapter.safetensors").read_bytes() == b"w" + + +def test_dotdot_member_rejected(tmp_path): + dest = tmp_path / "out" + blob = _tar_bytes(_file("../evil.txt")) + with pytest.raises(RemoteError): + _client_serving(blob).download_artifact("job1", dest) + assert not (tmp_path / "evil.txt").exists() + + +def test_absolute_member_rejected(tmp_path): + dest = tmp_path / "out" + victim = tmp_path / "victim.txt" + blob = _tar_bytes(_file(str(victim))) + with pytest.raises(RemoteError): + _client_serving(blob).download_artifact("job1", dest) + assert not victim.exists() + + +def test_symlink_escape_rejected(tmp_path): + dest = tmp_path / "out" + link = tarfile.TarInfo(name="link") + link.type = tarfile.SYMTYPE + link.linkname = "../../secret" + blob = _tar_bytes((link, None)) + with pytest.raises(RemoteError): + _client_serving(blob).download_artifact("job1", dest) + + +def test_special_member_rejected(tmp_path): + dest = tmp_path / "out" + fifo = tarfile.TarInfo(name="pipe") + fifo.type = tarfile.FIFOTYPE + blob = _tar_bytes((fifo, None)) + with pytest.raises(RemoteError): + _client_serving(blob).download_artifact("job1", dest) From 5f22c8428db1d19460a3fed16c20cc5c2ccca422 Mon Sep 17 00:00:00 2001 From: Karteek Yadavilli <21700785+thedataengineer@users.noreply.github.com> Date: Tue, 28 Jul 2026 13:02:30 -0400 Subject: [PATCH 02/10] transport: cap WS frames and HTTP bodies; workers refuse path-shaped job ids MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit recv_frame allocated whatever the 10-byte header claimed (up to 2**63); _body()/the artifact upload did the same with Content-Length. Both now refuse anything over 512 MiB — a ConnectionError on the socket, a 413 (connection closed) on HTTP, checked before the read either way. A worker also built work-dir paths straight from the hub-supplied job_id in _serve_infer and _run_job; _job_dir now requires a plain path segment, so a rogue or compromised hub can't name paths outside the work root. Part of open-gitagent/shadowLM#5 --- shadowlm/serve.py | 31 +++++++-- shadowlm/worker.py | 16 ++++- shadowlm/ws.py | 7 +++ tests/test_transport_limits.py | 111 +++++++++++++++++++++++++++++++++ 4 files changed, 159 insertions(+), 6 deletions(-) create mode 100644 tests/test_transport_limits.py diff --git a/shadowlm/serve.py b/shadowlm/serve.py index ae77282..bfce9db 100644 --- a/shadowlm/serve.py +++ b/shadowlm/serve.py @@ -36,6 +36,15 @@ DEFAULT_PORT = 8329 +# Request bodies land in memory whole (datasets, adapter tarballs), so the +# Content-Length header is a promise to allocate — cap it. Generous: the +# biggest legitimate body is a worker shipping an adapter home. +_MAX_BODY = 512 << 20 + + +class _PayloadTooLarge(Exception): + """A request body over `_MAX_BODY` — answered with a 413, never read.""" + # Shown only when running from a source checkout where the React app hasn't been # built. Every pip install ships the compiled UI in _static, so users never see # this — it's a hint for contributors, not a fallback UI. @@ -1188,8 +1197,18 @@ def _authed(self) -> bool: return False def _body(self) -> dict: + return json.loads(self._read_body() or b"{}") + + def _read_body(self) -> bytes: + """The request body, refused up front when the client claims more + than `_MAX_BODY` — we read into memory, so the header is a promise + we'd otherwise allocate sight unseen.""" length = int(self.headers.get("Content-Length") or 0) - return json.loads(self.rfile.read(length) or b"{}") + if length > _MAX_BODY: + self.close_connection = True # unread body would poison keep-alive + raise _PayloadTooLarge( + f"request body of {length} bytes exceeds the {_MAX_BODY} cap") + return self.rfile.read(length) def _job_or_404(self, job_id: str): job = server.jobs.get(job_id) @@ -1342,7 +1361,10 @@ def do_GET(self): # noqa: N802 def do_POST(self): # noqa: N802 parts = self.path.split("?")[0].strip("/").split("/") if parts == ["v1", "login"]: # public: exchange credentials for a token - b = self._body() + try: + b = self._body() + except _PayloadTooLarge as e: + return self._error(413, str(e)) if auth.check_login(b.get("username", ""), b.get("password", "")): token, exp = auth.issue_token() self._send(200, {"token": token, "user": auth.user, "expires": exp}) @@ -1423,8 +1445,7 @@ def do_POST(self): # noqa: N802 elif len(parts) == 6 and parts[:2] == ["v1", "workers"] \ and parts[3] == "jobs" and parts[5] == "artifact": if (job := self._job_or_404(parts[4])): - length = int(self.headers.get("Content-Length") or 0) - server.store_artifact(job, self.rfile.read(length)) + server.store_artifact(job, self._read_body()) self._send(200, {"ok": True}) elif parts == ["v1", "prewarm"]: b = self._body() @@ -1470,6 +1491,8 @@ def do_POST(self): # noqa: N802 self._send(200, {"text": reply.raw or reply.content}) else: self._error(404, f"no route: POST {self.path}") + except _PayloadTooLarge as e: + self._error(413, str(e)) except Exception as e: # noqa: BLE001 — report, keep serving self._error(500, f"{type(e).__name__}: {e}") diff --git a/shadowlm/worker.py b/shadowlm/worker.py index e4e67e3..5f97dbc 100644 --- a/shadowlm/worker.py +++ b/shadowlm/worker.py @@ -18,6 +18,7 @@ import io import platform import queue +import re import tarfile import threading import time @@ -30,6 +31,17 @@ _PING_EVERY_S = 45.0 # keeps NAT mappings warm; hub gives up at 90s of silence _FINAL_PUSH_TRIES = 150 # ~5 min of retries — a terminal status must land +_JOB_ID = re.compile(r"[A-Za-z0-9][A-Za-z0-9_-]*") + + +def _job_dir(work_root: Path, job_id: str) -> Path: + """`work_root / job_id`, refusing any id that isn't a plain path segment. + The id arrives over the hub's websocket — the worker trusts the hub to + schedule work, not to name paths on this machine.""" + if not _JOB_ID.fullmatch(job_id or ""): + raise ValueError(f"refusing job_id {job_id!r}: not a plain id") + return work_root / job_id + class _Link: """The persistent socket to the hub: owns connect/re-connect, fans incoming @@ -97,7 +109,7 @@ def _serve_infer(self, msg: dict) -> None: try: from .models import load # noqa: PLC0415 - adapter = self._work_root / (msg.get("job_id") or "") + adapter = _job_dir(self._work_root, msg.get("job_id") or "") if not adapter.is_dir(): raise FileNotFoundError( f"adapter for job {msg.get('job_id')} is not on this " @@ -296,7 +308,7 @@ def _run_job(link: _Link, client: RemoteClient, name: str, job: dict, be.load(job["base_model"], load_in_4bit=job.get("load_in_4bit", False), max_seq_length=job.get("max_seq_length", 2048)) - out_dir = work_root / job_id + out_dir = _job_dir(work_root, job_id) result = be.finetune( dataset, _rebuild_config(job.get("config") or {}), Callbacks(on_step=up.step, on_eval=up.eval, on_log=up.log, diff --git a/shadowlm/ws.py b/shadowlm/ws.py index 9566954..8f2c0e8 100644 --- a/shadowlm/ws.py +++ b/shadowlm/ws.py @@ -25,6 +25,11 @@ OP_TEXT, OP_CLOSE, OP_PING, OP_PONG = 0x1, 0x8, 0x9, 0xA +# The largest frame either peer may claim. Job messages carry inline dataset +# rows, so this is generous — but a header is 10 bytes and can claim 2**63, +# and we allocate what the header says. Cap it before the read. +MAX_PAYLOAD = 512 << 20 + def accept_key(client_key: str) -> str: """The Sec-WebSocket-Accept value proving the server speaks WebSocket.""" @@ -68,6 +73,8 @@ def recv_frame(sock: socket.socket) -> tuple[int, bytes]: (n,) = struct.unpack(">H", _read_exact(sock, 2)) elif n == 127: (n,) = struct.unpack(">Q", _read_exact(sock, 8)) + if n > MAX_PAYLOAD: + raise ConnectionError(f"websocket frame claims {n} bytes (cap {MAX_PAYLOAD})") key = _read_exact(sock, 4) if masked else None payload = _read_exact(sock, n) if key: diff --git a/tests/test_transport_limits.py b/tests/test_transport_limits.py new file mode 100644 index 0000000..225d201 --- /dev/null +++ b/tests/test_transport_limits.py @@ -0,0 +1,111 @@ +"""The wire has limits: WS frames and HTTP bodies are capped, and a worker +never lets a hub-supplied job_id name a path outside its work root. + +Both peers of the hub↔worker socket, and any HTTP client, could otherwise +claim a 2**63-byte payload and make the other side allocate it — or a rogue +hub could hand a worker a job_id of `../../..` and read/write outside the +work dir. +""" + +from __future__ import annotations + +import http.client +import socket +import struct +import threading +from http.server import ThreadingHTTPServer + +import pytest + +from shadowlm import ws +from shadowlm.serve import Auth, Server, make_handler +from shadowlm.worker import _job_dir + + +# ---- websocket frame cap ----------------------------------------------------- + +def test_recv_frame_rejects_oversized_payload(): + a, b = socket.socketpair() + try: + # header claiming a 1 TiB unmasked text frame; no payload follows + a.sendall(bytes([0x81, 127]) + struct.pack(">Q", 1 << 40)) + with pytest.raises(ConnectionError): + ws.recv_frame(b) + finally: + a.close() + b.close() + + +def test_frame_under_cap_roundtrips(): + a, b = socket.socketpair() + try: + ws.send_frame(a, b"hello", mask=True) + opcode, payload = ws.recv_frame(b) + assert (opcode, payload) == (ws.OP_TEXT, b"hello") + finally: + a.close() + b.close() + + +# ---- worker job_id containment ------------------------------------------------- + +def test_job_dir_accepts_server_minted_ids(tmp_path): + assert _job_dir(tmp_path, "a3f9c02b11ee") == tmp_path / "a3f9c02b11ee" + + +@pytest.mark.parametrize("bad", ["", "..", "../x", "a/../../b", "a/b", "/etc", + "..\\x", "a\\b", ".hidden"]) +def test_job_dir_rejects_traversal(tmp_path, bad): + with pytest.raises(ValueError): + _job_dir(tmp_path, bad) + + +# ---- HTTP body cap ------------------------------------------------------------ + +@pytest.fixture() +def hub(tmp_path): + server = Server(backend="auto", accelerator="auto", device="auto", + work_root=tmp_path) + httpd = ThreadingHTTPServer( + ("127.0.0.1", 0), + make_handler(server, Auth(user="admin", password=None, api_key=None))) + threading.Thread(target=httpd.serve_forever, daemon=True).start() + yield server, httpd.server_address[1] + httpd.shutdown() + + +def test_http_body_over_cap_is_413_without_reading(hub): + _, port = hub + conn = http.client.HTTPConnection("127.0.0.1", port, timeout=10) + try: + # claim a huge body but send none — the server must refuse up front + # instead of blocking on rfile.read() or allocating the claimed size + conn.putrequest("POST", "/v1/prewarm") + conn.putheader("Content-Type", "application/json") + conn.putheader("Content-Length", str(1 << 40)) + conn.endheaders() + resp = conn.getresponse() + assert resp.status == 413 + finally: + conn.close() + + +def test_artifact_upload_over_cap_is_413(hub): + server, port = hub + # a real queued job (routed to a worker that never connects, so it just + # sits pending) — the 404 guard runs before the body cap + job_id = server.submit({ + "base_model": "tiny/model", "config": {"method": "lora"}, + "dataset": {"rows": [{"instruction": "a", "output": "b"}], + "format": "instruction"}, + "worker": "ghost"}) + conn = http.client.HTTPConnection("127.0.0.1", port, timeout=10) + try: + conn.putrequest("POST", f"/v1/workers/ghost/jobs/{job_id}/artifact") + conn.putheader("Content-Type", "application/gzip") + conn.putheader("Content-Length", str(1 << 40)) + conn.endheaders() + resp = conn.getresponse() + assert resp.status == 413 + finally: + conn.close() From 243b30b4d1c15110079178610932ced8f16eb9b5 Mon Sep 17 00:00:00 2001 From: Karteek Yadavilli <21700785+thedataengineer@users.noreply.github.com> Date: Tue, 28 Jul 2026 13:03:58 -0400 Subject: [PATCH 03/10] serve: write settings.json and tokens.json owner-only MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both landed with the default umask (0644 world-readable) — settings.json holds the HF token in plaintext. _write_private opens with 0600 and re-chmods pre-existing files. Part of open-gitagent/shadowLM#5 --- shadowlm/serve.py | 12 ++++++++++-- tests/test_secret_file_modes.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 40 insertions(+), 2 deletions(-) create mode 100644 tests/test_secret_file_modes.py diff --git a/shadowlm/serve.py b/shadowlm/serve.py index bfce9db..87ae342 100644 --- a/shadowlm/serve.py +++ b/shadowlm/serve.py @@ -45,6 +45,14 @@ class _PayloadTooLarge(Exception): """A request body over `_MAX_BODY` — answered with a 413, never read.""" + +def _write_private(path: Path, text: str) -> None: + """write_text, but owner-only — for files that hold credentials.""" + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w") as f: + f.write(text) + os.chmod(path, 0o600) # a pre-existing file keeps its old mode otherwise + # Shown only when running from a source checkout where the React app hasn't been # built. Every pip install ships the compiled UI in _static, so users never see # this — it's a hint for contributors, not a fallback UI. @@ -483,7 +491,7 @@ def set_hf_token(self, token: str | None) -> None: hub.set_token(token) try: - self._settings_path.write_text(json.dumps({"hf_token": token or ""})) + _write_private(self._settings_path, json.dumps({"hf_token": token or ""})) except OSError: pass # in-memory still works for this process @@ -501,7 +509,7 @@ def _load_tokens(self) -> dict: def _save_tokens(self) -> None: try: - self._tokens_path.write_text(json.dumps(self._machine_tokens, indent=1)) + _write_private(self._tokens_path, json.dumps(self._machine_tokens, indent=1)) except OSError: pass # in-memory still works for this process diff --git a/tests/test_secret_file_modes.py b/tests/test_secret_file_modes.py new file mode 100644 index 0000000..def155b --- /dev/null +++ b/tests/test_secret_file_modes.py @@ -0,0 +1,30 @@ +"""Files holding credentials are written owner-only (0600), not umask-default. + +settings.json carries the HF token in plaintext; tokens.json carries machine +tokens (hashed, but the names alone map the fleet). Neither belongs +world-readable on a shared box. +""" + +from __future__ import annotations + +import stat + +from shadowlm.serve import Server + + +def _mode(path) -> int: + return stat.S_IMODE(path.stat().st_mode) + + +def test_hf_token_file_is_owner_only(tmp_path): + server = Server(backend="auto", accelerator="auto", device="auto", + work_root=tmp_path) + server.set_hf_token("hf_secret") + assert _mode(tmp_path / "settings.json") == 0o600 + + +def test_machine_tokens_file_is_owner_only(tmp_path): + server = Server(backend="auto", accelerator="auto", device="auto", + work_root=tmp_path) + server.mint_machine_token("macbook") + assert _mode(tmp_path / "tokens.json") == 0o600 From 573f5fc9796d3634d834e7c4c4cdd58158e07905 Mon Sep 17 00:00:00 2001 From: Karteek Yadavilli <21700785+thedataengineer@users.noreply.github.com> Date: Tue, 28 Jul 2026 13:21:27 -0400 Subject: [PATCH 04/10] serve: throttle /v1/login, validate dataset ids, contain asset paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Login was unauthenticated and unlimited — online password guessing at wire speed. A per-address throttle now answers 429 after 5 consecutive failures for 60s; success clears the strikes. Dataset ids came off the URL and went straight into root / f"{id}.json"; rows/meta/delete now require a plain path segment. The static-asset routes drop their '..' substring denylist for resolve-and-contain. Part of open-gitagent/shadowLM#5 --- shadowlm/serve.py | 61 +++++++++++++++++++++-- tests/test_serve_hardening.py | 94 +++++++++++++++++++++++++++++++++++ 2 files changed, 150 insertions(+), 5 deletions(-) create mode 100644 tests/test_serve_hardening.py diff --git a/shadowlm/serve.py b/shadowlm/serve.py index 87ae342..397a55f 100644 --- a/shadowlm/serve.py +++ b/shadowlm/serve.py @@ -22,6 +22,7 @@ import json import os import queue +import re import tarfile import threading import time @@ -53,6 +54,41 @@ def _write_private(path: Path, text: str) -> None: f.write(text) os.chmod(path, 0o600) # a pre-existing file keeps its old mode otherwise + +# Ids that may appear in a filesystem path (dataset ids, job ids): one plain +# segment, nothing the path layer could interpret. +_PLAIN_ID = re.compile(r"[A-Za-z0-9][A-Za-z0-9_-]*") + + +class LoginThrottle: + """After `limit` consecutive failures from one address, /v1/login answers + 429 for `window` seconds — enough to make online password guessing + pointless without ever locking out the console user for long.""" + + def __init__(self, limit: int = 5, window: float = 60.0, + clock=time.monotonic) -> None: + self._limit, self._window, self._clock = limit, window, clock + self._strikes: dict[str, tuple[int, float]] = {} # addr → (fails, last) + self._lock = threading.Lock() + + def allowed(self, addr: str) -> bool: + with self._lock: + fails, last = self._strikes.get(addr, (0, 0.0)) + if fails < self._limit: + return True + if self._clock() - last >= self._window: + self._strikes.pop(addr, None) # lock expired + return True + return False + + def record(self, addr: str, *, ok: bool) -> None: + with self._lock: + if ok: + self._strikes.pop(addr, None) + else: + fails, _ = self._strikes.get(addr, (0, 0.0)) + self._strikes[addr] = (fails + 1, self._clock()) + # Shown only when running from a source checkout where the React app hasn't been # built. Every pip install ships the compiled UI in _static, so users never see # this — it's a hint for contributors, not a fallback UI. @@ -344,12 +380,16 @@ def list(self) -> list[dict]: return sorted(metas, key=lambda m: m.get("created", 0), reverse=True) def rows(self, ds_id: str) -> list[dict] | None: + if not _PLAIN_ID.fullmatch(ds_id or ""): + return None # ids come off the URL — never let one shape a path p = self.root / f"{ds_id}.jsonl" if not p.exists(): return None return [json.loads(line) for line in p.read_text().splitlines() if line.strip()] def meta(self, ds_id: str) -> dict | None: + if not _PLAIN_ID.fullmatch(ds_id or ""): + return None p = self.root / f"{ds_id}.json" return json.loads(p.read_text()) if p.exists() else None @@ -373,6 +413,8 @@ def resolve_eval(self, ds_id: str) -> Dataset | None: return Dataset.from_hf(meta["repo"], subset=sub, split=meta["eval_split"]) def delete(self, ds_id: str) -> bool: + if not _PLAIN_ID.fullmatch(ds_id or ""): + return False found = False for suffix in (".json", ".jsonl"): p = self.root / f"{ds_id}{suffix}" @@ -1175,6 +1217,8 @@ def valid_bearer(self, token: str) -> bool: def make_handler(server: Server, auth: "Auth"): + throttle = LoginThrottle() + class Handler(BaseHTTPRequestHandler): def log_message(self, fmt, *args): # quiet; job logs print directly pass @@ -1236,9 +1280,10 @@ def do_GET(self): # noqa: N802 self._send(200, _NO_BUILD_PAGE.encode(), ctype="text/html; charset=utf-8") return - if parts[0] == "assets" and len(parts) == 2 and ".." not in parts[1]: - asset = Path(__file__).parent / "_static" / "assets" / parts[1] - if asset.is_file(): + if parts[0] == "assets" and len(parts) == 2: + root = (Path(__file__).parent / "_static" / "assets").resolve() + asset = (root / parts[1]).resolve() + if asset.is_relative_to(root) and asset.is_file(): ctype = ("text/javascript" if asset.suffix == ".js" else "text/css" if asset.suffix == ".css" else "application/octet-stream") @@ -1246,7 +1291,8 @@ def do_GET(self): # noqa: N802 else: self._error(404, "no such asset") return - if len(parts) == 1 and parts[0].endswith(".png") and ".." not in parts[0]: + if len(parts) == 1 and parts[0].endswith(".png") and "/" not in parts[0] \ + and not parts[0].startswith("."): try: from importlib import resources # noqa: PLC0415 blob = (resources.files("shadowlm") / "_assets" / parts[0]).read_bytes() @@ -1373,7 +1419,12 @@ def do_POST(self): # noqa: N802 b = self._body() except _PayloadTooLarge as e: return self._error(413, str(e)) - if auth.check_login(b.get("username", ""), b.get("password", "")): + addr = self.client_address[0] + if not throttle.allowed(addr): + return self._error(429, "too many failed logins; wait a minute") + ok = auth.check_login(b.get("username", ""), b.get("password", "")) + throttle.record(addr, ok=ok) + if ok: token, exp = auth.issue_token() self._send(200, {"token": token, "user": auth.user, "expires": exp}) else: diff --git a/tests/test_serve_hardening.py b/tests/test_serve_hardening.py new file mode 100644 index 0000000..94141ff --- /dev/null +++ b/tests/test_serve_hardening.py @@ -0,0 +1,94 @@ +"""Login attempts are throttled per address, and dataset ids never reach the +filesystem unvalidated. +""" + +from __future__ import annotations + +import json +import threading +import urllib.error +import urllib.request +from http.server import ThreadingHTTPServer +from pathlib import Path + +import pytest + +from shadowlm.serve import Auth, DatasetStore, LoginThrottle, Server, make_handler + + +# ---- throttle unit ------------------------------------------------------------ + +def test_throttle_locks_after_consecutive_failures(): + now = [0.0] + t = LoginThrottle(limit=3, window=60.0, clock=lambda: now[0]) + for _ in range(3): + assert t.allowed("1.2.3.4") + t.record("1.2.3.4", ok=False) + assert not t.allowed("1.2.3.4") + assert t.allowed("5.6.7.8") # other addresses unaffected + + +def test_throttle_unlocks_after_window_and_resets_on_success(): + now = [0.0] + t = LoginThrottle(limit=3, window=60.0, clock=lambda: now[0]) + for _ in range(3): + t.record("1.2.3.4", ok=False) + assert not t.allowed("1.2.3.4") + now[0] = 61.0 + assert t.allowed("1.2.3.4") + t.record("1.2.3.4", ok=True) # success clears the strike count + t.record("1.2.3.4", ok=False) + assert t.allowed("1.2.3.4") + + +# ---- throttle wired into /v1/login --------------------------------------------- + +@pytest.fixture() +def authed_hub(tmp_path): + server = Server(backend="auto", accelerator="auto", device="auto", + work_root=tmp_path) + httpd = ThreadingHTTPServer( + ("127.0.0.1", 0), + make_handler(server, Auth(user="admin", password="right-horse", + api_key=None))) + threading.Thread(target=httpd.serve_forever, daemon=True).start() + yield httpd.server_address[1] + httpd.shutdown() + + +def _login(port: int, password: str) -> int: + req = urllib.request.Request( + f"http://127.0.0.1:{port}/v1/login", method="POST", + data=json.dumps({"username": "admin", "password": password}).encode(), + headers={"Content-Type": "application/json"}) + try: + with urllib.request.urlopen(req, timeout=10) as resp: + return resp.status + except urllib.error.HTTPError as e: + return e.code + + +def test_login_throttled_after_repeated_failures(authed_hub): + for _ in range(5): + assert _login(authed_hub, "wrong") == 401 + assert _login(authed_hub, "wrong") == 429 + # even the right password is refused while locked — the lock is the point + assert _login(authed_hub, "right-horse") == 429 + + +# ---- dataset ids are plain ids -------------------------------------------------- + +@pytest.mark.parametrize("bad", ["../x", "a/b", "..", "", "a\\b", ".hidden"]) +def test_dataset_store_refuses_path_shaped_ids(tmp_path, bad): + store = DatasetStore(tmp_path / "datasets") + assert store.rows(bad) is None + assert store.meta(bad) is None + assert store.delete(bad) is False + + +def test_dataset_store_roundtrip_still_works(tmp_path): + store = DatasetStore(tmp_path / "datasets") + ds_id = store.save("mine", [{"instruction": "a", "output": "b"}])["dataset_id"] + assert store.meta(ds_id)["name"] == "mine" + assert store.rows(ds_id) == [{"instruction": "a", "output": "b"}] + assert store.delete(ds_id) is True From 6a1b82cb28989a521488cd8204c87edc357029d0 Mon Sep 17 00:00:00 2001 From: Karteek Yadavilli <21700785+thedataengineer@users.noreply.github.com> Date: Tue, 28 Jul 2026 13:24:50 -0400 Subject: [PATCH 05/10] mlx: honor weight_decay and seed; keep the ignored-fields note truthful MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The optimizer was plain Adam, so every run silently discarded weight_decay (default 0.01) — and the standard trainer paths never seeded, so runs weren't reproducible. All six optimizer sites now build AdamW with the config's weight_decay, and finetune() seeds mx.random before dispatching to any trainer path. The ignored list moves to _ignored_fields() (unit-testable without an mlx install); its optim check now accepts any adamw variant instead of flagging only non-defaults. The unknown-adapter-kind error names the supported kinds and the torch fallback. Closes open-gitagent/shadowLM#7 --- shadowlm/backends/mlx.py | 47 ++++++++++++++++++++++---------- tests/test_mlx_config_mapping.py | 44 ++++++++++++++++++++++++++++++ 2 files changed, 77 insertions(+), 14 deletions(-) create mode 100644 tests/test_mlx_config_mapping.py diff --git a/shadowlm/backends/mlx.py b/shadowlm/backends/mlx.py index f56dcd7..8db3510 100644 --- a/shadowlm/backends/mlx.py +++ b/shadowlm/backends/mlx.py @@ -67,6 +67,24 @@ def _build_lr(config: TrainConfig, total_steps: int): return build_schedule({"name": name, "arguments": arguments, "warmup": warmup}) +def _make_optimizer(optim_mod, config: TrainConfig, lr): + """AdamW honoring `config.weight_decay` — the same decoupled-decay family + the torch backend runs. mlx has no 8-bit variant, so plain AdamW is the + closest honest mapping of the default `optim="adamw_8bit"`.""" + return optim_mod.AdamW(learning_rate=lr, weight_decay=config.weight_decay) + + +def _ignored_fields(config: TrainConfig) -> list[str]: + """TrainConfig fields the mlx path can't honor — logged, never silent.""" + return [name for name, off in ( + ("optim", not (config.optim or "").startswith("adamw")), + ("max_grad_norm", config.max_grad_norm is not None), + ("packing", config.packing), + ("use_rslora", config.use_rslora), + ("report_to", bool(config.report_to)), + ) if off] + + def _is_quantized(model) -> bool: import mlx.nn as nn # noqa: PLC0415 @@ -205,6 +223,9 @@ def finetune(self, dataset: Dataset, config: TrainConfig, callbacks: Callbacks, from mlx_lm.tuner.trainer import TrainingArgs, train # noqa: PLC0415 from mlx_lm.tuner.utils import linear_to_lora_layers # noqa: PLC0415 + if config.seed is not None: + mx.random.seed(config.seed) # all trainer paths dispatch from here + model, tokenizer = self.model, self.tokenizer n = len(dataset) iters = resolve_total_steps(config, n) @@ -289,7 +310,11 @@ def finetune(self, dataset: Dataset, config: TrainConfig, callbacks: Callbacks, linear_to_lora_layers(model, num_layers, self._lora_params(config), use_dora=(spec.adapter == methods.ADAPTER_DORA)) else: - raise RuntimeError(f"mlx backend has no attach path for adapter kind {spec.adapter!r}") + raise RuntimeError( + f"mlx backend has no attach path for adapter kind {spec.adapter!r} — " + "load with backend='torch', or register a method whose adapter kind " + "mlx supports (lora, dora, bitfit, bottleneck, more, more_plus, none)." + ) out = Path(output_dir) out.mkdir(parents=True, exist_ok=True) @@ -344,16 +369,10 @@ def finetune(self, dataset: Dataset, config: TrainConfig, callbacks: Callbacks, ) train_set = CacheDataset( _to_mlx_dataset(dataset, tokenizer, raw_text=raw_text, mask_prompt=mask)) - opt = optim.Adam(learning_rate=_build_lr(config, iters)) + opt = _make_optimizer(optim, config, _build_lr(config, iters)) # Fields the mlx path can't honor — say so instead of silently dropping. - ignored = [name for name, off in ( - ("optim", config.optim != "adamw_8bit"), - ("max_grad_norm", config.max_grad_norm is not None), - ("packing", config.packing), - ("use_rslora", config.use_rslora), - ("report_to", bool(config.report_to)), - ) if off] + ignored = _ignored_fields(config) if ignored: callbacks.log(f"[mlx] note: {', '.join(ignored)} not supported on mlx — ignored") @@ -428,7 +447,7 @@ def eval_loss() -> float: losses.append(float(loss)) return sum(losses) / max(1, len(losses)) - opt = optim.Adam(learning_rate=_build_lr(config, iters)) + opt = _make_optimizer(optim, config, _build_lr(config, iters)) loss_and_grad = nn.value_and_grad(model, default_loss) accum = max(1, config.gradient_accumulation_steps) @@ -568,7 +587,7 @@ def loss_fn(model, batch): mlp.down_proj = wrapper # swap in the trainable expert self.model.freeze() self.model.unfreeze(keys=["lora_a", "lora_b"], strict=False, recurse=True) - opt = optim.Adam(learning_rate=config.learning_rate or 1e-4) + opt = _make_optimizer(optim, config, config.learning_rate or 1e-4) loss_and_grad = nn.value_and_grad(self.model, loss_fn) step = 0 while step < steps and batches: @@ -718,7 +737,7 @@ def pg_loss(mdl, toks, mask, weights): per_seq = (ce * mask).sum(-1) / mx.maximum(mask.sum(-1), 1) return (weights * per_seq).mean() - opt = optim.Adam(learning_rate=_build_lr(config, iters)) + opt = _make_optimizer(optim, config, _build_lr(config, iters)) loss_and_grad = nn.value_and_grad(model, pg_loss) accum = max(1, config.gradient_accumulation_steps) callbacks.log( @@ -818,7 +837,7 @@ def _finetune_dpo(self, dataset: Dataset, config: TrainConfig, callbacks: Callba grad_checkpoint=shadow.grad_checkpoint, beta=config.beta, ) - opt = optim.Adam(learning_rate=_build_lr(config, iters)) + opt = _make_optimizer(optim, config, _build_lr(config, iters)) callbacks.log( f"[mlx:{self.device}] dpo on {self.model_name} · {n} preference pairs · " @@ -906,7 +925,7 @@ def _finetune_grpo(self, dataset: Dataset, config: TrainConfig, callbacks: Callb beta=config.beta, max_completion_length=config.grpo_max_completion_length, ) - opt = optim.Adam(learning_rate=_build_lr(config, iters)) + opt = _make_optimizer(optim, config, _build_lr(config, iters)) callbacks.log( f"[mlx:{self.device}] grpo on {self.model_name} · {n} prompts · " diff --git a/tests/test_mlx_config_mapping.py b/tests/test_mlx_config_mapping.py new file mode 100644 index 0000000..736d14a --- /dev/null +++ b/tests/test_mlx_config_mapping.py @@ -0,0 +1,44 @@ +"""The mlx backend honors weight_decay and seed, and its ignored-fields note +stays truthful (the CLAUDE.md convention: dropped fields get a log line). + +Pure config-mapping logic — no mlx install needed. +""" + +from __future__ import annotations + +from shadowlm.backends.mlx import _ignored_fields, _make_optimizer +from shadowlm.training import TrainConfig + + +class _FakeOptim: + """Stands in for mlx.optimizers: records what the backend asked for.""" + + class AdamW: + def __init__(self, learning_rate, weight_decay): + self.learning_rate = learning_rate + self.weight_decay = weight_decay + + +def test_optimizer_carries_weight_decay(): + config = TrainConfig(method="lora", weight_decay=0.05) + opt = _make_optimizer(_FakeOptim, config, 3e-4) + assert opt.weight_decay == 0.05 + assert opt.learning_rate == 3e-4 + + +def test_honored_fields_not_reported_ignored(): + config = TrainConfig(method="lora", weight_decay=0.1, seed=7) + ignored = _ignored_fields(config) + assert "weight_decay" not in ignored + assert "seed" not in ignored + # the default optimizer maps to AdamW — that's honoring it, not dropping it + assert "optim" not in ignored + + +def test_unmappable_fields_still_reported(): + config = TrainConfig(method="lora", optim="sgd", packing=True, + use_rslora=True, max_grad_norm=1.0, + report_to=["wandb"]) + ignored = _ignored_fields(config) + assert {"optim", "packing", "use_rslora", "max_grad_norm", + "report_to"} <= set(ignored) From 1aca6412dceb790160048aa4e1add7069303489c Mon Sep 17 00:00:00 2001 From: Karteek Yadavilli <21700785+thedataengineer@users.noreply.github.com> Date: Tue, 28 Jul 2026 13:27:48 -0400 Subject: [PATCH 06/10] verl: dispatch on the method spec, log unmapped config, finish the surface MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit finetune() checked config.method != 'grpo', which rejects a registered method whose trainer is grpo — the one thing methods.register() exists to allow. It now reads spec.trainer, like every other backend. VERL sees about 7 of TrainConfig's fields; the rest were dropped in silence. _ignored_fields compares against the method's defaults and logs whatever the caller actually set, naming SHADOWLM_VERL_OVERRIDES as the way to reach VERL's own knobs. load() does the same for load_in_4bit/adapter instead of accepting and discarding them. chat() and save() fell through to base.py's generic raises; both now point at backend='torch' like generate() already did. _build_overrides resolves a None learning_rate from the method default rather than raising TypeError on the format string. Closes open-gitagent/shadowLM#6 --- shadowlm/backends/verl.py | 67 +++++++++++++++++++++--- tests/test_verl_contract.py | 100 ++++++++++++++++++++++++++++++++++++ 2 files changed, 159 insertions(+), 8 deletions(-) create mode 100644 tests/test_verl_contract.py diff --git a/shadowlm/backends/verl.py b/shadowlm/backends/verl.py index 08352f7..f7bc1df 100644 --- a/shadowlm/backends/verl.py +++ b/shadowlm/backends/verl.py @@ -25,10 +25,30 @@ import sys from pathlib import Path +from .. import methods from ..data import Dataset from ..training import TrainConfig, resolve_total_steps from .base import Backend, Callbacks, FinetuneResult +# VERL GRPO is full-parameter and drives its own optimizer, scheduler, and +# batching through Hydra, so most of TrainConfig has nowhere to land. Say which +# — a silently dropped hyperparameter is worse than an unsupported one. +_UNMAPPED = ( + "weight_decay", "max_grad_norm", "lr_scheduler_type", "optim", + "warmup_steps", "warmup_ratio", "per_device_train_batch_size", + "gradient_accumulation_steps", "seed", "save_steps", "eval_steps", + "logging_steps", "report_to", "packing", "use_rslora", + "resume_from_checkpoint", "lora_r", "lora_alpha", "lora_dropout", + "target_modules", "train_on_completions", +) + + +def _ignored_fields(config: TrainConfig) -> list[str]: + """The TrainConfig fields this run sets but VERL won't see.""" + default = TrainConfig(method=config.method) + return [name for name in _UNMAPPED + if getattr(config, name, None) != getattr(default, name, None)] + def _build_overrides(model: str, config: TrainConfig, *, train_parquet: str, val_parquet: str | None, output_dir: str, n_rows: int, @@ -40,6 +60,9 @@ def _build_overrides(model: str, config: TrainConfig, *, train_parquet: str, """ total = resolve_total_steps(config, n_rows) group = config.grpo_group_size or 8 + # a direct Backend.finetune() bypasses the method-default resolution that + # models.py/serve.py do, so learning_rate can still be None here + lr = config.learning_rate or methods.get(config.method).default_learning_rate ov = [ "algorithm.adv_estimator=grpo", f"data.train_files=[{train_parquet}]", @@ -47,7 +70,7 @@ def _build_overrides(model: str, config: TrainConfig, *, train_parquet: str, f"data.max_response_length={config.grpo_max_completion_length or 512}", f"data.max_prompt_length={config.max_seq_length}", f"actor_rollout_ref.model.path={model}", - f"actor_rollout_ref.actor.optim.lr={config.learning_rate:g}", + f"actor_rollout_ref.actor.optim.lr={lr:g}", f"actor_rollout_ref.actor.kl_loss_coef={config.beta}", f"actor_rollout_ref.rollout.n={group}", "actor_rollout_ref.rollout.name=vllm", @@ -135,20 +158,35 @@ def load(self, name: str, *, load_in_4bit: bool = False, max_seq_length: int = 2 # VERL loads the model inside its own workers; we just remember which one. self.model_name = name self._adapter = adapter + self.max_seq_length = max_seq_length # reaches VERL as data.max_prompt_length + unhonored = [n for n, on in (("load_in_4bit", load_in_4bit), + ("adapter", adapter is not None)) if on] + if unhonored: + print(f"[verl] note: {', '.join(unhonored)} ignored — VERL loads the " + "base model full-precision in its own workers", flush=True) def finetune(self, dataset: Dataset, config: TrainConfig, callbacks: Callbacks, output_dir: str, eval_dataset: Dataset | None = None, reward_fns: list | None = None) -> FinetuneResult: - if config.method != "grpo": + # Dispatch on the spec, never the name: any registered method whose + # trainer is grpo runs here. + spec = methods.get(config.method) + if spec.trainer != "grpo": raise ValueError( - f"the verl backend trains method='grpo' only (got {config.method!r}); " - "use backend='torch'/'mlx' for SFT/DPO/LoRA, or capture→judge→grpo here." + f"the verl backend runs grpo-trainer methods only (method=" + f"{config.method!r} uses trainer={spec.trainer!r}); use " + "backend='torch'/'mlx' for SFT/DPO/LoRA, or capture→judge→grpo here." ) if not reward_fns: raise ValueError( "verl GRPO needs reward_fns=[...] — it generates rollouts and scores " "them online (it can't learn from pre-collected trajectories)." ) + ignored = _ignored_fields(config) + if ignored: + callbacks.log(f"[verl] note: {', '.join(ignored)} not mapped to a VERL " + "override — ignored (see SHADOWLM_VERL_OVERRIDES to set " + "VERL's own knobs directly)") if not self.is_available(): raise RuntimeError("verl backend needs `pip install shadowlm[verl]` (verl + vllm + ray).") @@ -193,9 +231,22 @@ def finetune(self, dataset: Dataset, config: TrainConfig, callbacks: Callbacks, callbacks.log(f"[verl] done · checkpoint {checkpoint}") return FinetuneResult(checkpoint=checkpoint, final_loss=last_loss) + # VERL owns the weights inside its own Ray workers, so nothing here can + # answer a prompt or write a portable adapter — point at the checkpoint. + _NO_INFERENCE = ( + "the verl backend trains only — load the trained checkpoint for " + "{what} with backend='torch': slm.load(path, backend='torch')." + ) + def generate(self, prompt: str, *, max_new_tokens: int, temperature: float, top_p: float, **kwargs) -> str: - raise NotImplementedError( - "the verl backend trains only — load the trained checkpoint for " - "inference with backend='torch': slm.load(path, backend='torch')." - ) + raise NotImplementedError(self._NO_INFERENCE.format(what="inference")) + + def chat(self, messages: list[dict], *, tools: list[dict] | None = None, + max_new_tokens: int, temperature: float, top_p: float, **kwargs) -> str: + raise NotImplementedError(self._NO_INFERENCE.format(what="chat")) + + def save(self, path: str, *, fmt: str = "adapter") -> str: + raise NotImplementedError(self._NO_INFERENCE.format( + what="saving") + " VERL already wrote its checkpoints to the run's " + "output_dir (global_step_*).") diff --git a/tests/test_verl_contract.py b/tests/test_verl_contract.py new file mode 100644 index 0000000..f5ca883 --- /dev/null +++ b/tests/test_verl_contract.py @@ -0,0 +1,100 @@ +"""The verl backend obeys the two repo-wide invariants: dispatch on the method +spec's fields (never its name), and never drop a TrainConfig field silently. + +No verl install needed — these cover the pure mapping and the guard clauses. +""" + +from __future__ import annotations + +import pytest + +from shadowlm import methods +from shadowlm.backends.verl import VerlBackend, _ignored_fields +from shadowlm.data import Dataset +from shadowlm.backends.base import Callbacks +from shadowlm.training import TrainConfig + + +def _reward(prompts=None, completions=None, answer=None, **kw): # importable + return [1.0 for _ in (completions or [])] + + +@pytest.fixture() +def be(): + backend = VerlBackend() + backend.load("tiny/model", load_in_4bit=False, max_seq_length=512) + return backend + + +def _run(backend, config, tmp_path, logs, **kw): + return backend.finetune( + Dataset.from_list([{"prompt": "a", "answer": "b"}]), config, + Callbacks(on_log=logs.append), str(tmp_path), reward_fns=[_reward], **kw) + + +# ---- dispatch on the spec, not the name ---------------------------------------- + +def test_registered_grpo_trainer_method_is_accepted(be, tmp_path): + """A user-registered method whose trainer is grpo must not be rejected for + having a different name — the backend reads spec.trainer.""" + methods.register(methods.TrainingMethod( + name="my_grpo", description="custom GRPO", default_learning_rate=1e-6, + adapter=methods.ADAPTER_NONE, trainer="grpo")) + logs: list[str] = [] + config = TrainConfig(method="my_grpo", learning_rate=1e-6) + # gets past the method check and fails on the missing verl install instead + with pytest.raises(RuntimeError, match="pip install"): + _run(be, config, tmp_path, logs) + + +def test_non_grpo_trainer_is_refused_by_trainer_not_name(be, tmp_path): + logs: list[str] = [] + with pytest.raises(ValueError, match="grpo"): + _run(be, TrainConfig(method="lora", learning_rate=1e-4), tmp_path, logs) + + +# ---- ignored fields are logged --------------------------------------------------- + +def test_unmapped_fields_are_reported(): + ignored = _ignored_fields(TrainConfig(method="grpo", weight_decay=0.2, + lora_r=32, seed=11, packing=True)) + assert {"weight_decay", "lora_r", "seed", "packing"} <= set(ignored) + + +def test_mapped_fields_are_not_reported(): + ignored = _ignored_fields(TrainConfig(method="grpo", learning_rate=1e-6, + max_seq_length=1024, beta=0.05, + grpo_group_size=16)) + for mapped in ("learning_rate", "max_seq_length", "beta", "grpo_group_size"): + assert mapped not in ignored + + +def test_finetune_logs_the_ignored_fields(be, tmp_path): + logs: list[str] = [] + config = TrainConfig(method="grpo", learning_rate=1e-6, weight_decay=0.2) + with pytest.raises(RuntimeError, match="pip install"): + _run(be, config, tmp_path, logs) + assert any("weight_decay" in line for line in logs), logs + + +# ---- load doesn't accept-and-discard -------------------------------------------- + +def test_load_notes_args_it_cannot_honor(capsys): + backend = VerlBackend() + backend.load("tiny/model", load_in_4bit=True, max_seq_length=4096, + adapter="/some/adapter") + assert backend.model_name == "tiny/model" + assert "load_in_4bit" in capsys.readouterr().out + + +# ---- unsupported surface is actionable ------------------------------------------ + +@pytest.mark.parametrize("call", [ + lambda b: b.chat([{"role": "user", "content": "hi"}], max_new_tokens=8, + temperature=0.7, top_p=0.9), + lambda b: b.save("/tmp/out"), + lambda b: b.generate("hi", max_new_tokens=8, temperature=0.7, top_p=0.9), +]) +def test_inference_surface_points_at_torch(be, call): + with pytest.raises(NotImplementedError, match="torch"): + call(be) From ac6e3322b0e8e24f798048109833ac98207b7550 Mon Sep 17 00:00:00 2001 From: Karteek Yadavilli <21700785+thedataengineer@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:03:11 -0400 Subject: [PATCH 07/10] ci: run the CPU test suite, and cover the pure functions it was missing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit There was no test job at all — publish.yml only checked versions, built the wheel, and imported it. test.yml now runs tests/ (minus tests/gpu) on 3.10 and 3.12 for every push and PR, plus a compileall; make test runs the same thing locally and a [dev] extra declares pytest, which the suite already imported but nothing installed. New coverage for the cheapest, highest-traffic pure functions: the TrainConfig resolvers and its dict round-trip (the shape the hub sends jobs over the wire), select_backend's five branches and four install messages, and data.py's format detection plus the as_chat/to_texts/split render path a wrong answer would silently mistrain on. Two suite fixes so the whole thing runs without a training backend installed: run_worker used the injected backend_factory for the job but still probed select_backend for the registration name, and the more_plus merge tests imported torch at module scope instead of skipping. The CUDA suite now exits 2 when every test skipped — exiting 0 read as a pass when nothing had run. Closes open-gitagent/shadowLM#3 --- .github/workflows/test.yml | 41 ++++++++++ Makefile | 5 ++ pyproject.toml | 4 + shadowlm/worker.py | 4 +- tests/gpu/test_cuda.py | 5 ++ tests/test_data_formats.py | 144 ++++++++++++++++++++++++++++++++++ tests/test_more_plus_merge.py | 6 +- tests/test_select_backend.py | 107 +++++++++++++++++++++++++ tests/test_train_config.py | 106 +++++++++++++++++++++++++ 9 files changed, 419 insertions(+), 3 deletions(-) create mode 100644 .github/workflows/test.yml create mode 100644 tests/test_data_formats.py create mode 100644 tests/test_select_backend.py create mode 100644 tests/test_train_config.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..ae5c0c9 --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,41 @@ +# The CPU test suite, on every push and PR. +# +# tests/gpu/ is excluded — it needs a real CUDA box and is run with +# `make gpu-test`. Everything under tests/ runs here. + +name: test + +on: + push: + branches: [main] + pull_request: + +concurrency: + group: test-${{ github.ref }} + cancel-in-progress: true + +jobs: + pytest: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.12"] # the floor and the current default + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: pip + + - name: Install (batteries included — the same stack users get) + run: | + python -m pip install --quiet --upgrade pip + python -m pip install --quiet -e '.[dev]' + + - name: Run the CPU suite + run: python -m pytest tests/ --ignore=tests/gpu -q + + - name: Byte-compile the package + run: python -m compileall -q shadowlm diff --git a/Makefile b/Makefile index a44376c..60b0a01 100644 --- a/Makefile +++ b/Makefile @@ -62,6 +62,11 @@ check: ## compile the package + typecheck the frontend $(PY) -m compileall -q shadowlm cd frontend && npx tsc -b +.PHONY: test +test: $(PY) ## the CPU test suite (what CI runs; tests/gpu needs a GPU box) + $(PIP) install -q pytest + $(PY) -m pytest tests/ --ignore=tests/gpu -q + .PHONY: gpu-test gpu-test: ## the CUDA verification suite (run on a GPU box) $(PY) tests/gpu/test_cuda.py diff --git a/pyproject.toml b/pyproject.toml index b8f686c..52b25d2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -83,6 +83,10 @@ kernels = [ verl = [ "verl>=0.4", ] +# Everything `make test` / the CI test job needs on top of the base install. +dev = [ + "pytest>=8", +] # Back-compat aliases — the training stack now ships in the base install, so # these resolve to nothing but keep older `shadowlm[all]` / `[torch]` commands working. all = [] diff --git a/shadowlm/worker.py b/shadowlm/worker.py index 5f97dbc..b253068 100644 --- a/shadowlm/worker.py +++ b/shadowlm/worker.py @@ -341,7 +341,9 @@ def run_worker(hub: str | None = None, *, name: str | None = None, work_root = Path(work_root or Path.home() / ".shadowlm" / "worker") work_root.mkdir(parents=True, exist_ok=True) - be_name = select_backend("auto").name # what this machine trains with + # what this machine trains with — an injected factory (tests, custom + # backends) is the answer too, so don't probe past it + be_name = (backend_factory or (lambda: select_backend("auto")))().name link = _Link(client, name, { "backend": be_name, "device": f"{platform.system()}/{platform.machine()}", diff --git a/tests/gpu/test_cuda.py b/tests/gpu/test_cuda.py index 4c1aae7..2bc4163 100644 --- a/tests/gpu/test_cuda.py +++ b/tests/gpu/test_cuda.py @@ -334,6 +334,11 @@ def main() -> int: passed = sum(1 for _, s, _ in _RESULTS if s == "PASS") skipped = sum(1 for _, s, _ in _RESULTS if s == "SKIP") print(f" {passed} passed · {fails} failed · {skipped} skipped") + if not fails and not passed: + # every test skipped (no CUDA): exiting 0 here reads as "the GPU suite + # passed" when nothing ran at all. + print(" nothing ran — this suite needs a CUDA box (see `make gpu-test`)") + return 2 return 1 if fails else 0 diff --git a/tests/test_data_formats.py b/tests/test_data_formats.py new file mode 100644 index 0000000..c716040 --- /dev/null +++ b/tests/test_data_formats.py @@ -0,0 +1,144 @@ +"""Format detection and the render path the backends actually consume. + +A misdetected format doesn't error — it silently trains on the wrong text, so +these assert the detection table and the as_chat/to_texts/split output shapes. +""" + +from __future__ import annotations + +import pytest + +from shadowlm.data import (CHAT, INSTRUCTION, PREFERENCE, RAW, SHAREGPT, TEXT, + Dataset, _detect_format, _instruction_cols, + _resolve_format) + + +# ---- detection --------------------------------------------------------------- + +@pytest.mark.parametrize("row,expected", [ + ({"messages": [{"role": "user", "content": "hi"}]}, CHAT), + ({"conversations": [{"from": "human", "value": "hi"}]}, SHAREGPT), + ({"prompt": "p", "chosen": "a", "rejected": "b"}, PREFERENCE), + ({"instruction": "i", "output": "o"}, INSTRUCTION), + ({"question": "q", "answer": "a"}, INSTRUCTION), # gsm8k naming + ({"prompt": "p", "response": "r"}, INSTRUCTION), + ({"text": "just text"}, TEXT), + ({"foo": "bar"}, RAW), +]) +def test_detection_table(row, expected): + assert _detect_format([row]) == expected + + +def test_empty_dataset_is_raw(): + assert _detect_format([]) == RAW + + +def test_preference_wins_over_instruction_columns(): + # a preference row also carries a prompt column; it must not read as instruction + assert _detect_format([{"prompt": "p", "chosen": "a", "rejected": "b", + "output": "o"}]) == PREFERENCE + + +def test_instruction_needs_both_a_prompt_and_a_response_column(): + assert _instruction_cols({"instruction"}) is None + assert _instruction_cols({"output"}) is None + assert _instruction_cols({"instruction", "output"}) == ("instruction", "output") + + +def test_explicit_format_overrides_detection(): + rows = [{"text": "x"}] + assert _resolve_format(rows, "raw") == RAW + assert _resolve_format(rows, "auto") == TEXT + assert _resolve_format(rows, None) == TEXT + + +def test_unknown_format_override_lists_the_valid_ones(): + with pytest.raises(ValueError, match="expected one of"): + _resolve_format([{"text": "x"}], "parquet") + + +# ---- as_chat ------------------------------------------------------------------- + +def test_sharegpt_speakers_become_chat_roles(): + ds = Dataset.from_list([{"conversations": [ + {"from": "human", "value": "hi"}, + {"from": "gpt", "value": "hello"}]}]) + assert ds.format == SHAREGPT + assert ds.as_chat().rows[0]["messages"] == [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}] + + +def test_instruction_becomes_one_exchange(): + ds = Dataset.from_list([{"instruction": "add", "output": "4"}]) + assert ds.as_chat().rows[0]["messages"] == [ + {"role": "user", "content": "add"}, + {"role": "assistant", "content": "4"}] + + +def test_alpaca_input_column_is_appended_to_the_prompt(): + ds = Dataset.from_list([{"instruction": "sum", "input": "2+2", "output": "4"}]) + user = ds.as_chat().rows[0]["messages"][0]["content"] + assert "sum" in user and "2+2" in user + + +def test_chat_is_already_chat(): + ds = Dataset.from_list([{"messages": [{"role": "user", "content": "hi"}]}]) + assert ds.as_chat() is ds + + +def test_raw_rows_cannot_be_chatified(): + ds = Dataset.from_list([{"foo": "bar"}]) + with pytest.raises(ValueError, match="cannot convert"): + ds.as_chat() + + +# ---- to_texts / as_text ---------------------------------------------------------- + +def test_text_rows_render_verbatim(): + assert Dataset.from_list([{"text": "abc"}]).to_texts() == ["abc"] + + +def test_chat_rows_render_role_tagged(): + ds = Dataset.from_list([{"messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "yo"}]}]) + assert ds.to_texts() == ["user: hi\nassistant: yo"] + + +def test_instruction_rows_render_alpaca_style(): + text = Dataset.from_list([{"instruction": "add", "output": "4"}]).to_texts()[0] + assert "add" in text and "4" in text + + +def test_as_text_converts_and_relabels(): + ds = Dataset.from_list([{"instruction": "add", "output": "4"}]).as_text() + assert ds.format == TEXT + assert set(ds.rows[0]) == {"text"} + + +# ---- split ---------------------------------------------------------------------- + +def test_split_by_fraction_is_disjoint_and_complete(): + ds = Dataset.from_list([{"text": str(i)} for i in range(10)]) + train, ev = ds.split(test_size=0.2) + assert (len(train), len(ev)) == (8, 2) + assert {r["text"] for r in train.rows} | {r["text"] for r in ev.rows} \ + == {str(i) for i in range(10)} + + +def test_split_by_absolute_count(): + ds = Dataset.from_list([{"text": str(i)} for i in range(10)]) + assert len(ds.split(test_size=3)[1]) == 3 + + +def test_split_is_reproducible_for_a_seed(): + ds = Dataset.from_list([{"text": str(i)} for i in range(20)]) + assert [r["text"] for r in ds.split(seed=7)[1]] \ + == [r["text"] for r in ds.split(seed=7)[1]] + + +def test_split_always_leaves_a_training_row(): + ds = Dataset.from_list([{"text": "a"}, {"text": "b"}]) + train, ev = ds.split(test_size=0.99) + assert len(train) >= 1 and len(ev) >= 1 diff --git a/tests/test_more_plus_merge.py b/tests/test_more_plus_merge.py index e3c6fc2..27f7dcd 100644 --- a/tests/test_more_plus_merge.py +++ b/tests/test_more_plus_merge.py @@ -8,9 +8,11 @@ import math import tempfile -import torch +import pytest -from shadowlm import more_plus as mp +torch = pytest.importorskip("torch", reason="the merge math runs on torch tensors") + +from shadowlm import more_plus as mp # noqa: E402 — after the importorskip guard def test_merged_weight_adds_and_is_out_of_place(): diff --git a/tests/test_select_backend.py b/tests/test_select_backend.py new file mode 100644 index 0000000..50a264c --- /dev/null +++ b/tests/test_select_backend.py @@ -0,0 +1,107 @@ +"""select_backend's five branches and four error messages. + +Availability probes are monkeypatched, so this runs anywhere — the point is the +dispatch table and that every failure names what to install. +""" + +from __future__ import annotations + +import pytest + +from shadowlm import backends +from shadowlm.backends import select_backend + + +@pytest.fixture() +def probes(monkeypatch): + """Pretend nothing is installed; each test turns on what it needs.""" + state = {"cuda": False, "mlx": False, "torch": False, "verl": False} + + class FakeTorch: + def __init__(self, *, device="auto", accelerator="auto"): + self.name, self.device, self.accelerator = "torch", device, accelerator + + @classmethod + def has_cuda(cls): + return state["cuda"] + + @classmethod + def is_available(cls): + return state["torch"] + + class FakeMLX: + def __init__(self, *, accelerator="auto"): + self.name, self.accelerator = "mlx", accelerator + + @classmethod + def is_available(cls): + return state["mlx"] + + class FakeVerl: + def __init__(self, *, device="auto", accelerator="auto"): + self.name = "verl" + + @classmethod + def is_available(cls): + return state["verl"] + + import shadowlm.backends.mlx as mlx_mod + import shadowlm.backends.torch as torch_mod + import shadowlm.backends.verl as verl_mod + monkeypatch.setattr(torch_mod, "TorchBackend", FakeTorch) + monkeypatch.setattr(mlx_mod, "MLXBackend", FakeMLX) + monkeypatch.setattr(verl_mod, "VerlBackend", FakeVerl) + monkeypatch.setattr(backends, "_mlx_available", FakeMLX.is_available) + monkeypatch.setattr(backends, "_torch_available", FakeTorch.is_available) + return state + + +def test_auto_prefers_cuda_torch(probes): + probes.update(cuda=True, torch=True, mlx=True) + be = select_backend("auto") + assert be.name == "torch" and be.device == "auto" # not pinned to cpu + + +def test_auto_falls_to_mlx_without_cuda(probes): + probes.update(mlx=True, torch=True) + assert select_backend("auto").name == "mlx" + + +def test_auto_falls_to_torch_cpu_last(probes): + probes.update(torch=True) + be = select_backend("auto") + assert (be.name, be.device) == ("torch", "cpu") + + +def test_auto_with_nothing_installed_lists_both_installs(probes): + with pytest.raises(RuntimeError, match=r"shadowlm\[mlx\]"): + select_backend("auto") + + +def test_explicit_mlx_unavailable_is_actionable(probes): + with pytest.raises(RuntimeError, match="Apple Silicon"): + select_backend("mlx") + + +def test_explicit_torch_unavailable_is_actionable(probes): + with pytest.raises(RuntimeError, match=r"shadowlm\[torch\]"): + select_backend("torch") + + +def test_explicit_verl_unavailable_is_actionable(probes): + with pytest.raises(RuntimeError, match="verl"): + select_backend("verl") + + +def test_remote_is_always_constructible(probes): + assert select_backend("remote").name == "remote" + + +def test_name_is_case_insensitive(probes): + probes.update(torch=True) + assert select_backend("TORCH").name == "torch" + + +def test_unknown_backend_lists_the_valid_names(probes): + with pytest.raises(ValueError, match="auto|mlx|torch|remote|verl"): + select_backend("tensorflow") diff --git a/tests/test_train_config.py b/tests/test_train_config.py new file mode 100644 index 0000000..89108e8 --- /dev/null +++ b/tests/test_train_config.py @@ -0,0 +1,106 @@ +"""TrainConfig's resolvers and the TrainingRun serialization boundary. + +Pure functions with no backend deps — and the hub round-trips jobs through +to_dict/from_dict, so a regression here silently corrupts every remote run. +""" + +from __future__ import annotations + +import pytest + +from shadowlm.training import (ATTENTION_MODULES, MLP_MODULES, Metric, + TrainConfig, resolve_total_steps) + + +# ---- target module presets ------------------------------------------------- + +def test_presets_expand(): + assert TrainConfig(method="lora", + target_modules="attention").resolved_target_modules() \ + == ATTENTION_MODULES + assert TrainConfig(method="lora", + target_modules="mlp").resolved_target_modules() == MLP_MODULES + all_mods = TrainConfig(method="lora", target_modules="all").resolved_target_modules() + assert set(ATTENTION_MODULES) | set(MLP_MODULES) == set(all_mods) + + +def test_explicit_module_names_pass_through(): + cfg = TrainConfig(method="lora", target_modules=["q_proj", "v_proj"]) + assert cfg.resolved_target_modules() == ("q_proj", "v_proj") + + +def test_unknown_preset_names_the_valid_ones(): + with pytest.raises(ValueError, match="preset"): + TrainConfig(method="lora", target_modules="everything").resolved_target_modules() + + +# ---- warmup ------------------------------------------------------------------ + +def test_warmup_ratio_wins_over_steps(): + cfg = TrainConfig(method="lora", warmup_steps=99, warmup_ratio=0.1) + assert cfg.resolved_warmup(200) == 20 + + +def test_warmup_steps_used_when_no_ratio(): + assert TrainConfig(method="lora", warmup_steps=7).resolved_warmup(200) == 7 + + +def test_warmup_never_negative(): + assert TrainConfig(method="lora", warmup_steps=-5).resolved_warmup(100) == 0 + + +# ---- eval interval ------------------------------------------------------------- + +def test_eval_steps_defaults_to_quarter_of_the_run(): + assert TrainConfig(method="lora").resolved_eval_steps(100) == 25 + + +def test_eval_steps_fraction_scales_with_total(): + assert TrainConfig(method="lora", eval_steps=0.5).resolved_eval_steps(80) == 40 + + +def test_eval_steps_absolute_is_kept(): + assert TrainConfig(method="lora", eval_steps=13).resolved_eval_steps(80) == 13 + + +def test_eval_steps_never_zero(): + assert TrainConfig(method="lora").resolved_eval_steps(1) == 1 + assert TrainConfig(method="lora", eval_steps=0.001).resolved_eval_steps(10) == 1 + + +# ---- total steps ---------------------------------------------------------------- + +def test_max_steps_wins_over_epochs(): + cfg = TrainConfig(method="lora", max_steps=42, num_train_epochs=10) + assert resolve_total_steps(cfg, n_examples=1000) == 42 + + +def test_epochs_scale_with_rows_and_batching(): + cfg = TrainConfig(method="lora", max_steps=None, num_train_epochs=2, + per_device_train_batch_size=2, + gradient_accumulation_steps=1) + assert resolve_total_steps(cfg, n_examples=10) == 10 # 10/2 per epoch × 2 + + +def test_total_steps_is_at_least_one(): + cfg = TrainConfig(method="lora", max_steps=None, num_train_epochs=1, + per_device_train_batch_size=64) + assert resolve_total_steps(cfg, n_examples=1) == 1 + + +# ---- round trip ------------------------------------------------------------------- + +def test_config_dict_roundtrip_preserves_fields(): + cfg = TrainConfig(method="dpo", learning_rate=5e-6, lora_r=32, + target_modules="attention", report_to=["wandb"], beta=0.2) + back = TrainConfig(**cfg.to_dict()) + assert back.method == "dpo" + assert back.learning_rate == 5e-6 + assert back.lora_r == 32 + assert back.beta == 0.2 + assert back.resolved_target_modules() == ATTENTION_MODULES + + +def test_metric_survives_the_wire(): + m = Metric(step=3, loss=1.5, lr=1e-4, grad_norm=0.9, epoch=0.5) + assert Metric(**vars(m)) == m From 7eac2d2430fca129d535a61be3f24a6c49156b12 Mon Sep 17 00:00:00 2001 From: Karteek Yadavilli <21700785+thedataengineer@users.noreply.github.com> Date: Tue, 28 Jul 2026 16:26:12 -0400 Subject: [PATCH 08/10] docs: point at files that exist, stop calling shipped features unbuilt MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit README told people to run examples/quickstart.py and linked a Colab notebook — neither file is in the repo. Both now point at examples/, which does exist. SHADOWLM_EXTRAS=cli was documented as a lighter UI-only install in README and install.sh, but [cli] resolves to nothing; _cli_entry.py went further and recommended it as the fix for a missing Typer, which it cannot be. eval.py, traces.py, and the worker fleet had no README presence at all while the roadmap listed eval gates and a GPU job queue as future work. They're in the feature table now and the roadmap says what's actually left. CLAUDE.md gains the worker/ws tier (it still described serve.py as a one-job reference server), traces.py, eval.py, the three auth modes and machine tokens, the remote pool, and the two new UI pages. numpy is imported directly by more.py and more_plus.py and was riding in transitively through torch — declared now. models.py's method list was missing more_plus, and six CLI help strings plus load()'s docstring listed only auto|mlx|torch when remote and verl are accepted too. Closes open-gitagent/shadowLM#10 --- CLAUDE.md | 42 ++++++++++++++++++++++++++++++++---------- README.md | 25 +++++++++++++++---------- install.sh | 2 +- pyproject.toml | 3 +++ shadowlm/_cli_entry.py | 13 ++++++------- shadowlm/cli.py | 12 ++++++------ shadowlm/models.py | 6 ++++-- 7 files changed, 67 insertions(+), 36 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index eb2e274..f908f0d 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -28,19 +28,22 @@ make check # compileall the package + `tsc -b` the frontend make build # build the frontend, then the wheel+sdist, then twine check make release # bump patch (or BUMP=minor/major, V=x.y.z), build, tag, push -pytest # the CPU test suite (tests/, excludes gpu/) +make test # the CPU test suite — exactly what CI runs pytest tests/test_more_plus_router.py # a single test file pytest tests/test_more_plus_router.py::test_name -x # a single test, stop on first fail make gpu-test # the CUDA verification suite — run on a GPU box only ``` `tests/gpu/` (CUDA, real GPU) is **not** part of the default `pytest` run — it's -invoked explicitly via `make gpu-test` or `python tests/gpu/test_cuda.py`. +invoked explicitly via `make gpu-test` or `python tests/gpu/test_cuda.py`, and +it exits non-zero if every test skipped (no CUDA) rather than reading as a pass. -Releases publish to PyPI via `.github/workflows/publish.yml` on a `v*` tag. The -CI gate requires the version to match in **three** places: the git tag, -`pyproject.toml`, and `shadowlm/__init__.py`. `make bump`/`make release` keep -the latter two in sync — never edit only one. +Two workflows: `.github/workflows/test.yml` runs the CPU suite on 3.10 + 3.12 +for every push and PR; `.github/workflows/publish.yml` builds on every push to +main and publishes to PyPI on a `v*` tag. The publish gate requires the version +to match in **three** places: the git tag, `pyproject.toml`, and +`shadowlm/__init__.py`. `make bump`/`make release` keep the latter two in sync — +never edit only one. ## Architecture @@ -88,10 +91,16 @@ touches one file and no others. records an unmodified agent's traffic, reconstructing message-level trajectories (calls that extend a prior call's message prefix merge into one episode; use an `x-session-id` header to disambiguate interleaved conversations). +- `traces.py` — the offline sibling of `capture.py`: ingests OpenTelemetry GenAI + spans (OTLP JSON, event/spec/OpenInference/raw-wire shapes), groups them into + conversations, and `to_dataset()`s them. For agents already instrumented — + no proxy in the path. - `rl.py` — `Trajectory` / `TrajectoryGroup` / `judge_group` (LLM-judge scoring), fed into `method="grpo"`. - `apo.py` — `optimize_prompt()`: optimize the prompt instead of weights, same capture/judge front end, no GPU. +- `eval.py` — `slm.evaluate()` / `shadowlm eval`: score a model on a held-out set + (exact / contains / numeric / JSON / LLM-judge scorers). ### Signature methods (MoRE) @@ -105,14 +114,27 @@ routing; its run progress is one step per unit (see `resolve_total_steps`). - `serve.py` — `python -m shadowlm.serve` / `shadowlm serve`. Pure-stdlib (`http.server` + threads) reference server: trains on **this machine's real backend** (no mock), streams metrics, ships adapters as tar.gz, serves the - built React UI from `_static`. One job at a time — honest reference tier. + built React UI from `_static`. It runs one local job at a time, and doubles as + the **hub** for the worker tier below (per-worker inboxes, job routing, + inference forwarding). +- `worker.py` + `ws.py` — the fleet tier. `shadowlm worker --hub ` makes a + NAT'd machine dial *out* and hold one websocket open: jobs arrive on it, logs + and metrics stream back up it, cancels land mid-step, and playground + chat/generate is proxied to whichever worker holds the adapter. The trained + adapter ships home over plain HTTP (`POST /v1/workers//jobs//artifact`). + `ws.py` is a minimal RFC 6455 implementation — no fragmentation, no + extensions, capped frame size. - `remote.py` — the typed client for that JSON protocol (`/v1/finetunes`, …). - Same protocol backs `backend="remote"` and ShadowLM Studio. + Same protocol backs `backend="remote"` and ShadowLM Studio. `SHADOWLM_API_URL` + may list several servers; `pick()` binds to the least-busy reachable one for + the session (client-side routing, deliberately not a scheduler). - `frontend/` — React 19 + Vite + Tailwind v4 studio. `npm run build` outputs to `../shadowlm/_static` (the wheel ships the compiled UI; end users never need node). `frontend/src/api.ts` is the typed mirror of the remote protocol. The - pages (Datasets → Models → Train → Runs → Playground) are the capture→train→own - loop as a UI. Auth: studio routes are gated by username/password. + pages (Dashboard · Datasets → Models → Train → Runs → Playground · Machines) + are the capture→train→own loop as a UI. Auth has three modes — `password`, + `apikey`, or `none` (`GET /v1/auth` reports which) — plus long-lived, hashed, + individually-revocable **machine tokens** that workers authenticate with. ### The shadow accelerator (`accel.py`) diff --git a/README.md b/README.md index 3af37a9..3d558ee 100644 --- a/README.md +++ b/README.md @@ -72,6 +72,7 @@ The whole **capture → judge → train → own a shadowLM** loop runs on these: | Block | What it does | API | |-------|--------------|-----| | **Capture proxy** | drop-in OpenAI endpoint that records your agent's traffic into trajectories — agent unchanged | `slm.capture()` | +| **Trace ingestion** | already have OpenTelemetry GenAI spans? turn an OTLP dump into a training set, no proxy | `slm.traces.to_dataset()` | | **13 methods** | LoRA · QLoRA · DoRA · full · CPT · DPO · GRPO · MoRE · MoRE+ · BitFit · prompt · p-tuning · adapter | `method=` | | **Judge → train** | score episodes with an LLM judge, train with trajectory-GRPO or DPO | `judge_group` | | **APO** | optimize the *prompt* instead of weights — same capture/judge front end, no GPU | `slm.optimize_prompt()` | @@ -80,7 +81,9 @@ The whole **capture → judge → train → own a shadowLM** loop runs on these: | **Any hardware** | CUDA · TPU · Trainium · Intel · Apple · CPU (whatever HF accelerate targets) | `device=` | | **Shadow accelerator** | 4-bit, grad checkpointing, flash-attn, fused optimizer, optional Liger kernels — logged, never silent | `accelerator="shadow"` | | **Checkpoints** | save every N steps, then load or A/B any version — `step 200` vs `final` — in the playground | `save_steps=` · `run.checkpoint_at(step)` | +| **Eval** | score a model on a held-out set — exact/contains/numeric/JSON/LLM-judge scorers | `slm.evaluate()` · `shadowlm eval` | | **Remote + server** | train on a GPU box or fleet over one JSON protocol; metrics stream back | `backend="remote"` · `shadowlm serve` | +| **Worker fleet** | a NAT'd GPU box dials the hub over one websocket and takes jobs; the studio routes to it | `shadowlm worker --hub …` | | **Studio** | datasets → models → guided train → live runs (charts + console) → playground compare | `shadowlm serve` → `/` | | **CLI** | finetune / runs / plot / chat / export / methods from the shell | `shadowlm …` | | **Own the weights** | adapter/merged export, run records that survive restarts, nothing leaves your box | `model.save()` | @@ -139,8 +142,9 @@ curl -fsSL https://install.shadowlm.sh | sh It detects your hardware and installs the matching stack — Apple Silicon → mlx, NVIDIA → torch + Liger fused kernels, otherwise torch CPU — into an isolated env in `~/.shadowlm/venv`, then launches `shadowlm serve` at `http://127.0.0.1:8329`. -Re-run any time to upgrade. Override with `SHADOWLM_EXTRAS=cli` (UI only), -`SHADOWLM_PORT=…`, or `SHADOWLM_NO_SERVE=1` (install without launching). +Re-run any time to upgrade. Override with `SHADOWLM_EXTRAS=…` (pip extras — +NVIDIA gets `kernels` by default), `SHADOWLM_PORT=…`, or `SHADOWLM_NO_SERVE=1` +(install without launching). Or with pip — `pip install shadowlm` ships the full training stack (torch + HuggingFace, retrieval, CLI). On Apple Silicon the mlx dev backend is pulled in @@ -154,12 +158,12 @@ automatically. Two extras stay opt-in for specialized hardware: ```bash git clone https://github.com/open-gitagent/shadowLM && cd shadowLM python3 -m venv .venv && source .venv/bin/activate && pip install -e . -python examples/quickstart.py # datasets → finetune → inference, end to end +python examples/mlx/lora.py # datasets → finetune → inference, end to end ``` -No hardware handy? Test-drive the whole thing — checkpoints, faiss MoRE, APO — -on a free Colab GPU: -[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/open-gitagent/shadowLM/blob/main/examples/colab_test_drive.ipynb) +`examples/` has one runnable script per (backend × method) — `mlx/` runs a 0.5B +model in seconds, `torch/` and `remote/` use Qwen3-8B. See +[examples/README.md](./examples/README.md) for the full matrix. Run output (mlx, a 0.5B model, ~3.5s): @@ -204,16 +208,17 @@ The engine ships first; **ShadowLM Studio** (the hosted tier) wraps this exact API — nothing reimplemented — to turn the blocks into a one-click migration: - **Decision inbox** — captured traces surfaced for human approve/correct into chosen-vs-rejected pairs (today: auto-scored by an LLM judge). -- **Eval gates** — advance only when quality holds *and* savings beat cost: task-level evals + cost-per-task on the run records. +- **Cost gates** — `slm.evaluate` already scores a model on a held-out set; the gate that ties quality to cost-per-task and blocks the switch is Studio's. - **Shadow router** — the capture proxy evolved: run the shadowLM in parallel behind the live agent, then shift traffic % frontier → owned. -- **Fleet + teams** — GPU job queue, shared run history, dataset/adapter registry. +- **Teams** — shared run history, dataset/adapter registry, org auth on top of the fleet below. ``` [x] SDK — datasets → finetune → inference on mlx / torch / remote [x] 13 methods incl. MoRE, MoRE+ (decoupled MoE), trajectory GRPO, judge rewards -[x] Capture proxy · shadow accelerator · any-hardware +[x] Capture proxy · OTLP trace ingestion · shadow accelerator · any-hardware [x] Remote backend + reference server + the studio dashboard + CLI -[ ] Studio orchestration — decision inbox · eval gates · shadow router · switch +[x] Eval scorers (`slm.evaluate`, `shadowlm eval`) · worker fleet (`shadowlm worker`) +[ ] Studio orchestration — decision inbox · cost gates · shadow router · switch ``` ## Contributing diff --git a/install.sh b/install.sh index 94b265c..fc916b0 100755 --- a/install.sh +++ b/install.sh @@ -6,7 +6,7 @@ # Installs shadowlm into an isolated venv (~/.shadowlm/venv), picks the right # training backend for your machine, and opens the studio. Re-run any time to # upgrade. Override with env vars: -# SHADOWLM_EXTRAS=cli # lighter: just the UI/CLI, add a backend later +# SHADOWLM_EXTRAS=kernels # pip extras to add (NVIDIA gets this by default) # SHADOWLM_PORT=8329 # studio port # SHADOWLM_NO_SERVE=1 # install only, don't launch set -eu diff --git a/pyproject.toml b/pyproject.toml index 52b25d2..2e4a5d8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,9 @@ dependencies = [ # mixture of retrieval experts (method="more") — embeddings + faiss index "sentence-transformers>=3.0", "faiss-cpu>=1.8", + # more.py / more_plus.py use it directly (index math, routing scores) rather + # than relying on it arriving via torch/faiss + "numpy>=1.24", # the shadowlm CLI (Typer + Rich) "typer>=0.12", "rich>=13", diff --git a/shadowlm/_cli_entry.py b/shadowlm/_cli_entry.py index 751d921..9cbe9b8 100644 --- a/shadowlm/_cli_entry.py +++ b/shadowlm/_cli_entry.py @@ -1,9 +1,8 @@ """Console-script entry point — pure stdlib, always importable. -The real CLI (`shadowlm.cli`) is built on Typer + Rich, which ship in the -`[cli]` extra (and in `[all]` / `[mlx-all]`). This shim lets the bare -`shadowlm` command degrade to a friendly message instead of a traceback when -those aren't installed. +The real CLI (`shadowlm.cli`) is built on Typer + Rich, which the base install +ships. This shim lets the bare `shadowlm` command degrade to a friendly message +instead of a traceback if they're somehow missing (a partial or pruned install). """ from __future__ import annotations @@ -17,9 +16,9 @@ def main() -> int: except ModuleNotFoundError as e: if e.name in ("typer", "rich", "click", "click.core"): print( - "The shadowlm CLI needs Typer + Rich:\n\n" - " pip install 'shadowlm[cli]'\n\n" - "(included in 'shadowlm[all]' and 'shadowlm[mlx-all]')", + f"The shadowlm CLI needs Typer + Rich ({e.name} is missing).\n\n" + " pip install --upgrade --force-reinstall shadowlm\n\n" + "They ship in the base install, so this means a partial one.", file=sys.stderr, ) return 1 diff --git a/shadowlm/cli.py b/shadowlm/cli.py index d339ff7..50feb74 100644 --- a/shadowlm/cli.py +++ b/shadowlm/cli.py @@ -208,7 +208,7 @@ def finetune( eval: Annotated[Optional[str], typer.Option(help='eval dataset, or "auto"/"15%"/0.15 to hold out a fraction')] = None, save: Annotated[Optional[str], typer.Option(help="also export the adapter to DIR")] = None, output_dir: Annotated[Optional[str], typer.Option("--output-dir")] = None, - backend: Annotated[Optional[str], typer.Option(help="auto | mlx | torch")] = None, + backend: Annotated[Optional[str], typer.Option(help="auto | mlx | torch | remote | verl")] = None, accelerator: Annotated[Optional[str], typer.Option(help="auto | shadow | none")] = None, device: Annotated[Optional[str], typer.Option(help="auto | cuda | cpu (torch)")] = None, load_in_4bit: Annotated[Optional[bool], typer.Option("--load-in-4bit/--no-load-in-4bit")] = None, @@ -302,7 +302,7 @@ def generate( prompt: Annotated[str, typer.Argument(help="the prompt to complete")], model: Annotated[Optional[str], typer.Option("--model", "-m", help="base model override for adapter dirs")] = None, - backend: Annotated[str, typer.Option(help="auto | mlx | torch")] = "auto", + backend: Annotated[str, typer.Option(help="auto | mlx | torch | remote | verl")] = "auto", load_in_4bit: Annotated[bool, typer.Option("--load-in-4bit")] = False, temperature: Annotated[float, typer.Option()] = 0.7, max_new_tokens: Annotated[int, typer.Option("--max-new-tokens")] = 256, @@ -320,7 +320,7 @@ def chat( target: Annotated[str, typer.Argument(help="model name, or an adapter directory")], model: Annotated[Optional[str], typer.Option("--model", "-m", help="base model override for adapter dirs")] = None, - backend: Annotated[str, typer.Option(help="auto | mlx | torch")] = "auto", + backend: Annotated[str, typer.Option(help="auto | mlx | torch | remote | verl")] = "auto", load_in_4bit: Annotated[bool, typer.Option("--load-in-4bit")] = False, system: Annotated[Optional[str], typer.Option(help="system prompt")] = None, temperature: Annotated[float, typer.Option()] = 0.7, @@ -366,7 +366,7 @@ def export( help="adapter | merged")] = "adapter", model: Annotated[Optional[str], typer.Option("--model", "-m", help="base model override")] = None, - backend: Annotated[str, typer.Option(help="auto | mlx | torch")] = "auto", + backend: Annotated[str, typer.Option(help="auto | mlx | torch | remote | verl")] = "auto", hf_token: Annotated[Optional[str], typer.Option("--hf-token", envvar="HF_TOKEN")] = None, ): """Export a trained adapter — re-save it, or merge it into the base weights.""" @@ -391,7 +391,7 @@ def evaluate_cmd( show: Annotated[int, typer.Option(help="how many worst examples to show")] = 5, model: Annotated[Optional[str], typer.Option("--model", "-m", help="base model override for adapter dirs")] = None, - backend: Annotated[str, typer.Option(help="auto | mlx | torch")] = "auto", + backend: Annotated[str, typer.Option(help="auto | mlx | torch | remote | verl")] = "auto", load_in_4bit: Annotated[bool, typer.Option("--load-in-4bit")] = False, max_new_tokens: Annotated[int, typer.Option("--max-new-tokens")] = 256, hf_token: Annotated[Optional[str], typer.Option("--hf-token", envvar="HF_TOKEN")] = None, @@ -494,7 +494,7 @@ def plot( def serve( host: Annotated[str, typer.Option()] = "127.0.0.1", port: Annotated[int, typer.Option()] = 8329, - backend: Annotated[str, typer.Option(help="auto | mlx | torch")] = "auto", + backend: Annotated[str, typer.Option(help="auto | mlx | torch | remote | verl")] = "auto", accelerator: Annotated[str, typer.Option(help="auto | shadow | none")] = "auto", device: Annotated[str, typer.Option(help="auto | cuda | cpu (torch)")] = "auto", work_dir: Annotated[Optional[str], typer.Option("--work-dir")] = None, diff --git a/shadowlm/models.py b/shadowlm/models.py index a49ff82..1895728 100644 --- a/shadowlm/models.py +++ b/shadowlm/models.py @@ -148,7 +148,8 @@ def finetune( """Finetune on `dataset`, returning a `TrainingRun`. method: any registered training method — built-ins are lora, qlora, - dora, full, cpt, dpo, grpo, more, bitfit, prompt, ptuning, and adapter. + dora, full, cpt, dpo, grpo, more, more_plus, bitfit, prompt, ptuning, + and adapter. See `shadowlm.methods` to list or register methods. The default learning rate comes from the method's spec. @@ -303,7 +304,8 @@ def load( ) -> Model: """Load a model. - backend: "auto" (CUDA→torch, Apple→mlx, else torch-CPU) | "mlx" | "torch". + backend: "auto" (CUDA→torch, Apple→mlx, else torch-CPU) | "mlx" | "torch" | + "remote" (a ShadowLM server) | "verl" (multi-GPU GRPO). accelerator: the shadow optimization layer — "auto" | "shadow" | "none". device: "auto" | "cuda" | "cpu" (torch backend; mlx is always the Metal GPU). adapter: path to a previously-trained LoRA checkpoint to attach. From 204ee33e57aac1567929fbd1169ee1f47d1edb62 Mon Sep 17 00:00:00 2001 From: Karteek Yadavilli <21700785+thedataengineer@users.noreply.github.com> Date: Tue, 28 Jul 2026 16:29:09 -0400 Subject: [PATCH 09/10] errors: stop swallowing failures the caller needs to see MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The capture proxy read Content-Length and called model.chat() with no handling, so bad JSON or a failed generation killed the connection. An agent under capture is a real OpenAI client: it retries a JSON error and surfaces the message, but a dropped socket reads as a transport fault and usually takes the run down. Both now return OpenAI-shaped error bodies (400 / 500) and the proxy keeps serving. It also records the OS-assigned port when constructed with port=0, so base_url is right. do_POST's blanket except turned client typos into 500 KeyError: 'model' — /v1/generate, /v1/chat and /v1/prewarm now 422 naming the missing field, and chat validates that messages is a list. torch.has_cuda() swallowed a broken CUDA install and silently ran on CPU; it says so now. The two cuda.empty_cache() swallows in serve.py log what threw, and clear_vram reports the error instead of returning a clean-sweep result it didn't achieve. Closes open-gitagent/shadowLM#8 --- shadowlm/backends/torch.py | 6 +- shadowlm/capture.py | 26 +++++++- shadowlm/serve.py | 35 ++++++++-- tests/test_capture_errors.py | 117 +++++++++++++++++++++++++++++++++ tests/test_serve_validation.py | 61 +++++++++++++++++ 5 files changed, 237 insertions(+), 8 deletions(-) create mode 100644 tests/test_capture_errors.py create mode 100644 tests/test_serve_validation.py diff --git a/shadowlm/backends/torch.py b/shadowlm/backends/torch.py index d8f9f42..30ce794 100644 --- a/shadowlm/backends/torch.py +++ b/shadowlm/backends/torch.py @@ -71,7 +71,11 @@ def has_cuda(cls) -> bool: try: import torch # noqa: PLC0415 return torch.cuda.is_available() - except Exception: + except Exception as e: # noqa: BLE001 — a broken CUDA install must not + # abort backend selection, but falling back to CPU in silence is how + # a "why is this so slow" afternoon starts. Say what happened. + print(f"[torch] CUDA probe failed ({type(e).__name__}: {e}) — " + "treating this box as CPU-only", flush=True) return False def load(self, name, *, load_in_4bit=False, max_seq_length=2048, adapter=None, **kwargs) -> None: diff --git a/shadowlm/capture.py b/shadowlm/capture.py index ad00e8e..eae9560 100644 --- a/shadowlm/capture.py +++ b/shadowlm/capture.py @@ -106,12 +106,33 @@ def do_GET(self): # noqa: N802 — clients health-check /v1/models else: self.send_error(404) + def _error(self, code: int, message: str, kind: str) -> None: + """An OpenAI-shaped error body. Harnesses under capture are real + OpenAI clients: they handle a JSON error and retry, but a dropped + connection reads as a transport fault and usually kills the run.""" + data = json.dumps({"error": {"message": message, "type": kind, + "param": None, "code": None}}).encode() + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + def do_POST(self): # noqa: N802 (http.server API) if not self.path.rstrip("/").endswith("/chat/completions"): self.send_error(404) return - body = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) - reply = proxy._serve(body, self.headers.get("x-session-id")) + try: + length = int(self.headers.get("Content-Length") or 0) + body = json.loads(self.rfile.read(length) or b"") + if not isinstance(body, dict): + raise ValueError("request body must be a JSON object") + except ValueError as e: + return self._error(400, str(e), "invalid_request_error") + try: + reply = proxy._serve(body, self.headers.get("x-session-id")) + except Exception as e: # noqa: BLE001 — the agent gets the error, not a dead socket + return self._error(500, f"{type(e).__name__}: {e}", "server_error") if body.get("stream"): # synthetic provider-shaped stream from the completed # response — most real harnesses request streaming @@ -131,6 +152,7 @@ def do_POST(self): # noqa: N802 (http.server API) self.wfile.write(data) self._server = ThreadingHTTPServer((self.host, self.port), Handler) + self.port = self._server.server_address[1] # port=0 → whatever the OS gave us self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) self._thread.start() return self diff --git a/shadowlm/serve.py b/shadowlm/serve.py index 397a55f..747481f 100644 --- a/shadowlm/serve.py +++ b/shadowlm/serve.py @@ -844,8 +844,11 @@ def _load(): if torch.cuda.is_available(): torch.cuda.empty_cache() - except Exception: # noqa: BLE001 - pass + except Exception as e: # noqa: BLE001 — the retry below is the real + # recovery; but if the cache never cleared, the retry's OOM is a + # consequence of this, so leave a trail. + print(f"[serve] VRAM release failed ({type(e).__name__}: {e}) — " + "retrying the load anyway", flush=True) m = _load() if len(self._infer_cache) >= self._infer_cache_cap: self._infer_cache.pop(next(iter(self._infer_cache))) # oldest out @@ -871,6 +874,7 @@ def clear_vram(self) -> dict: import gc # noqa: PLC0415 before = self._gpu_used_mb() + freed_error: str | None = None with self._model_lock: n = len(self._infer_cache) self._infer_cache.clear() @@ -880,10 +884,14 @@ def clear_vram(self) -> dict: if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.synchronize() - except Exception: # noqa: BLE001 - pass + except Exception as e: # noqa: BLE001 — report, don't fail the request + freed_error = f"{type(e).__name__}: {e}" + print(f"[serve] VRAM release failed ({freed_error})", flush=True) gc.collect() - return {"unloaded": n, "before_mb": before, "after_mb": self._gpu_used_mb()} + out = {"unloaded": n, "before_mb": before, "after_mb": self._gpu_used_mb()} + if freed_error: # don't report a clean sweep when the allocator threw + out["error"] = freed_error + return out # ---- operations ------------------------------------------------------------ def submit(self, payload: dict) -> str: @@ -1262,6 +1270,15 @@ def _read_body(self) -> bytes: f"request body of {length} bytes exceeds the {_MAX_BODY} cap") return self.rfile.read(length) + def _require(self, body: dict, *fields: str) -> bool: + """422 naming the first missing/empty field. Without this the route + table's blanket except turns a client typo into `500 KeyError`.""" + for f in fields: + if not body.get(f): + self._error(422, f"missing field {f!r}") + return False + return True + def _job_or_404(self, job_id: str): job = server.jobs.get(job_id) if job is None: @@ -1508,10 +1525,14 @@ def do_POST(self): # noqa: N802 self._send(200, {"ok": True}) elif parts == ["v1", "prewarm"]: b = self._body() + if not self._require(b, "model"): + return self._send(200, server.prewarm( b["model"], b.get("adapter"), b.get("checkpoint"))) elif parts == ["v1", "generate"]: b = self._body() + if not self._require(b, "prompt", "model"): + return # a worker-trained shadow answers on its own machine — the # adapter format matches that backend, not necessarily ours wjob = server.jobs.get(b.get("adapter") or "") @@ -1532,6 +1553,10 @@ def do_POST(self): # noqa: N802 self._send(200, {"text": text}) elif parts == ["v1", "chat"]: b = self._body() + if not self._require(b, "messages", "model"): + return + if not isinstance(b["messages"], list): + return self._error(422, "field 'messages' must be a list") wjob = server.jobs.get(b.get("adapter") or "") if wjob is not None and wjob.worker: self._send(200, {"text": server.worker_infer(wjob, { diff --git a/tests/test_capture_errors.py b/tests/test_capture_errors.py new file mode 100644 index 0000000..4cd4405 --- /dev/null +++ b/tests/test_capture_errors.py @@ -0,0 +1,117 @@ +"""The capture proxy answers with an OpenAI-shaped error body instead of +dropping the connection. + +Agents under capture are real OpenAI clients: they retry a 4xx/5xx JSON error +and surface a useful message, but a killed connection reads as a transport +fault and usually takes the agent down with it. +""" + +from __future__ import annotations + +import http.client +import json + +import pytest + +from shadowlm.capture import CaptureProxy + + +class _StubModel: + """Stands in for a loaded Model — chat() either answers or explodes.""" + + name = "stub/model" + + def __init__(self, boom: Exception | None = None) -> None: + self._boom = boom + + def chat(self, messages, **kw): + if self._boom: + raise self._boom + + class _Reply: + content = "hi" + raw = "hi" + tool_calls = None + + def to_message(self): + return {"role": "assistant", "content": "hi"} + + return _Reply() + + +@pytest.fixture() +def proxy_for(): + live = [] + + def _make(model): + p = CaptureProxy(model, port=0).start() + live.append(p) + return p + + yield _make + for p in live: + p.stop() + + +def _post(proxy, body: bytes | str, *, path="/v1/chat/completions", + content_length=None): + host, port = proxy.host, proxy.port + conn = http.client.HTTPConnection(host, port, timeout=10) + try: + raw = body if isinstance(body, bytes) else body.encode() + headers = {"Content-Type": "application/json"} + if content_length is not None: + headers["Content-Length"] = str(content_length) + conn.request("POST", path, body=raw, headers=headers) + resp = conn.getresponse() + return resp.status, resp.read() + finally: + conn.close() + + +def test_happy_path_still_answers(proxy_for): + proxy = proxy_for(_StubModel()) + status, payload = _post(proxy, json.dumps({"messages": [ + {"role": "user", "content": "hi"}]})) + assert status == 200 + assert json.loads(payload)["choices"][0]["message"]["content"] == "hi" + + +def test_model_failure_becomes_a_json_error_not_a_dropped_socket(proxy_for): + proxy = proxy_for(_StubModel(boom=RuntimeError("CUDA out of memory"))) + status, payload = _post(proxy, json.dumps({"messages": [ + {"role": "user", "content": "hi"}]})) + assert status == 500 + body = json.loads(payload) + assert "CUDA out of memory" in body["error"]["message"] + assert body["error"]["type"] + + +def test_malformed_json_is_a_400(proxy_for): + proxy = proxy_for(_StubModel()) + status, payload = _post(proxy, "{not json") + assert status == 400 + assert json.loads(payload)["error"]["message"] + + +def test_missing_content_length_is_a_400(proxy_for): + proxy = proxy_for(_StubModel()) + status, _ = _post(proxy, b"", content_length=None) + assert status == 400 + + +def test_a_failed_call_is_not_recorded_as_a_trajectory(proxy_for): + proxy = proxy_for(_StubModel(boom=RuntimeError("boom"))) + _post(proxy, json.dumps({"messages": [{"role": "user", "content": "hi"}]})) + assert proxy.trajectories() == [] + + +def test_the_proxy_keeps_serving_after_a_failure(proxy_for): + model = _StubModel(boom=RuntimeError("boom")) + proxy = proxy_for(model) + _post(proxy, json.dumps({"messages": [{"role": "user", "content": "hi"}]})) + model._boom = None # the transient failure clears + status, payload = _post(proxy, json.dumps({"messages": [ + {"role": "user", "content": "hi"}]})) + assert status == 200 + assert json.loads(payload)["choices"][0]["message"]["content"] == "hi" diff --git a/tests/test_serve_validation.py b/tests/test_serve_validation.py new file mode 100644 index 0000000..5c154a2 --- /dev/null +++ b/tests/test_serve_validation.py @@ -0,0 +1,61 @@ +"""Missing request fields are 422s naming the field, not 500 KeyErrors. + +The whole do_POST route table sits inside one `except Exception` that renders +as a 500 — so a client typo used to come back as `500 KeyError: 'model'`, which +reads like a server bug and tells the caller nothing. +""" + +from __future__ import annotations + +import json +import threading +import urllib.error +import urllib.request +from http.server import ThreadingHTTPServer + +import pytest + +from shadowlm.serve import Auth, Server, make_handler + + +@pytest.fixture() +def hub(tmp_path): + server = Server(backend="auto", accelerator="auto", device="auto", + work_root=tmp_path) + httpd = ThreadingHTTPServer( + ("127.0.0.1", 0), + make_handler(server, Auth(user="admin", password=None, api_key=None))) + threading.Thread(target=httpd.serve_forever, daemon=True).start() + yield httpd.server_address[1] + httpd.shutdown() + + +def _post(port: int, path: str, body: dict) -> tuple[int, dict]: + req = urllib.request.Request( + f"http://127.0.0.1:{port}{path}", method="POST", + data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}) + try: + with urllib.request.urlopen(req, timeout=10) as resp: + return resp.status, json.loads(resp.read() or b"{}") + except urllib.error.HTTPError as e: + return e.code, json.loads(e.read() or b"{}") + + +@pytest.mark.parametrize("path,body,missing", [ + ("/v1/generate", {}, "prompt"), + ("/v1/generate", {"prompt": "hi"}, "model"), + ("/v1/chat", {}, "messages"), + ("/v1/chat", {"messages": [{"role": "user", "content": "hi"}]}, "model"), + ("/v1/prewarm", {}, "model"), +]) +def test_missing_field_is_a_422_naming_it(hub, path, body, missing): + status, payload = _post(hub, path, body) + assert status == 422, payload + assert missing in payload.get("error", "") + + +def test_chat_messages_must_be_a_list(hub): + status, payload = _post(hub, "/v1/chat", {"model": "m", "messages": "hi"}) + assert status == 422 + assert "messages" in payload.get("error", "") From acb29b2670decccd9b2d51e74973bd6fc26b94eb Mon Sep 17 00:00:00 2001 From: Karteek Yadavilli <21700785+thedataengineer@users.noreply.github.com> Date: Tue, 28 Jul 2026 16:38:29 -0400 Subject: [PATCH 10/10] protocol: close the client/server gaps in remote.py and api.ts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit submit_finetune never sent `name`, so every SDK-submitted run showed up nameless in the studio's run list, and it had no dataset_id path even though the server accepts one — you couldn't train on a dataset the studio already holds. Both are parameters now, with a check that one of rows/dataset_id is present instead of a server-side 422. The remote backend passes a name built from the model and method. pick() ranked on running+pending+gpus and ignored the `workers` count capacity() reports, so a hub with four idle workers tied with a lone box. Workers now break the tie ahead of GPU count. api.ts: getHealth was typed {ok, backend, version} while the server also returns gpus/running/pending/workers; submitFinetune took `object`, the one POST body worth typing; login dropped user/expires; MethodInfo's adapter comment was missing more_plus. Train.tsx narrows model/ds explicitly — `ready` already guaranteed both, but not to tsc. Closes open-gitagent/shadowLM#9 --- frontend/src/api.ts | 37 ++++++++++++-- frontend/src/pages/Train.tsx | 2 +- shadowlm/backends/remote.py | 2 + shadowlm/remote.py | 27 +++++++--- shadowlm/serve.py | 2 +- tests/test_remote_protocol.py | 92 +++++++++++++++++++++++++++++++++++ 6 files changed, 149 insertions(+), 13 deletions(-) create mode 100644 tests/test_remote_protocol.py diff --git a/frontend/src/api.ts b/frontend/src/api.ts index f271648..ebd0659 100644 --- a/frontend/src/api.ts +++ b/frontend/src/api.ts @@ -38,7 +38,8 @@ export interface MethodInfo { description: string; default_lr: number; trainer: string; - adapter: string; // lora | dora | more | bottleneck | bitfit | prompt | ptuning | none + // lora | dora | more | more_plus | bottleneck | bitfit | prompt | ptuning | none + adapter: string; } export interface JobSummary { @@ -101,6 +102,8 @@ export interface AuthInfo { auth_required: boolean; mode: "password" | "apikey" export const getAuthInfo = () => fetch("/v1/auth").then((r) => r.json() as Promise); +export interface LoginResponse { token: string; user: string; expires: number; } + export async function login(username: string, password: string): Promise { const r = await fetch("/v1/login", { method: "POST", headers: { "Content-Type": "application/json" }, @@ -110,14 +113,23 @@ export async function login(username: string, password: string): Promise { const d = await r.json().catch(() => ({} as { error?: string })); throw new Error(d.error || "login failed"); } - const { token } = (await r.json()) as { token: string }; + const { token } = (await r.json()) as LoginResponse; apiKey.set(token); } export const logout = () => apiKey.clear(); -export const getHealth = () => - api<{ ok: boolean; backend: string; version: string }>("/v1/health"); +// the server spreads capacity() into this — fleet depth, not just liveness +export interface Health { + ok: boolean; + backend: string; + version: string; + gpus: number; + running: number; + pending: number; + workers: number; +} +export const getHealth = () => api("/v1/health"); export const getDatasets = () => api<{ datasets: DatasetMeta[] }>("/v1/datasets"); export const getDataset = (id: string) => api(`/v1/datasets/${id}`); @@ -194,7 +206,22 @@ export const getCheckpoints = (id: string) => api<{ checkpoints: Checkpoint[] }>(`/v1/finetunes/${id}/checkpoints`); export const cancelJob = (id: string) => api<{ ok: boolean }>(`/v1/finetunes/${id}/cancel`, { method: "POST" }); -export const submitFinetune = (body: object) => +export interface DatasetPayload { rows: unknown[]; format?: string } + +// what POST /v1/finetunes accepts — either `dataset` rows or a `dataset_id`. +// eval_dataset is rows, a holdout ("20%" or a 0–1 fraction), or null for none. +export interface FinetuneRequest { + base_model: string; + config: Record; + name?: string; + dataset?: DatasetPayload | null; + dataset_id?: string | null; + eval_dataset?: DatasetPayload | string | number | null; + worker?: string | null; + load_in_4bit?: boolean; + max_seq_length?: number; +} +export const submitFinetune = (body: FinetuneRequest) => api<{ job_id: string }>("/v1/finetunes", { method: "POST", body: JSON.stringify(body) }); export const prewarm = (model: string, adapter: string | null, checkpoint: number | null = null) => api<{ ready: boolean; error?: string }>("/v1/prewarm", { diff --git a/frontend/src/pages/Train.tsx b/frontend/src/pages/Train.tsx index b6d8f46..20a9676 100644 --- a/frontend/src/pages/Train.tsx +++ b/frontend/src/pages/Train.tsx @@ -215,7 +215,7 @@ export default function Train({ methods }: { methods: MethodInfo[] }) { } async function start() { - if (!ready || busy) return; + if (!ready || busy || !model || !ds) return; // `ready` covers this; tsc can't see it setErr(""); setBusy(true); try { const out = await submitFinetune({ diff --git a/shadowlm/backends/remote.py b/shadowlm/backends/remote.py index 9aca73c..2cb8652 100644 --- a/shadowlm/backends/remote.py +++ b/shadowlm/backends/remote.py @@ -102,6 +102,8 @@ def finetune(self, dataset: Dataset, config: TrainConfig, callbacks: Callbacks, eval_dataset=_serialize(eval_dataset), load_in_4bit=self.load_in_4bit, max_seq_length=self.max_seq_length, + # so an SDK-submitted run isn't nameless in the studio's run list + name=f"{self.model_name} · {config.method}", ) self._last_job = job_id gpus = f", {self._server_gpus} gpu" if self._server_gpus else "" diff --git a/shadowlm/remote.py b/shadowlm/remote.py index 3dd761b..05a44c6 100644 --- a/shadowlm/remote.py +++ b/shadowlm/remote.py @@ -5,9 +5,10 @@ hosted tier). The SDK's remote backend speaks to either — `backend="remote"` plus `SHADOWLM_API_URL` / `SHADOWLM_API_KEY` is all a caller configures. ` -Endpoints (JSON over HTTP, optional Bearer auth): +Endpoints this client speaks (JSON over HTTP, optional Bearer auth): - GET /v1/health → {ok, backend, version, gpus, running, pending} + GET /v1/health → {ok, backend, version, gpus, running, + pending, workers} POST /v1/finetunes → {job_id} GET /v1/finetunes/ → {status, error, checkpoint, final_loss} GET /v1/finetunes//metrics → {steps: [...], evals: [...]} @@ -20,6 +21,10 @@ GET /v1/workers//socket → websocket upgrade (see worker.py) POST /v1/workers//jobs//artifact ← a worker ships its adapter home +The server also serves the studio's own routes — datasets, models, settings, +vram, methods, tokens, prewarm, checkpoints, auth/login. They're driven from +`frontend/src/api.ts`, not from here; add a method when the SDK needs one. + `SHADOWLM_API_URL` may name several servers, comma-separated — `pick()` binds the client to the least-busy reachable one and stays there for the session. That is the whole "cluster": more boxes, each already using all of its own GPUs (the torch @@ -143,8 +148,10 @@ def pick(self) -> dict: h = self.health(base=url) except RemoteError: continue # a dead box in the pool is skipped, not fatal + # queue depth first; break ties on capacity — a hub with workers + # attached can absorb more than a lone box with the same GPU count ranked.append(((h.get("running", 0) + h.get("pending", 0), - -h.get("gpus", 0)), url, h)) + -h.get("workers", 0), -h.get("gpus", 0)), url, h)) if not ranked: raise RemoteError( f"no reachable ShadowLM server in the pool: {', '.join(self.pool)}") @@ -152,13 +159,21 @@ def pick(self) -> dict: _, self.api_url, health = ranked[0] return health - def submit_finetune(self, *, base_model: str, config: dict, dataset: dict, - eval_dataset: dict | None, load_in_4bit: bool, - max_seq_length: int, worker: str | None = None) -> str: + def submit_finetune(self, *, base_model: str, config: dict, + dataset: dict | None, eval_dataset: dict | None, + load_in_4bit: bool, max_seq_length: int, + worker: str | None = None, name: str | None = None, + dataset_id: str | None = None) -> str: + """Queue a run. Send inline `dataset` rows, or a `dataset_id` already + stored on the server (what the studio's Datasets page holds).""" + if not dataset and not dataset_id: + raise ValueError("pass dataset={'rows': [...]} or dataset_id='...'") out = self._request("POST", "/v1/finetunes", { "base_model": base_model, "config": config, + "name": name, # what the studio's run list shows "dataset": dataset, + "dataset_id": dataset_id, "eval_dataset": eval_dataset, "load_in_4bit": load_in_4bit, "max_seq_length": max_seq_length, diff --git a/shadowlm/serve.py b/shadowlm/serve.py index 747481f..e0e69a1 100644 --- a/shadowlm/serve.py +++ b/shadowlm/serve.py @@ -905,7 +905,7 @@ def submit(self, payload: dict) -> str: with self._lock: self.jobs[job_id] = job self._persist(job) # visible (as pending) the instant it's queued - # target="" routes the job to that machine's inbox instead + # worker="" routes the job to that machine's inbox instead # of this box's own training queue. Submitting ahead of the worker # connecting is fine — the inbox waits. if payload.get("worker"): diff --git a/tests/test_remote_protocol.py b/tests/test_remote_protocol.py new file mode 100644 index 0000000..7499cdf --- /dev/null +++ b/tests/test_remote_protocol.py @@ -0,0 +1,92 @@ +"""The SDK client sends what the server accepts, and ranks the pool on the +capacity the server actually reports. +""" + +from __future__ import annotations + +import pytest + +from shadowlm.remote import RemoteClient, RemoteError + + +class _Recorder(RemoteClient): + """Captures requests instead of making them; replays canned responses.""" + + def __init__(self, url="http://one", responses=None): + super().__init__(api_url=url) + self.sent: list[tuple[str, str, dict | None, str | None]] = [] + self._responses = responses or {} + + def _request(self, method, path, body=None, *, raw=False, timeout=None, + retries=2, base=None, blob=None): + self.sent.append((method, path, body, base)) + target = base or self.api_url + canned = self._responses.get((path, target), self._responses.get(path)) + if isinstance(canned, Exception): + raise canned + return canned if canned is not None else {"job_id": "j1"} + + +# ---- submit_finetune ------------------------------------------------------------ + +def test_submit_sends_the_run_name(): + c = _Recorder() + c.submit_finetune(base_model="m", config={}, dataset={"rows": []}, + eval_dataset=None, load_in_4bit=False, max_seq_length=512, + name="nightly shadow") + assert c.sent[0][2]["name"] == "nightly shadow" + + +def test_submit_can_reference_a_server_side_dataset(): + c = _Recorder() + c.submit_finetune(base_model="m", config={}, dataset=None, eval_dataset=None, + load_in_4bit=False, max_seq_length=512, dataset_id="ds123") + assert c.sent[0][2]["dataset_id"] == "ds123" + + +def test_submit_needs_rows_or_a_dataset_id(): + c = _Recorder() + with pytest.raises(ValueError, match="dataset"): + c.submit_finetune(base_model="m", config={}, dataset=None, + eval_dataset=None, load_in_4bit=False, + max_seq_length=512) + + +# ---- pick() --------------------------------------------------------------------- + +def _health(running=0, pending=0, gpus=1, workers=0): + return {"ok": True, "backend": "torch", "version": "0", + "running": running, "pending": pending, "gpus": gpus, + "workers": workers} + + +def test_pick_prefers_the_shallower_queue(): + c = _Recorder("http://busy,http://idle", responses={ + ("/v1/health", "http://busy"): _health(running=3), + ("/v1/health", "http://idle"): _health(running=0)}) + c.pick() + assert c.api_url == "http://idle" + + +def test_pick_counts_idle_workers_as_capacity(): + """Equal queues: the hub with workers attached can absorb more.""" + c = _Recorder("http://solo,http://fleet", responses={ + ("/v1/health", "http://solo"): _health(gpus=1, workers=0), + ("/v1/health", "http://fleet"): _health(gpus=1, workers=4)}) + c.pick() + assert c.api_url == "http://fleet" + + +def test_pick_skips_a_dead_box(): + c = _Recorder("http://dead,http://alive", responses={ + ("/v1/health", "http://dead"): RemoteError("refused"), + ("/v1/health", "http://alive"): _health()}) + c.pick() + assert c.api_url == "http://alive" + + +def test_pick_raises_when_the_whole_pool_is_down(): + c = _Recorder("http://a,http://b", responses={ + "/v1/health": RemoteError("refused")}) + with pytest.raises(RemoteError, match="no reachable"): + c.pick()