colabctl 0.1.0__py3-none-any.whl

colabctl/__init__.py

@@ -0,0 +1,71 @@
+"""colabctl — programmatic control of Google Colab for developers and AI agents.
+
+Public API is intentionally small and stable; everything else is an
+implementation detail behind the transport/provider abstractions.
+"""
+
+from __future__ import annotations
+
+from colabctl.drive import DriveSync, drive_checkpoint_hooks
+from colabctl.errors import (
+    AcceleratorUnavailableError,
+    AllocationError,
+    AuthError,
+    ColabctlError,
+    ConfigurationError,
+    ExecutionError,
+    FileTransferError,
+    KeepAliveError,
+    KernelError,
+    QuotaExceededError,
+    SecretStoreError,
+    TooManyAssignmentsError,
+    TransportError,
+)
+from colabctl.lifecycle import RuntimeLifecycleManager
+from colabctl.models import (
+    Accelerator,
+    Assignment,
+    ExecutionResult,
+    MachineShape,
+    RuntimeProxyInfo,
+    RuntimeSpec,
+    SessionInfo,
+    SessionStatus,
+    Variant,
+)
+from colabctl.sdk import ColabClient, ColabSession, remote
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Accelerator",
+    "AcceleratorUnavailableError",
+    "AllocationError",
+    "Assignment",
+    "AuthError",
+    "ColabClient",
+    "ColabSession",
+    "ColabctlError",
+    "ConfigurationError",
+    "DriveSync",
+    "ExecutionError",
+    "ExecutionResult",
+    "FileTransferError",
+    "KeepAliveError",
+    "KernelError",
+    "MachineShape",
+    "QuotaExceededError",
+    "RuntimeLifecycleManager",
+    "RuntimeProxyInfo",
+    "RuntimeSpec",
+    "SecretStoreError",
+    "SessionInfo",
+    "SessionStatus",
+    "TooManyAssignmentsError",
+    "TransportError",
+    "Variant",
+    "__version__",
+    "drive_checkpoint_hooks",
+    "remote",
+]

colabctl/auth/__init__.py

@@ -0,0 +1,21 @@
+"""Authentication providers for the Colab backend."""
+
+from __future__ import annotations
+
+from colabctl.auth.adc import ADCAuthProvider
+from colabctl.auth.base import (
+    ADC_LOGIN_SCOPES,
+    COLAB_SCOPES,
+    AuthProvider,
+    StaticTokenProvider,
+    TokenCallable,
+)
+
+__all__ = [
+    "ADC_LOGIN_SCOPES",
+    "COLAB_SCOPES",
+    "ADCAuthProvider",
+    "AuthProvider",
+    "StaticTokenProvider",
+    "TokenCallable",
+]

colabctl/auth/adc.py

@@ -0,0 +1,80 @@
+"""ADC auth provider — the Phase 0-verified working path.
+
+Uses Application Default Credentials (``gcloud auth application-default login
+--scopes=…colaboratory``). ``google.auth`` is sync, so refreshes run in a thread;
+a lock serializes concurrent refreshes. ``google-auth`` is imported lazily.
+
+Setup (once, by the user):
+
+    gcloud auth application-default login \\
+        --scopes=openid,https://www.googleapis.com/auth/cloud-platform,\\
+    https://www.googleapis.com/auth/userinfo.email,\\
+    https://www.googleapis.com/auth/colaboratory,\\
+    https://www.googleapis.com/auth/drive.file
+"""
+
+from __future__ import annotations
+
+import asyncio
+import warnings
+from typing import Any
+
+from colabctl.auth.base import COLAB_SCOPES, AuthProvider
+from colabctl.errors import AuthError
+
+
+class ADCAuthProvider(AuthProvider):
+    """Bearer tokens from Application Default Credentials with the Colab scopes."""
+
+    def __init__(self, *, scopes: tuple[str, ...] = COLAB_SCOPES) -> None:
+        self._scopes = list(scopes)
+        self._creds: Any | None = None
+        self._lock = asyncio.Lock()
+
+    async def token(self) -> str:
+        async with self._lock:
+            return await asyncio.to_thread(self._sync_token)
+
+    async def email(self) -> str | None:
+        async with self._lock:
+            creds = self._creds
+        # ADC user creds expose the signer email for service accounts; user creds
+        # usually don't, so this is best-effort.
+        return getattr(creds, "service_account_email", None) if creds else None
+
+    def _sync_token(self) -> str:
+        try:
+            import google.auth
+            from google.auth.transport.requests import Request
+        except ImportError as exc:  # pragma: no cover - only without the extra
+            raise AuthError(
+                "google-auth is not installed. Install with `pip install 'colabctl[native]'`."
+            ) from exc
+
+        # Typed Any: google-auth credential subclasses differ structurally
+        # (only scopable creds have with_scopes; refresh is untyped upstream).
+        creds: Any = self._creds
+        if creds is None:
+            # ADC user creds emit a noisy "no quota project" UserWarning on every
+            # call; it's irrelevant here (Colab calls pin their own project), so
+            # suppress just that one message.
+            with warnings.catch_warnings():
+                warnings.filterwarnings(
+                    "ignore",
+                    message=r"Your application has authenticated using end user credentials.*",
+                    category=UserWarning,
+                )
+                creds, _ = google.auth.default(scopes=self._scopes)
+        # User creds ignore scopes= in default(); re-apply when supported.
+        if getattr(creds, "requires_scopes", False):
+            creds = creds.with_scopes(self._scopes)
+        if not creds.valid:
+            try:
+                creds.refresh(Request())
+            except Exception as exc:
+                raise AuthError(f"Failed to refresh ADC credentials: {exc}") from exc
+        self._creds = creds
+        token = getattr(creds, "token", None)
+        if not token:
+            raise AuthError("ADC credentials produced no access token.")
+        return str(token)

colabctl/auth/base.py

@@ -0,0 +1,63 @@
+"""Authentication contract.
+
+An :class:`AuthProvider` yields a fresh OAuth bearer token (with the Colab scopes)
+on demand. Transports depend only on this; concrete providers (ADC, OAuth2-loopback,
+static) are swappable. The Colab scope set is the one verified in Phase 0 — note
+that ``colaboratory`` is not third-party-grantable, so ADC (gcloud's client) is the
+working path; ``cloud-platform`` is additionally required by gcloud itself.
+"""
+
+from __future__ import annotations
+
+import abc
+from collections.abc import Awaitable, Callable
+
+#: OAuth scopes the Colab backend requires (verified from CLI ``PUBLIC_SCOPES``).
+COLAB_SCOPES: tuple[str, ...] = (
+    "openid",
+    "https://www.googleapis.com/auth/userinfo.profile",
+    "https://www.googleapis.com/auth/userinfo.email",
+    "https://www.googleapis.com/auth/colaboratory",
+    "https://www.googleapis.com/auth/drive.file",
+)
+
+#: ADC via ``gcloud`` additionally requires these (gcloud refuses otherwise).
+ADC_LOGIN_SCOPES: tuple[str, ...] = (
+    "openid",
+    "https://www.googleapis.com/auth/cloud-platform",
+    "https://www.googleapis.com/auth/userinfo.email",
+    "https://www.googleapis.com/auth/colaboratory",
+    "https://www.googleapis.com/auth/drive.file",
+)
+
+TokenCallable = Callable[[], Awaitable[str]]
+
+
+class AuthProvider(abc.ABC):
+    """Yields fresh bearer tokens for the Colab backend."""
+
+    @abc.abstractmethod
+    async def token(self) -> str:
+        """Return a currently-valid bearer token, refreshing if needed."""
+
+    async def email(self) -> str | None:
+        """Return the authenticated account email if known (best-effort)."""
+        return None
+
+    def as_token_callable(self) -> TokenCallable:
+        """Adapt to the ``TokenProvider`` callable the native client expects."""
+        return self.token
+
+
+class StaticTokenProvider(AuthProvider):
+    """An :class:`AuthProvider` wrapping a fixed token — for tests/injection."""
+
+    def __init__(self, token: str, *, email: str | None = None) -> None:
+        self._token = token
+        self._email = email
+
+    async def token(self) -> str:
+        return self._token
+
+    async def email(self) -> str | None:
+        return self._email

colabctl/backends/__init__.py

@@ -0,0 +1,42 @@
+"""Provider abstraction: pluggable batch backends + capability-based routing.
+
+- :class:`Backend` — the submit/status/logs/result/cancel contract.
+- :class:`ColabBackend` — Colab via an interactive transport (sanctioned default).
+- :class:`ModalBackend` — gVisor GPU sandboxes (best for agent code).
+- :class:`VertexBackend` — sanctioned, headless, deadline-bound GPU jobs.
+- :class:`BackendRouter` — selects a backend by capability and fails over on infra errors.
+
+HF Jobs / Kaggle / IaaS are registered-but-deferred (Phase 4).
+"""
+
+from __future__ import annotations
+
+from colabctl.backends.base import (
+    Backend,
+    BackendCapabilities,
+    JobInfo,
+    JobResult,
+    JobSpec,
+    JobState,
+)
+from colabctl.backends.colab import ColabBackend
+from colabctl.backends.hf_backend import HFJobsBackend
+from colabctl.backends.kaggle_backend import KaggleBackend
+from colabctl.backends.modal_backend import ModalBackend
+from colabctl.backends.router import BackendRouter
+from colabctl.backends.vertex_backend import VertexBackend
+
+__all__ = [
+    "Backend",
+    "BackendCapabilities",
+    "BackendRouter",
+    "ColabBackend",
+    "HFJobsBackend",
+    "JobInfo",
+    "JobResult",
+    "JobSpec",
+    "JobState",
+    "KaggleBackend",
+    "ModalBackend",
+    "VertexBackend",
+]

colabctl/backends/base.py

@@ -0,0 +1,147 @@
+"""Provider abstraction: the batch-`Backend` contract.
+
+Two complementary abstractions exist in colabctl:
+
+- :class:`~colabctl.transport.base.TransportAdapter` — *interactive* runtimes
+  (allocate a warm GPU, run cells on a live kernel). Colab's native shape.
+- :class:`Backend` (this module) — *batch jobs* (submit code → poll → fetch result →
+  cancel). The natural shape for Modal, Vertex, HF Jobs, etc.
+
+The :class:`~colabctl.backends.router.BackendRouter` selects a backend by capability
+and fails over between them, so a Colab outage/quota/ban degrades to Modal or Vertex
+instead of failing. Colab is also exposed as a batch backend
+(:class:`~colabctl.backends.colab.ColabBackend`) so callers can use one job API
+across every provider.
+"""
+
+from __future__ import annotations
+
+import abc
+import enum
+
+from pydantic import BaseModel, model_validator
+
+from colabctl.models import Accelerator
+
+
+class JobState(enum.StrEnum):
+    PENDING = "PENDING"
+    RUNNING = "RUNNING"
+    SUCCEEDED = "SUCCEEDED"
+    FAILED = "FAILED"
+    CANCELLED = "CANCELLED"
+    UNKNOWN = "UNKNOWN"
+
+    @property
+    def is_terminal(self) -> bool:
+        return self in {JobState.SUCCEEDED, JobState.FAILED, JobState.CANCELLED}
+
+
+class JobSpec(BaseModel):
+    """What to run on a backend. Provide exactly one of ``code`` or ``script_path``."""
+
+    code: str | None = None
+    script_path: str | None = None
+    requirements: list[str] = []
+    accelerator: Accelerator = Accelerator.T4
+    env: dict[str, str] = {}
+    timeout: int | None = None
+    name: str | None = None
+
+    @model_validator(mode="after")
+    def _exactly_one_source(self) -> JobSpec:
+        if bool(self.code) == bool(self.script_path):
+            raise ValueError("Provide exactly one of `code` or `script_path`.")
+        return self
+
+    def resolved_code(self) -> str:
+        """Return the code to run (reads ``script_path`` if that's what was given)."""
+        if self.code is not None:
+            return self.code
+        from pathlib import Path
+
+        return Path(self.script_path or "").read_text()
+
+
+class JobInfo(BaseModel):
+    """A lightweight view of a submitted job."""
+
+    id: str
+    backend: str
+    state: JobState
+    accelerator: Accelerator = Accelerator.NONE
+    detail: str | None = None
+
+
+class JobResult(BaseModel):
+    """The outcome of a job."""
+
+    id: str
+    backend: str
+    state: JobState
+    exit_code: int | None = None
+    stdout: str = ""
+    stderr: str = ""
+    error: str | None = None
+
+    @property
+    def ok(self) -> bool:
+        return self.state is JobState.SUCCEEDED
+
+
+class BackendCapabilities(BaseModel):
+    """What a backend can do — used for routing and honest disclosure."""
+
+    name: str
+    accelerators: list[str] = []  # supported accelerator values; empty = any/unknown
+    interactive: bool = False
+    streaming_logs: bool = False
+    persistent: bool = False  # runtime survives between calls
+    max_runtime_seconds: int | None = None
+    requires_account: bool = True
+    tos_posture: str = "sanctioned"  # "sanctioned" | "gray-area" | "prohibited"
+    notes: list[str] = []
+
+    def supports(self, accelerator: Accelerator) -> bool:
+        if not self.accelerators:
+            return True
+        return accelerator.value in self.accelerators or accelerator is Accelerator.NONE
+
+
+class Backend(abc.ABC):
+    """A pluggable batch-execution backend (Colab, Modal, Vertex, ...)."""
+
+    name: str = "backend"
+
+    @property
+    @abc.abstractmethod
+    def capabilities(self) -> BackendCapabilities: ...
+
+    @abc.abstractmethod
+    async def submit(self, spec: JobSpec) -> JobInfo:
+        """Start a job and return immediately with its handle."""
+
+    @abc.abstractmethod
+    async def status(self, job_id: str) -> JobInfo:
+        """Return the job's current state."""
+
+    @abc.abstractmethod
+    async def logs(self, job_id: str) -> str:
+        """Return the job's logs so far (best-effort)."""
+
+    @abc.abstractmethod
+    async def result(self, job_id: str) -> JobResult:
+        """Wait for the job to finish and return its result."""
+
+    @abc.abstractmethod
+    async def cancel(self, job_id: str) -> None:
+        """Cancel a running job."""
+
+    async def run(self, spec: JobSpec) -> JobResult:
+        """Convenience: submit and wait for the result."""
+        info = await self.submit(spec)
+        return await self.result(info.id)
+
+    async def aclose(self) -> None:
+        """Release backend-level resources (default no-op; override if needed)."""
+        return None

colabctl/backends/colab.py

@@ -0,0 +1,166 @@
+"""Colab as a batch :class:`Backend` (wraps a TransportAdapter).
+
+Presents Colab's interactive transport through the unified job API: ``submit``
+launches an in-process asyncio task that allocates a runtime, optionally pip-installs
+requirements, runs the code, captures output, and releases the runtime. State/logs/
+result read the in-memory job record.
+
+Limitation: job tracking is in-process (if the host process dies, the record is lost).
+Cross-process durability is the runtime-lifecycle manager's concern (checkpoint to
+Drive/GCS); for the interactive, warm-GPU workflow use the SDK's ``ColabSession`` directly.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import uuid
+from dataclasses import dataclass, field
+
+from colabctl.backends.base import (
+    Backend,
+    BackendCapabilities,
+    JobInfo,
+    JobResult,
+    JobSpec,
+    JobState,
+)
+from colabctl.errors import ColabctlError
+from colabctl.models import RuntimeSpec
+from colabctl.transport.base import TransportAdapter
+
+
+def _install_code(requirements: list[str]) -> str:
+    pkgs = ", ".join(repr(r) for r in requirements)
+    return (
+        "import subprocess, sys\n"
+        f"subprocess.run([sys.executable, '-m', 'pip', 'install', '-q', {pkgs}], check=True)\n"
+    )
+
+
+@dataclass
+class _Job:
+    info: JobInfo
+    spec: JobSpec
+    task: asyncio.Task[None] | None = None
+    session: str | None = None
+    stdout: str = ""
+    stderr: str = ""
+    exit_code: int | None = None
+    error: str | None = None
+    logbuf: list[str] = field(default_factory=list)
+
+
+class ColabBackend(Backend):
+    """Run batch jobs on Colab via an interactive transport."""
+
+    name = "colab"
+
+    def __init__(self, transport: TransportAdapter) -> None:
+        self._transport = transport
+        self._jobs: dict[str, _Job] = {}
+
+    @property
+    def capabilities(self) -> BackendCapabilities:
+        caps = self._transport.capabilities
+        return BackendCapabilities(
+            name=self.name,
+            accelerators=["T4", "L4", "G4", "A100", "H100"],
+            interactive=caps.interactive,
+            streaming_logs=False,
+            persistent=True,
+            requires_account=True,
+            tos_posture="sanctioned" if self._transport.name == "cli" else "gray-area",
+            notes=[f"via the {self._transport.name!r} transport", *caps.caveats],
+        )
+
+    async def submit(self, spec: JobSpec) -> JobInfo:
+        job_id = f"colab-{uuid.uuid4().hex[:10]}"
+        info = JobInfo(
+            id=job_id, backend=self.name, state=JobState.PENDING, accelerator=spec.accelerator
+        )
+        job = _Job(info=info, spec=spec)
+        self._jobs[job_id] = job
+        job.task = asyncio.create_task(self._execute(job))
+        return info
+
+    async def status(self, job_id: str) -> JobInfo:
+        return self._require(job_id).info
+
+    async def logs(self, job_id: str) -> str:
+        return "".join(self._require(job_id).logbuf)
+
+    async def result(self, job_id: str) -> JobResult:
+        job = self._require(job_id)
+        if job.task is not None:
+            with contextlib.suppress(asyncio.CancelledError):
+                await job.task
+        return JobResult(
+            id=job_id,
+            backend=self.name,
+            state=job.info.state,
+            exit_code=job.exit_code,
+            stdout=job.stdout,
+            stderr=job.stderr,
+            error=job.error,
+        )
+
+    async def cancel(self, job_id: str) -> None:
+        job = self._require(job_id)
+        if job.task is not None and not job.task.done():
+            job.task.cancel()
+        job.info.state = JobState.CANCELLED
+
+    async def aclose(self) -> None:
+        for job in list(self._jobs.values()):
+            if job.task is not None and not job.task.done():
+                job.task.cancel()
+        await self._transport.aclose()
+
+    # -- internals ----------------------------------------------------------
+
+    def _require(self, job_id: str) -> _Job:
+        job = self._jobs.get(job_id)
+        if job is None:
+            raise ColabctlError(f"No such job: {job_id!r}")
+        return job
+
+    async def _execute(self, job: _Job) -> None:
+        job.info.state = JobState.RUNNING
+        try:
+            session = await self._transport.allocate(
+                RuntimeSpec(accelerator=job.spec.accelerator, name=job.spec.name)
+            )
+            job.session = session.name
+            if job.spec.requirements:
+                install = await self._transport.execute(
+                    session.name, _install_code(job.spec.requirements), timeout=job.spec.timeout
+                )
+                if not install.ok:
+                    job.error = "pip install failed: " + (install.stderr or install.text)[:400]
+                    job.info.state = JobState.FAILED
+                    return
+            result = await self._transport.execute(
+                session.name, job.spec.resolved_code(), timeout=job.spec.timeout
+            )
+            job.stdout = result.text
+            job.stderr = result.stderr
+            job.logbuf.append(result.text)
+            if result.ok:
+                job.info.state = JobState.SUCCEEDED
+                job.exit_code = 0
+            else:
+                job.info.state = JobState.FAILED
+                job.exit_code = 1
+                if result.error is not None:
+                    job.error = f"{result.error.ename}: {result.error.evalue}"
+        except asyncio.CancelledError:
+            job.info.state = JobState.CANCELLED
+            raise
+        except ColabctlError as exc:
+            job.info.state = JobState.FAILED
+            job.error = str(exc)
+        finally:
+            if job.session is not None:
+                with contextlib.suppress(ColabctlError):
+                    await self._transport.stop(job.session)

colabctl/backends/factory.py

@@ -0,0 +1,65 @@
+"""Backend construction by name — the seam the CLI and MCP server build on.
+
+Keeps backend wiring in one place so ``colabctl job run --backend modal`` and the
+MCP ``run_job`` tool construct backends identically. The Colab backend is built over
+the chosen transport (cli/native); Modal/Vertex read their own env/config.
+"""
+
+from __future__ import annotations
+
+from colabctl.backends.base import Backend
+from colabctl.backends.modal_backend import ModalBackend
+from colabctl.backends.router import BackendRouter
+from colabctl.backends.vertex_backend import VertexBackend
+from colabctl.errors import ConfigurationError
+
+#: Backends available for selection.
+BACKEND_NAMES: tuple[str, ...] = ("colab", "modal", "vertex", "hf", "kaggle")
+
+
+def build_backend(
+    name: str,
+    *,
+    transport_name: str = "cli",
+    auth_mode: str = "adc",
+    colab_bin: str = "colab",
+) -> Backend:
+    """Construct a backend by name. Colab uses the chosen transport."""
+    key = name.lower()
+    if key == "colab":
+        from colabctl.backends.colab import ColabBackend
+        from colabctl.sdk.client import ColabClient
+
+        client = ColabClient(
+            transport_name=transport_name, auth_mode=auth_mode, colab_bin=colab_bin
+        )
+        return ColabBackend(client.transport)
+    if key == "modal":
+        return ModalBackend()
+    if key == "vertex":
+        return VertexBackend()
+    if key == "hf":
+        from colabctl.backends.hf_backend import HFJobsBackend
+
+        return HFJobsBackend()
+    if key == "kaggle":
+        from colabctl.backends.kaggle_backend import KaggleBackend
+
+        return KaggleBackend()
+    raise ConfigurationError(f"Unknown backend {name!r}. Choose from: {', '.join(BACKEND_NAMES)}.")
+
+
+def build_router(
+    names: list[str] | None = None,
+    *,
+    transport_name: str = "cli",
+    auth_mode: str = "adc",
+    colab_bin: str = "colab",
+) -> BackendRouter:
+    """Build a router over the named backends (default: all), in failover order."""
+    selected = names or list(BACKEND_NAMES)
+    backends = [
+        build_backend(n, transport_name=transport_name, auth_mode=auth_mode, colab_bin=colab_bin)
+        for n in selected
+    ]
+    return BackendRouter(backends, order=selected)