bladekit 4.5.0__py3-none-any.whl

bladekit/_batch.py

@@ -0,0 +1,150 @@
+"""``@bladekit.batched`` — server-side dynamic batching (SPEC: batching.md).
+
+Dynamic batching auto-collects individual calls into a *batch* and invokes the
+decorated function ONCE per batch, then fans the per-call results back to each
+caller. This is distinct from ``@app.concurrent`` (which forwards N concurrent
+requests to a long-lived server, each handled independently): a batched function
+is written to take *lists* of arguments and return a *list* of results.
+
+    @app.function(gpu="RTX-A4500", image="fabric/torch")
+    @bladekit.batched(max_batch_size=32, wait_ms=50)
+    def embed(prompts: list[str]) -> list[list[float]]:
+        # `prompts` is the collected batch; return one result per prompt.
+        return model.encode(prompts).tolist()
+
+Callers still invoke ONE item at a time and get ONE result back::
+
+    vec = embed.remote("hello")          # batched with concurrent callers
+    handle = embed.spawn("world")        # non-blocking
+
+The batch fills when it reaches ``max_batch_size`` inputs OR ``wait_ms``
+milliseconds elapse since the batch's first input — whichever comes first
+(mirrors ``modal.batched``). The runner owns the accumulation/dispatch loop
+(``forge_runner.batch``); this module is the SDK surface: the decorator, its
+validation, and the wire ``batch_config`` carried on the run/deploy manifest.
+
+CONTRACT: the collected argument is passed *positionally*. A batched function
+takes exactly the positional parameters the individual call passes, each widened
+to a list of the per-call values in batch order. Keyword arguments a caller
+passes are batched into a list under the same keyword. Every returned list MUST
+have the same length as the batch, else the runner raises
+``BATCH_RESULT_LENGTH_MISMATCH`` and fails every call in the batch.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Optional
+
+from . import exceptions as exc
+from ._constants import BATCH_MAX_SIZE_LIMIT, BATCH_WAIT_MS_LIMIT
+
+
+def _validate_batch_config(max_batch_size: int, wait_ms: int) -> None:
+    """Validate the batching knobs at decoration time (before any network I/O)."""
+    if (
+        not isinstance(max_batch_size, int)
+        or isinstance(max_batch_size, bool)
+        or max_batch_size < 1
+        or max_batch_size > BATCH_MAX_SIZE_LIMIT
+    ):
+        raise exc.InvalidArgumentError(
+            "BATCH_CONFIG_INVALID",
+            f"max_batch_size must be an integer between 1 and "
+            f"{BATCH_MAX_SIZE_LIMIT}; got {max_batch_size!r}",
+        )
+    if (
+        not isinstance(wait_ms, int)
+        or isinstance(wait_ms, bool)
+        or wait_ms < 0
+        or wait_ms > BATCH_WAIT_MS_LIMIT
+    ):
+        raise exc.InvalidArgumentError(
+            "BATCH_CONFIG_INVALID",
+            f"wait_ms must be an integer between 0 and {BATCH_WAIT_MS_LIMIT}; "
+            f"got {wait_ms!r}",
+        )
+
+
+class BatchedFunction:
+    """Wraps a batch-shaped user function, carrying its batch configuration.
+
+    ``@app.function`` detects this wrapper (via ``__forge_batched__``) and folds
+    the config into the ``Function`` so the run/deploy manifest carries it. The
+    object is transparently callable — it delegates to the user function — so
+    ``.local()`` (which passes lists directly) and any introspection keep working
+    exactly as if the plain function had been decorated.
+    """
+
+    def __init__(self, fn: Callable[..., Any], max_batch_size: int, wait_ms: int) -> None:
+        _validate_batch_config(max_batch_size, wait_ms)
+        self._fn = fn
+        self.max_batch_size = max_batch_size
+        self.wait_ms = wait_ms
+        # The wire config; the shape mirrors the runner's parser and the server's
+        # EndpointSpec.batch_config (batching.md §4).
+        self.__forge_batched__ = {
+            "max_batch_size": max_batch_size,
+            "wait_ms": wait_ms,
+        }
+        # Keep the wrapped function introspectable / importable by the runner:
+        # FunctionRef.qualname points at the module attribute, which is THIS
+        # object, so it must carry the user function's identity.
+        self.__name__ = getattr(fn, "__name__", "fn")
+        self.__qualname__ = getattr(fn, "__qualname__", self.__name__)
+        self.__module__ = getattr(fn, "__module__", "__main__")
+        self.__doc__ = getattr(fn, "__doc__", None)
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """Invoke the underlying batch function directly (lists in, list out).
+
+        This is the path the runner takes once it has assembled a batch, and the
+        path ``.local()`` takes for local testing with already-batched inputs.
+        """
+        return self._fn(*args, **kwargs)
+
+    def __repr__(self) -> str:
+        return (
+            f"<bladekit.batched {self.__qualname__} "
+            f"max_batch_size={self.max_batch_size} wait_ms={self.wait_ms}>"
+        )
+
+
+def batched(
+    _fn: Optional[Callable[..., Any]] = None,
+    *,
+    max_batch_size: int,
+    wait_ms: int,
+) -> Any:
+    """Decorator: auto-collect individual calls into server-side batches.
+
+    Stack it UNDER ``@app.function`` (Modal order), so the function decorator
+    wraps the batched function::
+
+        @app.function(gpu="RTX-A4500", image="fabric/torch")
+        @bladekit.batched(max_batch_size=32, wait_ms=50)
+        def embed(prompts: list[str]) -> list[list[float]]:
+            return model.encode(prompts).tolist()
+
+    Args:
+        max_batch_size: the maximum number of individual calls fused into one
+            invocation (1..1000). A batch is dispatched as soon as it fills.
+        wait_ms: how long the runner waits, after the batch's first input, for
+            more inputs to arrive before dispatching a partial batch (0..600000).
+            ``0`` dispatches as soon as the event loop yields — useful when you
+            only want to fuse calls that are already queued.
+
+    Returns:
+        A :class:`BatchedFunction` wrapping the user function.
+    """
+
+    def wrap(fn: Callable[..., Any]) -> BatchedFunction:
+        if isinstance(fn, BatchedFunction):
+            raise exc.InvalidArgumentError(
+                "BATCH_CONFIG_INVALID",
+                "@bladekit.batched cannot wrap an already-batched function",
+            )
+        return BatchedFunction(fn, max_batch_size, wait_ms)
+
+    if _fn is not None:
+        return wrap(_fn)
+    return wrap

bladekit/_bucket.py

@@ -0,0 +1,211 @@
+"""``bladekit.CloudBucketMount`` — mount an external object-store bucket as a path.
+
+A ``CloudBucketMount`` mounts an S3 / GCS / R2 / MinIO bucket (anything that
+speaks the S3 API) as a POSIX filesystem path inside a function, deployed
+replica, or sandbox, using a FUSE mount helper (``s3fs``/``goofys``/``rclone``)
+that the runner sets up on container start and tears down on teardown.
+
+Unlike :class:`bladekit.Volume` — an *AI Fabric-managed* network volume with generations
+and ``commit()`` — a bucket mount is a live view of a bucket **you own** in an
+**external** cloud account. AI Fabric never stores the objects; it only mounts them.
+
+Usage — pass it in the same ``volumes=`` map as a ``Volume`` (both are
+"filesystem things mounted at a path"), keyed by mount path::
+
+    creds = bladekit.Secret.from_name("aws-credentials")
+
+    @app.function(
+        image=img,
+        volumes={
+            "/mnt/data": bladekit.CloudBucketMount(
+                "my-data-bucket",
+                secret=creds,
+                read_only=True,
+                key_prefix="raw/2024/",
+            ),
+        },
+    )
+    def process():
+        import pathlib
+        for p in pathlib.Path("/mnt/data").glob("*.parquet"):
+            ...
+
+Security invariant (audit-security-idor.md, conventions.md §11): the bucket
+credentials come **only** from a :class:`bladekit.Secret`, never inline. The mount
+spec on the wire carries the *secret name*, never a value; the control plane
+resolves the secret and injects its values as container env vars (the same path
+as ``secrets=[...]``), and the runner reads those env vars at mount time. A
+credential value never appears in the SDK object, the run body, or a log line.
+
+``read_only=True`` (the default is read-write) mounts the bucket without write
+access — the runner passes the mount helper's read-only flag and the mount is
+refused any object PUT/DELETE.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+from . import exceptions as exc
+
+# S3 bucket names are 3–63 chars, lowercase letters/digits/dots/hyphens. We keep
+# the check permissive (providers vary: GCS allows underscores, R2 is looser)
+# but reject the obvious footguns — empty, whitespace, a leading scheme — so a
+# typo fails in the SDK instead of as an opaque mount error in the container.
+_MAX_BUCKET_LEN = 255
+
+
+class CloudBucketMount:
+    """Reference to an external object-store bucket mounted at a path.
+
+    Construct directly (there is no ``from_name`` — a bucket is not an AI Fabric
+    object, it is your external bucket)::
+
+        bladekit.CloudBucketMount("my-bucket", secret=bladekit.Secret.from_name("aws-credentials"))
+
+    Parameters
+    ----------
+    bucket_name:
+        The external bucket to mount (S3/GCS/R2/MinIO).
+    secret:
+        A :class:`bladekit.Secret` holding the credentials. Required — credentials
+        are never passed inline. The secret's env keys are provider-specific
+        (e.g. ``AWS_ACCESS_KEY_ID`` / ``AWS_SECRET_ACCESS_KEY`` for S3).
+    read_only:
+        Mount without write access (default ``False`` = read-write).
+    key_prefix:
+        Mount only the objects under this key prefix (must end with ``/``); the
+        prefix is stripped from the visible paths. Omit to mount the whole
+        bucket.
+    bucket_endpoint_url:
+        Override the S3 endpoint for non-AWS stores (R2, MinIO, GCS-S3,
+        Scaleway). Omit for AWS S3.
+    requester_pays:
+        Set for buckets configured with S3 Requester Pays.
+    """
+
+    __slots__ = (
+        "bucket_name",
+        "_secret",
+        "_read_only",
+        "_key_prefix",
+        "_endpoint_url",
+        "_requester_pays",
+    )
+
+    def __init__(
+        self,
+        bucket_name: str,
+        *,
+        secret: Any,
+        read_only: bool = False,
+        key_prefix: Optional[str] = None,
+        bucket_endpoint_url: Optional[str] = None,
+        requester_pays: bool = False,
+    ) -> None:
+        # Import here (not at module top) to avoid any import cycle and keep the
+        # isinstance check against the concrete class.
+        from ._secret import Secret
+
+        if not isinstance(bucket_name, str) or not bucket_name.strip():
+            raise exc.InvalidArgumentError(
+                "INVALID_ARGUMENT", "bucket_name must be a non-empty string"
+            )
+        if len(bucket_name) > _MAX_BUCKET_LEN:
+            raise exc.InvalidArgumentError(
+                "INVALID_ARGUMENT",
+                f"bucket_name too long ({len(bucket_name)} > {_MAX_BUCKET_LEN})",
+            )
+        if "://" in bucket_name or "/" in bucket_name:
+            raise exc.InvalidArgumentError(
+                "INVALID_ARGUMENT",
+                f"bucket_name {bucket_name!r} must be a bare bucket name, not a "
+                "URL or path (use key_prefix= for a sub-path)",
+            )
+        # Credentials come ONLY from a Secret, never inline (security invariant).
+        if not isinstance(secret, Secret):
+            raise exc.InvalidArgumentError(
+                "INVALID_ARGUMENT",
+                "CloudBucketMount requires secret=bladekit.Secret.from_name(...); "
+                "credentials are never passed inline",
+            )
+        if key_prefix is not None:
+            if not isinstance(key_prefix, str) or not key_prefix:
+                raise exc.InvalidArgumentError(
+                    "INVALID_ARGUMENT", "key_prefix must be a non-empty string"
+                )
+            if key_prefix.startswith("/"):
+                raise exc.InvalidArgumentError(
+                    "INVALID_ARGUMENT",
+                    "key_prefix must be a bucket-relative key, not start with '/'",
+                )
+            if not key_prefix.endswith("/"):
+                raise exc.InvalidArgumentError(
+                    "INVALID_ARGUMENT",
+                    f"key_prefix {key_prefix!r} must end with '/'",
+                )
+
+        self.bucket_name = bucket_name
+        self._secret = secret
+        self._read_only = bool(read_only)
+        self._key_prefix = key_prefix
+        self._endpoint_url = bucket_endpoint_url
+        self._requester_pays = bool(requester_pays)
+
+    @property
+    def secret_name(self) -> str:
+        """Name of the Secret that provides this bucket's credentials."""
+        return self._secret.name
+
+    @property
+    def read_only(self) -> bool:
+        return self._read_only
+
+    def _spec(self, mount_path: str) -> Dict[str, Any]:
+        """Serialize to the wire spec (openapi BucketMountSpec).
+
+        Carries the secret *name* only — never a credential value. Optional
+        fields are emitted only when set so a minimal mount body stays minimal.
+        """
+        spec: Dict[str, Any] = {
+            "bucket_name": self.bucket_name,
+            "mount_path": mount_path,
+            "secret": self._secret.name,
+        }
+        if self._read_only:
+            spec["read_only"] = True
+        if self._key_prefix is not None:
+            spec["key_prefix"] = self._key_prefix
+        if self._endpoint_url is not None:
+            spec["endpoint_url"] = self._endpoint_url
+        if self._requester_pays:
+            spec["requester_pays"] = True
+        return spec
+
+    def __repr__(self) -> str:
+        ro = ", read_only=True" if self._read_only else ""
+        return f"<bladekit.CloudBucketMount bucket={self.bucket_name!r}{ro}>"
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, CloudBucketMount):
+            return NotImplemented
+        return (
+            self.bucket_name == other.bucket_name
+            and self._secret == other._secret
+            and self._read_only == other._read_only
+            and self._key_prefix == other._key_prefix
+            and self._endpoint_url == other._endpoint_url
+            and self._requester_pays == other._requester_pays
+        )
+
+    def __hash__(self) -> int:
+        return hash(
+            (
+                self.bucket_name,
+                self._secret,
+                self._read_only,
+                self._key_prefix,
+                self._endpoint_url,
+                self._requester_pays,
+            )
+        )