jambu-gpu 0.1.0__py3-none-any.whl

jambu_gpu/providers/base.py

@@ -0,0 +1,125 @@
+"""The provider contract (spec section 6).
+
+Every provider implements exactly this interface. Provider-specific APIs,
+SDKs, identifiers, pricing structures and authentication MUST NOT leak
+through it. The core branches on ``capabilities``, never on ``name``.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+from ..core.config import RuntimeConfig
+from ..core.credentials import CredentialResolver
+from ..core.errors import UnsupportedCapabilityError
+from ..core.models import (
+    ComputeSpec,
+    Execution,
+    Instance,
+    InstanceStatus,
+    LogStream,
+    Offer,
+    ProviderCapabilities,
+    ValidationResult,
+    Workload,
+)
+
+
+@dataclass
+class WatchdogSpec:
+    """What the adapter contributes to the remote lifecycle watchdog.
+
+    The *policy* lives in the core; the adapter only supplies the provider
+    calls the watchdog needs to shut its own machine down (spec section 14).
+    """
+
+    # Python source executed inside the watchdog. Must define
+    #   def provider_stop(ctx) -> None
+    #   def provider_destroy(ctx) -> None
+    # using the stdlib only.
+    actions_source: str
+    # Environment variables the instance needs for those calls (credentials).
+    env: dict[str, str] = field(default_factory=dict)
+    # Static values available to the actions through ctx["provider"].
+    context: dict[str, Any] = field(default_factory=dict)
+
+
+class GPUProvider(ABC):
+    """Stable internal contract implemented by every adapter."""
+
+    name: str = "abstract"
+    display_name: str = "Abstract"
+    required_credentials: list[str] = []
+    capabilities: ProviderCapabilities = ProviderCapabilities()
+
+    def __init__(self, config: RuntimeConfig, credentials: CredentialResolver) -> None:
+        self.config = config
+        self.credentials = credentials
+
+    # -- contract -----------------------------------------------------------
+
+    @abstractmethod
+    def validate(self, config: RuntimeConfig) -> ValidationResult:
+        """Check credentials, capabilities and the requested spec. No side effects."""
+
+    @abstractmethod
+    def setup(self, spec: ComputeSpec) -> Instance:
+        """Provision compute matching ``spec`` and return a normalized handle."""
+
+    @abstractmethod
+    def start(self, instance: Instance) -> Instance:
+        """Bring a stopped instance back to running."""
+
+    @abstractmethod
+    def run(self, instance: Instance, workload: Workload) -> Execution:
+        """Execute a workload on the instance (requires ``capabilities.remote_exec``)."""
+
+    @abstractmethod
+    def status(self, instance: Instance) -> InstanceStatus:
+        """Observed provider-side state, normalized."""
+
+    @abstractmethod
+    def stop(self, instance: Instance) -> None:
+        """Stop compute. MUST be idempotent (spec section 23)."""
+
+    @abstractmethod
+    def destroy(self, instance: Instance) -> None:
+        """Destroy the remote resource. MUST be idempotent."""
+
+    @abstractmethod
+    def logs(self, instance: Instance) -> LogStream:
+        """Provider-side logs for the instance."""
+
+    # -- optional, capability-gated ----------------------------------------
+
+    def find_offers(self, spec: ComputeSpec, limit: int = 10) -> list[Offer]:
+        """Discover purchasable capacity (requires ``capabilities.pricing_query``)."""
+        raise UnsupportedCapabilityError(f"{self.name} cannot query offers")
+
+    def watchdog_spec(self, spec_or_instance: Any) -> Optional[WatchdogSpec]:
+        """Provider calls the remote watchdog uses to stop/destroy itself."""
+        return None
+
+    def refresh(self, instance: Instance) -> Instance:
+        """Re-read provider data into the normalized handle (ports, ip, state)."""
+        status = self.status(instance)
+        return status.instance or instance
+
+    # -- helpers ------------------------------------------------------------
+
+    def require_capability(self, name: str) -> None:
+        if not getattr(self.capabilities, name, False):
+            raise UnsupportedCapabilityError(
+                f"provider '{self.name}' does not support capability '{name}'"
+            )
+
+    def credential_status(self) -> dict[str, bool]:
+        return {name: bool(self.credentials.get(name)) for name in self.required_credentials}
+
+    def is_configured(self) -> bool:
+        return all(self.credential_status().values()) if self.required_credentials else True
+
+    def __repr__(self) -> str:  # pragma: no cover - debug helper
+        return f"<{type(self).__name__} name={self.name!r}>"

jambu_gpu/providers/registry.py

@@ -0,0 +1,63 @@
+"""Provider registry (spec section 20).
+
+Adding a provider means adding ``providers/<name>/`` and registering the class.
+No lifecycle, execution or CLI code changes.
+"""
+
+from __future__ import annotations
+
+from typing import Iterator, Optional, Type
+
+from ..core.config import RuntimeConfig
+from ..core.credentials import CredentialResolver
+from ..core.errors import ConfigError
+from .base import GPUProvider
+
+
+class ProviderRegistry:
+    def __init__(self) -> None:
+        self._providers: dict[str, Type[GPUProvider]] = {}
+
+    def register(self, provider_cls: Type[GPUProvider]) -> Type[GPUProvider]:
+        name = provider_cls.name
+        if not name or name == "abstract":
+            raise ValueError(f"{provider_cls!r} must define a provider name")
+        self._providers[name] = provider_cls
+        return provider_cls
+
+    def names(self) -> list[str]:
+        return sorted(self._providers)
+
+    def classes(self) -> Iterator[Type[GPUProvider]]:
+        for name in self.names():
+            yield self._providers[name]
+
+    def get_class(self, name: str) -> Type[GPUProvider]:
+        try:
+            return self._providers[name]
+        except KeyError:
+            known = ", ".join(self.names()) or "<none>"
+            raise ConfigError(f"unknown provider '{name}' (available: {known})") from None
+
+    def get(
+        self,
+        config: RuntimeConfig,
+        credentials: Optional[CredentialResolver] = None,
+        name: Optional[str] = None,
+    ) -> GPUProvider:
+        provider_cls = self.get_class(name or config.provider.name)
+        resolver = credentials or CredentialResolver(
+            config.project_dir, profile=config.provider.profile
+        )
+        return provider_cls(config, resolver)
+
+
+registry = ProviderRegistry()
+
+
+def load_builtin_providers() -> ProviderRegistry:
+    """Import and register every adapter shipped with the CLI."""
+    from .vast.provider import VastProvider  # noqa: WPS433 (deferred import by design)
+
+    registry.register(VastProvider)
+    return registry

jambu_gpu/providers/vast/__init__.py

@@ -0,0 +1,3 @@
+from .provider import VastProvider
+
+__all__ = ["VastProvider"]

jambu_gpu/providers/vast/client.py

@@ -0,0 +1,230 @@
+"""Thin HTTP client for the Vast.ai REST API.
+
+Only this module knows Vast's wire format. It raises normalized errors so
+nothing Vast-shaped escapes the adapter.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any, Optional
+
+import httpx
+
+from ...core.errors import (
+    AuthenticationError,
+    CapacityUnavailableError,
+    ProviderError,
+    ProviderTimeoutError,
+    ProvisioningError,
+)
+from .config import API_BASE
+
+
+class VastNotFound(ProviderError):
+    """The requested Vast object does not exist (any more)."""
+
+
+class VastClient:
+    def __init__(
+        self,
+        api_key: str,
+        api_base: str = API_BASE,
+        timeout: float = 60.0,
+        client: Optional[httpx.Client] = None,
+    ) -> None:
+        self.api_key = api_key
+        self.api_base = api_base.rstrip("/")
+        self.timeout = timeout
+        self._client = client
+
+    # -- transport ----------------------------------------------------------
+
+    @property
+    def client(self) -> httpx.Client:
+        if self._client is None:
+            self._client = httpx.Client(
+                timeout=self.timeout,
+                headers={
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Accept": "application/json",
+                    "Content-Type": "application/json",
+                    "User-Agent": "jambu-gpu/0.1",
+                },
+            )
+        return self._client
+
+    def close(self) -> None:
+        if self._client is not None:
+            self._client.close()
+            self._client = None
+
+    def request(
+        self,
+        method: str,
+        path: str,
+        *,
+        json_body: Optional[dict] = None,
+        params: Optional[dict] = None,
+    ) -> Any:
+        url = path if path.startswith("http") else f"{self.api_base}{path}"
+        try:
+            response = self.client.request(method, url, json=json_body, params=params)
+        except httpx.TimeoutException as exc:
+            raise ProviderTimeoutError(f"vast.ai timed out on {method} {path}") from exc
+        except httpx.HTTPError as exc:
+            raise ProviderError(f"vast.ai request failed ({method} {path}): {exc}") from exc
+
+        # These two are reliable regardless of what's in the body.
+        if response.status_code == 404:
+            raise VastNotFound(f"vast.ai object not found: {method} {path}")
+        if response.status_code in (401, 403):
+            raise AuthenticationError(
+                "vast.ai rejected the API key (check VAST_API_KEY at "
+                "https://cloud.vast.ai/account/)"
+            )
+
+        # Vast returns its own {"success": false, "msg": "..."} error shape
+        # on other 4xx responses just as often as on 200 - the old behavior
+        # (raise a blunt ProviderError for any >= 400 BEFORE ever looking at
+        # the body) intercepted real, specifically classifiable errors like
+        # "no_such_ask ... is not available" and turned them into a generic
+        # ProviderError that setup()'s offer-retry loop can't walk past.
+        # Parse the body first and classify by its content when there is one;
+        # only fall back to the bare status code when there isn't.
+        data: Any = None
+        if response.content:
+            try:
+                data = response.json()
+            except (ValueError, json.JSONDecodeError):
+                data = None
+
+        if isinstance(data, dict) and data.get("success") is False:
+            message = str(data.get("msg") or data.get("error") or data)
+            # Vast spells the same "this offer is gone" failure differently
+            # across endpoints/versions ("no_such_ask" underscored, "no such
+            # ask" spaced, "is not available" two words, "unavailable" one) -
+            # normalize separators before matching so none of them slip
+            # through as a generic ProviderError.
+            normalized = message.lower().replace("_", " ").replace("-", " ")
+            lowered = message.lower()
+            if (
+                "no such ask" in normalized
+                or "not available" in normalized
+                or "unavailable" in normalized
+                or "taken" in normalized
+                or "no longer" in normalized
+            ):
+                raise CapacityUnavailableError(f"vast.ai: {message}")
+            if "insufficient" in lowered or "credit" in lowered or "balance" in lowered:
+                raise ProvisioningError(f"vast.ai: {message}")
+            raise ProviderError(f"vast.ai: {message}")
+
+        if response.status_code == 429:
+            raise ProviderError("vast.ai rate limit hit; retry in a moment")
+        if response.status_code >= 400:
+            raise ProviderError(
+                f"vast.ai error {response.status_code} on {method} {path}: "
+                f"{response.text[:400]}"
+            )
+
+        if data is None:
+            return {"raw": response.text} if response.content else {}
+        return data
+
+    # -- account ------------------------------------------------------------
+
+    def whoami(self) -> dict:
+        data = self.request("GET", "/users/current/")
+        return data if isinstance(data, dict) else {}
+
+    # -- offers -------------------------------------------------------------
+
+    def search_offers(self, query: dict) -> list[dict]:
+        """Vast has shipped three shapes of this endpoint; try them in order.
+
+        Stops at the first shape that *answers in the expected shape*
+        (a dict with an "offers" key), even when that list is legitimately
+        empty - a correct "no matches" MUST NOT fall through to the next
+        (possibly incompatible) shape, or a real empty result gets replaced
+        by a confusing error from a fallback shape that was never going to
+        work. Only a shape that outright fails (raises) is skipped.
+        """
+        attempts = (
+            ("PUT", "/search/asks/", {"q": query}, None),
+            ("PUT", "/search/asks/", query, None),
+            ("GET", "/bundles/", None, {"q": json.dumps(query)}),
+        )
+        last_error: Optional[Exception] = None
+        for method, path, body, params in attempts:
+            try:
+                data = self.request(method, path, json_body=body, params=params)
+            except (ProviderError, VastNotFound) as exc:
+                last_error = exc
+                continue
+            if isinstance(data, dict) and "offers" in data:
+                return list(data["offers"] or [])
+        if last_error is not None and not isinstance(last_error, VastNotFound):
+            raise last_error
+        return []
+
+    # -- instances ----------------------------------------------------------
+
+    def create_instance(self, ask_id: str, payload: dict) -> dict:
+        data = self.request("PUT", f"/asks/{ask_id}/", json_body=payload)
+        if not isinstance(data, dict) or not data.get("new_contract"):
+            raise ProvisioningError(f"vast.ai did not return an instance id: {data}")
+        return data
+
+    def list_instances(self) -> list[dict]:
+        data = self.request("GET", "/instances/", params={"owner": "me"})
+        if isinstance(data, dict):
+            return list(data.get("instances") or [])
+        return []
+
+    def get_instance(self, instance_id: str) -> Optional[dict]:
+        try:
+            data = self.request("GET", f"/instances/{instance_id}/")
+        except VastNotFound:
+            return None
+        if isinstance(data, dict):
+            found = data.get("instances")
+            if isinstance(found, list):
+                return found[0] if found else None
+            if isinstance(found, dict):
+                return found or None
+        # Some deployments answer the single-instance route with a bare object.
+        for candidate in self.list_instances():
+            if str(candidate.get("id")) == str(instance_id):
+                return candidate
+        return None
+
+    def set_instance_state(self, instance_id: str, state: str) -> dict:
+        data = self.request("PUT", f"/instances/{instance_id}/", json_body={"state": state})
+        return data if isinstance(data, dict) else {}
+
+    def destroy_instance(self, instance_id: str) -> dict:
+        try:
+            data = self.request("DELETE", f"/instances/{instance_id}/")
+        except VastNotFound:
+            return {"success": True, "already_gone": True}
+        return data if isinstance(data, dict) else {}
+
+    def request_logs(self, instance_id: str, tail: int = 500) -> Optional[str]:
+        data = self.request(
+            "PUT",
+            f"/instances/request_logs/{instance_id}/",
+            json_body={"tail": str(tail)},
+        )
+        if not isinstance(data, dict):
+            return None
+        return data.get("result_url")
+
+    def fetch_url(self, url: str) -> str:
+        try:
+            response = httpx.get(url, timeout=self.timeout)
+        except httpx.HTTPError as exc:
+            raise ProviderError(f"could not download vast.ai logs: {exc}") from exc
+        if response.status_code >= 400:
+            raise ProviderError(f"vast.ai log download failed ({response.status_code})")
+        return response.text

jambu_gpu/providers/vast/config.py

@@ -0,0 +1,51 @@
+"""Vast.ai adapter settings.
+
+Everything here is provider-local. Nothing in this module may be referenced
+by the core.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+API_BASE = "https://console.vast.ai/api/v0"
+API_KEY_ENV = "VAST_API_KEY"
+
+# runtype that gives us sshd plus our own onstart script (the image entrypoint
+# is replaced by Vast's supervisor, so we launch the model server ourselves).
+DEFAULT_RUNTYPE = "ssh_direc ssh_proxy"
+
+DEFAULT_SEARCH_LIMIT = 24
+DEFAULT_MIN_RELIABILITY = 0.95
+DEFAULT_MIN_INET_DOWN_MBPS = 100.0
+DEFAULT_MIN_CUDA = 12.1
+
+
+@dataclass
+class VastOptions:
+    """`provider.options` in jambu.yaml (gpu_runtime.provider.options), all optional."""
+
+    api_base: str = API_BASE
+    runtype: str = DEFAULT_RUNTYPE
+    verified_only: bool = True
+    min_reliability: float = DEFAULT_MIN_RELIABILITY
+    min_inet_down_mbps: float = DEFAULT_MIN_INET_DOWN_MBPS
+    min_cuda: float = DEFAULT_MIN_CUDA
+    search_limit: int = DEFAULT_SEARCH_LIMIT
+    order_by: str = "dph_total"  # dph_total | score | inet_down
+    order_dir: str = "asc"
+    image_login: Optional[str] = None
+    # Extra raw filters merged into the offer query, escape hatch for power users.
+    search_filters: Optional[dict[str, Any]] = None
+
+    @classmethod
+    def from_mapping(cls, data: Optional[dict[str, Any]]) -> "VastOptions":
+        data = data or {}
+        known = {f for f in cls.__dataclass_fields__}  # type: ignore[attr-defined]
+        unknown = set(data) - known
+        if unknown:
+            raise ValueError(
+                "unknown provider.options for vast: " + ", ".join(sorted(unknown))
+            )
+        return cls(**{k: v for k, v in data.items() if k in known})