Renewable GPU leases: heartbeat + generic lease-exec wrapper

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-<random>)
+  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 <program-args...>")
+        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())

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.
 

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}"}

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.
 

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

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