commoncompute 0.1.0__py3-none-any.whl

commoncompute/aws_batch.py

@@ -0,0 +1,589 @@
+"""boto3-shaped AWS Batch shim that talks to Common Compute.
+
+Lets a codebase that calls ``boto3.client("batch")`` migrate to Common
+Compute by changing one import line. The methods accept the same
+keyword arguments and return dicts with the same key names that AWS
+Batch returns, so downstream code that reads ``response["jobId"]``,
+``response["jobs"][0]["status"]``, etc. keeps working.
+
+Quick start::
+
+    # before
+    import boto3
+    batch = boto3.client("batch", region_name="us-east-1")
+
+    # after
+    from commoncompute.aws_batch import client
+    batch = client()                       # reads CC_API_KEY
+
+    resp = batch.submit_job(
+        jobName="nightly-embed",
+        jobQueue="ml-gpu",                 # routed by us, name is informational
+        jobDefinition="bge-large:1",       # parsed as workload_id[:model_id]
+        containerOverrides={
+            "environment": [
+                {"name": "INPUT_PAYLOAD", "value": "{...json...}"},
+            ],
+        },
+    )
+    print(resp["jobId"])
+
+Mapping rules:
+
+* ``jobDefinition`` is parsed as ``"<workload_id>[:<model_id>]"``. A
+  bare ``"bge-large"`` is treated as ``model_id`` only and the workload
+  is inferred (currently ``coreml_embed``).
+* ``containerOverrides.environment`` is collected into a dict and passed
+  through as the job ``payload``. The well-known key ``INPUT_PAYLOAD``
+  may carry a JSON string that is parsed and used as-is.
+* ``schedulingPriority`` (AWS) or ``shareIdentifier=='priority'`` maps
+  to Common Compute ``priority='priority'``; otherwise we default to
+  ``'batch'`` (cheaper and matches the typical AWS Batch use case).
+* AWS Batch state strings (``SUBMITTED``, ``PENDING``, ``RUNNABLE``,
+  ``STARTING``, ``RUNNING``, ``SUCCEEDED``, ``FAILED``) are mapped from
+  Common Compute states (``queued``, ``leased``, ``running``,
+  ``completed``, ``failed``, ``dead_letter``, ``cancelled``).
+
+What this shim does NOT do (out of scope, by design):
+
+* It does not call AWS. It does not require boto3 to be installed.
+* It does not implement compute environments or job queues as
+  configurable resources — those are synthesised as defaults by the
+  CommonCompute backend.
+
+For full boto3 wire-protocol parity (so existing scripts using
+``boto3.client('batch', endpoint_url='https://aws.commoncompute.ai')``
+work without any code change), see the AWS-shim Worker. It accepts
+SigV4-signed JSON-1.1 traffic and supports RegisterJobDefinition,
+SubmitJob (including arrayProperties), DescribeJobs, ListJobs,
+TerminateJob, CancelJob, and the synthesised Describe* read-only set.
+Issue access keys at https://commoncompute.ai/app/aws-keys.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import time
+from datetime import datetime, timezone
+from typing import Any, Iterable, Mapping, Optional
+
+from .client import Client
+
+
+# ─── State + priority translation ───────────────────────────────────────
+
+# Common Compute → AWS Batch terminal/transient state strings. Multiple
+# CC states fold into a single AWS state on purpose — AWS Batch consumers
+# only care about a handful of values.
+_CC_TO_AWS_STATE = {
+    "queued": "SUBMITTED",
+    "leased": "STARTING",
+    "running": "RUNNING",
+    "completed": "SUCCEEDED",
+    "failed": "FAILED",
+    "dead_letter": "FAILED",
+    "cancelled": "FAILED",
+}
+
+
+def _aws_state(cc_state: Optional[str]) -> str:
+    return _CC_TO_AWS_STATE.get(cc_state or "", "SUBMITTED")
+
+
+def _to_iso(ms: Optional[int]) -> Optional[str]:
+    if not ms:
+        return None
+    return datetime.fromtimestamp(ms / 1000, tz=timezone.utc).isoformat()
+
+
+# ─── jobDefinition parsing ──────────────────────────────────────────────
+
+# Heuristic: anything containing one of these workload ids gets routed
+# directly. Otherwise we default to ``coreml_embed`` (the most common
+# batch workload migrated from AWS Batch + bge-* embeddings pipelines).
+# Workload identity comes from the generated catalog mirror so this shim
+# resolves jobDefinitions identically to the wire shim (apps/aws-shim) and the
+# Node SDK. The previous hand-maintained set had drifted to 12 of the catalog's
+# 24 workloads. Regenerate with:
+#   npx tsx packages/workloads/codegen/emit-sdk-workloads.ts
+from ._generated_workloads import IMAGE_PATTERNS as _IMAGE_PATTERN_SOURCES
+from ._generated_workloads import KNOWN_WORKLOAD_IDS as _KNOWN_WORKLOADS
+
+_DEF_RE = re.compile(r"^([A-Za-z0-9_\-]+)(?::(\d+|\$LATEST))?$")
+
+# Compile the generated (source, workload) patterns case-insensitively.
+_IMAGE_PATTERNS: list[tuple[re.Pattern, str]] = [
+    (re.compile(src, re.I), workload) for src, workload in _IMAGE_PATTERN_SOURCES
+]
+
+
+def _infer_workload_from_image(image: str) -> Optional[str]:
+    for pattern, workload in _IMAGE_PATTERNS:
+        if pattern.search(image):
+            return workload
+    return None
+
+
+def _parse_job_definition(job_definition: str) -> tuple[str, Optional[str]]:
+    """Returns (workload_id, model_id) from a jobDefinition like
+    ``"bge-large:1"`` or ``"coreml_embed:bge-large"`` or
+    ``"my-pipeline-rev3"``."""
+    m = _DEF_RE.match(job_definition or "")
+    if m:
+        head, _rev = m.group(1), m.group(2)
+        if head in _KNOWN_WORKLOADS:
+            return head, None
+        return "coreml_embed", head
+    if ":" in job_definition:
+        left, right = job_definition.split(":", 1)
+        if left in _KNOWN_WORKLOADS:
+            return left, right
+    return "coreml_embed", job_definition or None
+
+
+def _payload_from_overrides(overrides: Optional[Mapping[str, Any]]) -> Any:
+    """AWS Batch passes user data through environment variables. We
+    collect them into a dict; if INPUT_PAYLOAD is a JSON string we use
+    it directly so the user can carry arbitrary structure across the
+    boundary without inventing a new convention."""
+    if not overrides:
+        return None
+    env_list = overrides.get("environment") or []
+    env: dict[str, str] = {}
+    for item in env_list:
+        name = item.get("name")
+        value = item.get("value")
+        if isinstance(name, str) and isinstance(value, str):
+            env[name] = value
+    if "INPUT_PAYLOAD" in env:
+        raw = env.pop("INPUT_PAYLOAD")
+        try:
+            return json.loads(raw)
+        except json.JSONDecodeError:
+            return {"input": raw, "env": env}
+    return env if env else None
+
+
+def _priority_from_args(*, scheduling_priority: Any, share_identifier: Any) -> str:
+    if isinstance(share_identifier, str) and share_identifier.lower() == "priority":
+        return "priority"
+    if isinstance(scheduling_priority, (int, float)) and scheduling_priority >= 50:
+        return "priority"
+    return "batch"
+
+
+# ─── Public surface ─────────────────────────────────────────────────────
+
+class BatchClient:
+    """boto3.client('batch')-shaped facade over a CommonCompute client.
+
+    Implements the AWS Batch methods that account for >95% of usage in
+    inference pipelines: submit_job, describe_jobs, list_jobs,
+    cancel_job, terminate_job. Any unimplemented method raises
+    ``NotImplementedError`` with a link to the migration guide.
+    """
+
+    def __init__(
+        self,
+        cc: Optional[Client] = None,
+        *,
+        api_key: Optional[str] = None,
+        base_url: Optional[str] = None,
+    ) -> None:
+        if cc is not None:
+            self._cc = cc
+        else:
+            kwargs: dict[str, Any] = {}
+            if api_key:
+                kwargs["api_key"] = api_key
+            if base_url:
+                kwargs["base_url"] = base_url
+            self._cc = Client(**kwargs)
+        # Process-local jobName cache so describe_jobs/list_jobs can echo
+        # back the original AWS-Batch jobName the customer chose. The
+        # CC backend doesn't store user-supplied job names today, so
+        # without this cache the only way to identify a job in describe
+        # output is by jobId. Hash-keyed by jobId.
+        self._job_names: dict[str, str] = {}
+
+    # ── submit_job ─────────────────────────────────────────────────────
+
+    def submit_job(
+        self,
+        *,
+        jobName: str,
+        jobQueue: Optional[str] = None,
+        jobDefinition: str = "",
+        containerOverrides: Optional[Mapping[str, Any]] = None,
+        nodeOverrides: Optional[Mapping[str, Any]] = None,
+        retryStrategy: Optional[Mapping[str, Any]] = None,
+        timeout: Optional[Mapping[str, Any]] = None,
+        propagateTags: Optional[bool] = None,
+        tags: Optional[Mapping[str, str]] = None,
+        schedulingPriorityOverride: Any = None,
+        shareIdentifier: Any = None,
+        dependsOn: Optional[Iterable[Mapping[str, Any]]] = None,
+        arrayProperties: Optional[Mapping[str, Any]] = None,
+        parameters: Optional[Mapping[str, Any]] = None,
+        # AWS Batch accepts **kwargs we don't model — silently ignore
+        # for forward-compat with new boto3 fields.
+        **_extra: Any,
+    ) -> dict[str, Any]:
+        if arrayProperties:
+            return self._submit_array_job(
+                jobName=jobName,
+                jobDefinition=jobDefinition,
+                containerOverrides=containerOverrides,
+                parameters=parameters,
+                arrayProperties=arrayProperties,
+                schedulingPriorityOverride=schedulingPriorityOverride,
+                shareIdentifier=shareIdentifier,
+            )
+        if dependsOn:
+            raise NotImplementedError(
+                "Job dependencies aren't supported by this shim. Sequence "
+                "submissions in your application code or open an issue with "
+                "your dependency pattern."
+            )
+
+        workload_id, model_id = _parse_job_definition(jobDefinition)
+        payload = _payload_from_overrides(containerOverrides)
+        priority = _priority_from_args(
+            scheduling_priority=schedulingPriorityOverride,
+            share_identifier=shareIdentifier,
+        )
+
+        # max_spend_usd from `timeout.attemptDurationSeconds` would be a
+        # very loose proxy; leave it unset so the user gets a quote-shape
+        # error and learns the correct field if they want a cap.
+        idempotency_key = jobName  # AWS Batch dedupes on jobName per queue.
+
+        job = self._cc.jobs.submit(
+            workload_id=workload_id,
+            model_id=model_id,
+            priority=priority,
+            payload=payload,
+            idempotency_key=idempotency_key,
+        )
+        job_id = job.get("id") or ""
+        if job_id:
+            self._job_names[job_id] = jobName
+        return {
+            "jobName": jobName,
+            "jobId": job_id,
+            "jobArn": f"cc:job:{job_id}",  # AWS callers expect an ARN-ish string
+        }
+
+    # ── array jobs ────────────────────────────────────────────────────
+
+    def _submit_array_job(
+        self,
+        *,
+        jobName: str,
+        jobDefinition: str,
+        containerOverrides: Optional[Mapping[str, Any]],
+        parameters: Optional[Mapping[str, Any]],
+        arrayProperties: Mapping[str, Any],
+        schedulingPriorityOverride: Any,
+        shareIdentifier: Any,
+    ) -> dict[str, Any]:
+        """Translate AWS arrayProperties into a CommonCompute batch.
+
+        Item source priority:
+          1. ``parameters['array_items_json']`` — literal JSON array
+          2. ``parameters['array_input_uri']`` — JSONL/JSON file URI
+          3. fan-out: N copies, each with ``payload['array_index'] = i``
+        """
+        size = int(arrayProperties.get("size") or 0)
+        if size <= 0:
+            raise ValueError("arrayProperties.size must be a positive integer")
+        workload_id, model_id = _parse_job_definition(jobDefinition)
+        priority = _priority_from_args(
+            scheduling_priority=schedulingPriorityOverride,
+            share_identifier=shareIdentifier,
+        )
+        env_payload = _payload_from_overrides(containerOverrides)
+
+        params = dict(parameters or {})
+        items: list[dict[str, Any]] = []
+        if "array_items_json" in params:
+            arr = json.loads(params["array_items_json"])
+            for i, raw in enumerate(arr[:size]):
+                if isinstance(raw, dict) and "payload" in raw:
+                    items.append({"custom_id": raw.get("custom_id") or f"{jobName}-{i}", "payload": raw["payload"]})
+                else:
+                    items.append({"custom_id": f"{jobName}-{i}", "payload": raw})
+        elif "array_input_uri" in params:
+            raise NotImplementedError(
+                "array_input_uri requires the wire-protocol shim at "
+                "aws.commoncompute.ai (the in-process shim can't fetch R2 URIs). "
+                "Set endpoint_url=https://aws.commoncompute.ai on boto3 to use it."
+            )
+        else:
+            base_payload = env_payload if isinstance(env_payload, dict) else {}
+            for i in range(size):
+                items.append({
+                    "custom_id": f"{jobName}-{i}",
+                    "payload": {**base_payload, "array_index": i},
+                })
+
+        resp = self._cc._transport.request(  # type: ignore[attr-defined]
+            "POST", "/v1/batches",
+            json_body={
+                "name": jobName,
+                "workload_id": workload_id,
+                "model_id": model_id,
+                "priority": priority,
+                "items": items,
+            },
+        )
+        batch_id = resp.get("id") or ""
+        if batch_id:
+            self._job_names[batch_id] = jobName
+        return {
+            "jobName": jobName,
+            "jobId": batch_id,
+            "jobArn": f"arn:cc:batch:::job/{batch_id}",
+            "arrayProperties": {"size": resp.get("total", size), "statusSummary": {}},
+        }
+
+    # ── register_job_definition ───────────────────────────────────────
+
+    def register_job_definition(
+        self,
+        *,
+        jobDefinitionName: str,
+        type: str = "container",
+        containerProperties: Optional[Mapping[str, Any]] = None,
+        parameters: Optional[Mapping[str, Any]] = None,
+        tags: Optional[Mapping[str, str]] = None,
+        **_extra: Any,
+    ) -> dict[str, Any]:
+        """In-process shim: persist a JobDefinition by recording the
+        resolved workload locally so subsequent ``submit_job`` calls
+        can reference it by name. The wire-protocol shim at
+        aws.commoncompute.ai persists this server-side; this shim
+        keeps a per-process registry only.
+        """
+        if not hasattr(self, "_local_defs"):
+            self._local_defs: dict[str, dict[str, Any]] = {}
+        # Resolve a workload from the image hint or the cc-workload tag.
+        cc_workload = (tags or {}).get("cc-workload") or (parameters or {}).get("cc_workload")
+        if cc_workload and cc_workload in _KNOWN_WORKLOADS:
+            workload_id = cc_workload
+        else:
+            image = (containerProperties or {}).get("image", "") or ""
+            workload_id = _infer_workload_from_image(image)
+            if not workload_id:
+                raise ValueError(
+                    f"Couldn't infer a CommonCompute workload from image={image!r}. "
+                    "Set tags['cc-workload'] to one of: " + ", ".join(sorted(_KNOWN_WORKLOADS))
+                )
+        revision = 1 + sum(1 for d in self._local_defs.values() if d["name"] == jobDefinitionName)
+        arn = f"arn:cc:batch:::job-definition/{jobDefinitionName}:{revision}"
+        self._local_defs[arn] = {
+            "name": jobDefinitionName, "revision": revision, "workload_id": workload_id,
+            "model_id": (tags or {}).get("cc-model"), "containerProperties": containerProperties or {},
+            "parameters": parameters or {}, "tags": tags or {},
+        }
+        # Also key by name (latest) and name:rev for boto3 lookup parity.
+        self._local_defs[jobDefinitionName] = self._local_defs[arn]
+        self._local_defs[f"{jobDefinitionName}:{revision}"] = self._local_defs[arn]
+        return {
+            "jobDefinitionName": jobDefinitionName,
+            "jobDefinitionArn": arn,
+            "revision": revision,
+        }
+
+    # ── describe_jobs ──────────────────────────────────────────────────
+
+    def describe_jobs(self, *, jobs: Iterable[str], **_extra: Any) -> dict[str, Any]:
+        """Mirror AWS describe_jobs — returns ``{"jobs": [...]}``.
+
+        Just-submitted jobs may not be queryable for ~100ms while the
+        underlying D1 row propagates across the worker's read replicas.
+        We absorb a single 404 by retrying once after 250ms, which is
+        enough to mask the eventual-consistency window in practice."""
+        out: list[dict[str, Any]] = []
+        for jid in jobs:
+            row = None
+            for attempt in range(2):
+                try:
+                    row = self._cc.jobs.get(jid)
+                    break
+                except Exception:
+                    if attempt == 0:
+                        time.sleep(0.25)
+                        continue
+            if row is None:
+                continue
+            shape = _aws_job_shape(row)
+            # Patch in the original jobName from our process-local cache,
+            # so describe_jobs is round-trip-stable for in-process callers.
+            cached_name = self._job_names.get(jid)
+            if cached_name:
+                shape["jobName"] = cached_name
+            out.append(shape)
+        return {"jobs": out}
+
+    # ── list_jobs ──────────────────────────────────────────────────────
+
+    def list_jobs(
+        self,
+        *,
+        jobQueue: Optional[str] = None,
+        arrayJobId: Optional[str] = None,
+        multiNodeJobId: Optional[str] = None,
+        jobStatus: Optional[str] = None,
+        maxResults: Optional[int] = None,
+        nextToken: Optional[str] = None,
+        filters: Any = None,
+        **_extra: Any,
+    ) -> dict[str, Any]:
+        """List jobs from /v1/me/submissions and translate each to AWS shape."""
+        rows = self._cc._transport.request("GET", "/v1/me/submissions").get("jobs", [])  # type: ignore[attr-defined]
+        if maxResults:
+            rows = rows[: int(maxResults)]
+        target = jobStatus.upper() if isinstance(jobStatus, str) else None
+        summaries = []
+        for row in rows:
+            shape = _aws_job_shape(row)
+            if target and shape["status"] != target:
+                continue
+            summaries.append({
+                "jobId": shape["jobId"],
+                "jobName": shape["jobName"],
+                "createdAt": shape.get("createdAt"),
+                "startedAt": shape.get("startedAt"),
+                "stoppedAt": shape.get("stoppedAt"),
+                "status": shape["status"],
+                "statusReason": shape.get("statusReason"),
+            })
+        return {"jobSummaryList": summaries}
+
+    # ── cancel_job / terminate_job ─────────────────────────────────────
+
+    def cancel_job(self, *, jobId: str, reason: str = "", **_extra: Any) -> dict[str, Any]:
+        self._cc._transport.request("POST", f"/v1/jobs/{jobId}/cancel", json_body={"reason": reason})  # type: ignore[attr-defined]
+        return {}
+
+    def terminate_job(self, *, jobId: str, reason: str = "", **_extra: Any) -> dict[str, Any]:
+        # AWS distinguishes cancel (queued) vs terminate (running). On
+        # Common Compute the same endpoint handles both — the router
+        # picks the right path based on current state.
+        return self.cancel_job(jobId=jobId, reason=reason)
+
+    # ── Convenience: poll-to-completion ───────────────────────────────
+
+    def wait_until_done(
+        self,
+        *,
+        jobId: str,
+        timeout: float = 300.0,
+        poll_interval: float = 2.0,
+    ) -> dict[str, Any]:
+        """Block until job is SUCCEEDED or FAILED. Not part of boto3 API
+        but commonly needed at migration sites that previously used
+        ``waiter = batch.get_waiter('jobs_complete')``."""
+        deadline = time.monotonic() + timeout
+        while True:
+            row = self._cc.jobs.get(jobId)
+            state = _aws_state(row.get("state"))
+            if state in ("SUCCEEDED", "FAILED"):
+                shape = _aws_job_shape(row)
+                cached_name = self._job_names.get(jobId)
+                if cached_name:
+                    shape["jobName"] = cached_name
+                return shape
+            if time.monotonic() > deadline:
+                raise TimeoutError(f"Job {jobId} not done in {timeout}s (state={state})")
+            time.sleep(poll_interval)
+
+    # ── Result download (not in boto3 — commonly needed for CC) ────────
+
+    def download_result(self, *, jobId: str, dest: Optional[str] = None) -> bytes:
+        """Download a completed job's result body. AWS Batch customers
+        typically read S3 themselves; on Common Compute the API streams
+        the result directly. Returns the raw bytes; if ``dest`` is given,
+        writes to that path and returns the bytes anyway.
+
+        Raises if the job isn't in a terminal state yet."""
+        body = self._cc._transport.request("GET", f"/v1/jobs/{jobId}/result")
+        # The transport may return parsed JSON or raw bytes depending on
+        # content-type — coerce to bytes for the file-write path.
+        if isinstance(body, (dict, list)):
+            data = json.dumps(body).encode("utf-8")
+        elif isinstance(body, str):
+            data = body.encode("utf-8")
+        elif isinstance(body, (bytes, bytearray)):
+            data = bytes(body)
+        else:
+            data = json.dumps(body).encode("utf-8")
+        if dest:
+            with open(dest, "wb") as f:
+                f.write(data)
+        return data
+
+    # ── Unimplemented stubs for legibility ────────────────────────────
+
+    def __getattr__(self, name: str) -> Any:
+        # Catch any boto3 method we haven't implemented and return a
+        # helpful error rather than AttributeError.
+        if name.startswith("_"):
+            raise AttributeError(name)
+
+        def _missing(*_a: Any, **_kw: Any) -> Any:
+            raise NotImplementedError(
+                f"BatchClient.{name} is not implemented by the Common Compute "
+                "AWS Batch shim. The shim covers submit_job, describe_jobs, "
+                "list_jobs, cancel_job, terminate_job, and wait_until_done — "
+                "AWS Batch concepts like compute environments and job queues "
+                "are obviated by the Common Compute architecture. See "
+                "https://commoncompute.ai/docs/migrate-from-aws-batch."
+            )
+
+        return _missing
+
+
+def _aws_job_shape(cc_row: Mapping[str, Any]) -> dict[str, Any]:
+    """Translate a Common Compute job row into the AWS Batch describe_jobs
+    shape that downstream code expects."""
+    state = cc_row.get("state")
+    return {
+        "jobId": cc_row.get("id"),
+        "jobName": (cc_row.get("metadata") or {}).get("jobName") or cc_row.get("type"),
+        "jobArn": f"cc:job:{cc_row.get('id', '')}",
+        "status": _aws_state(state),
+        "statusReason": cc_row.get("error") or None,
+        "createdAt": _to_iso(cc_row.get("created_at")),
+        "startedAt": _to_iso(cc_row.get("leased_at") or cc_row.get("started_at")),
+        "stoppedAt": _to_iso(cc_row.get("completed_at")),
+        "attempts": cc_row.get("attempts"),
+        "container": {
+            "exitCode": 0 if state == "completed" else None,
+        },
+        # Common Compute extension fields — not in AWS Batch but useful
+        # to callers who learn the shape and want extra signal.
+        "_cc_workload_id": cc_row.get("type"),
+        "_cc_device_id": cc_row.get("device_id"),
+        "_cc_result_uri": cc_row.get("result_uri"),
+    }
+
+
+def client(
+    *,
+    api_key: Optional[str] = None,
+    base_url: Optional[str] = None,
+) -> BatchClient:
+    """Construct a boto3-shaped client. Mirrors ``boto3.client('batch')``.
+
+    ``region_name``, ``aws_access_key_id``, etc. are accepted by the
+    boto3 caller convention but ignored here (no AWS region applies).
+    Use ``api_key`` to override the ``CC_API_KEY`` env var.
+    """
+    return BatchClient(api_key=api_key, base_url=base_url)
+
+
+__all__ = ["BatchClient", "client"]