dataspires 0.1.0__py3-none-any.whl

afrilink/__init__.py

@@ -0,0 +1,205 @@
+"""
+DataSpires SDK — Train, finetune, and generate on GPUs from your notebook.
+
+Quick start:
+    from dataspires import DataSpiresClient
+    client = DataSpiresClient()
+    client.authenticate()
+
+    ft = client.finetune(model="qwen2.5-0.5b", training_mode="low", data=my_dataset, gpus=1)
+    result = ft.run(wait=True)
+
+User guide:
+    import dataspires
+    dataspires.docs("help")
+"""
+
+__version__ = "0.1.0"
+
+# Credentials provisioning (auto-runs in Colab)
+from .credentials import (
+    provision_credentials,
+    get_cineca_credentials,
+    CredentialProvider,
+    OrgCredentials,
+)
+
+# Core authentication
+from .auth import AfriLinkAuth, authenticate, AuthResult
+from .cineca_auth import CinecaDirectAuth, authenticate_cineca, CinecaAuthResult
+
+# Telemetry
+from .telemetry import TelemetryClient, JobUsageTracker, track_cineca_jobs
+
+# Finetuning API
+from .finetune import finetune, FinetuneJob, FinetuneJobSpec, TrainingMode, TrainingConfig, SUPPORTED_BACKENDS, DEFAULT_BACKEND
+
+# Generate API (frozen run)
+from .generate import generate, GenerateJob, GenerateJobSpec, GENERATE_RUNNER_MODELS
+
+# Pretrain API (managed recipes + architecture scratch)
+from .pretrain import pretrain, PretrainJob, PretrainJobSpec
+
+# Training API (deprecated — use pretrain)
+from .train import train, TrainJob, TrainJobSpec, CONTAINER_REGISTRY
+
+# Routing helpers
+from .routing import KIND_REGISTRY, ARCHITECTURE_CATALOG, ArchitectureClass, get_kind
+
+# SLURM job management
+from .slurm import SlurmJobManager, SlurmConfig
+
+# Distributed training
+from .distributed import DistributedConfig, DistributedStrategy, get_optimal_strategy
+
+# Model/dataset registry
+from .registry import (
+    ModelRegistry,
+    ModelInfo,
+    DatasetInfo,
+    ModelSize,
+    ModelType,
+    get_registry,
+    list_models,
+    list_datasets,
+)
+
+# Vouchers
+from .vouchers import generate_voucher, validate_voucher, generate_batch
+
+# Data transfer
+from .transfer import DataTransferManager, create_transfer_manager
+
+# Session watchdog & recovery
+from .watchdog import SessionWatchdog, EmailNotifier, JobRecoveryInfo
+
+# HuggingFace inference
+from .inference import HFInferenceClient, InferenceResult
+
+# Custom-image build service (Cloud Build + Artifact Registry)
+from .build import BuildSpec, ImageBuilder, ModelSource, PRESETS as BUILD_PRESETS
+from .gcp_auth import GcpServiceAccountAuth
+
+# Main client (recommended entry point)
+from .client import AfriLinkClient, ClientConfig, create_client
+
+# Public name for the same client class.
+DataSpiresClient = AfriLinkClient
+
+# Help / manual system — prefer: import afrilink; afrilink.docs("help")
+from .help import help_system as _help_system, _resolve_topic as _resolve_help_topic
+
+import sys as _sys
+import types as _types
+
+
+class _AfriLinkModule(_types.ModuleType):
+    """Module subclass that implements afrilink / \"topic\" slash help."""
+
+    def __truediv__(self, topic):
+        _help_system._show(_resolve_help_topic(topic))
+        return self
+
+    def docs(self, topic: str = "help"):
+        """Print an inline help page. Preferred over slash syntax."""
+        _help_system._show(topic)
+        return None
+
+    def __repr__(self):
+        return (
+            f"<module 'afrilink' v{__version__}> — "
+            'run  afrilink.docs("help")  for the user guide.'
+        )
+
+
+# Replace this module with the instrumented subclass so that
+# `afrilink / "specs"` and `afrilink.docs("specs")` work at module level.
+_current = _sys.modules[__name__]
+_upgraded = _AfriLinkModule(__name__, __doc__)
+_upgraded.__dict__.update({k: v for k, v in _current.__dict__.items()
+                            if not k.startswith("_upgraded")})
+_sys.modules[__name__] = _upgraded
+
+
+__all__ = [
+    "__version__",
+    # Main client
+    "AfriLinkClient",
+    "DataSpiresClient",
+    "ClientConfig",
+    "create_client",
+    # Credentials
+    "provision_credentials",
+    "get_cineca_credentials",
+    "CredentialProvider",
+    "OrgCredentials",
+    # Authentication
+    "AfriLinkAuth",
+    "authenticate",
+    "AuthResult",
+    "CinecaDirectAuth",
+    "authenticate_cineca",
+    "CinecaAuthResult",
+    # Finetuning
+    "finetune",
+    "FinetuneJob",
+    "FinetuneJobSpec",
+    "TrainingMode",
+    "TrainingConfig",
+    "SUPPORTED_BACKENDS",
+    "DEFAULT_BACKEND",
+    # Generate (frozen run)
+    "generate",
+    "GenerateJob",
+    "GenerateJobSpec",
+    "GENERATE_RUNNER_MODELS",
+    # Pretrain (managed recipes + architecture scratch)
+    "pretrain",
+    "PretrainJob",
+    "PretrainJobSpec",
+    "KIND_REGISTRY",
+    "ARCHITECTURE_CATALOG",
+    "ArchitectureClass",
+    "get_kind",
+    # Training (deprecated — use pretrain)
+    "train",
+    "TrainJob",
+    "TrainJobSpec",
+    "CONTAINER_REGISTRY",
+    # SLURM
+    "SlurmJobManager",
+    "SlurmConfig",
+    # Distributed
+    "DistributedConfig",
+    "DistributedStrategy",
+    "get_optimal_strategy",
+    # Registry
+    "ModelRegistry",
+    "ModelInfo",
+    "DatasetInfo",
+    "ModelSize",
+    "ModelType",
+    "get_registry",
+    "list_models",
+    "list_datasets",
+    # Vouchers
+    "generate_voucher",
+    "validate_voucher",
+    "generate_batch",
+    # Transfer
+    "DataTransferManager",
+    "create_transfer_manager",
+    # Telemetry
+    "TelemetryClient",
+    "JobUsageTracker",
+    "track_cineca_jobs",
+    # Watchdog & recovery
+    "SessionWatchdog",
+    "EmailNotifier",
+    "JobRecoveryInfo",
+    # HuggingFace inference
+    "HFInferenceClient",
+    "InferenceResult",
+    # Help system
+    "help_system",
+]

afrilink/adapt.py

@@ -0,0 +1,60 @@
+"""
+Deprecated adapt() — use pretrain(kind=..., pretrained=True).
+
+Kept as a thin wrapper for one release; removed in 1.0.0.
+"""
+
+from __future__ import annotations
+
+import warnings
+from typing import Any, Dict, Optional
+
+from .train import TrainJob, TrainJobSpec
+
+AdaptJob = TrainJob
+AdaptJobSpec = TrainJobSpec
+
+
+def adapt(
+    recipe: Optional[str] = None,
+    data: Any = None,
+    weights: Optional[str] = None,
+    model: Optional[str] = None,
+    data_config: Optional[str] = None,
+    params: Optional[Dict[str, Any]] = None,
+    gpus: int = 1,
+    time_limit: str = "04:00:00",
+    backend: str = None,
+    output_dir: str = None,
+    script: Optional[str] = None,
+    script_content: Optional[str] = None,
+    container: str = "afrilink-yolo",
+    **kwargs,
+) -> AdaptJob:
+    """Deprecated — use ``pretrain(kind=..., pretrained=True)``."""
+    warnings.warn(
+        "adapt() is deprecated in AfriLink SDK 0.9.0. "
+        "Use pretrain(kind=..., weights=..., data=..., pretrained=True) instead. "
+        "adapt() will be removed in 1.0.0.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    from .pretrain import pretrain
+
+    return pretrain(
+        kind=recipe,
+        data=data,
+        weights=weights,
+        model=model,
+        data_config=data_config,
+        params=params,
+        pretrained=True,
+        gpus=gpus,
+        time_limit=time_limit,
+        backend=backend,
+        output_dir=output_dir,
+        script=script,
+        script_content=script_content,
+        container=container,
+        **kwargs,
+    )

afrilink/api_key_auth.py

@@ -0,0 +1,255 @@
+"""
+AfriLink API-Key Auth — stateless SDK auth via Cloudflare Worker exchange.
+
+Users mint API keys at https://dataspires.com/dashboard/profile and pass them
+to the SDK either as the `api_key=` arg to `client.authenticate()` or via the
+`AFRILINK_API_KEY` env var (Colab/Kaggle "Add a secret" UI).
+
+The SDK hits `https://api.dataspires.com/v1/sdk-auth/exchange` with the key
+in `Authorization: Bearer ...`, and receives back everything needed for the
+session: Supabase JWT (for RLS-protected writes against `sessions` and
+`deduct_credits`), Supabase project URL + anon key, and the shared A100 SSH
+key (base64) which the SDK lands in a per-process tempfile at 0600.
+
+Nothing persists. Kill the process and the SSH key tempfile + JWT in memory
+are gone with it.
+
+NOTE: v1 uses a single shared `dataspires` SSH key for the A100 across all
+users. Per-user keys are planned as a follow-up so the A100's `last login`
+audit reflects the actual DataSpires identity.
+"""
+
+import base64
+import json
+import os
+import urllib.error
+import urllib.request
+from dataclasses import dataclass
+from typing import Optional
+
+
+# The afrilink-sdk-auth Cloudflare Worker. Default points at the working
+# .workers.dev URL; override via AFRILINK_AUTH_URL once api.dataspires.com
+# is wired to its Custom Domain in the Cloudflare dashboard.
+DEFAULT_EXCHANGE_URL = (
+    "https://afrilink-sdk-auth.sweet-star-2474.workers.dev/v1/sdk-auth/exchange"
+)
+
+
+@dataclass
+class ApiKeyExchangeResult:
+    """Everything the SDK needs to run jobs after an API-key exchange."""
+    success: bool
+    error: Optional[str] = None
+    user_id: Optional[str] = None
+    email: Optional[str] = None
+    balance_usd: float = 0.0
+    supabase_jwt: Optional[str] = None
+    supabase_url: Optional[str] = None
+    supabase_anon_key: Optional[str] = None
+    a100_host: Optional[str] = None
+    a100_user: Optional[str] = None
+    a100_ssh_key_material: Optional[str] = None  # decoded PEM, not base64
+    # GHCR creds for pulling private curated containers (org-owned)
+    ghcr_token: Optional[str] = None
+    ghcr_username: Optional[str] = None
+    # GCP build service bundle (Cloud Build + Artifact Registry on anadrome)
+    gcp_sa_key_json: Optional[str] = None       # single-line JSON of the SA key
+    gcp_project: Optional[str] = None
+    gcp_region: Optional[str] = None
+    gcp_ar_repo: Optional[str] = None
+    gcp_build_bucket: Optional[str] = None
+    expires_at: Optional[int] = None
+
+
+def _post_json(url: str, body: dict, headers: dict, timeout: int = 15) -> dict:
+    """Minimal HTTP POST with stdlib fallback so the SDK works without `requests`."""
+    try:
+        import requests
+        resp = requests.post(url, json=body, headers=headers, timeout=timeout)
+        try:
+            return {"status": resp.status_code, "body": resp.json()}
+        except Exception:
+            return {"status": resp.status_code, "body": {"raw": resp.text[:500]}}
+    except ImportError:
+        req = urllib.request.Request(
+            url,
+            data=json.dumps(body).encode("utf-8") if body else b"",
+            headers={**headers, "Content-Type": "application/json"},
+            method="POST",
+        )
+        try:
+            with urllib.request.urlopen(req, timeout=timeout) as r:
+                return {"status": r.status, "body": json.loads(r.read().decode("utf-8"))}
+        except urllib.error.HTTPError as e:
+            try:
+                return {"status": e.code, "body": json.loads(e.read().decode("utf-8"))}
+            except Exception:
+                return {"status": e.code, "body": {"error": str(e)}}
+
+
+def resolve_api_key(explicit: Optional[str] = None) -> Optional[str]:
+    """
+    Resolution order:
+      1. `explicit` argument
+      2. API_KEY env / notebook secret
+      3. AFRILINK_API_KEY env / notebook secret (legacy)
+      4. None — caller must prompt
+    """
+    return resolve_secret("API_KEY", explicit=explicit) or resolve_secret(
+        "AFRILINK_API_KEY", explicit=None
+    )
+
+
+def resolve_secret(name: str, explicit: Optional[str] = None) -> Optional[str]:
+    """
+    Generic secret resolver. Resolution order:
+      1. `explicit` arg
+      2. `os.environ[name]`
+      3. Colab notebook secret (google.colab.userdata)
+      4. Kaggle notebook secret (kaggle_secrets.UserSecretsClient)
+      5. None
+
+    Used for AFRILINK_API_KEY, HUGGINGFACE_TOKEN, and other third-party
+    tokens users set as notebook secrets without exporting to os.environ.
+    """
+    if explicit:
+        return explicit.strip()
+    env = os.environ.get(name)
+    if env:
+        return env.strip()
+
+    # Colab
+    try:
+        from google.colab import userdata  # type: ignore
+        try:
+            v = userdata.get(name)
+            if v:
+                # Mirror into os.environ so downstream SDK paths (subprocesses,
+                # docker -e flags, dataset loaders) see it without re-resolving.
+                os.environ[name] = v.strip()
+                return v.strip()
+        except Exception:
+            pass
+    except ImportError:
+        pass
+
+    # Kaggle
+    try:
+        from kaggle_secrets import UserSecretsClient  # type: ignore
+        try:
+            v = UserSecretsClient().get_secret(name)
+            if v:
+                os.environ[name] = v.strip()
+                return v.strip()
+        except Exception:
+            # Kaggle raises if the secret isn't attached to the notebook;
+            # we want to fall through to the explicit-prompt path, not crash.
+            pass
+    except ImportError:
+        pass
+
+    return None
+
+
+def exchange_api_key(
+    api_key: str,
+    exchange_url: str = DEFAULT_EXCHANGE_URL,
+) -> ApiKeyExchangeResult:
+    """
+    Exchange a plaintext API key for a session bundle.
+
+    Network failure → ApiKeyExchangeResult(success=False, error="…")
+    401 → "Invalid or revoked API key"
+    Anything else → the Worker's `error` field
+    """
+    if not api_key or not api_key.startswith(("afk_live_", "afk_test_")):
+        return ApiKeyExchangeResult(
+            success=False,
+            error="Malformed API key (expected to start with afk_live_ or afk_test_)",
+        )
+    try:
+        resp = _post_json(
+            exchange_url,
+            body={},
+            headers={"Authorization": f"Bearer {api_key}"},
+        )
+    except Exception as e:
+        return ApiKeyExchangeResult(success=False, error=f"Network error: {e}")
+
+    status = resp.get("status", 0)
+    data = resp.get("body") or {}
+    if status != 200 or not data.get("ok"):
+        return ApiKeyExchangeResult(
+            success=False,
+            error=(data.get("error") or f"HTTP {status}"),
+        )
+
+    # Decode the SSH key from base64 into PEM material
+    key_b64 = data.get("a100_ssh_key_b64") or ""
+    try:
+        key_material = base64.b64decode(key_b64).decode("utf-8")
+    except Exception:
+        key_material = None
+
+    return ApiKeyExchangeResult(
+        success=True,
+        user_id=data.get("user_id"),
+        email=data.get("email"),
+        balance_usd=float(data.get("balance_usd") or 0),
+        supabase_jwt=data.get("supabase_jwt"),
+        supabase_url=data.get("supabase_url"),
+        supabase_anon_key=data.get("supabase_anon_key"),
+        a100_host=data.get("a100_host"),
+        a100_user=data.get("a100_user"),
+        a100_ssh_key_material=key_material,
+        ghcr_token=data.get("ghcr_token"),
+        ghcr_username=data.get("ghcr_username"),
+        gcp_sa_key_json=data.get("gcp_sa_key_json"),
+        gcp_project=data.get("gcp_project"),
+        gcp_region=data.get("gcp_region"),
+        gcp_ar_repo=data.get("gcp_ar_repo"),
+        gcp_build_bucket=data.get("gcp_build_bucket"),
+        expires_at=data.get("expires_at"),
+    )
+
+
+# ── Convenience: a minimal AuthClient stand-in that the rest of the SDK
+# (telemetry.SupabaseTelemetryWriter, JobUsageTracker) treats just like
+# AfriLinkAuth. Same attribute surface, populated from the exchange.
+
+class ApiKeyAuthClient:
+    """
+    Drop-in replacement for AfriLinkAuth when authenticating via API key.
+
+    Exposes the same attributes the rest of the SDK reads:
+        user_id, access_token, supabase_url, supabase_key, balance_usd,
+        is_authenticated
+    """
+
+    def __init__(self, result: ApiKeyExchangeResult, api_key: str = None, exchange_url: str = None):
+        self._result = result
+        self._api_key = api_key
+        self._exchange_url = exchange_url or DEFAULT_EXCHANGE_URL
+        self.user_id = result.user_id
+        self.email = result.email
+        self.access_token = result.supabase_jwt
+        self.supabase_url = result.supabase_url
+        self.supabase_key = result.supabase_anon_key
+        self.balance_usd = result.balance_usd
+
+    def refresh_access_token(self) -> bool:
+        """Re-exchange API key for a fresh Supabase JWT (billing writes)."""
+        if not self._api_key:
+            return False
+        result = exchange_api_key(self._api_key, exchange_url=self._exchange_url)
+        if not result.success or not result.supabase_jwt:
+            return False
+        self._result = result
+        self.access_token = result.supabase_jwt
+        self.balance_usd = result.balance_usd
+        return True
+
+    @property
+    def is_authenticated(self) -> bool:
+        return self._result.success and bool(self._result.supabase_jwt)