Pass large worker tensors through sealed memfds

docs/index.md

@@ -13,6 +13,7 @@ services can build kernels for a separate GPU server.
 | Write a client program | [Write a client program](client-guide/writing-a-program.md), [KCoral Protocol](client-guide/protocol.md) |
 | Run Python, checks, profilers, or shell commands | [Builtin CLI Tools](client-guide/builtin-cli-tools.md) |
 | Operate a service | [Launch the server](server-guide/launch-the-server.md), [Logging](server-guide/logging.md) and [Launch the router](server-guide/router.md) |
+| Understand large-tensor transfer cost | [Worker Tensor Transfer](server-guide/tensor-transfer.md) |
 | Compile remotely, then execute the returned library in another request | [Remote Compilation](tutorials/remote-compilation.md) |
 | Measure or automate a workload | [Benchmark a Kernel with KCoral](tutorials/benchmark-kernel.md) and [Agent Integration Guide](tutorials/agent-integration.md) |
 | Look up Python interfaces | [Python API](python-api/index.rst) |
@@ -52,6 +53,7 @@ client-guide/builtin-cli-tools
 server-guide/launch-the-server
 server-guide/router
 server-guide/logging
+server-guide/tensor-transfer
 ```
 
 ```{toctree}

docs/server-guide/tensor-transfer.md

@@ -0,0 +1,49 @@
+# Worker Tensor Transfer
+
+A KCoral server runs each request in a worker process and moves the request's
+tensors across the process boundary. For small inputs this is cheap. For large
+tensors the copy dominates the request, so the server passes big blobs through
+shared memory instead of serializing them over the worker control socket. This
+page explains what crosses the boundary, when the shared-memory path is used,
+and what it saves.
+
+## What crosses the boundary
+
+The server and each worker communicate over a control socket. By default a
+request's inputs and results are pickled and written over that socket, including
+on cache hits. Pickling and copying a few kilobytes is negligible, but a
+64 MiB tensor is copied several times on the way in and again on the way out,
+and that copying, not the GPU work, sets the request's wall time.
+
+## The shared-memory path
+
+Blobs of at least 1 MiB — inputs and results alike — are written into a sealed
+[`memfd`](https://man7.org/linux/man-pages/man2/memfd_create.2.html) and the
+file descriptor is handed to the worker over the same socket. The receiver
+checks the seals and the exact size, then maps the buffer read-only; no bytes
+are copied over the socket. Smaller blobs, non-Linux hosts, and Python builds
+without `memfd_create` keep the pickling path, so the change is transparent to
+clients.
+
+## Measured effect
+
+Vector add with 64 MiB tensors, medians of 30 trials in balanced order,
+sandbox off, measured with the benchmark harness on two consumer GPUs:
+
+| GPU | worker | serialized | shared memory |
+| --- | --- | ---: | ---: |
+| RTX 5090 | reused | 997 ms | 643 ms (−36%) |
+| RTX 5090 | fresh | 1335 ms | 940 ms (−30%) |
+| RTX 2070 | reused | 982 ms | 657 ms (−33%) |
+| RTX 2070 | fresh | 1310 ms | 940 ms (−28%) |
+
+A 4 MiB reused-worker request drops 17–28%. Requests below the 1 MiB threshold
+do not use the shared-memory path and are unchanged.
+
+## Reproduce
+
+The numbers above come from a vector-add workload driven through the server over
+loopback. Follow [Benchmark a Kernel with KCoral](../tutorials/benchmark-kernel.md)
+to run a kernel end to end, and vary the tensor size and
+`--max-requests-per-worker` (reused vs. fresh workers) to see the transfer cost
+and the shared-memory path at each size.

python/kcoral/client/program.py

@@ -4,6 +4,7 @@
 
 import os
 import sys
+from concurrent.futures import Future, ThreadPoolExecutor
 from dataclasses import dataclass, field
 from typing import Any
 
@@ -19,6 +20,9 @@
 )
 from kcoral.schemas import DTYPE_ITEM_SIZES
 
+_HASH_POOL = ThreadPoolExecutor(max_workers=4, thread_name_prefix="kcoral-hash")
+_HASH_THRESHOLD = 1 << 20
+
 
 @dataclass(frozen=True)
 class Register:
@@ -50,7 +54,8 @@ class Program:
     """
 
     _instructions: list[dict[str, Any]] = field(default_factory=list, init=False)
-    _blobs: dict[str, bytes] = field(default_factory=dict, init=False)
+    _blob_data: dict[str, bytes] = field(default_factory=dict, init=False)
+    _pending_blobs: dict[Future, bytes] = field(default_factory=dict, init=False, repr=False)
     _ids: set[str] = field(default_factory=set, init=False)
     _next_id: int = field(default=0, init=False)
     _return_keys: set[str] = field(default_factory=set, init=False)
@@ -63,8 +68,35 @@ def instructions(self) -> list[dict[str, Any]]:
         The contained dictionaries are shared with this program. Treat them
         as read-only; use the builder methods to add instructions.
         """
+        self._finish_blobs()
         return list(self._instructions)
 
+    @property
+    def _blobs(self) -> dict[str, bytes]:
+        self._finish_blobs()
+        return self._blob_data
+
+    def _queue_blob(self, data: bytes):
+        if len(data) >= _HASH_THRESHOLD:
+            digest = _HASH_POOL.submit(compute_blob_hash, data)
+            self._pending_blobs[digest] = data
+            return digest
+        digest = compute_blob_hash(data)
+        self._blob_data.setdefault(digest, data)
+        return digest
+
+    def _finish_blobs(self) -> None:
+        if not self._pending_blobs:
+            return
+        digests = {future: future.result() for future in self._pending_blobs}
+        for future, data in self._pending_blobs.items():
+            self._blob_data.setdefault(digests[future], data)
+        for instruction in self._instructions:
+            digest = instruction.get("blob")
+            if isinstance(digest, Future):
+                instruction["blob"] = digests[digest]
+        self._pending_blobs.clear()
+
     def upload_file(self, *, blob: Any, path: str, id: str | None = None) -> Register:
         """Snapshot bytes-like data as a file in the request workspace.
 
@@ -167,8 +199,7 @@ def upload(
             if source is not None:
                 raise TypeError("tensor upload does not accept 'source'")
             tensor_dtype, tensor_shape, raw = _tensor_fields(value, dtype=dtype, shape=shape)
-            blob_hash = compute_blob_hash(raw)
-            self._blobs.setdefault(blob_hash, raw)
+            blob_hash = self._queue_blob(raw)
             instruction = {
                 "op": "upload",
                 "kind": "tensor",
@@ -185,8 +216,7 @@ def upload(
                 raw = value if isinstance(value, bytes) else bytes(memoryview(value))
             except TypeError as exc:
                 raise TypeError("bytes upload requires a bytes-like 'value'") from exc
-            blob_hash = compute_blob_hash(raw)
-            self._blobs.setdefault(blob_hash, raw)
+            blob_hash = self._queue_blob(raw)
             instruction = {
                 "op": "upload",
                 "kind": "bytes",
@@ -198,8 +228,7 @@ def upload(
             if dtype is not None or shape is not None:
                 raise TypeError("library upload does not accept tensor fields")
             raw = value if isinstance(value, bytes) else bytes(memoryview(value))
-            blob_hash = compute_blob_hash(raw)
-            self._blobs.setdefault(blob_hash, raw)
+            blob_hash = self._queue_blob(raw)
             instruction = {
                 "op": "upload",
                 "kind": "library",

python/kcoral/runtime/engine.py

@@ -39,7 +39,7 @@ class Runtime(Protocol):
     def load_module(self, source: str) -> Any: ...
     def load_library(self, data: bytes) -> Any: ...
     def get_function(self, module: Any, name: str) -> Any: ...
-    def load_tensor(self, data: bytes, dtype: str, shape: list[int]) -> Any: ...
+    def load_tensor(self, data: bytes | memoryview, dtype: str, shape: list[int]) -> Any: ...
     def export_tensor(self, value: Any) -> tuple[str, list[int], bytes] | None: ...
     def forbid_gpu(self) -> AbstractContextManager[None]: ...
     def synchronize(self) -> None: ...
@@ -116,7 +116,7 @@ def _execute_in_workspace(
                         _materialize_file(
                             workspace_dir,
                             instruction.path,
-                            program.blob_bytes[instruction.blob],
+                            bytes(program.blob_bytes[instruction.blob]),
                         )
                         env[instruction.id] = instruction.path
                     elif isinstance(instruction, Upload):
@@ -125,11 +125,11 @@ def _execute_in_workspace(
                             env[instruction.id] = runtime.load_module(instruction.source)
                         elif instruction.kind == "bytes":
                             assert instruction.blob is not None
-                            env[instruction.id] = program.blob_bytes[instruction.blob]
+                            env[instruction.id] = bytes(program.blob_bytes[instruction.blob])
                         elif instruction.kind == "library":
                             assert instruction.blob is not None
                             env[instruction.id] = runtime.load_library(
-                                program.blob_bytes[instruction.blob]
+                                bytes(program.blob_bytes[instruction.blob])
                             )
                         else:
                             assert (

python/kcoral/runtime/transfer.py

@@ -0,0 +1,121 @@
+"""Request-local sealed buffers on the existing worker control socket."""
+
+import array
+import ctypes
+import mmap
+import os
+import socket
+import sys
+from contextlib import ExitStack, contextmanager
+from dataclasses import dataclass
+from multiprocessing import reduction
+
+if sys.platform == "linux":
+    import fcntl
+
+MEMFD_THRESHOLD = 1 << 20
+_SEAL_NAMES = ("SEAL", "SHRINK", "GROW", "WRITE")
+
+
+def _seals():
+    return sum(getattr(fcntl, "F_SEAL_" + name, 1 << i) for i, name in enumerate(_SEAL_NAMES))
+
+
+def _memfd_create():
+    create = getattr(os, "memfd_create", None)
+    fd = create("blob", 3) if create else ctypes.CDLL(None, use_errno=True).memfd_create(b"blob", 3)
+    if fd < 0:
+        raise OSError(ctypes.get_errno(), "memfd_create failed")
+    return fd
+
+
+class HandleTimeout(OSError):
+    pass
+
+
+class ResultSizeError(OSError):
+    pass
+
+
+@dataclass(frozen=True)
+class BlobHandle:
+    size: int
+
+
+@contextmanager
+def pack_blobs(blobs, cache=None):
+    fds, packed = [], blobs.copy()
+    with ExitStack() as cleanup:
+        try:
+            if sys.platform == "linux":
+                for key, data in blobs.items():
+                    if len(data) < MEMFD_THRESHOLD:
+                        continue
+                    cached_fd = cache.borrow_memfd(key) if cache is not None else None
+                    if cached_fd is not None:
+                        cleanup.callback(os.close, cached_fd)
+                        packed[key] = BlobHandle(len(data))
+                        fds.append(cached_fd)
+                        continue
+                    fd = _memfd_create()  # MFD_CLOEXEC | MFD_ALLOW_SEALING
+                    cleanup.callback(os.close, fd)
+                    with os.fdopen(fd, "wb", closefd=False) as stream:
+                        stream.write(data)
+                    fcntl.fcntl(fd, getattr(fcntl, "F_ADD_SEALS", 1033), _seals())
+                    packed[key] = BlobHandle(len(data))
+                    fds.append(fd)
+        except (OSError, AttributeError, NotImplementedError):
+            cleanup.close()
+            fds, packed = [], blobs
+        yield packed, fds
+
+
+def _recv_handle(conn):
+    fds = array.array("i")
+    try:
+        with socket.fromfd(conn.fileno(), socket.AF_UNIX, socket.SOCK_STREAM) as sock:
+            data, ancillary, flags, _ = sock.recvmsg(
+                1, socket.CMSG_SPACE(256 * fds.itemsize), socket.MSG_CMSG_CLOEXEC
+            )
+            for level, kind, payload in ancillary:
+                if level == socket.SOL_SOCKET and kind == socket.SCM_RIGHTS:
+                    fds.frombytes(payload[: len(payload) - len(payload) % fds.itemsize])
+            if len(fds) != 1 or data != bytes([1]) or flags & socket.MSG_CTRUNC:
+                raise OSError("invalid memfd handle message")
+            if getattr(reduction, "ACKNOWLEDGE", True):
+                sock.sendall(b"A")
+        return fds.pop()
+    finally:
+        for fd in fds:
+            os.close(fd)
+
+
+@contextmanager
+def mapped_blobs(conn, blobs, max_bytes=None, ready=lambda: True):
+    descriptors = [(key, value) for key, value in blobs.items() if isinstance(value, BlobHandle)]
+    if not descriptors:
+        yield blobs
+        return
+    if max_bytes is not None:
+        sizes = [p.size if isinstance(p, BlobHandle) else len(p) for p in blobs.values()]
+        if any(type(n) is not int or n < 0 for n in sizes) or sum(sizes) > max_bytes:
+            raise ResultSizeError("results exceed the configured response-size limit")
+    mapped = blobs.copy()
+    with ExitStack() as cleanup:
+        for key, descriptor in descriptors:
+            if not ready():
+                raise HandleTimeout("worker result handles exceeded the execution timeout")
+            fd = _recv_handle(conn)
+            try:
+                if (
+                    descriptor.size < MEMFD_THRESHOLD
+                    or fcntl.fcntl(fd, getattr(fcntl, "F_GET_SEALS", 1034)) & _seals() != _seals()
+                    or os.fstat(fd).st_size != descriptor.size
+                ):
+                    raise OSError("invalid memfd size or seals")
+                mapping = mmap.mmap(fd, descriptor.size, access=mmap.ACCESS_READ)
+            finally:
+                os.close(fd)
+            cleanup.callback(mapping.close)
+            mapped[key] = cleanup.enter_context(memoryview(mapping))
+        yield mapped