jambu-gpu 0.1.0__py3-none-any.whl

jambu_gpu/cli/run.py

@@ -0,0 +1,37 @@
+"""`gpu run` - execute a workload against a guaranteed runtime (spec section 16)."""
+
+from __future__ import annotations
+
+from typing import List
+
+import typer
+
+from ..core.errors import JambuError, WorkloadFailed
+from ._common import get_context, handle, print_json
+
+
+def run_command(
+    ctx: typer.Context,
+    command: List[str] = typer.Argument(
+        ...,
+        help=(
+            "Command to execute, e.g. `python experiments/run.py`. Flags are passed "
+            "through; put `--` first if one collides with a gpu option."
+        ),
+    ),
+    label: str = typer.Option("", "--label", help="Human label recorded with the workload."),
+) -> None:
+    """Ensure the runtime is READY, then run the command with heartbeat tracking."""
+    context = get_context(ctx)
+    try:
+        engine = context.engine()
+        engine.run(list(command), label=label)
+    except WorkloadFailed as exc:
+        if context.json:
+            print_json({"ok": False, "returncode": exc.returncode})
+        raise typer.Exit(exc.returncode) from exc
+    except JambuError as exc:
+        handle(exc)
+        return
+    if context.json:
+        print_json({"ok": True, "returncode": 0})

jambu_gpu/cli/setup.py

@@ -0,0 +1,165 @@
+"""`gpu setup` / `init` / `offers` - provisioning entry points."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import typer
+from rich.table import Table
+from rich.text import Text
+
+from ..core.config import (
+    DEFAULT_CONFIG_TEMPLATE,
+    DEFAULT_CONFIG_TEMPLATE_FLAT,
+    CONFIG_FILENAMES,
+)
+from ..core.errors import JambuError
+from ._common import console, fmt_money, get_context, handle, print_json
+
+
+def init_command(
+    ctx: typer.Context,
+    model: str = typer.Option("empero-ai/Qwythos-9B-v2", "--model", help="HuggingFace model id."),
+    force: bool = typer.Option(False, "--force", help="Overwrite an existing config file."),
+    flat: bool = typer.Option(
+        False,
+        "--flat",
+        help=(
+            "Write a legacy, gpu-runtime-only config.yml (no gpu_runtime: "
+            "wrapper) instead of the shared jambu.yaml."
+        ),
+    ),
+) -> None:
+    """Write a starter jambu.yaml (and .env.example) in the current directory."""
+    if flat:
+        target = Path.cwd() / "config.yml"
+        content = DEFAULT_CONFIG_TEMPLATE_FLAT.format(model_id=model)
+    else:
+        target = Path.cwd() / CONFIG_FILENAMES[0]  # jambu.yaml
+        content = DEFAULT_CONFIG_TEMPLATE.format(model_id=model)
+
+    if target.exists() and not force:
+        console.print(Text(f"! {target} already exists (use --force)", style="yellow"))
+        raise typer.Exit(1)
+    target.write_text(content)
+    console.print(Text(f"+ wrote {target}", style="green"))
+    if not flat:
+        console.print(
+            Text(
+                "  gpu-runtime config lives under the gpu_runtime: key - other "
+                "Jambu Lab tools can share this file",
+                style="dim",
+            )
+        )
+
+    env_example = Path.cwd() / ".env.example"
+    if not env_example.exists():
+        env_example.write_text(
+            "# Credentials never belong in jambu.yaml / config.yml.\n"
+            "VAST_API_KEY=\n"
+            "# Optional, for gated HuggingFace models:\n"
+            "HF_TOKEN=\n"
+        )
+        console.print(Text(f"+ wrote {env_example}", style="green"))
+    console.print("next: export VAST_API_KEY=... && gpu validate")
+
+
+def setup_command(
+    ctx: typer.Context,
+    recreate: bool = typer.Option(
+        False, "--recreate", help="Replace an existing instance that no longer matches jambu.yaml."
+    ),
+    allow_cost_override: bool = typer.Option(
+        False, "--allow-cost-override", help="Provision even above budget.max_hourly_cost_usd."
+    ),
+    no_wait: bool = typer.Option(
+        False, "--no-wait", help="Return as soon as the instance is created."
+    ),
+) -> None:
+    """Provision the instance, install the watchdog and start the model server."""
+    context = get_context(ctx)
+    try:
+        engine = context.engine()
+        state = engine.setup(
+            recreate=recreate, allow_cost_override=allow_cost_override, wait=not no_wait
+        )
+    except JambuError as exc:
+        handle(exc)
+        return
+
+    payload = {
+        "instance_id": state.instance_id,
+        "state": state.state,
+        "endpoint": state.endpoint,
+        "model": state.model_id,
+        "hourly_cost_usd": state.hourly_cost_usd,
+        "watchdog": state.watchdog.to_dict(),
+    }
+    if context.json:
+        print_json(payload)
+        return
+    if state.endpoint:
+        console.print(
+            f"\n  endpoint  [bold]{state.endpoint}[/bold]\n"
+            f"  openai    [bold]{state.endpoint.rstrip('/')}/v1[/bold]\n"
+            f"  model     [bold]{state.model_id}[/bold]"
+        )
+    console.print("\nnext: gpu run python your_experiment.py")
+
+
+def offers_command(
+    ctx: typer.Context,
+    limit: int = typer.Option(10, "--limit", "-n", help="How many offers to show."),
+) -> None:
+    """Show matching capacity and its price before committing to it."""
+    context = get_context(ctx)
+    try:
+        engine = context.engine()
+        engine.provider.require_capability("pricing_query")
+        spec = engine.build_spec()
+        offers = engine.provider.find_offers(spec, limit=limit)
+    except JambuError as exc:
+        handle(exc)
+        return
+
+    rows = [
+        {
+            "id": o.id,
+            "gpu": o.gpu_name,
+            "count": o.gpu_count,
+            "vram_gb": o.gpu_vram_gb,
+            "disk_gb": o.disk_gb,
+            "hourly_cost_usd": o.hourly_cost_usd,
+            "region": o.region,
+        }
+        for o in offers
+    ]
+    if context.json:
+        print_json({"offers": rows})
+        return
+
+    if not rows:
+        console.print(Text("no matching offers", style="yellow"))
+        return
+
+    ceiling = engine.config.budget.max_hourly_cost_usd
+    table = Table(title=f"{engine.provider.display_name} offers", title_justify="left")
+    for column in ("OFFER", "GPU", "N", "VRAM", "DISK", "$/H", "REGION"):
+        table.add_column(column)
+    for row in rows:
+        over = ceiling is not None and row["hourly_cost_usd"] > float(ceiling)
+        table.add_row(
+            row["id"],
+            row["gpu"],
+            str(row["count"]),
+            f"{row['vram_gb']:.0f} GB",
+            f"{row['disk_gb']:.0f} GB",
+            Text(
+                fmt_money(row["hourly_cost_usd"]) or "-",
+                style="red" if over else "green",
+            ),
+            row["region"] or "-",
+        )
+    console.print(table)
+    if ceiling:
+        console.print(Text(f"  budget ceiling: ${float(ceiling):.3f}/h", style="dim"))

jambu_gpu/cli/status.py

@@ -0,0 +1,214 @@
+"""`gpu status` / `logs` / `events` - reconciled observability."""
+
+from __future__ import annotations
+
+import time
+
+import typer
+from rich.table import Table
+from rich.text import Text
+
+from ..core.durations import format_duration
+from ..core.errors import JambuError
+from ._common import (
+    console,
+    fmt_money,
+    fmt_time,
+    get_context,
+    handle,
+    kv_table,
+    print_json,
+    state_text,
+)
+
+
+def status_command(
+    ctx: typer.Context,
+    watch: float = typer.Option(
+        0, "--watch", "-w", help="Refresh every N seconds (0 = once)."
+    ),
+) -> None:
+    """Reconcile local and provider state, then report the normalized result."""
+    context = get_context(ctx)
+    while True:
+        try:
+            payload = context.engine().status()
+        except JambuError as exc:
+            handle(exc)
+            return
+
+        if context.json:
+            print_json(payload)
+            return
+
+        _render(payload)
+        if not watch:
+            return
+        time.sleep(max(2.0, watch))
+        console.rule(style="dim")
+
+
+def _render(payload: dict) -> None:
+    decision = payload.get("decision") or {}
+    watchdog = payload.get("watchdog") or {}
+    rows = [
+        ("state", state_text(payload.get("state", "unknown"))),
+        ("provider", payload.get("provider")),
+        ("instance", payload.get("instance_id") or "-"),
+        ("model", payload.get("model")),
+        ("endpoint", payload.get("endpoint") or "-"),
+        ("created", fmt_time(payload.get("created_at"))),
+        (
+            "age",
+            format_duration(payload.get("age_seconds"))
+            if payload.get("age_seconds") is not None
+            else None,
+        ),
+        ("idle", format_duration(payload.get("idle_seconds"))),
+        ("cost", _cost(payload)),
+        ("locked", _lock(payload)),
+        (
+            "watchdog",
+            _watchdog(watchdog),
+        ),
+        ("guard", _guard(decision, payload)),
+    ]
+    console.print(kv_table("Runtime", rows))
+
+    workloads = [w for w in payload.get("workloads", []) if w["state"] in ("running", "queued")]
+    recent = sorted(
+        payload.get("workloads", []), key=lambda w: w.get("started_at") or 0, reverse=True
+    )[:5]
+    if recent:
+        table = Table(title="Workloads", title_justify="left", box=None)
+        for column in ("ID", "STATE", "LABEL", "STARTED", "RC"):
+            table.add_column(column)
+        for record in recent:
+            table.add_row(
+                record["id"],
+                Text(
+                    record["state"],
+                    style="cyan" if record["state"] == "running" else "dim",
+                ),
+                (record.get("label") or "")[:48],
+                fmt_time(record.get("started_at")) or "-",
+                "-" if record.get("returncode") is None else str(record["returncode"]),
+            )
+        console.print(table)
+    if workloads:
+        console.print(
+            Text(f"  {len(workloads)} active workload(s) hold the instance open", style="cyan")
+        )
+
+
+def _cost(payload: dict) -> str:
+    hourly = payload.get("hourly_cost_usd")
+    session = payload.get("session_cost_usd")
+    if hourly is None and session is None:
+        return "-"
+    parts = []
+    if hourly is not None:
+        parts.append(f"{fmt_money(hourly)}/h")
+    if session:
+        parts.append(f"{fmt_money(session)} this session")
+    return "  ".join(parts) or "-"
+
+
+def _lock(payload: dict) -> Text:
+    if not payload.get("locked"):
+        return Text("no", style="dim")
+    until = fmt_time(payload.get("locked_until"))
+    return Text(f"yes (until {until})" if until else "yes (indefinite)", style="yellow")
+
+
+def _watchdog(watchdog: dict) -> Text:
+    if not watchdog.get("enabled"):
+        return Text("disabled - shutdown depends on this CLI staying alive", style="yellow")
+    if watchdog.get("reachable"):
+        return Text(f"live at {watchdog.get('url')}", style="green")
+    if watchdog.get("installed"):
+        return Text(
+            f"unreachable ({watchdog.get('last_error') or 'no answer'})", style="yellow"
+        )
+    return Text("not installed", style="yellow")
+
+
+def _guard(decision: dict, payload: dict) -> Text:
+    action = decision.get("action", "none")
+    reason = decision.get("reason") or "-"
+    if action == "none":
+        policy = payload.get("policy") or {}
+        idle_timeout = policy.get("idle_timeout_s")
+        idle = payload.get("idle_seconds") or 0
+        if idle_timeout and reason == "within_policy":
+            remaining = max(0.0, float(idle_timeout) - float(idle))
+            return Text(
+                f"ok ({reason}); auto-stop in {format_duration(remaining)} if idle",
+                style="green",
+            )
+        return Text(f"ok ({reason})", style="green")
+    return Text(f"{action} pending: {reason}", style="bold yellow")
+
+
+def logs_command(
+    ctx: typer.Context,
+    lines: int = typer.Option(200, "--lines", "-n", help="How many lines to show."),
+) -> None:
+    """Show provider-side container logs for the instance."""
+    context = get_context(ctx)
+    try:
+        output = context.engine().logs(limit=lines)
+    except JambuError as exc:
+        handle(exc)
+        return
+    if context.json:
+        print_json({"lines": output})
+        return
+    for line in output:
+        console.print(line, markup=False, highlight=False)
+
+
+def events_command(
+    ctx: typer.Context,
+    limit: int = typer.Option(40, "--limit", "-n", help="How many events to show."),
+    event: str = typer.Option("", "--event", help="Filter by event name."),
+) -> None:
+    """Show the structured lifecycle event log (spec section 24)."""
+    context = get_context(ctx)
+    try:
+        engine = context.engine()
+    except JambuError as exc:
+        handle(exc)
+        return
+    records = engine.events.tail(limit=limit, event=event or None)
+    if context.json:
+        print_json({"events": records})
+        return
+    if not records:
+        console.print(Text("no events recorded yet", style="dim"))
+        return
+    table = Table(title="Lifecycle events", title_justify="left", box=None)
+    for column in ("TS", "EVENT", "DETAIL"):
+        table.add_column(column)
+    for record in records:
+        detail = {
+            k: v
+            for k, v in record.items()
+            if k not in ("ts", "event", "pid", "provider", "source")
+        }
+        table.add_row(
+            str(record.get("ts", "")),
+            Text(str(record.get("event", "")), style=_event_style(record.get("event", ""))),
+            ", ".join(f"{k}={v}" for k, v in detail.items())[:110],
+        )
+    console.print(table)
+
+
+def _event_style(event: str) -> str:
+    if "error" in event or "failed" in event:
+        return "red"
+    if "guard" in event or "stopped" in event or "destroyed" in event:
+        return "yellow"
+    if "ready" in event or "created" in event:
+        return "green"
+    return "white"

jambu_gpu/cli/validate.py

@@ -0,0 +1,57 @@
+"""`gpu validate` - fail before provisioning (spec sections 5, 15)."""
+
+from __future__ import annotations
+
+import typer
+from rich.text import Text
+
+from ..core.errors import JambuError
+from ._common import console, get_context, handle, print_json
+
+
+def validate_command(
+    ctx: typer.Context,
+    offline: bool = typer.Option(
+        False, "--offline", help="Skip provider API checks (schema and policy only)."
+    ),
+) -> None:
+    """Validate jambu.yaml, credentials, capabilities and available capacity."""
+    context = get_context(ctx)
+    try:
+        engine = context.engine()
+        result = engine.validate(check_provider=not offline)
+    except JambuError as exc:
+        if context.json:
+            print_json({"ok": False, "errors": [{"message": str(exc)}]})
+            raise typer.Exit(getattr(exc, "exit_code", 1)) from exc
+        handle(exc)
+        return
+
+    payload = {
+        "ok": result.ok,
+        "provider": engine.provider.name,
+        "model": engine.config.model.id,
+        "runtime": engine.config.runtime.engine,
+        "fingerprint": engine.config.fingerprint(),
+        "errors": [{"field": i.field, "message": i.message} for i in result.errors],
+        "warnings": [{"field": i.field, "message": i.message} for i in result.warnings],
+    }
+    if context.json:
+        print_json(payload)
+        raise typer.Exit(0 if result.ok else 2)
+
+    console.print(
+        f"config: [bold]{engine.config.source_path}[/bold]\n"
+        f"provider: [bold]{engine.provider.display_name}[/bold]  "
+        f"runtime: [bold]{engine.runtime.describe()}[/bold]"
+    )
+    for issue in result.warnings:
+        console.print(Text(f"! {issue.field or 'config'}: {issue.message}", style="yellow"))
+    for issue in result.errors:
+        console.print(Text(f"x {issue.field or 'config'}: {issue.message}", style="bold red"))
+
+    if result.ok:
+        console.print(Text("+ configuration is valid", style="bold green"))
+    else:
+        console.print(Text(f"{len(result.errors)} error(s) must be fixed", style="bold red"))
+        raise typer.Exit(2)

jambu_gpu/core/beacon.py

@@ -0,0 +1,124 @@
+"""Client for the remote lifecycle watchdog.
+
+The watchdog is the authority on activity: it holds the heartbeat state and
+performs the shutdown. The CLI only reports activity to it and reads it back.
+Every call is best-effort - losing the watchdog must never break a workload,
+and an unreachable watchdog is surfaced as a warning, not a crash.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+import httpx
+
+from .errors import JambuError
+
+
+class WatchdogUnavailable(JambuError):
+    """The remote watchdog could not be reached."""
+
+
+class WatchdogClient:
+    def __init__(self, base_url: str, token: Optional[str], timeout: float = 10.0) -> None:
+        self.base_url = base_url.rstrip("/")
+        self.token = token
+        self.timeout = timeout
+
+    @property
+    def headers(self) -> dict[str, str]:
+        return {"X-Jambu-Token": self.token} if self.token else {}
+
+    # -- transport ----------------------------------------------------------
+
+    def _request(self, method: str, path: str, payload: Optional[dict] = None) -> dict[str, Any]:
+        url = f"{self.base_url}{path}"
+        try:
+            response = httpx.request(
+                method,
+                url,
+                json=payload,
+                headers=self.headers,
+                timeout=self.timeout,
+            )
+        except httpx.HTTPError as exc:
+            raise WatchdogUnavailable(f"watchdog unreachable at {url}: {exc}") from exc
+        if response.status_code == 401:
+            raise WatchdogUnavailable("watchdog rejected the control token")
+        if response.status_code >= 400:
+            raise WatchdogUnavailable(
+                f"watchdog returned {response.status_code}: {response.text[:200]}"
+            )
+        try:
+            return response.json()
+        except ValueError:
+            return {}
+
+    def try_request(self, method: str, path: str, payload: Optional[dict] = None) -> Optional[dict]:
+        try:
+            return self._request(method, path, payload)
+        except WatchdogUnavailable:
+            return None
+
+    # -- routes -------------------------------------------------------------
+
+    def ping(self) -> bool:
+        try:
+            self._request("GET", "/health")
+            return True
+        except WatchdogUnavailable:
+            return False
+
+    def state(self) -> dict[str, Any]:
+        return self._request("GET", "/state")
+
+    def try_state(self) -> Optional[dict[str, Any]]:
+        return self.try_request("GET", "/state")
+
+    def events(self, limit: int = 100) -> list[dict[str, Any]]:
+        data = self.try_request("GET", f"/events?limit={limit}") or {}
+        return data.get("events", [])
+
+    def identify(
+        self, instance_id: str, hourly_cost_usd: Optional[float] = None
+    ) -> Optional[dict]:
+        return self.try_request(
+            "POST",
+            "/identify",
+            {"instance_id": instance_id, "hourly_cost_usd": hourly_cost_usd},
+        )
+
+    def activity(self) -> Optional[dict]:
+        return self.try_request("POST", "/activity")
+
+    def mark_starting(self) -> Optional[dict]:
+        return self.try_request("POST", "/starting")
+
+    def mark_ready(self) -> Optional[dict]:
+        return self.try_request("POST", "/ready")
+
+    def workload_start(self, workload_id: str, label: str = "") -> Optional[dict]:
+        return self.try_request(
+            "POST", "/workload/start", {"id": workload_id, "label": label}
+        )
+
+    def workload_heartbeat(self, workload_id: str) -> Optional[dict]:
+        return self.try_request("POST", "/workload/heartbeat", {"id": workload_id})
+
+    def workload_end(
+        self, workload_id: str, returncode: Optional[int], state: Optional[str] = None
+    ) -> Optional[dict]:
+        return self.try_request(
+            "POST",
+            "/workload/end",
+            {"id": workload_id, "returncode": returncode, "state": state},
+        )
+
+    def lock(self, ttl_seconds: Optional[float], reason: str = "") -> dict:
+        return self._request("POST", "/lock", {"ttl_seconds": ttl_seconds, "reason": reason})
+
+    def unlock(self, reason: str = "manual") -> dict:
+        return self._request("POST", "/unlock", {"reason": reason})
+
+    def shutdown(self, action: str = "stop", reason: str = "manual") -> Optional[dict]:
+        return self.try_request("POST", "/shutdown", {"action": action, "reason": reason})