strata-pool 0.7.0__py3-none-any.whl

strata_pool/__init__.py

@@ -0,0 +1,31 @@
+"""Worker pool for dispatching Strata jobs to ephemeral machines."""
+
+from strata_pool.backend import Backend, ProvisionedWorker
+from strata_pool.backends import DockerBackend, RunPodBackend
+from strata_pool.pool import Pool
+from strata_pool.store import PoolStore
+from strata_pool.types import (
+    Job,
+    JobState,
+    MachineType,
+    UsageEvent,
+    Worker,
+    WorkerState,
+)
+
+__all__ = [
+    "Backend",
+    "DockerBackend",
+    "Job",
+    "JobState",
+    "MachineType",
+    "Pool",
+    "PoolStore",
+    "ProvisionedWorker",
+    "RunPodBackend",
+    "UsageEvent",
+    "Worker",
+    "WorkerState",
+]
+
+__version__ = "0.1.0"

strata_pool/api.py

@@ -0,0 +1,217 @@
+"""HTTP surface, for running the pool as a service.
+
+Needs the `server` extra. The pool is usable as a library without it, and the
+proxy that composes it may well bring its own framework, so a web server is
+not something `import strata_pool` should pull in.
+
+Job payloads and results are opaque bytes, so they travel as raw request and
+response bodies rather than being wedged into JSON. Everything else — status,
+fleet, usage — is JSON.
+
+The caller is trusted for tenant identity. It presents the pool's API token
+and asserts a tenant in a header; the pool does not authenticate end users and
+has no idea who they are. That is the same trusted-proxy model the Strata
+server uses, and it means the pool must never be reachable from anywhere but
+the proxy.
+"""
+
+import logging
+from dataclasses import asdict
+from typing import Annotated
+
+from fastapi import Depends, FastAPI, Header, HTTPException, Query, Request, Response
+from fastapi.responses import JSONResponse
+
+from strata_pool.pool import Pool
+from strata_pool.types import Job, JobState, UsageEvent, Worker
+
+logger = logging.getLogger(__name__)
+
+TENANT_HEADER = "X-Strata-Tenant"
+
+
+def _job_json(job: Job) -> dict:
+    """A job without its payload or result, which are bytes and often large."""
+    fields = asdict(job)
+    fields.pop("payload")
+    fields.pop("result")
+    fields["has_result"] = job.result is not None
+    return fields
+
+
+def _worker_json(worker: Worker) -> dict:
+    """A machine without its credential.
+
+    Built by hand rather than from asdict: `repr=False` keeps the token out of
+    logs but not out of a dict, and this response is the one place it would
+    otherwise be handed to whoever asked.
+    """
+    fields = asdict(worker)
+    fields.pop("auth_token")
+    return fields
+
+
+def _usage_json(event: UsageEvent) -> dict:
+    return asdict(event)
+
+
+def create_app(
+    pool: Pool,
+    *,
+    api_token: str | None = None,
+    scaler_interval_seconds: float = 10.0,
+) -> FastAPI:
+    """Build the pool's HTTP app.
+
+    Args:
+        api_token: Bearer token every route except `/health` requires. None
+            disables the check, which is a local-development choice: an open
+            submit endpoint runs arbitrary payloads on machines you pay for.
+        scaler_interval_seconds: How often idle machines are reaped. The app
+            starts the scaler itself, so a deployment cannot forget to.
+    """
+    if api_token is None:
+        logger.warning("pool API starting with no token; anyone who can reach it can run jobs")
+
+    async def require_token(authorization: Annotated[str | None, Header()] = None) -> None:
+        if api_token is None:
+            return
+        if authorization != f"Bearer {api_token}":
+            raise HTTPException(status_code=401, detail="invalid or missing API token")
+
+    async def tenant(request: Request) -> str:
+        value = request.headers.get(TENANT_HEADER)
+        if not value:
+            raise HTTPException(status_code=400, detail=f"{TENANT_HEADER} is required")
+        return value
+
+    async def lifespan(app: FastAPI):
+        # Reconcile before serving: machines from a previous process are
+        # either still reachable or still billing.
+        await pool.recover()
+        pool.start_scaler(scaler_interval_seconds)
+        yield
+        await pool.aclose()
+
+    app = FastAPI(title="strata-pool", lifespan=lifespan)
+    guard = [Depends(require_token)]
+
+    @app.get("/health")
+    async def health() -> dict:
+        """Deliberately outside the token check, for load balancers."""
+        workers = pool.store.list_workers()
+        counts: dict[str, int] = {}
+        for worker in workers:
+            counts[worker.state.value] = counts.get(worker.state.value, 0) + 1
+        return {"status": "ok", "workers": counts, "machine_types": list(pool.machine_types)}
+
+    @app.post("/v1/jobs", status_code=202, dependencies=guard)
+    async def submit_job(
+        request: Request,
+        machine_type: str,
+        tenant_id: Annotated[str, Depends(tenant)],
+        priority: int = 0,
+        session_id: str | None = None,
+        timeout_seconds: float | None = None,
+    ) -> JSONResponse:
+        """Queue a job. The request body is the payload, verbatim."""
+        job = await _submit(
+            pool,
+            tenant_id=tenant_id,
+            machine_type=machine_type,
+            payload=await request.body(),
+            priority=priority,
+            session_id=session_id,
+            timeout_seconds=timeout_seconds,
+        )
+        return JSONResponse(_job_json(job), status_code=202)
+
+    @app.post("/v1/jobs/sync", dependencies=guard)
+    async def submit_and_wait(
+        request: Request,
+        machine_type: str,
+        tenant_id: Annotated[str, Depends(tenant)],
+        priority: int = 0,
+        session_id: str | None = None,
+        timeout_seconds: float | None = None,
+        wait_seconds: float = 300.0,
+    ) -> Response:
+        """Queue a job and block until it finishes.
+
+        What the notebook's executor protocol wants, since it expects one
+        synchronous response. `wait_seconds` bounds how long the caller waits,
+        not how long the job may run — a job that outlives it keeps going and
+        can be collected by ID.
+        """
+        job = await _submit(
+            pool,
+            tenant_id=tenant_id,
+            machine_type=machine_type,
+            payload=await request.body(),
+            priority=priority,
+            session_id=session_id,
+            timeout_seconds=timeout_seconds,
+        )
+        try:
+            done = await pool.wait(job.id, timeout=wait_seconds)
+        except TimeoutError:
+            # 202: it is still running, and the ID is how you find it. Re-read
+            # for the freshest state, falling back to the submitted snapshot
+            # rather than pretending a row we just wrote could be missing.
+            latest = pool.store.get_job(job.id) or job
+            return JSONResponse(_job_json(latest), status_code=202)
+        return _terminal_response(done)
+
+    @app.get("/v1/jobs/{job_id}", dependencies=guard)
+    async def get_job(job_id: str) -> dict:
+        job = pool.store.get_job(job_id)
+        if job is None:
+            raise HTTPException(status_code=404, detail=f"no such job: {job_id}")
+        return _job_json(job)
+
+    @app.get("/v1/jobs/{job_id}/result", dependencies=guard)
+    async def get_job_result(job_id: str) -> Response:
+        """The raw result bytes, once there are any."""
+        job = pool.store.get_job(job_id)
+        if job is None:
+            raise HTTPException(status_code=404, detail=f"no such job: {job_id}")
+        if job.state not in (JobState.COMPLETED, JobState.FAILED, JobState.TIMED_OUT):
+            raise HTTPException(status_code=409, detail=f"job is {job.state.value}")
+        return _terminal_response(job)
+
+    @app.get("/v1/machine-types", dependencies=guard)
+    async def list_machine_types() -> list[dict]:
+        """What a caller may ask for. The catalogue an annotation resolves against."""
+        return [asdict(spec) for spec in pool.machine_types.values()]
+
+    @app.get("/v1/workers", dependencies=guard)
+    async def list_workers() -> list[dict]:
+        return [_worker_json(worker) for worker in pool.store.list_workers()]
+
+    @app.get("/v1/usage", dependencies=guard)
+    async def list_usage(tenant_id: str | None = Query(default=None)) -> list[dict]:
+        """The billing feed. One event per terminal job, monotonic duration."""
+        return [_usage_json(event) for event in pool.store.list_usage(tenant_id)]
+
+    return app
+
+
+async def _submit(pool: Pool, **kwargs) -> Job:
+    try:
+        return await pool.submit(**kwargs)
+    except ValueError as exc:
+        # An unknown machine type is the caller's mistake, not a pool failure.
+        raise HTTPException(status_code=400, detail=str(exc)) from exc
+
+
+def _terminal_response(job: Job) -> Response:
+    """Map a finished job onto a status code.
+
+    A failure on the worker is reported as a failure of the job, not of the
+    pool: the caller needs to tell "your code raised" from "we could not run
+    it", and a 500 would blur the two.
+    """
+    if job.state is JobState.COMPLETED:
+        return Response(content=job.result or b"", media_type="application/octet-stream")
+    status = 504 if job.state is JobState.TIMED_OUT else 502
+    return JSONResponse({"state": job.state.value, "error": job.error}, status_code=status)

strata_pool/backend.py

@@ -0,0 +1,56 @@
+"""The seam between the pool and whatever actually runs machines."""
+
+from dataclasses import dataclass, field
+from typing import Protocol
+
+from strata_pool.types import MachineType
+
+
+@dataclass
+class ProvisionedWorker:
+    """What a backend hands back from `start()`."""
+
+    backend_id: str
+    """Cloud-specific resource ID, used for stop()."""
+
+    endpoint: str
+    """HTTP base URL where the worker will accept jobs.
+
+    Required at start time. A backend that only learns the address later
+    should block until it knows it: the pool has no other way to health-check
+    a booting machine, and a worker it cannot reach is a worker it cannot
+    stop billing for.
+    """
+
+    region: str | None = None
+    metadata: dict[str, str] = field(default_factory=dict)
+
+
+class Backend(Protocol):
+    """A cloud provider or local runtime that can start and stop machines."""
+
+    name: str
+
+    async def start(
+        self,
+        spec: MachineType,
+        env: dict[str, str] | None = None,
+    ) -> ProvisionedWorker:
+        """Provision and start a worker.
+
+        Takes the whole machine type rather than a name and an image: a
+        backend needs the resource limits too, and each will grow its own
+        provider-specific fields. A backend ignores what it cannot express.
+
+        Returns once the machine is booting; it does not need to be ready.
+        The pool polls `health()` until it is.
+        """
+        ...
+
+    async def stop(self, backend_id: str) -> None:
+        """Stop and deallocate a worker. Must be idempotent."""
+        ...
+
+    async def health(self, endpoint: str) -> bool:
+        """Report whether a worker is ready to accept jobs."""
+        ...

strata_pool/backends/__init__.py

@@ -0,0 +1,6 @@
+"""Backends that provision machines."""
+
+from strata_pool.backends.docker import DockerBackend, DockerError
+from strata_pool.backends.runpod import RunPodBackend, RunPodError
+
+__all__ = ["DockerBackend", "DockerError", "RunPodBackend", "RunPodError"]

strata_pool/backends/docker.py

@@ -0,0 +1,188 @@
+"""Run workers as containers on a local Docker daemon.
+
+The first backend, chosen first so that dispatch, boot handling, and the
+metering path are all exercised by CI before any of them produces a number
+someone pays for.
+
+It talks to the Docker Engine API over its UNIX socket with httpx rather than
+pulling in the Docker SDK: the pool's dependency list stays at one entry, and
+every request shape is testable against `httpx.MockTransport` without a
+daemon.
+
+The backend does not care what runs inside the container. It starts an image,
+finds the host port Docker published, and reports the endpoint; whether that
+image is `strata-worker` or something else is the caller's business.
+"""
+
+import logging
+
+import httpx
+
+from strata_pool.backend import ProvisionedWorker
+from strata_pool.types import MachineType
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_SOCKET = "/var/run/docker.sock"
+
+
+class DockerError(RuntimeError):
+    """The Docker daemon refused a request."""
+
+
+class DockerBackend:
+    """Provisions workers as containers on the local Docker daemon."""
+
+    name = "docker"
+
+    def __init__(
+        self,
+        *,
+        socket_path: str = DEFAULT_SOCKET,
+        worker_port: int = 8080,
+        command: list[str] | None = None,
+        stop_timeout_seconds: int = 5,
+        api: httpx.AsyncClient | None = None,
+        probe: httpx.AsyncClient | None = None,
+    ):
+        """
+        Args:
+            worker_port: The port the image listens on inside the container.
+                Docker publishes it on an arbitrary host port, which is what
+                the pool connects to.
+            command: Overrides the image's entrypoint. For images that host
+                more than one, and for tests.
+            api: Client for the Docker Engine API. Defaults to one bound to
+                the daemon socket.
+            probe: Client for worker health checks, which go over TCP to the
+                published port rather than through the daemon.
+        """
+        self.worker_port = worker_port
+        self.command = command
+        self.stop_timeout_seconds = stop_timeout_seconds
+        self._owns_clients = api is None and probe is None
+        self._api = api if api is not None else _socket_client(socket_path)
+        self._probe = probe if probe is not None else httpx.AsyncClient()
+
+    async def aclose(self) -> None:
+        """Release the HTTP clients. The pool does not own the backend."""
+        if self._owns_clients:
+            await self._api.aclose()
+            await self._probe.aclose()
+
+    async def start(
+        self,
+        spec: MachineType,
+        env: dict[str, str] | None = None,
+    ) -> ProvisionedWorker:
+        port_key = f"{self.worker_port}/tcp"
+        host_config: dict[str, object] = {
+            # Empty HostPort means "pick a free one". Binding to loopback
+            # keeps a worker off the network: the pool is the only thing that
+            # should be able to reach it.
+            "PortBindings": {port_key: [{"HostIp": "127.0.0.1", "HostPort": ""}]},
+        }
+        # Unset means the container may consume the whole host, which is one
+        # tenant's job able to starve every other container on the box.
+        if spec.cpus is not None:
+            host_config["NanoCpus"] = int(spec.cpus * 1_000_000_000)
+        if spec.memory_mb is not None:
+            host_config["Memory"] = spec.memory_mb * 1024 * 1024
+
+        create = await self._api.post(
+            "/containers/create",
+            json={
+                "Image": spec.image,
+                "Env": [f"{key}={value}" for key, value in (env or {}).items()],
+                "Cmd": self.command,
+                "ExposedPorts": {port_key: {}},
+                "HostConfig": host_config,
+                # A machine the pool forgot about can still be found by hand,
+                # and gives a later reconcile pass something to match on.
+                "Labels": {"strata.pool.machine-type": spec.name},
+            },
+        )
+        if create.status_code >= 400:
+            raise DockerError(f"could not create a {spec.image} container: {_message(create)}")
+        container_id = create.json()["Id"]
+
+        started = await self._api.post(f"/containers/{container_id}/start")
+        if started.status_code >= 400:
+            # The container exists and would sit there costing disk, so take
+            # it back out before reporting the failure.
+            await self.stop(container_id)
+            raise DockerError(f"could not start {container_id}: {_message(started)}")
+
+        port = await self._published_port(container_id, port_key)
+        return ProvisionedWorker(
+            backend_id=container_id,
+            endpoint=f"http://127.0.0.1:{port}",
+            region="local",
+            metadata={"machine_type": spec.name},
+        )
+
+    async def stop(self, backend_id: str) -> None:
+        """Stop and remove a container. Idempotent."""
+        stopped = await self._api.post(
+            f"/containers/{backend_id}/stop",
+            params={"t": self.stop_timeout_seconds},
+        )
+        # 304 is "already stopped" and 404 is "already gone"; both are the
+        # state this method exists to reach.
+        if stopped.status_code >= 400 and stopped.status_code != 404:
+            logger.warning(
+                "docker refused to stop a container; removing it anyway",
+                extra={"container_id": backend_id, "detail": _message(stopped)},
+            )
+
+        removed = await self._api.delete(f"/containers/{backend_id}", params={"force": "true"})
+        if removed.status_code >= 400 and removed.status_code != 404:
+            raise DockerError(f"could not remove {backend_id}: {_message(removed)}")
+
+    async def health(self, endpoint: str) -> bool:
+        """Report whether the worker answers on its published port.
+
+        Never raises: a refused connection is the normal state of a container
+        that is still booting, and the pool polls this in a loop.
+        """
+        try:
+            response = await self._probe.get(f"{endpoint}/health", timeout=2.0)
+        except httpx.HTTPError:
+            return False
+        return response.status_code < 400
+
+    async def _published_port(self, container_id: str, port_key: str) -> int:
+        inspect = await self._api.get(f"/containers/{container_id}/json")
+        if inspect.status_code >= 400:
+            raise DockerError(f"could not inspect {container_id}: {_message(inspect)}")
+
+        bindings = inspect.json().get("NetworkSettings", {}).get("Ports") or {}
+        published = bindings.get(port_key) or []
+        if not published:
+            raise DockerError(
+                f"{container_id} published no host port for {port_key}. The image must "
+                f"listen on {self.worker_port}, or the backend needs a different worker_port."
+            )
+        return int(published[0]["HostPort"])
+
+
+def _socket_client(socket_path: str) -> httpx.AsyncClient:
+    """A client bound to the daemon socket.
+
+    The host in the URL is ignored for a UNIX-socket transport but httpx still
+    requires one, hence the placeholder.
+    """
+    return httpx.AsyncClient(
+        transport=httpx.AsyncHTTPTransport(uds=socket_path),
+        base_url="http://docker",
+        timeout=30.0,
+    )
+
+
+def _message(response: httpx.Response) -> str:
+    """The daemon's own error text, which is JSON when it is well behaved."""
+    try:
+        payload = response.json()
+    except ValueError:
+        return f"HTTP {response.status_code}: {response.text[:200]}"
+    return f"HTTP {response.status_code}: {payload.get('message', response.text[:200])}"