bladekit 4.5.0__py3-none-any.whl

bladekit/_map.py

@@ -0,0 +1,555 @@
+"""Function.map() / .starmap() — batch fan-out over the GPU fleet.
+
+Implements the client side of the map protocol (map-protocol.md):
+
+1. Create a map invocation (POST /v1/maps) — state=open.
+2. Push inputs in batches of <= 1000 (POST /v1/maps/{id}/inputs) — streaming,
+   never loads all items into memory at once.
+3. Finalize (POST /v1/maps/{id}/finalize) — state=finalized.
+4. Stream results as a generator (GET /v1/maps/{id}/results, paginated by
+   cursor) — yields values in order=index by default, optionally order=arrival.
+
+Key design decisions (from the contract):
+- return_exceptions=False (default): the first failed input raises immediately
+  and stops the generator; the batch continues server-side.
+- return_exceptions=True: exceptions are yielded in-place; the generator never
+  raises.
+- Reconnection: if the generator is interrupted and restarted, passing the
+  cursor from the previous run resumes from the last seen result. The CLI's
+  --detach/resume pattern uses this mechanism.
+- Input chunks are <= 1000 per POST (conventions.md §6 limits).
+- Results are paged with limit <= 200 (conventions.md §6).
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any, Generator, Iterable, List, Optional, Tuple
+
+from . import exceptions as exc
+from ._constants import (
+    MAP_INPUT_BATCH_SIZE,
+    MAP_RESULTS_PAGE_LIMIT,
+    MAP_POLL_INTERVAL,
+    MAP_POLL_MAX_INTERVAL,
+)
+
+
+def _make_input_payload(args: tuple, kwargs: dict) -> Tuple[dict, Optional[bytes]]:
+    """Serialize a single (args, kwargs) pair into a map input args payload."""
+    from ._serialize import dumps_args, make_payload_in
+
+    raw = dumps_args(args, kwargs)
+    payload_in, object_bytes = make_payload_in(raw)
+    return payload_in, object_bytes
+
+
+class MapHandle:
+    """Handle for an in-flight or completed map invocation.
+
+    Returned by Function.map() and Function.starmap() when called with
+    ``detach=True``, or by ``blade map resume <id>``. The generator is lazy:
+    no network calls until iteration begins.
+
+    Attributes:
+        map_id: The server-assigned map invocation ID (``map_<ULID>``).
+    """
+
+    def __init__(self, client: Any, map_id: str) -> None:
+        self._client = client
+        self.map_id = map_id
+
+    def __repr__(self) -> str:
+        return f"MapHandle(map_id={self.map_id!r})"
+
+    def results(
+        self,
+        *,
+        order: str = "index",
+        return_exceptions: bool = False,
+        cursor: Optional[str] = None,
+    ) -> Generator[Any, None, None]:
+        """Yield results, resumable from ``cursor``.
+
+        Args:
+            order: ``"index"`` (default, like ``list(map(f, xs))``) or
+                ``"arrival"`` (throughput-max order).
+            return_exceptions: if False (default) the first failed input raises
+                immediately; if True exceptions are yielded in-place.
+            cursor: opaque resume cursor (from a previous interrupted iteration).
+        """
+        yield from _stream_results(
+            self._client,
+            self.map_id,
+            order=order,
+            return_exceptions=return_exceptions,
+            cursor=cursor,
+        )
+
+    def cancel(self) -> None:
+        """Request cancellation of the batch (MAP_NOT_CANCELLABLE if terminal)."""
+        self._client.cancel_map(self.map_id)
+
+    def status(self) -> dict:
+        """Return the current map invocation object from the server."""
+        return self._client.get_map(self.map_id)
+
+
+def _create_map_batch(
+    client: Any,
+    *,
+    function_ref: dict,
+    bundle_id: str,
+    resource_spec: dict,
+    runtime: dict,
+    image: Optional[str],
+    image_build_id: Optional[str],
+    app_name: Optional[str],
+    inputs: Iterable[Tuple[tuple, dict]],
+    retries: Optional[dict],
+    input_timeout_seconds: Optional[int],
+    priority: Optional[int],
+    namespace: Optional[str],
+    max_concurrency: Optional[int],
+    gpu_class: str,
+    best_effort: bool,
+    max_interruptions: Optional[int],
+) -> str:
+    """Create a map invocation, push inputs, finalize; return its ``map_id``.
+
+    Shared by :func:`run_map` (which then streams results) and
+    :func:`spawn_map` (which returns a handle without blocking). Inputs are
+    pushed in chunks of up to 1000 per request so a 1M-item iterable never sits
+    in memory. On any error during push the batch is best-effort cancelled so the
+    server does not keep workers for an abandoned batch.
+    """
+    # 1. Create the map invocation.
+    body: dict[str, Any] = {
+        "function": function_ref,
+        "bundle_id": bundle_id,
+        "resources": resource_spec,
+        "runtime": runtime,
+    }
+    if image is not None:
+        body["image"] = image
+    if image_build_id is not None:
+        body["image_build_id"] = image_build_id
+    if app_name:
+        body["app_name"] = app_name
+    if retries is not None:
+        body["retries"] = retries
+    if input_timeout_seconds is not None:
+        body["input_timeout_seconds"] = input_timeout_seconds
+    if priority is not None:
+        body["priority"] = priority
+    if namespace is not None:
+        body["namespace"] = namespace
+    if max_concurrency is not None:
+        body["max_concurrency"] = max_concurrency
+    # capacity class + preemption budget (omitted when default so a default
+    # map body is byte-for-byte unchanged; m8-spot.md §13.6).
+    if gpu_class != "spot":
+        body["gpu_class"] = gpu_class
+    if best_effort:
+        body["best_effort"] = True
+    if max_interruptions is not None:
+        body["max_interruptions"] = max_interruptions
+
+    map_resp = client.create_map(body)
+    map_id = map_resp["map_id"]
+
+    # Emit any creation warnings (GPU_QUOTA_REACHED, GPU_TYPE_UNAVAILABLE).
+    _emit_map_warnings(map_resp)
+
+    # 2. Push inputs in streaming chunks.
+    try:
+        _push_inputs(client, map_id, inputs)
+    except Exception:
+        # Best-effort cancel so the server doesn't keep workers for an
+        # abandoned batch.
+        try:
+            client.cancel_map(map_id)
+        except exc.FabricError:
+            pass
+        raise
+
+    # 3. Finalize.
+    client.finalize_map(map_id)
+    return map_id
+
+
+def run_map(
+    client: Any,
+    *,
+    function_ref: dict,
+    bundle_id: str,
+    resource_spec: dict,
+    runtime: dict,
+    image: Optional[str] = None,
+    image_build_id: Optional[str] = None,
+    app_name: Optional[str] = None,
+    inputs: Iterable[Tuple[tuple, dict]],
+    order: str = "index",
+    return_exceptions: bool = False,
+    retries: Optional[dict] = None,
+    input_timeout_seconds: Optional[int] = None,
+    priority: Optional[int] = None,
+    namespace: Optional[str] = None,
+    max_concurrency: Optional[int] = None,
+    gpu_class: str = "spot",
+    best_effort: bool = False,
+    max_interruptions: Optional[int] = None,
+) -> Generator[Any, None, None]:
+    """Create a map batch, push inputs, finalize, stream results.
+
+    This is the implementation backing ``Function.map()`` / ``Function.starmap()``.
+
+    The function is a generator: it creates the batch server-side, then yields
+    results as they become available. Inputs are pushed in chunks of up to 1000
+    per request so that 1M-item iterables never sit in memory.
+    """
+    map_id = _create_map_batch(
+        client,
+        function_ref=function_ref,
+        bundle_id=bundle_id,
+        resource_spec=resource_spec,
+        runtime=runtime,
+        image=image,
+        image_build_id=image_build_id,
+        app_name=app_name,
+        inputs=inputs,
+        retries=retries,
+        input_timeout_seconds=input_timeout_seconds,
+        priority=priority,
+        namespace=namespace,
+        max_concurrency=max_concurrency,
+        gpu_class=gpu_class,
+        best_effort=best_effort,
+        max_interruptions=max_interruptions,
+    )
+
+    # Stream results.
+    yield from _stream_results(
+        client,
+        map_id,
+        order=order,
+        return_exceptions=return_exceptions,
+    )
+
+
+def spawn_map(
+    client: Any,
+    *,
+    function_ref: dict,
+    bundle_id: str,
+    resource_spec: dict,
+    runtime: dict,
+    image: Optional[str] = None,
+    image_build_id: Optional[str] = None,
+    app_name: Optional[str] = None,
+    inputs: Iterable[Tuple[tuple, dict]],
+    retries: Optional[dict] = None,
+    input_timeout_seconds: Optional[int] = None,
+    priority: Optional[int] = None,
+    namespace: Optional[str] = None,
+    max_concurrency: Optional[int] = None,
+    gpu_class: str = "spot",
+    best_effort: bool = False,
+    max_interruptions: Optional[int] = None,
+) -> "MapHandle":
+    """Fire an async fan-out and return a :class:`MapHandle` WITHOUT blocking.
+
+    This is the implementation backing ``Function.spawn_map()`` /
+    ``Function.spawn_starmap()``. It creates the batch, pushes all inputs, and
+    finalizes it server-side — then returns immediately. The batch is persisted
+    server-side, so results can be collected later (even from another process,
+    via ``blade map resume <map_id>``) by iterating ``handle.results()``.
+    """
+    map_id = _create_map_batch(
+        client,
+        function_ref=function_ref,
+        bundle_id=bundle_id,
+        resource_spec=resource_spec,
+        runtime=runtime,
+        image=image,
+        image_build_id=image_build_id,
+        app_name=app_name,
+        inputs=inputs,
+        retries=retries,
+        input_timeout_seconds=input_timeout_seconds,
+        priority=priority,
+        namespace=namespace,
+        max_concurrency=max_concurrency,
+        gpu_class=gpu_class,
+        best_effort=best_effort,
+        max_interruptions=max_interruptions,
+    )
+    return MapHandle(client, map_id)
+
+
+def _push_inputs(
+    client: Any,
+    map_id: str,
+    inputs: Iterable[Tuple[tuple, dict]],
+) -> None:
+    """Push inputs to the server in batches of MAP_INPUT_BATCH_SIZE."""
+    batch: List[dict] = []
+    idx = 0
+
+    for args, kwargs in inputs:
+        payload_in, object_bytes = _make_input_payload(args, kwargs)
+        if object_bytes is not None:
+            # Large payload: upload to object store first.
+            decl = client.create_payload(payload_in["size_bytes"], payload_in["sha256"])
+            client.upload_to_presigned(
+                decl["put_url"], object_bytes, "application/octet-stream"
+            )
+            payload_in["payload_id"] = decl["payload_id"]
+
+        # client_input_id for idempotent push (conventions.md §14.1):
+        # a unique string per input so a retry of a failed POST batch
+        # does not produce duplicate inputs.
+        batch.append({
+            "args": payload_in,
+            "client_input_id": f"sdk-{map_id}-{idx}",
+        })
+        idx += 1
+
+        if len(batch) >= MAP_INPUT_BATCH_SIZE:
+            client.push_map_inputs(map_id, batch)
+            batch = []
+
+    # Flush the last partial batch.
+    if batch:
+        client.push_map_inputs(map_id, batch)
+
+
+_MAP_TERMINAL_STATES = frozenset({"completed", "cancelled"})
+
+
+def _stream_results(
+    client: Any,
+    map_id: str,
+    *,
+    order: str = "index",
+    return_exceptions: bool = False,
+    cursor: Optional[str] = None,
+) -> Generator[Any, None, None]:
+    """Poll GET /v1/maps/{id}/results pages until the batch is done.
+
+    Termination signal (openapi.yaml MapResultPage):
+      - ``next_cursor`` is the *only* contractual termination signal:
+        ``null`` ⟺ all currently-available results have been read AND the
+        batch is terminal (map_state in {completed, cancelled}).
+      - ``map_state`` is read to confirm the batch is terminal when
+        ``next_cursor`` is null (defensive belt-and-suspenders check).
+
+    Uses exponential backoff between empty non-terminal pages (the server may
+    not have results yet if workers are still starting up).
+    """
+    poll_interval = MAP_POLL_INTERVAL
+
+    while True:
+        page = client.get_map_results(
+            map_id,
+            order=order,
+            return_exceptions=return_exceptions,
+            cursor=cursor,
+            limit=MAP_RESULTS_PAGE_LIMIT,
+        )
+
+        items = page.get("items") or []
+        next_cursor = page.get("next_cursor")
+        map_state = page.get("map_state", "")
+
+        for item in items:
+            outcome = item.get("outcome")
+            if outcome == "succeeded":
+                yield _decode_result(client, item)
+                poll_interval = MAP_POLL_INTERVAL  # reset after progress
+            elif outcome in ("exception", "error"):
+                if return_exceptions:
+                    yield _decode_exception_or_error(client, item)
+                    poll_interval = MAP_POLL_INTERVAL
+                else:
+                    # Raise the exception and stop the generator.
+                    # Attach next_cursor so callers can resume (S1: openapi.yaml:2026-2030).
+                    error = _build_map_error(client, item, map_id)
+                    if next_cursor is not None:
+                        # error may be an arbitrary reconstructed user exception,
+                        # so attach details dynamically (only FabricError defines it).
+                        err_any: Any = error
+                        if getattr(err_any, "details", None) is None:
+                            err_any.details = {}
+                        err_any.details["resume_cursor"] = next_cursor
+                    raise error
+            else:
+                # Unknown outcome: treat as an error to be safe.
+                if return_exceptions:
+                    yield exc.RunError(
+                        f"unexpected map result outcome: {outcome!r}",
+                        code="INTERNAL",
+                    )
+                else:
+                    raise exc.RunError(
+                        f"unexpected map result outcome: {outcome!r}",
+                        code="INTERNAL",
+                    )
+
+        # Advance cursor with whatever the server returned (opaque, persistable).
+        if next_cursor is not None:
+            cursor = next_cursor
+
+        # Contractual termination: next_cursor=null ⟺ drained AND terminal.
+        # (openapi.yaml:5209-5213: "null ⟺ all results read AND batch terminal")
+        if next_cursor is None and map_state in _MAP_TERMINAL_STATES:
+            return
+
+        # next_cursor=null but map not yet terminal (e.g. finalized, not completed):
+        # keep current cursor and back off.
+        if not items:
+            time.sleep(poll_interval)
+            poll_interval = min(poll_interval * 1.5, MAP_POLL_MAX_INTERVAL)
+
+
+def _decode_result(client: Any, item: dict) -> Any:
+    """Deserialize a succeeded map result."""
+    from ._serialize import decode_payload_out, loads
+
+    result = item.get("result")
+    if result is None:
+        return None
+    raw = decode_payload_out(result, fetch=client.download_presigned)
+    return loads(raw)
+
+
+def _decode_exception_or_error(client: Any, item: dict) -> BaseException:
+    """Return the exception object for a failed input (return_exceptions=True)."""
+    outcome = item.get("outcome")
+    if outcome == "exception":
+        exc_info = item.get("exception")
+        if exc_info:
+            return _build_user_exception(client, exc_info)
+    # outcome == "error": infra error (RETRIES_EXHAUSTED, INPUT_TIMEOUT, PREEMPTED)
+    error = item.get("error") or {}
+    code = error.get("code", "INTERNAL")
+    message = error.get("message", code)
+    hint = _map_error_hint(code, error.get("details") or {})
+    if hint:
+        message = f"{message}\n{hint}"
+    return exc.RunError(message, code=code, details=error.get("details") or {})
+
+
+def _build_map_error(client: Any, item: dict, map_id: str) -> BaseException:
+    """Build the exception to raise when return_exceptions=False."""
+    outcome = item.get("outcome")
+    if outcome == "exception":
+        exc_info = item.get("exception")
+        if exc_info:
+            return _build_user_exception(client, exc_info)
+    # Infra error.
+    error = item.get("error") or {}
+    code = error.get("code", "INTERNAL")
+    message = error.get("message", code)
+    hint = _map_error_hint(code, error.get("details") or {})
+    if hint:
+        message = f"{message}\n{hint}"
+    return exc.RunError(
+        message, code=code, run_id=map_id, details=error.get("details") or {}
+    )
+
+
+def _build_user_exception(client: Any, info: dict) -> BaseException:
+    """Reconstruct (or wrap) a remote user exception."""
+    from ._serialize import decode_payload_out, loads
+
+    tb = info.get("traceback_str", "")
+    pickled = info.get("pickled")
+    if pickled:
+        try:
+            raw = decode_payload_out(pickled, fetch=client.download_presigned)
+            original = loads(raw)
+            if isinstance(original, BaseException):
+                original._forge_remote_traceback = tb  # type: ignore[attr-defined]
+                note = f"\n\n--- remote traceback ---\n{tb}"
+                if hasattr(original, "add_note"):
+                    try:
+                        original.add_note(note)
+                    except Exception:
+                        pass
+                return original
+        except Exception:
+            pass
+
+    exc_type = info.get("type", "Exception")
+    message = info.get("message", "")
+    return exc.RemoteFunctionError(exc_type, message, tb)
+
+
+def _emit_map_warnings(resp: dict) -> None:
+    """Print structured creation warnings to stderr."""
+    warnings = resp.get("warnings")
+    if not warnings:
+        return
+    try:
+        from ._render import err_console
+    except Exception:
+        return
+    for w in warnings:
+        code = w.get("code", "WARNING")
+        message = w.get("message", "")
+        err_console.print(f"[yellow]warning[/yellow] [bold]{code}[/bold]: {message}")
+        details = w.get("details") or {}
+        if code == "GPU_QUOTA_REACHED":
+            ns = details.get("namespace", "")
+            gpu_type = details.get("gpu_type", "")
+            limit = details.get("limit", "?")
+            in_use = details.get("in_use", "?")
+            err_console.print(
+                f"  namespace {ns!r} has used {in_use}/{limit} GPUs of type "
+                f"{gpu_type!r}; batch will queue."
+            )
+        elif code == "ORG_GPU_QUOTA_REACHED":
+            in_use = details.get("in_use", "?")
+            limit = details.get("limit", "?")
+            gpu_type = details.get("gpu_type", "*")
+            scope = "GPUs" if gpu_type == "*" else f"GPUs of type {gpu_type!r}"
+            err_console.print(
+                f"  organization is using {in_use}/{limit} {scope}; "
+                "the batch advances as usage drops below the limit."
+            )
+        elif code == "GPU_TYPE_UNAVAILABLE":
+            known = details.get("known_types") or []
+            if known:
+                err_console.print(
+                    f"  [yellow]available GPU types:[/yellow] {', '.join(known)}"
+                )
+
+
+def _map_error_hint(code: str, details: dict) -> str:
+    """Return a short actionable hint for map-specific error codes."""
+    if code == "RETRIES_EXHAUSTED":
+        last = details.get("last_error_code", "")
+        base = "The input exhausted all retry attempts."
+        if last:
+            base += f" Last error: {last}."
+        base += " Check your function for transient failures, or increase retries=."
+        return base
+    if code == "INPUT_TIMEOUT":
+        return (
+            "An input exceeded its per-input timeout. Increase input_timeout_seconds= "
+            "on the map() call, or optimize the function for speed."
+        )
+    if code == "PREEMPTED":
+        return (
+            "A worker was preempted by a higher-priority run. "
+            "The input will be automatically re-queued."
+        )
+    if code == "GPU_QUOTA_EXCEEDED":
+        ns = details.get("namespace", "")
+        gpu_type = details.get("gpu_type", "")
+        return (
+            f"Namespace {ns!r} has exceeded its GPU quota for {gpu_type!r}. "
+            "Wait for quota to free up, or ask an admin to raise the limit."
+        )
+    return ""

bladekit/_preempt.py

@@ -0,0 +1,94 @@
+"""``@app.on_preempt`` / ``@bladekit.on_preempt`` hook + checkpoint helpers.
+
+Phase A formalises the soft-preemption checkpoint hook that runner-protocol.md
+§15.2 left as "le hook checkpoint quand il existera". The SDK exposes:
+
+- a decorator — ``@app.on_preempt`` (or its alias ``@bladekit.on_preempt``) — that
+  registers ONE zero-arg callback to run when the container is being preempted
+  (SIGTERM + ``FABRIC_PREEMPT_GRACE_SECONDS``, runner-protocol.md §15.2.1);
+- ``bladekit.save_checkpoint(obj, path)`` / ``bladekit.load_checkpoint(path,
+  default=None)`` — the idempotent-resume primitives, writing/reading a
+  cloudpickle blob on a mounted Volume (the Volume owns durability).
+
+Hook discovery is the SAME mechanism as ``@bladekit.enter`` (``_cls.py``): the
+decorator marks the callable with the dunder attribute below, and the runner
+scans for it (the bundle module for a free function, the live instance for an
+``@app.cls`` method). That attribute is THE SDK↔runner contract — see also
+``forge_runner.preempt``.
+"""
+
+from __future__ import annotations
+
+import os
+import tempfile
+from typing import Any, Callable, Optional
+
+# SDK↔runner contract (runner-protocol.md §15.2.1): the decorator marks the
+# callback with this attribute; the runner discovers it the same way it finds
+# ``@bladekit.enter`` hooks (``__forge_enter__``). A method on an ``@app.cls`` class
+# is found on the live instance; a free function is found in the bundle module.
+FORGE_ON_PREEMPT_ATTR = "__forge_on_preempt__"
+
+
+def on_preempt(fn: Callable[..., Any]) -> Callable[..., Any]:
+    """Register ``fn`` to run when the container is being preempted.
+
+    Usable as ``@bladekit.on_preempt`` on a free function or on an ``@app.cls``
+    method (``self`` is supplied — the checkpoint can read model/optimizer state
+    off the live instance), and re-exported as ``@app.on_preempt``. The callback
+    takes no argument (besides ``self`` for a method) and is executed
+    best-effort exactly once on the SIGTERM grace window, BEFORE the process
+    exits (runner-protocol.md §15.2.1). The last decorated callback wins
+    (deterministic); the hook decides nothing about re-queue/retry — the control
+    plane stays the only judge.
+    """
+    fn.__forge_on_preempt__ = True  # type: ignore[attr-defined]
+    return fn
+
+
+def save_checkpoint(obj: Any, path: str) -> None:
+    """Serialize ``obj`` (cloudpickle) and write it atomically to ``path``.
+
+    ``path`` should point at a file on a mounted Volume: the Volume — not
+    this helper — owns durability. For a read-write volume the upperdir persists
+    for the run; ``volume.commit()`` (runner-protocol.md §11) persists it across
+    runs. The write is atomic (tmp file + ``os.replace``) so a preemption mid-
+    write never leaves a torn checkpoint that ``load_checkpoint`` would choke on.
+    """
+    from ._serialize import dumps
+
+    raw = dumps(obj)
+    directory = os.path.dirname(os.path.abspath(path)) or "."
+    os.makedirs(directory, exist_ok=True)
+    # Write to a temp file in the SAME directory (os.replace is only atomic
+    # within a filesystem) then rename over the target.
+    fd, tmp = tempfile.mkstemp(prefix=".forge-ckpt-", dir=directory)
+    try:
+        with os.fdopen(fd, "wb") as f:
+            f.write(raw)
+            f.flush()
+            os.fsync(f.fileno())
+        os.replace(tmp, path)
+    except BaseException:
+        # Best-effort cleanup of the temp file on any failure.
+        try:
+            os.unlink(tmp)
+        except OSError:
+            pass
+        raise
+
+
+def load_checkpoint(path: str, default: Optional[Any] = None) -> Any:
+    """Read+deserialize the checkpoint at ``path``, else return ``default``.
+
+    The idempotent-resume primitive: call it at function start. A missing file
+    (no checkpoint yet — the first attempt) returns ``default`` instead of
+    raising, so resume code reads "load the last checkpoint or start fresh".
+    """
+    from ._serialize import loads
+
+    if not os.path.exists(path):
+        return default
+    with open(path, "rb") as f:
+        raw = f.read()
+    return loads(raw)