jambu-gpu 0.1.0__py3-none-any.whl

jambu_gpu/core/errors.py

@@ -0,0 +1,68 @@
+"""Normalized error model (spec section 22).
+
+Provider adapters MUST translate their native failures into these types so the
+CLI behaves identically regardless of which provider is configured.
+"""
+
+from __future__ import annotations
+
+
+class JambuError(Exception):
+    """Base class for every error raised by the CLI."""
+
+    exit_code = 1
+
+
+class ConfigError(JambuError):
+    """jambu.yaml (or legacy config.yml) is missing, malformed or semantically invalid."""
+
+    exit_code = 2
+
+
+class StateError(JambuError):
+    """Local runtime state is missing or inconsistent."""
+
+    exit_code = 3
+
+
+class ProviderError(JambuError):
+    """Base class for every provider-originated failure."""
+
+    exit_code = 4
+
+
+class AuthenticationError(ProviderError):
+    exit_code = 5
+
+
+class CapacityUnavailableError(ProviderError):
+    exit_code = 6
+
+
+class ProvisioningError(ProviderError):
+    exit_code = 7
+
+
+class ProviderTimeoutError(ProviderError):
+    exit_code = 8
+
+
+class UnsupportedCapabilityError(ProviderError):
+    exit_code = 9
+
+
+class BudgetExceededError(JambuError):
+    exit_code = 10
+
+
+class RuntimeStartupError(JambuError):
+    exit_code = 11
+
+
+class WorkloadFailed(JambuError):
+    """The user workload exited non-zero. Carries the original exit code."""
+
+    def __init__(self, message: str, returncode: int) -> None:
+        super().__init__(message)
+        self.returncode = returncode
+        self.exit_code = returncode or 1

jambu_gpu/core/events.py

@@ -0,0 +1,115 @@
+"""Structured lifecycle events (spec section 24).
+
+Every lifecycle action appends one JSON object to ``.jambu/logs/events.jsonl``
+so unexpected GPU cost can always be traced back to a decision and a reason.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+import threading
+from pathlib import Path
+from typing import Any, Iterable, Optional
+
+from .state import to_iso, utc_now
+
+# Minimum event vocabulary required by the spec.
+PROVISION_REQUESTED = "provision_requested"
+INSTANCE_CREATED = "instance_created"
+RUNTIME_STARTING = "runtime_starting"
+RUNTIME_READY = "runtime_ready"
+WORKLOAD_STARTED = "workload_started"
+WORKLOAD_FINISHED = "workload_finished"
+INSTANCE_IDLE = "instance_idle"
+INSTANCE_LOCKED = "instance_locked"
+INSTANCE_UNLOCKED = "instance_unlocked"
+INSTANCE_STOPPED = "instance_stopped"
+INSTANCE_STARTED = "instance_started"
+INSTANCE_DESTROYED = "instance_destroyed"
+LIFECYCLE_GUARD_TRIGGERED = "lifecycle_guard_triggered"
+PROVIDER_ERROR = "provider_error"
+STATE_RECONCILED = "state_reconciled"
+SETUP_FAILED = "setup_failed"
+CLEANUP_PERFORMED = "cleanup_performed"
+BUDGET_REJECTED = "budget_rejected"
+
+
+class EventLog:
+    """Append-only JSONL sink. Never raises into the caller's path."""
+
+    def __init__(self, logs_dir: Path, echo: bool = False) -> None:
+        self.dir = Path(logs_dir)
+        self.path = self.dir / "events.jsonl"
+        self.echo = echo
+        self._lock = threading.Lock()
+        self._context: dict[str, Any] = {}
+
+    def bind(self, **context: Any) -> "EventLog":
+        self._context.update({k: v for k, v in context.items() if v is not None})
+        return self
+
+    def emit(self, event: str, **fields: Any) -> dict[str, Any]:
+        record: dict[str, Any] = {
+            "ts": to_iso(utc_now()),
+            "event": event,
+            "pid": os.getpid(),
+        }
+        record.update(self._context)
+        record.update({k: v for k, v in fields.items() if v is not None})
+        line = json.dumps(record, default=str)
+        with self._lock:
+            try:
+                self.dir.mkdir(parents=True, exist_ok=True)
+                with open(self.path, "a") as handle:
+                    handle.write(line + "\n")
+            except OSError:
+                pass
+            if self.echo:
+                print(line, file=sys.stderr)
+        return record
+
+    def tail(self, limit: int = 50, event: Optional[str] = None) -> list[dict[str, Any]]:
+        if not self.path.is_file():
+            return []
+        out: list[dict[str, Any]] = []
+        for line in self.path.read_text().splitlines():
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                record = json.loads(line)
+            except json.JSONDecodeError:
+                continue
+            if event and record.get("event") != event:
+                continue
+            out.append(record)
+        return out[-limit:]
+
+    def extend(self, records: Iterable[dict[str, Any]]) -> int:
+        """Merge records produced elsewhere (e.g. pulled from the remote watchdog)."""
+        count = 0
+        seen = {(r.get("ts"), r.get("event")) for r in self.tail(limit=2000)}
+        with self._lock:
+            try:
+                self.dir.mkdir(parents=True, exist_ok=True)
+                with open(self.path, "a") as handle:
+                    for record in records:
+                        key = (record.get("ts"), record.get("event"))
+                        if key in seen:
+                            continue
+                        seen.add(key)
+                        handle.write(json.dumps(record, default=str) + "\n")
+                        count += 1
+            except OSError:
+                pass
+        return count
+
+
+class NullEventLog(EventLog):
+    def __init__(self) -> None:  # noqa: D107
+        super().__init__(Path(os.devnull).parent, echo=False)
+
+    def emit(self, event: str, **fields: Any) -> dict[str, Any]:  # noqa: D102
+        return {"event": event, **fields}

jambu_gpu/core/execution.py

@@ -0,0 +1,169 @@
+"""Workload execution and heartbeat (spec sections 12, 13, 16).
+
+The CLI registers the workload with the remote watchdog, keeps a heartbeat
+alive while it runs, and records completion. If the heartbeat disappears
+(CLI crash, SSH drop, laptop sleep), the watchdog lets the grace period
+expire and the instance falls back to IDLE and the normal shutdown rules.
+"""
+
+from __future__ import annotations
+
+import os
+import signal
+import subprocess
+import threading
+import time
+from typing import Optional, Sequence
+
+from .beacon import WatchdogClient
+from .events import WORKLOAD_FINISHED, WORKLOAD_STARTED, EventLog
+from .models import Execution, Workload
+from .state import StateStore, WorkloadRecord, new_workload_id, utc_now
+
+
+class Heartbeat:
+    """Background thread that keeps the remote watchdog informed."""
+
+    def __init__(
+        self,
+        workload_id: str,
+        interval: float,
+        store: StateStore,
+        beacon: Optional[WatchdogClient] = None,
+    ) -> None:
+        self.workload_id = workload_id
+        self.interval = max(5.0, float(interval))
+        self.store = store
+        self.beacon = beacon
+        self._stop = threading.Event()
+        self._thread: Optional[threading.Thread] = None
+        self.failures = 0
+
+    def _beat(self) -> None:
+        if self.beacon is not None:
+            if self.beacon.workload_heartbeat(self.workload_id) is None:
+                self.failures += 1
+            else:
+                self.failures = 0
+        try:
+            with self.store.transaction(timeout=10) as state:
+                record = state.workloads.get(self.workload_id)
+                if record is not None:
+                    record.last_heartbeat = utc_now()
+                state.touch()
+        except Exception:  # noqa: BLE001 - a heartbeat must never kill a workload
+            pass
+
+    def start(self) -> "Heartbeat":
+        self._beat()
+
+        def _loop() -> None:
+            while not self._stop.wait(self.interval):
+                self._beat()
+
+        self._thread = threading.Thread(target=_loop, name="jambu-heartbeat", daemon=True)
+        self._thread.start()
+        return self
+
+    def stop(self) -> None:
+        self._stop.set()
+        if self._thread is not None:
+            self._thread.join(timeout=2.0)
+
+    def __enter__(self) -> "Heartbeat":
+        return self.start()
+
+    def __exit__(self, *exc_info) -> None:
+        self.stop()
+
+
+class ExecutionManager:
+    """Registers, runs and closes out workloads."""
+
+    def __init__(
+        self,
+        store: StateStore,
+        events: EventLog,
+        beacon: Optional[WatchdogClient] = None,
+        heartbeat_interval: float = 30.0,
+    ) -> None:
+        self.store = store
+        self.events = events
+        self.beacon = beacon
+        self.heartbeat_interval = heartbeat_interval
+
+    def register(self, command: Sequence[str], label: str = "") -> WorkloadRecord:
+        record = WorkloadRecord(
+            id=new_workload_id(),
+            label=label or " ".join(command)[:80],
+            command=list(command),
+            state="running",
+            pid=os.getpid(),
+        )
+        with self.store.transaction() as state:
+            state.workloads[record.id] = record
+            state.touch()
+            state.prune_workloads()
+        if self.beacon is not None:
+            self.beacon.workload_start(record.id, record.label)
+        self.events.emit(WORKLOAD_STARTED, workload_id=record.id, label=record.label)
+        return record
+
+    def finish(self, record: WorkloadRecord, returncode: Optional[int]) -> None:
+        state_name = "finished" if returncode == 0 else "failed"
+        with self.store.transaction() as state:
+            stored = state.workloads.get(record.id)
+            if stored is not None:
+                stored.state = state_name
+                stored.finished_at = utc_now()
+                stored.returncode = returncode
+            state.touch()
+        if self.beacon is not None:
+            self.beacon.workload_end(record.id, returncode, state_name)
+        self.events.emit(
+            WORKLOAD_FINISHED,
+            workload_id=record.id,
+            returncode=returncode,
+            state=state_name,
+            duration_seconds=round(utc_now() - record.started_at, 1),
+        )
+
+    def run_local(self, workload: Workload, record: WorkloadRecord) -> Execution:
+        """Run the command on the developer machine against the remote runtime."""
+        execution = Execution(
+            id=record.id, workload_id=record.id, started_at=time.time()
+        )
+        env = os.environ.copy()
+        env.update({k: str(v) for k, v in workload.env.items()})
+
+        process = subprocess.Popen(  # noqa: S603 - the user's own command, by design
+            list(workload.command),
+            cwd=workload.cwd or None,
+            env=env,
+        )
+
+        def _forward(signum, _frame):  # pragma: no cover - interactive path
+            try:
+                process.send_signal(signum)
+            except ProcessLookupError:
+                pass
+
+        previous = {}
+        for sig in (signal.SIGINT, signal.SIGTERM):
+            try:
+                previous[sig] = signal.signal(sig, _forward)
+            except ValueError:
+                pass
+        try:
+            returncode = process.wait()
+        finally:
+            for sig, handler in previous.items():
+                try:
+                    signal.signal(sig, handler)
+                except ValueError:
+                    pass
+
+        execution.finished_at = time.time()
+        execution.returncode = returncode
+        execution.state = "finished" if returncode == 0 else "failed"
+        return execution

jambu_gpu/core/health.py

@@ -0,0 +1,62 @@
+"""Model server readiness checks.
+
+A passing health check means the runtime is *available*; per spec section 12
+it is explicitly NOT workload activity and never refreshes the idle timer.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Callable, Optional
+
+import httpx
+
+from .errors import RuntimeStartupError
+
+
+def probe(url: str, timeout: float = 5.0) -> tuple[bool, str]:
+    try:
+        response = httpx.get(url, timeout=timeout)
+    except httpx.HTTPError as exc:
+        return False, type(exc).__name__
+    if response.status_code < 400:
+        return True, f"{response.status_code}"
+    return False, f"HTTP {response.status_code}"
+
+
+def wait_for_healthy(
+    url: str,
+    *,
+    timeout: float,
+    interval: float = 10.0,
+    on_attempt: Optional[Callable[[int, float, str], None]] = None,
+    should_abort: Optional[Callable[[], Optional[str]]] = None,
+) -> None:
+    """Poll ``url`` until it answers or ``timeout`` elapses.
+
+    ``should_abort`` lets the caller fail fast when the instance itself died
+    instead of waiting out the full startup timeout.
+    """
+    deadline = time.time() + timeout
+    attempt = 0
+    last_detail = "not started"
+    interval = max(2.0, min(interval, 30.0))
+
+    while time.time() < deadline:
+        attempt += 1
+        if should_abort is not None:
+            reason = should_abort()
+            if reason:
+                raise RuntimeStartupError(f"model server did not start: {reason}")
+        ok, last_detail = probe(url)
+        if ok:
+            return
+        if on_attempt is not None:
+            on_attempt(attempt, max(0.0, deadline - time.time()), last_detail)
+        time.sleep(interval)
+
+    raise RuntimeStartupError(
+        f"model server did not become healthy within {int(timeout)}s "
+        f"({url}, last probe: {last_detail}). "
+        "Check `gpu logs` - model downloads can be slow on first boot."
+    )