freesolo-flash-dev 0.2.25__py3-none-any.whl

flash/providers/runpod/auth.py

@@ -0,0 +1,37 @@
+"""RunPod credential handling for the managed Flash backend (operator-side).
+
+The Flash SDK authenticates via the ``RUNPOD_API_KEY`` environment variable, set by
+the **operator** on the control-plane host. End users never
+provide provider credentials — they authenticate to the control plane with an Flash
+key. Deliberately env-only: ``~/.flash/config.json`` holds the *Flash* key, which
+must never be mistaken for a RunPod key.
+"""
+
+from __future__ import annotations
+
+from .._auth import load_provider_key
+from . import keys as _keys
+
+_ENV_VAR = "RUNPOD_API_KEY"
+
+
+def load_api_key() -> str | None:
+    """API key from the environment (operator configuration)."""
+    return load_provider_key(_ENV_VAR)
+
+
+def ensure_auth() -> str:
+    """Select the active account and collapse ``RUNPOD_API_KEY`` to that single key.
+
+    ``RUNPOD_API_KEY`` may be a comma-separated pool of per-account keys (see
+    ``runpod.keys``). The runpod_flash SDK reads the raw env var, so a multi-key value
+    would be sent as one bearer token (a 401); ``select_active`` collapses it to the
+    single active key while the cached pool keeps the rest for failover. Raises if no
+    key is configured.
+    """
+    key = _keys.select_active()
+    if not key:
+        raise RuntimeError(
+            "no RunPod API key found; set RUNPOD_API_KEY on the control-plane host"
+        )
+    return key

flash/providers/runpod/cost.py

@@ -0,0 +1,57 @@
+"""Shape RunPod realized billing rows into a RealizedCost. Pure shaping (offline-testable);
+the HTTP call is isolated in ``api.billing_endpoints``."""
+
+from __future__ import annotations
+
+from datetime import UTC, datetime
+from typing import Any
+
+from flash.providers.realized import RealizedCost
+
+
+def _get(row: Any, key: str) -> Any:
+    if isinstance(row, dict):
+        return row.get(key)
+    return getattr(row, key, None)
+
+
+def _iso(ts: float) -> str:
+    return datetime.fromtimestamp(float(ts), UTC).isoformat()
+
+
+def shape_endpoint_cost(rows: list[Any], *, endpoint_id: str | None) -> RealizedCost:
+    """Sum realized USD across the billing rows for ``endpoint_id``.
+
+    Each row is {endpointId, amount (USD), timeBilledMs, ...}. When ``endpoint_id`` is given we
+    keep only its rows (the response may be grouped/unfiltered); otherwise we sum everything
+    returned. RunPod's ``amount`` already includes disk, so it's reported as a single resource.
+    """
+    total = 0.0
+    billed_ms = 0
+    for row in rows:
+        if endpoint_id is not None:
+            row_eid = _get(row, "endpointId") or _get(row, "endpoint_id")
+            if row_eid is not None and str(row_eid) != str(endpoint_id):
+                continue
+        total += float(_get(row, "amount") or 0)
+        billed_ms += int(_get(row, "timeBilledMs") or 0)
+    total = round(total, 6)
+    return RealizedCost(
+        provider="runpod",
+        realized_usd=total,
+        by_resource={"gpu": total},
+        wall_seconds=(billed_ms / 1000.0) if billed_ms else None,
+        source={"endpoint_id": endpoint_id},
+    )
+
+
+def realized_cost(endpoint_id: str | None, *, start: float, end: float) -> RealizedCost | None:
+    """Pull + shape this run's realized RunPod cost; None when there's no endpoint to query."""
+    if not endpoint_id:
+        return None
+    from flash.providers.runpod import api
+
+    rows = api.billing_endpoints(
+        start_time=_iso(start), end_time=_iso(end), endpoint_id=endpoint_id
+    )
+    return shape_endpoint_cost(rows, endpoint_id=endpoint_id)

flash/providers/runpod/gpus.py

@@ -0,0 +1,46 @@
+"""RunPod's GPU classes + the Flash-specific bits of the shared GPU table.
+
+The class table itself is provider-agnostic and lives in ``providers/base.py`` (one
+canonical row per friendly name). This module carves out RunPod's rows
+(``gpu_classes()`` == every class with a Flash ``enum_member``) and owns the
+RunPod-only translation: friendly name -> Flash ``GpuType``.
+"""
+
+from __future__ import annotations
+
+from flash.providers.base import (
+    GpuClass,
+    UnsupportedGpuError,
+    get_gpu_info,
+    providers_for,
+)
+
+
+# Lazy import so unit tests that only exercise the mapping don't pull the whole SDK
+# graph unless needed. ``runpod_flash`` is a hard dependency, so this import is safe.
+def _gpu_enum():
+    from runpod_flash import GpuType
+
+    return GpuType
+
+
+def gpu_classes() -> list[GpuClass]:
+    """The GPU classes RunPod Flash can provision (those with a ``GpuType`` member)."""
+    from flash.providers.base import GPU_INFO
+
+    return [g for g in GPU_INFO.values() if g.enum_member]
+
+
+def flash_gpu(name: str):
+    """Return the RunPod Flash ``GpuType`` for a friendly GPU name."""
+    info = get_gpu_info(name)
+    if not info.enum_member:
+        raise UnsupportedGpuError(
+            f"{info.name} is not available on RunPod (providers: {', '.join(providers_for(name))})"
+        )
+    return getattr(_gpu_enum(), info.enum_member)
+
+
+def gpu_api_id(name: str) -> str:
+    """RunPod API GPU id (the ``GpuType`` enum value, e.g. 'NVIDIA GeForce RTX 4090')."""
+    return flash_gpu(name).value