diff --git a/v2/assets/lease-exec.py b/v2/assets/lease-exec.py new file mode 100755 index 0000000..e95f70a --- /dev/null +++ b/v2/assets/lease-exec.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python3 +"""Ordo GPU lease wrapper — run any command under an ops-controller VRAM lease. + +Generic and stdlib-only: any GPU plugin bind-mounts this file (or prefixes its command with it) +to route GPU-heavy work through the scheduler. Nothing here is specific to any one service. +First consumer: the ai-toolkit plugin mounts it at /app/ai-toolkit/venv/bin/python — the path the +AI-toolkit UI prefers when spawning trainers — so every training run leases the GPU first. + +Contract (env): + OPS_CONTROLLER_URL scheduler base URL (required), e.g. http://ops-controller:9000 + OPS_CONTROLLER_TOKEN optional bearer token (sent when set) + ORDO_LEASE_VRAM_GB VRAM to lease (required); ~full card => exclusive lease + ORDO_LEASE_KIND job kind label (default "generic") + ORDO_LEASE_JOB_ID explicit job id (default: lease-$AITK_JOB_ID, else lease-) + ORDO_LEASE_ACQUIRE_TIMEOUT_S max wait for admission (default 3600) + ORDO_LEASE_POLL_S admission poll interval (default 5) + ORDO_LEASE_HEARTBEAT_S heartbeat interval (default 60) + +Behavior: POST /jobs → poll GET /status until admitted (fail on rejected/timeout) → run +`sys.executable argv[1:]` → heartbeat while it runs (re-acquiring if the controller lost the +lease, e.g. across a restart) → forward SIGINT/SIGTERM to the child → POST /jobs/complete on +exit, propagating the child's exit code. If the controller is unreachable up front, exit non-zero +WITHOUT running the command: GPU work must never run unleased (arbitration is not optional). +""" +from __future__ import annotations + +import json +import os +import signal +import subprocess +import sys +import threading +import time +import urllib.error +import urllib.request +import uuid + + +def _log(msg: str) -> None: + print(f"[ordo-lease] {msg}", file=sys.stderr, flush=True) + + +def _req(method: str, path: str, body: dict | None = None) -> dict: + base = os.environ["OPS_CONTROLLER_URL"].rstrip("/") + req = urllib.request.Request( + base + path, + data=json.dumps(body).encode() if body is not None else None, + method=method, + headers={"Content-Type": "application/json"}, + ) + token = os.environ.get("OPS_CONTROLLER_TOKEN", "").strip() + if token: + req.add_header("Authorization", f"Bearer {token}") + with urllib.request.urlopen(req, timeout=30) as r: + return json.loads(r.read() or b"{}") + + +def _gpu(payload: dict) -> dict: + # GET /status nests the scheduler block under "gpu"; POST /jobs* returns it bare. + return payload.get("gpu", payload) + + +def _acquire(job_id: str, vram_gb: float, kind: str) -> None: + poll_s = float(os.environ.get("ORDO_LEASE_POLL_S", "5")) + timeout_s = float(os.environ.get("ORDO_LEASE_ACQUIRE_TIMEOUT_S", "3600")) + st = _gpu(_req("POST", "/jobs", {"id": job_id, "vram_gb": vram_gb, "kind": kind})) + deadline = time.monotonic() + timeout_s + while True: + if job_id in [j.get("id") for j in st.get("running", [])]: + _log(f"lease '{job_id}' admitted ({vram_gb} GB, kind={kind})") + return + if job_id in st.get("rejected", []): + raise SystemExit( + f"[ordo-lease] job '{job_id}' REJECTED by scheduler ({vram_gb} GB does not fit)" + ) + if time.monotonic() >= deadline: + raise SystemExit(f"[ordo-lease] timed out after {timeout_s}s waiting for lease '{job_id}'") + time.sleep(poll_s) + st = _gpu(_req("GET", "/status")) + + +def _complete(job_id: str) -> None: + for attempt in range(5): + try: + _req("POST", "/jobs/complete", {"id": job_id}) + _log(f"lease '{job_id}' released") + return + except (urllib.error.URLError, OSError) as e: + _log(f"release attempt {attempt + 1}/5 failed: {e}") + time.sleep(2) + _log(f"WARNING: could not release lease '{job_id}' — scheduler sweep will reclaim it") + + +def _heartbeat_loop(job_id: str, vram_gb: float, kind: str, child: subprocess.Popen) -> None: + interval = float(os.environ.get("ORDO_LEASE_HEARTBEAT_S", "60")) + while child.poll() is None: + time.sleep(interval) + if child.poll() is not None: + return + try: + _req("POST", "/jobs/heartbeat", {"id": job_id}) + except urllib.error.HTTPError as e: + if e.code == 404: + # Controller lost the lease (e.g. restart wiped in-memory state). Re-acquire so + # the resident is re-evicted rather than restored into our VRAM. + _log(f"lease '{job_id}' lost (404) — re-acquiring") + try: + _req("POST", "/jobs", {"id": job_id, "vram_gb": vram_gb, "kind": kind}) + except (urllib.error.URLError, OSError) as e2: + _log(f"re-acquire failed: {e2}") + else: + _log(f"heartbeat failed: HTTP {e.code}") + except (urllib.error.URLError, OSError) as e: + _log(f"heartbeat failed: {e}") + + +def main() -> int: + if len(sys.argv) < 2: + _log("usage: lease-exec.py ") + return 2 + vram_env = os.environ.get("ORDO_LEASE_VRAM_GB", "").strip() + if not os.environ.get("OPS_CONTROLLER_URL", "").strip() or not vram_env: + _log("OPS_CONTROLLER_URL and ORDO_LEASE_VRAM_GB are required — refusing to run unleased") + return 2 + vram_gb = float(vram_env) + kind = os.environ.get("ORDO_LEASE_KIND", "generic") + job_id = os.environ.get("ORDO_LEASE_JOB_ID", "").strip() or ( + f"lease-{os.environ['AITK_JOB_ID']}" + if os.environ.get("AITK_JOB_ID") + else f"lease-{uuid.uuid4().hex[:8]}" + ) + + try: + _acquire(job_id, vram_gb, kind) + except (urllib.error.URLError, OSError) as e: + _log(f"cannot reach ops-controller: {e} — refusing to run unleased") + return 3 + + child = subprocess.Popen([sys.executable] + sys.argv[1:]) + + def _forward(signum: int, _frame) -> None: + if child.poll() is None: + child.send_signal(signum) + + signal.signal(signal.SIGINT, _forward) # AI-toolkit's Stop button sends SIGINT + signal.signal(signal.SIGTERM, _forward) + + threading.Thread(target=_heartbeat_loop, args=(job_id, vram_gb, kind, child), daemon=True).start() + try: + rc = child.wait() + finally: + _complete(job_id) + return rc + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/v2/ordo/broker.py b/v2/ordo/broker.py index 4058739..1c2015f 100644 --- a/v2/ordo/broker.py +++ b/v2/ordo/broker.py @@ -129,6 +129,10 @@ def complete(self, job_id: str) -> None: self.backend.stop(job_id) self.reconcile() + def heartbeat(self, job_id: str) -> bool: + """Renew a running job's lease (liveness-based). No reconcile — nothing starts or stops.""" + return self.scheduler.heartbeat(job_id) + def sweep_leases(self) -> list[str]: """Force-complete stranded leases (TTL elapsed) and reconcile — restores the resident. diff --git a/v2/ordo/control.py b/v2/ordo/control.py index 7761a21..9e497a2 100644 --- a/v2/ordo/control.py +++ b/v2/ordo/control.py @@ -114,6 +114,16 @@ def complete_job(self, body: dict[str, Any]) -> dict[str, Any]: self.broker.complete(job_id) return self.scheduler.status() + def heartbeat_job(self, body: dict[str, Any]) -> dict[str, Any]: + if not self.broker: + return self._error(503, "no broker configured") + job_id = str(body.get("id", "")).strip() + if not job_id: + return self._error(400, "body must include 'id'") + if not self.broker.heartbeat(job_id): + return self._error(404, f"no running job '{job_id}'") + return self.scheduler.status() + # --- routing (also pure) --- def route(self, method: str, path: str, body: dict[str, Any] | None = None) -> tuple[int, dict]: body = body or {} @@ -128,6 +138,8 @@ def route(self, method: str, path: str, body: dict[str, Any] | None = None) -> t return self._as_response(self.request_job(body)) if m == "POST" and path == "/jobs/complete": return self._as_response(self.complete_job(body)) + if m == "POST" and path == "/jobs/heartbeat": + return self._as_response(self.heartbeat_job(body)) if m == "GET" and path in ("/health", "/healthz"): return 200, {"ok": True} return 404, {"error": f"no route {method} {path}"} diff --git a/v2/ordo/scheduler.py b/v2/ordo/scheduler.py index 4d7743e..451309c 100644 --- a/v2/ordo/scheduler.py +++ b/v2/ordo/scheduler.py @@ -35,6 +35,12 @@ LEASE_TTL_EST_MULT = 2.0 # a job may legitimately run up to ~2x its estimate LEASE_TTL_MAX_SECONDS = 3600.0 # absolute ceiling — no lease outlives this +# Renewable leases: a client that HEARTBEATS proves liveness, so its lease is extended past the +# estimate-based caps above (those defend against clients that never call /jobs/complete — a +# heartbeating client that dies simply stops beating and is swept within HEARTBEAT_TTL_SECONDS). +# This is what lets a multi-hour job (LoRA training) hold a lease without weakening self-heal. +HEARTBEAT_TTL_SECONDS = 900.0 + @dataclasses.dataclass class Job: @@ -51,11 +57,13 @@ def __init__( cloud_fallback: bool = False, lease_ttl_default: float = DEFAULT_LEASE_TTL_SECONDS, lease_ttl_max: float = LEASE_TTL_MAX_SECONDS, + heartbeat_ttl: float = HEARTBEAT_TTL_SECONDS, ): self.total_vram_gb = float(total_vram_gb) self.cloud_fallback = bool(cloud_fallback) self.lease_ttl_default = float(lease_ttl_default) self.lease_ttl_max = float(lease_ttl_max) + self.heartbeat_ttl = float(heartbeat_ttl) self._queue: list[Job] = [] self._running: dict[str, Job] = {} self._elapsed: dict[str, float] = {} # running job id -> seconds elapsed @@ -203,6 +211,19 @@ def complete(self, job_id: str) -> None: self._elapsed.pop(job_id, None) self._deadline.pop(job_id, None) + def heartbeat(self, job_id: str) -> bool: + """Renew a running job's lease: deadline moves to now + heartbeat_ttl (liveness-based). + + Deliberately NOT capped by lease_ttl_max — the cap defends against clients that never + complete, and a heartbeating client re-proves liveness on every beat. Returns False (and + changes nothing) for a job that isn't running, so a client can detect a lost lease (e.g. + across an ops-controller restart) and re-acquire. + """ + if job_id not in self._running: + return False + self._deadline[job_id] = self._clock + self.heartbeat_ttl + return True + def sweep_expired_leases(self) -> list[str]: """Force-complete any running job whose lease TTL has elapsed. Returns the swept ids. diff --git a/v2/tests/test_broker.py b/v2/tests/test_broker.py index 281c86c..b743fcb 100644 --- a/v2/tests/test_broker.py +++ b/v2/tests/test_broker.py @@ -141,3 +141,17 @@ def test_abstract_lease_job_never_touched_only_resident_is(): b.complete("lease-test") # release the lease assert b.backend.started == ["llamacpp"] # resident restarted; lease-test never touched assert "lease-test" not in b.backend.stopped and "lease-test" not in b.backend.started + + +def test_broker_heartbeat_passes_through_without_reconcile(): + sched = Scheduler(32) + backend = MockBackend() + b = Broker(sched, backend) + b.request(Job("train", 30, "training")) + started_before = list(backend.started) + stopped_before = list(backend.stopped) + assert b.heartbeat("train") is True + assert b.heartbeat("ghost") is False + # a heartbeat never starts or stops anything — it only moves a deadline + assert backend.started == started_before + assert backend.stopped == stopped_before diff --git a/v2/tests/test_control.py b/v2/tests/test_control.py index b266393..9e1d8e5 100644 --- a/v2/tests/test_control.py +++ b/v2/tests/test_control.py @@ -115,3 +115,28 @@ def test_complete_job_restores_resident_over_control_plane(tmp_path): assert after["evicted_residents"] == {} assert after["idle_cached"] == {"llamacpp": 25.0} # re-armed as evictable assert "llamacpp" in cp.broker.backend.started # broker actually issued the restart + + +# ── renewable leases over the control plane ────────────────────────────────────────────────────── + +def test_heartbeat_route_renews_running_lease(tmp_path): + cp, _ = _cp(tmp_path) + cp.route("POST", "/jobs", {"id": "train", "vram_gb": 30, "kind": "training"}) + cp.scheduler.tick(1700) # near the 1800s no-estimate TTL + code, body = cp.route("POST", "/jobs/heartbeat", {"id": "train"}) + assert code == 200 + ttl = next(r["lease_ttl_s"] for r in body["running"] if r["id"] == "train") + assert ttl == 900.0 # renewed to HEARTBEAT_TTL_SECONDS + cp.scheduler.tick(200) # past the ORIGINAL deadline — but renewed + assert cp.scheduler.sweep_expired_leases() == [] + + +def test_heartbeat_unknown_job_is_404(tmp_path): + cp, _ = _cp(tmp_path) + code, body = cp.route("POST", "/jobs/heartbeat", {"id": "ghost"}) + assert code == 404 and "error" in body + + +def test_heartbeat_missing_id_is_400(tmp_path): + cp, _ = _cp(tmp_path) + assert cp.route("POST", "/jobs/heartbeat", {})[0] == 400 diff --git a/v2/tests/test_lease_exec.py b/v2/tests/test_lease_exec.py new file mode 100644 index 0000000..e01e083 --- /dev/null +++ b/v2/tests/test_lease_exec.py @@ -0,0 +1,165 @@ +"""lease-exec wrapper behavior against a stub control plane (real subprocess, real HTTP).""" +import json +import os +import signal +import subprocess +import sys +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +import pytest + +LEASE_EXEC = str(Path(__file__).resolve().parent.parent / "assets" / "lease-exec.py") + + +class _StubOps: + """Minimal ops-controller: admits after `admit_after_polls` GET /status calls.""" + + def __init__(self, admit_after_polls=0, reject=False): + self.admit_after_polls = admit_after_polls + self.reject = reject + self.jobs, self.heartbeats, self.completes = [], [], [] + self.polls = 0 + self.server = ThreadingHTTPServer(("127.0.0.1", 0), self._make_handler()) + self.url = f"http://127.0.0.1:{self.server.server_port}" + threading.Thread(target=self.server.serve_forever, daemon=True).start() + + def _status(self): + jid = self.jobs[-1] if self.jobs else None + if self.reject: + return {"running": [], "rejected": [jid]} + admitted = self.polls >= self.admit_after_polls + return {"running": [{"id": jid}] if (jid and admitted) else [], "rejected": []} + + def _make_handler(stub): # noqa: N805 — closure over the stub instance + class Handler(BaseHTTPRequestHandler): + def _send(self, status, payload): + data = json.dumps(payload).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + + def do_GET(self): + stub.polls += 1 + self._send(200, {"gpu": stub._status()}) # GET /status nests under "gpu" + + def do_POST(self): + length = int(self.headers.get("Content-Length") or 0) + body = json.loads(self.rfile.read(length) or b"{}") + if self.path == "/jobs": + stub.jobs.append(body["id"]) + self._send(200, stub._status()) # POST returns bare status + elif self.path == "/jobs/heartbeat": + stub.heartbeats.append(body["id"]) + self._send(200, stub._status()) + elif self.path == "/jobs/complete": + stub.completes.append(body["id"]) + self._send(200, stub._status()) + else: + self._send(404, {"error": "no route"}) + + def log_message(self, *_): + pass + + return Handler + + +def _env(stub_url, extra=None): + env = { + "OPS_CONTROLLER_URL": stub_url, + "ORDO_LEASE_VRAM_GB": "30", + "ORDO_LEASE_KIND": "training", + "ORDO_LEASE_JOB_ID": "train-test", + "ORDO_LEASE_POLL_S": "0.05", + "ORDO_LEASE_HEARTBEAT_S": "0.1", + "ORDO_LEASE_ACQUIRE_TIMEOUT_S": "5", + "PATH": os.environ["PATH"], + } + if os.environ.get("SYSTEMROOT"): # Windows: sockets need it + env["SYSTEMROOT"] = os.environ["SYSTEMROOT"] + env.update(extra or {}) + return env + + +def _run(stub_url, child_code, marker, timeout=30): + code = child_code.replace("MARKER", str(marker).replace("\\", "/")) + return subprocess.run( + [sys.executable, LEASE_EXEC, "-c", code], + env=_env(stub_url), capture_output=True, text=True, timeout=timeout, + ) + + +def test_runs_child_under_lease_and_completes(tmp_path): + stub = _StubOps() + marker = tmp_path / "ran.txt" + r = _run(stub.url, "open('MARKER','w').write('ok')", marker) + assert r.returncode == 0, r.stderr + assert marker.read_text() == "ok" + assert stub.jobs == ["train-test"] + assert stub.completes == ["train-test"] + + +def test_child_exit_code_passes_through_and_lease_still_released(tmp_path): + stub = _StubOps() + r = _run(stub.url, "import sys; sys.exit(7)", tmp_path / "x") + assert r.returncode == 7 + assert stub.completes == ["train-test"] + + +def test_waits_for_admission_before_running_child(tmp_path): + stub = _StubOps(admit_after_polls=3) + marker = tmp_path / "ran.txt" + r = _run(stub.url, "open('MARKER','w').write('ok')", marker) + assert r.returncode == 0, r.stderr + assert stub.polls >= 3 # it actually waited through the queue + assert marker.exists() + + +def test_rejected_job_fails_loudly_without_running_child(tmp_path): + stub = _StubOps(reject=True) + marker = tmp_path / "ran.txt" + r = _run(stub.url, "open('MARKER','w').write('ok')", marker) + assert r.returncode != 0 + assert not marker.exists() # GPU work must never run unleased + assert stub.completes == [] + + +def test_unreachable_controller_fails_loudly_without_running_child(tmp_path): + marker = tmp_path / "ran.txt" + r = _run("http://127.0.0.1:1", "open('MARKER','w').write('ok')", marker) + assert r.returncode != 0 + assert not marker.exists() + + +def test_heartbeats_while_child_runs(tmp_path): + stub = _StubOps() + r = _run(stub.url, "import time; time.sleep(0.6)", tmp_path / "x") + assert r.returncode == 0, r.stderr + assert len(stub.heartbeats) >= 2 # 0.6s child at 0.1s beat interval + + +@pytest.mark.skipif(sys.platform == "win32", reason="POSIX signal forwarding") +def test_sigint_is_forwarded_and_lease_released(tmp_path): + stub = _StubOps() + marker = tmp_path / "sig.txt" + child = ( + "import signal, sys, time\n" + "signal.signal(signal.SIGINT, lambda *a: (open('MARKER','w').write('int'), sys.exit(0)))\n" + "time.sleep(30)\n" + ) + p = subprocess.Popen( + [sys.executable, LEASE_EXEC, "-c", child.replace("MARKER", str(marker))], + env=_env(stub.url), stdout=subprocess.PIPE, stderr=subprocess.PIPE, + ) + deadline = time.time() + 10 + while not stub.jobs and time.time() < deadline: + time.sleep(0.05) + time.sleep(0.3) # let the child reach its sleep + p.send_signal(signal.SIGINT) # what the AI-toolkit Stop button sends + p.wait(timeout=10) + assert marker.read_text() == "int" # child received the forwarded SIGINT + assert stub.completes == ["train-test"] # and the lease was released diff --git a/v2/tests/test_scheduler.py b/v2/tests/test_scheduler.py index e0160d7..bb417c2 100644 --- a/v2/tests/test_scheduler.py +++ b/v2/tests/test_scheduler.py @@ -180,3 +180,51 @@ def test_status_surfaces_lease_and_resident_fields(): assert st["idle_cached"] == {} assert st["running"][0]["lease_ttl_s"] == 240.0 # 120 * 2 assert st["free_vram_gb"] == 32 - 18 # evicted resident holds no VRAM + + +# ── renewable leases: a live client heartbeats past any TTL; a dead one is still swept ──────────── + +def test_heartbeat_extends_lease_beyond_absolute_cap(): + s = Scheduler(total_vram_gb=32) + s.submit(Job("train", 30, "training")) # no est_seconds → default 1800s TTL + s.pump() + # simulate 2h of runtime with a beat every 60s — the lease never expires + for _ in range(120): + s.tick(60) + assert s.heartbeat("train") is True + assert s.sweep_expired_leases() == [] + assert "train" in s.running_ids + + +def test_dead_client_stops_beating_and_is_swept_within_heartbeat_ttl(): + s = Scheduler(total_vram_gb=32) + s.submit(Job("train", 30, "training")) + s.pump() + s.tick(60) + assert s.heartbeat("train") is True # last sign of life + s.tick(900) # HEARTBEAT_TTL elapses with no further beats + assert s.sweep_expired_leases() == ["train"] + assert s.running_ids == [] + + +def test_heartbeat_unknown_or_completed_job_is_false_and_harmless(): + s = Scheduler(total_vram_gb=32) + assert s.heartbeat("ghost") is False + s.submit(Job("j", 4, "chat")) + s.pump() + s.complete("j") + assert s.heartbeat("j") is False + + +def test_heartbeat_lease_restores_resident_after_sweep(): + # end-to-end lease semantics survive: a heartbeating training job that dies is swept + # and the evicted resident becomes restorable again (llamacpp can never be stranded down). + s = Scheduler(total_vram_gb=32) + s.cache_idle("llamacpp", 25) + s.submit(Job("train", 30, "training")) + s.pump() + assert "llamacpp" in s.evicted_residents + s.heartbeat("train") + s.tick(901) # client died right after its beat + assert s.sweep_expired_leases() == ["train"] + assert "llamacpp" in s.take_restorable()