Mimiry 0.2.1__py3-none-any.whl

mimiry/__init__.py

@@ -0,0 +1,35 @@
+"""mimiry — Python SDK for Mimiry GPU compute (softlaunch).
+
+v0.2.0 — wraps the existing /api/compute/v1/sessions API. See README.md.
+"""
+
+from mimiry._config import configure, get_config
+from mimiry.exceptions import (
+    MimiryError,
+    AuthError,
+    SessionError,
+    SessionFailed,
+    SessionTimeout,
+    ResultParseError,
+)
+from mimiry.function import function, Function
+from mimiry.image import Image
+from mimiry.run import run
+
+__version__ = "0.2.0"
+
+__all__ = [
+    "__version__",
+    "configure",
+    "get_config",
+    "function",
+    "Function",
+    "Image",
+    "run",
+    "MimiryError",
+    "AuthError",
+    "SessionError",
+    "SessionFailed",
+    "SessionTimeout",
+    "ResultParseError",
+]

mimiry/__main__.py

@@ -0,0 +1,8 @@
+"""Enables ``python -m mimiry ...`` (notably ``python -m mimiry setup``)."""
+
+import sys
+
+from mimiry._cli import main
+
+if __name__ == "__main__":
+    sys.exit(main())

mimiry/_auth.py

@@ -0,0 +1,188 @@
+"""SSH-JWT authentication against softlaunch.mimiry.com.
+
+Ports the algorithm from .claude/skills/mimiry-softlaunch/scripts/mimiry-auth.sh.
+We shell out to ``ssh-keygen`` for signing — implementing the SSH signature
+format in pure Python would mean re-deriving an OpenSSH-compatible format from
+the ``cryptography`` library, and ``ssh-keygen`` is universally available on
+the systems Mimiry users actually run on.
+
+Tokens last 1 hour (Mimiry default). The Token class refreshes itself when
+within 5 minutes of expiry.
+"""
+
+from __future__ import annotations
+
+import base64
+import os
+import secrets
+import shutil
+import subprocess
+import tempfile
+import time
+from dataclasses import dataclass
+from pathlib import Path
+
+import httpx
+
+from mimiry.exceptions import AuthError
+
+_TOKEN_REFRESH_BUFFER_SECONDS = 300  # refresh if < 5 min left
+_DEFAULT_TOKEN_TTL_SECONDS = 3600
+
+
+@dataclass
+class Token:
+    """A short-lived JWT issued by Mimiry. Self-refreshes when near expiry."""
+
+    access_token: str
+    expires_at: float  # unix seconds
+    fingerprint: str
+    ssh_key_path: Path
+    api_base: str
+
+    @property
+    def is_expired(self) -> bool:
+        return time.time() >= self.expires_at
+
+    @property
+    def is_near_expiry(self) -> bool:
+        return time.time() + _TOKEN_REFRESH_BUFFER_SECONDS >= self.expires_at
+
+    def refresh(self) -> "Token":
+        """Re-exchange the SSH signature for a fresh JWT. Mutates self."""
+        new = exchange_ssh_for_token(self.ssh_key_path, self.api_base)
+        self.access_token = new.access_token
+        self.expires_at = new.expires_at
+        return self
+
+    def get(self) -> str:
+        """Return the bearer string, refreshing if near expiry."""
+        if self.is_near_expiry:
+            self.refresh()
+        return self.access_token
+
+
+def _normalize_key_path(key_path: str | Path) -> Path:
+    """Accept either the private key path or the .pub path; return the private path."""
+    p = Path(key_path).expanduser()
+    if p.suffix == ".pub":
+        p = p.with_suffix("")
+    if not p.is_file():
+        raise AuthError(f"SSH private key not found: {p}")
+    if not p.with_suffix(p.suffix + ".pub").is_file() and not Path(f"{p}.pub").is_file():
+        raise AuthError(f"SSH public key not found: {p}.pub")
+    return p
+
+
+def _fingerprint(public_key_path: Path) -> str:
+    """Run ``ssh-keygen -lf`` to get the SHA256 fingerprint of a public key."""
+    if shutil.which("ssh-keygen") is None:
+        raise AuthError("ssh-keygen not found on PATH — required for Mimiry auth")
+    try:
+        out = subprocess.run(
+            ["ssh-keygen", "-lf", str(public_key_path)],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+    except subprocess.CalledProcessError as e:
+        raise AuthError(f"ssh-keygen -lf failed: {e.stderr.strip()}") from e
+    # Format: "<bits> <fingerprint> <comment> (<type>)"
+    parts = out.stdout.strip().split()
+    if len(parts) < 2:
+        raise AuthError(f"unexpected ssh-keygen output: {out.stdout!r}")
+    return parts[1]
+
+
+def _sign(message: bytes, private_key_path: Path) -> bytes:
+    """Sign ``message`` with the SSH key under the ``mimiry-auth`` namespace.
+
+    Returns the raw .sig file contents (OpenSSH SSHSIG format), as Mimiry expects.
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        msg_path = Path(tmpdir) / "msg"
+        msg_path.write_bytes(message)
+        try:
+            subprocess.run(
+                ["ssh-keygen", "-Y", "sign", "-f", str(private_key_path), "-n", "mimiry-auth",
+                 str(msg_path)],
+                capture_output=True,
+                check=True,
+            )
+        except subprocess.CalledProcessError as e:
+            raise AuthError(
+                f"ssh-keygen -Y sign failed: {e.stderr.decode(errors='replace').strip()}"
+            ) from e
+        sig_path = Path(f"{msg_path}.sig")
+        if not sig_path.is_file():
+            raise AuthError("ssh-keygen did not produce a .sig file")
+        return sig_path.read_bytes()
+
+
+def exchange_ssh_for_token(
+    ssh_key_path: str | Path,
+    api_base: str = "https://softlaunch.mimiry.com",
+) -> Token:
+    """Do the SSH-signature → JWT exchange. Returns a Token."""
+    api_base = api_base.rstrip("/")
+    priv = _normalize_key_path(ssh_key_path)
+    pub = Path(f"{priv}.pub")
+
+    fingerprint = _fingerprint(pub)
+    timestamp = str(int(time.time()))
+    nonce = secrets.token_hex(16)
+
+    message = f"{fingerprint}\n{timestamp}\n{nonce}".encode()
+    signature_bytes = _sign(message, priv)
+    signature_b64 = base64.b64encode(signature_bytes).decode()
+
+    try:
+        resp = httpx.post(
+            f"{api_base}/api/v1/auth/token",
+            headers={
+                "X-SSH-Fingerprint": fingerprint,
+                "X-SSH-Signature": signature_b64,
+                "X-SSH-Timestamp": timestamp,
+                "X-SSH-Nonce": nonce,
+                "Content-Type": "application/json",
+            },
+            json={"expires_in": _DEFAULT_TOKEN_TTL_SECONDS},
+            timeout=30.0,
+        )
+    except httpx.HTTPError as e:
+        raise AuthError(f"token exchange request failed: {e}") from e
+
+    if resp.status_code != 200:
+        raise AuthError(
+            f"token exchange returned {resp.status_code}: {resp.text[:500]}"
+        )
+
+    body = resp.json()
+    access_token = body.get("access_token")
+    if not access_token:
+        raise AuthError(f"token exchange response missing access_token: {body}")
+
+    expires_in = body.get("expires_in", _DEFAULT_TOKEN_TTL_SECONDS)
+    return Token(
+        access_token=access_token,
+        expires_at=time.time() + float(expires_in),
+        fingerprint=fingerprint,
+        ssh_key_path=priv,
+        api_base=api_base,
+    )
+
+
+def get_token(
+    ssh_key_path: str | Path | None = None,
+    api_base: str = "https://softlaunch.mimiry.com",
+) -> Token:
+    """Get a fresh Token. Falls back to ``MIMIRY_SSH_KEY`` env var when ``ssh_key_path`` is None."""
+    if ssh_key_path is None:
+        env_key = os.environ.get("MIMIRY_SSH_KEY")
+        if not env_key:
+            raise AuthError(
+                "no SSH key provided — pass ssh_key_path=, set MIMIRY_SSH_KEY env var, "
+                "or call mimiry.configure(ssh_key_path=...) first"
+            )
+        ssh_key_path = env_key
+    return exchange_ssh_for_token(ssh_key_path, api_base)

mimiry/_cli.py

@@ -0,0 +1,101 @@
+"""Minimal CLI for sanity-checking SDK auth + API connectivity.
+
+This is intentionally tiny in v1 — the SDK's primary surface is the Python
+decorator. The CLI exists so users can verify their auth setup before
+writing code:
+
+    $ mimiry balance
+    {"account_type": "user", "balance": 49.95, ...}
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+
+from mimiry._auth import get_token
+from mimiry._client import MimiryClient
+from mimiry._config import configure, get_config
+
+
+def _client() -> MimiryClient:
+    cfg = get_config()
+    token = get_token(cfg.ssh_key_path, cfg.api_base)
+    return MimiryClient(token)
+
+
+def cmd_balance(_: argparse.Namespace) -> int:
+    with _client() as c:
+        print(json.dumps(c.get_balance(), indent=2))
+    return 0
+
+
+def cmd_quota(_: argparse.Namespace) -> int:
+    with _client() as c:
+        print(json.dumps(c.get_quota(), indent=2))
+    return 0
+
+
+def cmd_availability(args: argparse.Namespace) -> int:
+    with _client() as c:
+        params = {}
+        if args.gpu_family:
+            params["gpu_family"] = args.gpu_family
+        print(json.dumps(c.get_availability(**params), indent=2))
+    return 0
+
+
+def cmd_token(_: argparse.Namespace) -> int:
+    """Print a fresh JWT — useful for piping into curl during debugging."""
+    cfg = get_config()
+    token = get_token(cfg.ssh_key_path, cfg.api_base)
+    print(token.access_token)
+    return 0
+
+
+def cmd_setup(args: argparse.Namespace) -> int:
+    """Interactive first-time auth wizard (also aliased as ``mimiry init``)."""
+    from mimiry._setup import run_setup_wizard
+
+    return run_setup_wizard(ssh_key_path=args.ssh_key, api_base=args.api_base)
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(prog="mimiry", description=__doc__.splitlines()[0])
+    parser.add_argument(
+        "--ssh-key",
+        help="Path to SSH private key (overrides MIMIRY_SSH_KEY env var).",
+    )
+    parser.add_argument("--api-base", help="API base URL (default: softlaunch.mimiry.com).")
+    subs = parser.add_subparsers(dest="cmd", required=True)
+
+    subs.add_parser("balance", help="Show account balance.").set_defaults(func=cmd_balance)
+    subs.add_parser("quota", help="Show account quota / usage.").set_defaults(func=cmd_quota)
+
+    avail = subs.add_parser("availability", help="Show GPU availability (no auth).")
+    avail.add_argument("--gpu-family", help="Filter, e.g. T4 or H100.")
+    avail.set_defaults(func=cmd_availability)
+
+    subs.add_parser("token", help="Print a fresh JWT (for debugging).").set_defaults(func=cmd_token)
+
+    setup_help = "Interactive auth setup wizard (generate/register SSH key, verify)."
+    subs.add_parser("setup", help=setup_help).set_defaults(func=cmd_setup)
+    subs.add_parser("init", help="Alias for `setup`.").set_defaults(func=cmd_setup)
+
+    args = parser.parse_args(argv)
+
+    if args.ssh_key or args.api_base:
+        configure(ssh_key_path=args.ssh_key, api_base=args.api_base)
+
+    try:
+        return args.func(args)
+    except KeyboardInterrupt:
+        return 130
+    except Exception as e:
+        print(f"error: {e}", file=sys.stderr)
+        return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

mimiry/_client.py

@@ -0,0 +1,122 @@
+"""Thin HTTP wrapper over the Mimiry compute API.
+
+Stateless except for the Token. Higher-level lifecycle helpers (polling for
+state transitions, scanning logs for sentinels) live in ``_session.py``.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import httpx
+
+from mimiry._auth import Token
+from mimiry.exceptions import SessionError
+
+
+class MimiryClient:
+    """Thin client for /api/v1/* and /api/compute/v1/*."""
+
+    def __init__(self, token: Token, http_timeout: float = 30.0) -> None:
+        self._token = token
+        self._compute_base = f"{token.api_base}/api/compute/v1"
+        self._http = httpx.Client(timeout=http_timeout)
+
+    def close(self) -> None:
+        self._http.close()
+
+    def __enter__(self) -> "MimiryClient":
+        return self
+
+    def __exit__(self, *exc: Any) -> None:
+        self.close()
+
+    @property
+    def token(self) -> Token:
+        return self._token
+
+    def _headers(self) -> dict[str, str]:
+        return {
+            "Authorization": f"Bearer {self._token.get()}",
+            "Content-Type": "application/json",
+        }
+
+    def _request(self, method: str, path: str, **kwargs: Any) -> httpx.Response:
+        url = f"{self._compute_base}{path}"
+        try:
+            resp = self._http.request(method, url, headers=self._headers(), **kwargs)
+        except httpx.HTTPError as e:
+            raise SessionError(f"{method} {path} request failed: {e}") from e
+        return resp
+
+    @staticmethod
+    def _json_or_raise(resp: httpx.Response) -> Any:
+        if resp.status_code >= 400:
+            try:
+                body = resp.json()
+            except Exception:
+                body = resp.text[:500]
+            raise SessionError(f"HTTP {resp.status_code} on {resp.request.url}: {body}")
+        return resp.json()
+
+    # ────────── balance / quota / availability ──────────
+
+    def get_balance(self) -> dict:
+        return self._json_or_raise(self._request("GET", "/balance"))
+
+    def get_quota(self) -> dict:
+        return self._json_or_raise(self._request("GET", "/quota"))
+
+    def get_availability(self, **params: Any) -> dict:
+        """Public endpoint — no auth required, but it's convenient to call from the client."""
+        return self._json_or_raise(self._http.get(f"{self._compute_base}/availability", params=params))
+
+    # ────────── sessions ──────────
+
+    def create_session(self, payload: dict) -> dict:
+        """POST /sessions. Returns the initial session object (state=submitted)."""
+        return self._json_or_raise(self._request("POST", "/sessions", json=payload))
+
+    def get_session(self, session_id: str, *, events_tail: int | None = None) -> dict:
+        params = {}
+        if events_tail is not None:
+            params["events_tail"] = events_tail
+        return self._json_or_raise(self._request("GET", f"/sessions/{session_id}", params=params))
+
+    def list_sessions(self, **params: Any) -> list[dict]:
+        body = self._json_or_raise(self._request("GET", "/sessions", params=params))
+        return body.get("sessions", body) if isinstance(body, dict) else body
+
+    def terminate_session(self, session_id: str) -> dict | None:
+        """DELETE /sessions/{id}. Returns the response body or None on 202/204."""
+        resp = self._request("DELETE", f"/sessions/{session_id}")
+        if resp.status_code in (202, 204):
+            return None
+        return self._json_or_raise(resp)
+
+    def get_logs(self, session_id: str, *, tail: int = 200, timestamps: bool = False) -> dict:
+        """GET /sessions/{id}/logs.
+
+        Returns a dict like ``{"logs": "<text>"}`` on 200, or
+        ``{"retry_after_seconds": N}`` on 503 (container still booting).
+
+        Caller is responsible for retry/backoff loops — see :func:`mimiry._session.wait_for_marker`.
+        """
+        resp = self._request(
+            "GET",
+            f"/sessions/{session_id}/logs",
+            params={"tail": tail, "timestamps": "true" if timestamps else "false"},
+        )
+        if resp.status_code == 503:
+            try:
+                body = resp.json()
+            except Exception:
+                body = {"retry_after_seconds": 5}
+            return {"_status": 503, **body}
+        if resp.status_code == 409:
+            try:
+                body = resp.json()
+            except Exception:
+                body = {"message": resp.text}
+            return {"_status": 409, **body}
+        return {"_status": 200, **self._json_or_raise(resp)}

mimiry/_config.py

@@ -0,0 +1,58 @@
+"""Process-global SDK configuration. Holds auth + API base URL.
+
+Users either set MIMIRY_SSH_KEY env var or call ``mimiry.configure(ssh_key_path=...)``
+once. Subsequent calls override.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+
+DEFAULT_API_BASE = "https://softlaunch.mimiry.com"
+DEFAULT_TIMEOUT_SECONDS = 1800  # 30 min — must accommodate ~2 min cold start + work
+
+
+@dataclass
+class Config:
+    ssh_key_path: Path | None = None
+    api_base: str = DEFAULT_API_BASE
+    timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS
+    poll_interval_seconds: float = 5.0
+    # Tighter than state polling — the marker-emission → auto_terminate window is
+    # only the wrapper's tail sleep (~20s), so we need ~2s cadence to catch it reliably.
+    log_poll_interval_seconds: float = 2.0
+    extras: dict = field(default_factory=dict)
+
+
+_config = Config()
+
+
+def configure(
+    ssh_key_path: str | Path | None = None,
+    api_base: str | None = None,
+    timeout_seconds: int | None = None,
+    poll_interval_seconds: float | None = None,
+) -> Config:
+    """Set SDK-wide configuration. Any argument left as ``None`` is unchanged."""
+    if ssh_key_path is not None:
+        _config.ssh_key_path = Path(ssh_key_path).expanduser()
+    if api_base is not None:
+        _config.api_base = api_base.rstrip("/")
+    if timeout_seconds is not None:
+        _config.timeout_seconds = timeout_seconds
+    if poll_interval_seconds is not None:
+        _config.poll_interval_seconds = poll_interval_seconds
+    return _config
+
+
+def get_config() -> Config:
+    """Return the active configuration. Auto-bootstraps from env vars on first read."""
+    if _config.ssh_key_path is None:
+        env_key = os.environ.get("MIMIRY_SSH_KEY")
+        if env_key:
+            _config.ssh_key_path = Path(env_key).expanduser()
+    if "MIMIRY_API_BASE" in os.environ:
+        _config.api_base = os.environ["MIMIRY_API_BASE"].rstrip("/")
+    return _config