jambu-gpu 0.1.0__py3-none-any.whl

jambu_gpu/__init__.py

@@ -0,0 +1,5 @@
+"""gpu: provider-agnostic GPU runtime provisioning and lifecycle CLI."""
+
+__version__ = "0.1.0"
+
+__all__ = ["__version__"]

jambu_gpu/__main__.py

@@ -0,0 +1,4 @@
+from .cli.main import main
+
+if __name__ == "__main__":
+    main()

jambu_gpu/cli/_common.py

@@ -0,0 +1,117 @@
+"""Shared CLI plumbing: config loading, output and error rendering."""
+
+from __future__ import annotations
+
+import json as jsonlib
+import sys
+from pathlib import Path
+from typing import Any, Optional
+
+import typer
+from rich.console import Console
+from rich.table import Table
+from rich.text import Text
+
+from ..core.config import RuntimeConfig, load_config
+from ..core.engine import Engine
+from ..core.errors import JambuError
+from ..core.state import to_iso
+
+console = Console()
+err_console = Console(stderr=True)
+
+STATE_STYLES = {
+    "ready": "bold green",
+    "busy": "bold cyan",
+    "idle": "yellow",
+    "starting": "cyan",
+    "provisioning": "cyan",
+    "stopped": "dim",
+    "destroyed": "dim",
+    "failed": "bold red",
+    "unknown": "magenta",
+    "none": "dim",
+}
+
+
+class Context:
+    """Per-invocation state carried on the typer context object."""
+
+    def __init__(self) -> None:
+        self.config_path: Optional[Path] = None
+        self.json: bool = False
+        self.verbose: bool = False
+        self._config: Optional[RuntimeConfig] = None
+        self._engine: Optional[Engine] = None
+
+    def config(self) -> RuntimeConfig:
+        if self._config is None:
+            self._config = load_config(self.config_path)
+        return self._config
+
+    def engine(self) -> Engine:
+        if self._engine is None:
+            self._engine = Engine(self.config(), echo=self.echo)
+        return self._engine
+
+    def echo(self, message: str) -> None:
+        if self.json:
+            return
+        style = ""
+        if message.startswith("!"):
+            style = "yellow"
+        elif message.startswith("+"):
+            style = "green"
+        elif message.startswith("-"):
+            style = "red"
+        elif message.startswith(">"):
+            style = "cyan"
+        console.print(Text(message, style=style) if style else message)
+
+
+def get_context(ctx: typer.Context) -> Context:
+    if ctx.obj is None:
+        ctx.obj = Context()
+    return ctx.obj
+
+
+def fail(message: str, code: int = 1) -> None:
+    err_console.print(Text(f"error: {message}", style="bold red"))
+    raise typer.Exit(code)
+
+
+def handle(exc: Exception) -> None:
+    """Render a JambuError the same way regardless of which provider raised it."""
+    if isinstance(exc, JambuError):
+        err_console.print(Text(f"error: {exc}", style="bold red"))
+        raise typer.Exit(getattr(exc, "exit_code", 1))
+    raise exc
+
+
+def print_json(payload: Any) -> None:
+    sys.stdout.write(jsonlib.dumps(payload, indent=2, default=str) + "\n")
+
+
+def state_text(state: str) -> Text:
+    return Text(state.upper(), style=STATE_STYLES.get(state, "white"))
+
+
+def kv_table(title: str, rows: list[tuple[str, Any]]) -> Table:
+    table = Table(title=title, show_header=False, box=None, title_justify="left", pad_edge=False)
+    table.add_column("field", style="dim", no_wrap=True)
+    table.add_column("value")
+    for key, value in rows:
+        if value is None:
+            continue
+        table.add_row(key, value if isinstance(value, Text) else str(value))
+    return table
+
+
+def fmt_time(ts: Optional[float]) -> Optional[str]:
+    return to_iso(ts) if ts else None
+
+
+def fmt_money(value: Optional[float], suffix: str = "") -> Optional[str]:
+    if value is None:
+        return None
+    return f"${value:.3f}{suffix}"

jambu_gpu/cli/inspect.py

@@ -0,0 +1,207 @@
+"""`gpu inspect-model` - where jambu.yaml's numbers actually come from.
+
+Fetches the two public HuggingFace documents that answer "what should
+compute.gpu / runtime look like for this model" (see
+jambu_gpu/core/model_inspect.py for exactly which fields), and prints a
+ready-to-paste gpu_runtime: snippet plus the reasoning behind each value.
+
+`--add-profile` closes the loop: instead of hand-writing a new full jambu.yaml
+per model (which only ever grows, never shrinks, and duplicates across every
+project directory), it writes the model-specific fragment straight into the
+shared jambu.models.yaml catalog under a name - `model_profile: <name>` in a
+project's jambu.yaml is then the only thing that ever needs to change.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional
+
+import typer
+import yaml
+from rich.text import Text
+
+from rich.table import Table
+
+from ..core.config import CATALOG_FILENAMES, find_catalog, load_catalog
+from ..core.credentials import CredentialResolver
+from ..core.errors import ConfigError, JambuError
+from ..core.model_inspect import inspect_model
+from ._common import console, get_context, handle, print_json
+
+
+def inspect_model_command(
+    ctx: typer.Context,
+    repo_id: str = typer.Argument(
+        ..., help="HuggingFace repo id to inspect, e.g. Qwen/Qwen2.5-7B-Instruct."
+    ),
+    revision: str = typer.Option("main", "--revision", help="Repo revision/branch/tag."),
+    hf_token_env: str = typer.Option(
+        "HF_TOKEN", "--hf-token-env", help="Env var holding a token for gated repos."
+    ),
+    add_profile: str = typer.Option(
+        "",
+        "--add-profile",
+        help=(
+            "Write the result straight into jambu.models.yaml under this name "
+            "instead of just printing a snippet (creates the catalog if none "
+            "exists yet). Use model_profile: <name> in jambu.yaml to consume it."
+        ),
+    ),
+    catalog_path_opt: Optional[Path] = typer.Option(
+        None,
+        "--catalog",
+        help=f"Catalog file to write to with --add-profile (default: nearest {CATALOG_FILENAMES[0]}, or a new one here).",
+    ),
+    force: bool = typer.Option(
+        False, "--force", help="With --add-profile, overwrite an existing profile of that name."
+    ),
+) -> None:
+    """Inspect a HuggingFace repo and suggest jambu.yaml values for it.
+
+    Reads exactly two public documents, cited in the output:
+
+    \b
+      https://huggingface.co/api/models/<repo>            (param count, dtype)
+      https://huggingface.co/<repo>/raw/<rev>/config.json  (architecture, context)
+
+    This is a starting point, not an oracle - the VRAM figure is weights plus
+    a flat margin for KV cache/activations (which depend on context length
+    and concurrency this command has no way to know). Confirm with
+    `gpu offers` and a real health check before trusting it in production.
+    """
+    context = get_context(ctx)
+    resolver = CredentialResolver(Path.cwd())
+    token = resolver.get(hf_token_env)
+
+    try:
+        result = inspect_model(repo_id, revision=revision, token=token)
+    except JambuError as exc:
+        handle(exc)
+        return
+
+    if add_profile:
+        try:
+            target = _write_profile(result, add_profile, catalog_path_opt, force=force)
+        except JambuError as exc:
+            handle(exc)
+            return
+        if context.json:
+            print_json({**result.to_dict(), "profile_name": add_profile, "catalog_path": str(target)})
+        else:
+            console.print(Text(f"+ wrote profile '{add_profile}' to {target}", style="green"))
+            console.print(f"  use it with: model_profile: {add_profile}")
+        return
+
+    if context.json:
+        print_json(result.to_dict())
+        return
+
+    console.print(f"[bold]{result.repo_id}[/bold] @ {result.revision}")
+    console.print(
+        f"  architecture   {', '.join(result.architectures) or '?'} "
+        f"(model_type: {result.model_type or '?'})"
+    )
+    if result.total_params:
+        console.print(
+            f"  parameters     {result.total_params / 1e9:.2f}B "
+            f"({result.param_dtype or '?'}"
+            + (f", quantization: {result.quant_method}" if result.quant_method else "")
+            + ")"
+        )
+        console.print(
+            f"  weight size    ~{result.weight_gb:.1f} GB "
+            f"-> suggested min_vram_gb: {result.suggested_min_vram_gb}"
+        )
+    else:
+        console.print(Text("  parameters     unknown (no safetensors metadata)", style="yellow"))
+    console.print(f"  context        native {result.max_position_embeddings or '?'}")
+    console.print(
+        f"  tool calling   {result.suggested_tool_call_parser or '(no parser guessed)'}"
+    )
+    if result.gated:
+        console.print(Text("  gated          yes - a real HF token is required", style="yellow"))
+    if result.is_hybrid_attention:
+        console.print(
+            Text("  hybrid attention  yes - needs Ampere-or-newer GPU", style="yellow")
+        )
+
+    for note in result.notes:
+        console.print(Text(f"  ! {note}", style="yellow"))
+
+    console.print("\n[bold]Sources[/bold]")
+    console.print(f"  https://huggingface.co/api/models/{repo_id}")
+    console.print(f"  https://huggingface.co/{repo_id}/raw/{revision}/config.json")
+
+    console.print("\n[bold]Suggested gpu_runtime: snippet[/bold]")
+    console.print(result.to_yaml_snippet())
+    console.print(
+        Text(
+            f"\nTip: `gpu inspect-model {repo_id} --add-profile <name>` writes this "
+            "straight into your model catalog instead of a one-off file.",
+            style="dim",
+        )
+    )
+
+
+def profiles_command(
+    ctx: typer.Context,
+    catalog_path_opt: Optional[Path] = typer.Option(
+        None, "--catalog", help=f"Catalog to list (default: nearest {CATALOG_FILENAMES[0]})."
+    ),
+) -> None:
+    """List the model profiles available in the model catalog (jambu.models.yaml)."""
+    context = get_context(ctx)
+    target = catalog_path_opt or find_catalog(Path.cwd())
+    if target is None or not target.is_file():
+        if context.json:
+            print_json({"catalog": None, "profiles": {}})
+        else:
+            console.print(
+                Text(
+                    f"no {CATALOG_FILENAMES[0]} found in this directory or any parent. "
+                    "Create one with `gpu inspect-model <repo> --add-profile <name>`.",
+                    style="yellow",
+                )
+            )
+        return
+
+    try:
+        catalog = load_catalog(target)
+    except JambuError as exc:
+        handle(exc)
+        return
+
+    if context.json:
+        print_json({"catalog": str(target), "profiles": catalog})
+        return
+
+    table = Table(title=f"Model profiles ({target})", title_justify="left")
+    for column in ("PROFILE", "MODEL", "MIN VRAM", "TOOL CALLING"):
+        table.add_column(column)
+    for name, fragment in sorted(catalog.items()):
+        if not isinstance(fragment, dict):
+            continue
+        model_id = (fragment.get("model") or {}).get("id", "?")
+        vram = ((fragment.get("compute") or {}).get("gpu") or {}).get("min_vram_gb", "-")
+        runtime = fragment.get("runtime") or {}
+        tool_calling = runtime.get("tool_call_parser") or runtime.get("tool_calling", "-")
+        table.add_row(name, str(model_id), str(vram), str(tool_calling))
+    console.print(table)
+    console.print(Text(f"  use with: model_profile: <name> in jambu.yaml", style="dim"))
+
+
+def _write_profile(result, name: str, catalog_path_opt: Optional[Path], *, force: bool) -> Path:
+    if catalog_path_opt:
+        target = catalog_path_opt
+    else:
+        target = find_catalog(Path.cwd()) or (Path.cwd() / CATALOG_FILENAMES[0])
+
+    catalog: dict = load_catalog(target) if target.is_file() else {}
+    if name in catalog and not force:
+        raise ConfigError(
+            f"profile '{name}' already exists in {target} (use --force to overwrite)"
+        )
+    catalog[name] = result.to_profile_fragment()
+    target.write_text(yaml.safe_dump(catalog, sort_keys=False))
+    return target

jambu_gpu/cli/lifecycle.py

@@ -0,0 +1,127 @@
+"""`gpu lock|unlock|stop|destroy|guard` - lifecycle control."""
+
+from __future__ import annotations
+
+import typer
+from rich.text import Text
+
+from ..core.durations import parse_duration
+from ..core.errors import ConfigError, JambuError
+from ..core.state import to_iso
+from ._common import console, get_context, handle, print_json
+
+
+def lock_command(
+    ctx: typer.Context,
+    duration: str = typer.Option(
+        "", "--for", help="Lock TTL, e.g. 2h. Omit to use lifecycle.lock.default_ttl."
+    ),
+    reason: str = typer.Option("", "--reason", help="Why the instance must stay up."),
+) -> None:
+    """Prevent automatic shutdown until the lock expires (spec section 11)."""
+    context = get_context(ctx)
+    try:
+        engine = context.engine()
+        if duration:
+            ttl = parse_duration(duration)
+        else:
+            ttl = engine.config.lock_default_ttl_s
+        result = engine.lock(ttl, reason=reason)
+    except ValueError as exc:
+        handle(ConfigError(str(exc)))
+        return
+    except JambuError as exc:
+        handle(exc)
+        return
+
+    if context.json:
+        print_json(result)
+        return
+    until = result.get("locked_until")
+    console.print(
+        Text(
+            f"+ locked until {to_iso(until)}" if until else "+ locked indefinitely",
+            style="yellow",
+        )
+    )
+
+
+def unlock_command(ctx: typer.Context) -> None:
+    """Resume normal lifecycle rules."""
+    context = get_context(ctx)
+    try:
+        result = context.engine().unlock()
+    except JambuError as exc:
+        handle(exc)
+        return
+    if context.json:
+        print_json(result)
+        return
+    console.print(Text("+ unlocked; idle and lifetime rules resume", style="green"))
+
+
+def stop_command(
+    ctx: typer.Context,
+    reason: str = typer.Option("manual", "--reason", help="Recorded with the stop event."),
+) -> None:
+    """Stop the compute resource (idempotent)."""
+    context = get_context(ctx)
+    try:
+        result = context.engine().stop(reason=reason)
+    except JambuError as exc:
+        handle(exc)
+        return
+    if context.json:
+        print_json(result)
+
+
+def destroy_command(
+    ctx: typer.Context,
+    yes: bool = typer.Option(False, "--yes", "-y", help="Skip the confirmation prompt."),
+    reason: str = typer.Option("manual", "--reason", help="Recorded with the destroy event."),
+) -> None:
+    """Destroy the remote resource (idempotent)."""
+    context = get_context(ctx)
+    try:
+        engine = context.engine()
+        if not yes and not context.json:
+            state = engine.store.load()
+            if state.instance_id:
+                confirm = typer.confirm(
+                    f"destroy instance {state.instance_id} on {engine.provider.display_name}?"
+                )
+                if not confirm:
+                    console.print("aborted")
+                    raise typer.Exit(1)
+        result = engine.destroy(reason=reason)
+    except JambuError as exc:
+        handle(exc)
+        return
+    if context.json:
+        print_json(result)
+
+
+def guard_command(
+    ctx: typer.Context,
+    dry_run: bool = typer.Option(
+        False, "--dry-run", help="Evaluate the policy without enforcing it."
+    ),
+) -> None:
+    """Evaluate the lifecycle rule now (external controller path, spec section 14).
+
+    The remote watchdog already does this on the instance. Use this from cron
+    or CI as a second, independent enforcement path.
+    """
+    context = get_context(ctx)
+    try:
+        result = context.engine().guard_tick(apply=not dry_run)
+    except JambuError as exc:
+        handle(exc)
+        return
+    if context.json:
+        print_json(result)
+        return
+    action = result.get("action")
+    reason = result.get("reason")
+    style = "yellow" if action != "none" else "green"
+    console.print(Text(f"{action}: {reason}", style=style))

jambu_gpu/cli/main.py

@@ -0,0 +1,103 @@
+"""gpu entrypoint."""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+from typing import Optional
+
+import typer
+
+from ..core.errors import JambuError
+from . import inspect as inspect_cli
+from . import lifecycle, providers, run, setup, status, validate
+from ._common import Context, err_console, get_context
+
+app = typer.Typer(
+    name="gpu",
+    help=(
+        "Provider-agnostic CLI for provisioning and managing temporary GPU runtimes.\n\n"
+        "jambu.yaml is the desired-state source of truth; .jambu/state.json is what "
+        "was last observed."
+    ),
+    no_args_is_help=True,
+    add_completion=False,
+    rich_markup_mode="rich",
+)
+
+
+@app.callback()
+def main_callback(
+    ctx: typer.Context,
+    config: Optional[Path] = typer.Option(
+        None, "--config", "-c", help="Path to jambu.yaml (default: nearest one upwards)."
+    ),
+    json_output: bool = typer.Option(
+        False, "--json", help="Machine-readable output on stdout."
+    ),
+    verbose: bool = typer.Option(False, "--verbose", "-v", help="Verbose diagnostics."),
+) -> None:
+    context = Context()
+    context.config_path = config
+    context.json = json_output
+    context.verbose = verbose
+    ctx.obj = context
+
+
+app.command("init")(setup.init_command)
+app.command("inspect-model")(inspect_cli.inspect_model_command)
+app.command("profiles")(inspect_cli.profiles_command)
+app.command("providers")(providers.providers_command)
+app.command("validate")(validate.validate_command)
+app.command("offers")(setup.offers_command)
+app.command("setup")(setup.setup_command)
+app.command(
+    "run",
+    context_settings={"allow_extra_args": True, "ignore_unknown_options": True},
+)(run.run_command)
+app.command("status")(status.status_command)
+app.command("logs")(status.logs_command)
+app.command("events")(status.events_command)
+app.command("lock")(lifecycle.lock_command)
+app.command("unlock")(lifecycle.unlock_command)
+app.command("stop")(lifecycle.stop_command)
+app.command("destroy")(lifecycle.destroy_command)
+app.command("guard")(lifecycle.guard_command)
+
+
+@app.command("endpoint")
+def endpoint_command(ctx: typer.Context) -> None:
+    """Print the model endpoint URL (handy for shell pipelines)."""
+    context = get_context(ctx)
+    try:
+        state = context.engine().store.load()
+    except JambuError as exc:
+        err_console.print(f"error: {exc}")
+        raise typer.Exit(getattr(exc, "exit_code", 1)) from exc
+    if not state.endpoint:
+        err_console.print("error: no endpoint; run `gpu setup` first")
+        raise typer.Exit(3)
+    print(state.endpoint)
+
+
+@app.command("version")
+def version_command() -> None:
+    """Print the CLI version."""
+    from .. import __version__
+
+    print(__version__)
+
+
+def main() -> None:
+    try:
+        app()
+    except JambuError as exc:  # safety net: errors are normalized, never tracebacks
+        err_console.print(f"error: {exc}")
+        sys.exit(getattr(exc, "exit_code", 1))
+    except KeyboardInterrupt:
+        err_console.print("interrupted")
+        sys.exit(130)
+
+
+if __name__ == "__main__":
+    main()

jambu_gpu/cli/providers.py

@@ -0,0 +1,96 @@
+"""`gpu providers` - capability matrix (spec section 7)."""
+
+from __future__ import annotations
+
+import typer
+from rich.table import Table
+from rich.text import Text
+
+from ..core.credentials import CredentialResolver
+from ..core.errors import ConfigError
+from ..providers.registry import load_builtin_providers
+from ._common import get_context, print_json
+
+
+def providers_command(ctx: typer.Context) -> None:
+    """List available providers and what each one can do."""
+    context = get_context(ctx)
+    registry = load_builtin_providers()
+
+    try:
+        config = context.config()
+        project_dir = config.project_dir
+        profile = config.provider.profile
+        selected = config.provider.name
+    except ConfigError:
+        config = None
+        project_dir = None
+        profile = "default"
+        selected = None
+
+    resolver = CredentialResolver(project_dir, profile=profile)
+    rows = []
+    for provider_cls in registry.classes():
+        missing = [n for n in provider_cls.required_credentials if not resolver.get(n)]
+        rows.append(
+            {
+                "provider": provider_cls.name,
+                "display_name": provider_cls.display_name,
+                "selected": provider_cls.name == selected,
+                "status": "configured" if not missing else "unavailable",
+                "missing_credentials": missing,
+                "required_credentials": list(provider_cls.required_credentials),
+                "capabilities": dict(provider_cls.capabilities.__dict__),
+            }
+        )
+
+    if context.json:
+        print_json({"providers": rows})
+        return
+
+    table = Table(title="Providers", title_justify="left")
+    table.add_column("PROVIDER")
+    table.add_column("STATUS")
+    table.add_column("GPU")
+    table.add_column("STOP")
+    table.add_column("DESTROY")
+    table.add_column("SSH")
+    table.add_column("PORTS")
+    table.add_column("SPOT")
+    table.add_column("PRICING")
+    for row in rows:
+        caps = row["capabilities"]
+        name = Text(row["provider"] + (" *" if row["selected"] else ""))
+        status = Text(
+            row["status"],
+            style="green" if row["status"] == "configured" else "yellow",
+        )
+        table.add_row(
+            name,
+            status,
+            _yn(caps["gpu_selection"]),
+            _yn(caps["stop"]),
+            _yn(caps["destroy"]),
+            _yn(caps["ssh"]),
+            _yn(caps["public_ports"]),
+            _yn(caps["spot_instances"]),
+            _yn(caps["pricing_query"]),
+        )
+    from ._common import console
+
+    console.print(table)
+    for row in rows:
+        if row["missing_credentials"]:
+            console.print(
+                Text(
+                    f"  {row['provider']}: set "
+                    + ", ".join(row["missing_credentials"]),
+                    style="yellow",
+                )
+            )
+    if selected:
+        console.print(Text(f"  * selected in jambu.yaml ({selected})", style="dim"))
+
+
+def _yn(value: bool) -> str:
+    return "yes" if value else "no"