[sandbox audit] Bound what a transfer or a command's output costs the client

docs/source/en/concepts/sandbox.md

@@ -243,7 +243,7 @@ This section is deliberately exhaustive rather than reassuring: if you are decid
 - **Host discovery starts from Job labels**, which any Job creator in the namespace can set — so a label match is a claim, not proof. By default (`adopt_hosts="own"`) the client only adopts hosts *this principal started*, per the Jobs API's `initiator`, and additionally checks the image, flavor, command and exposed URL. Setting `adopt_hosts="namespace"` restores cross-user host sharing and, with it, the ability of any Job creator in that namespace to publish a host your client will send a token to — only use it in a namespace whose members you trust.
 - **Landlock can degrade silently.** If Landlock is unavailable, or its ruleset cannot be built, the server currently falls back to uid-only isolation and creates the sandbox anyway — without telling the client. Under uid-only isolation, `/tmp`, `/dev/shm`, TCP bind and cross-home filesystem access are *not* denied. The server also accepts Landlock ABI 1, while the ✅ list above needs ABI 4 (TCP bind) and ABI 6 (abstract sockets); production kernels provide ABI 6, but a lower one would silently drop those two guarantees.
 - **Residual shared channels**, none of which Landlock or uid isolation closes: unrestricted outbound TCP; loopback access to the control server; **UDP bind is allowed** (Landlock has no UDP coverage); a sibling's `/proc/<pid>/cmdline` and `status` are readable (`environ` is not — that is the part that would leak credentials, and it is denied); `/proc` and `/sys` are readable and `/dev` is broadly readable and writable; kernel IPC and all machine resources are shared.
-- **No CPU, disk, FD or total-memory quotas.** Only per-process `RLIMIT_NPROC` and `RLIMIT_AS` are set; cgroup delegation is not available on Jobs. One sandbox can starve its neighbours. The `max_procs`/`max_mem_mb` values are caller-supplied and not clamped server-side.
+- **No CPU-share, disk or total-memory *quota*.** Per-process `RLIMIT_NPROC`, `RLIMIT_AS`, `RLIMIT_NOFILE`, `RLIMIT_FSIZE` and `RLIMIT_CPU` are set, and `max_procs`/`max_mem_mb` are clamped server-side — but cgroup delegation is not available on Jobs, so these are per-process ceilings rather than a share of the host. A sandbox running many processes within its limits can still crowd its neighbours.
 - **GPU flavors are untested in pool mode.** The client does not prevent one. Use `Sandbox.create` for GPU.
 
 ### Lifecycle and operational gaps

src/huggingface_hub/_sandbox.py

@@ -70,6 +70,12 @@
 
 DEFAULT_SANDBOXES_PER_HOST = 50
 
+# Ceiling on the output `run()` will accumulate for its result. Without one, a
+# runaway command (a loop printing to stdout, say) is an unbounded allocation in
+# the caller's process -- and it happened even when a callback was already
+# consuming the output. Raising beats being OOM-killed with no explanation.
+MAX_CAPTURED_OUTPUT_CHARS = 64 * 1024 * 1024
+
 # Which pool hosts `create()` may adopt from job labels. Labels are set by whoever
 # creates a Job, so "any job carrying our pool's labels" is not an identity claim --
 # see the `adopt_hosts` argument of [`SandboxPool`].
@@ -241,15 +247,36 @@ class SandboxFiles:
     PARALLEL_THRESHOLD = 2 * 1024 * 1024
     PARALLEL_CHUNK_SIZE = 1 * 1024 * 1024
     PARALLEL_MAX_WORKERS = 16
+    # Ceiling on what `read`/`read_text` will materialize in memory. A parallel
+    # read used to collect every chunk into a list and *then* join it, so a 2 GB
+    # file peaked at roughly twice its size before the caller saw a byte. Reading
+    # a file into memory is inherently bounded by the file; this makes the bound
+    # explicit and points at the streaming alternative instead of dying in an
+    # allocator.
+    MAX_READ_BYTES = 512 * 1024 * 1024
 
     def __init__(self, sandbox: "Sandbox") -> None:
         self._sandbox = sandbox
 
     def read(self, path: str) -> bytes:
-        """Read a file from the sandbox and return its content as bytes."""
+        """Read a file from the sandbox and return its content as bytes.
+
+        Raises [`SandboxError`] above `MAX_READ_BYTES`; use [`download`] for
+        anything that large, which streams to disk instead of buffering.
+        """
         size = self.stat(path).size
+        if size > self.MAX_READ_BYTES:
+            raise SandboxError(
+                f"{path} is {size} bytes, over the {self.MAX_READ_BYTES}-byte limit for reading into "
+                "memory. Use `files.download(path, local_path)`, which streams to disk."
+            )
         if size > self.PARALLEL_THRESHOLD:
-            return b"".join(self._read_ranges(path, size))
+            # Write each range into one preallocated buffer as it arrives, rather
+            # than collecting every chunk and joining: the join doubled peak memory.
+            buffer = bytearray(size)
+            for offset, part in self._read_ranges(path, size):
+                buffer[offset : offset + len(part)] = part
+            return bytes(buffer)
         response = self._sandbox._request("GET", "/files/read", params={"path": path})
         return response.content
 
@@ -279,8 +306,12 @@ def upload(self, local_path: str | Path, path: str, mode: str | None = None) ->
         whatever the path resolves to a moment later.
         """
         with open(local_path, "rb") as f:
-            if os.fstat(f.fileno()).st_size > self.PARALLEL_THRESHOLD:
-                self._write_ranges(path, f.read(), mode)
+            size = os.fstat(f.fileno()).st_size
+            if size > self.PARALLEL_THRESHOLD:
+                # Each worker reads its own range straight off the descriptor, so
+                # a large upload no longer means holding the whole file in memory
+                # (`f.read()` on a 10 GB file was 10 GB of resident bytes).
+                self._write_ranges_from_file(path, f, size, mode)
             else:
                 self.write(path, f, mode=mode)
 
@@ -294,7 +325,10 @@ def download(self, path: str, local_path: str | Path) -> None:
         size = self.stat(path).size
         with _open_download_target(Path(local_path)) as f:
             if size > self.PARALLEL_THRESHOLD:
-                for part in self._read_ranges(path, size):
+                # Seek-and-write per range as it completes, so peak memory is a
+                # few chunks rather than the whole file.
+                for offset, part in self._read_ranges(path, size):
+                    f.seek(offset)
                     f.write(part)
                 return
             with self._sandbox._stream("GET", "/files/read", params={"path": path}) as response:
@@ -315,30 +349,76 @@ def _parallel(self, items: List[Any], fn: Callable[[Any], Any]) -> List[Any]:
         with ThreadPoolExecutor(workers) as executor:
             return list(executor.map(fn, items))
 
-    def _read_ranges(self, path: str, size: int) -> List[bytes]:
-        def fetch(rng: tuple[int, int]) -> bytes:
+    def _read_ranges(self, path: str, size: int) -> Iterator[tuple[int, bytes]]:
+        """Fetch every range in parallel, yielding `(offset, bytes)` as they land.
+
+        Yielding rather than returning a list is the point: the caller places each
+        chunk and drops it, so peak memory is a few chunks instead of the whole
+        file plus a copy.
+        """
+
+        def fetch(rng: tuple[int, int]) -> tuple[int, bytes]:
             offset, length = rng
             response = self._sandbox._request(
                 "GET", "/files/read", params={"path": path, "offset": offset, "length": length}
             )
-            return response.content
+            return offset, response.content
 
-        return self._parallel(self._ranges(size), fetch)
+        ranges = self._ranges(size)
+        workers = min(self.PARALLEL_MAX_WORKERS, len(ranges))
+        with ThreadPoolExecutor(workers) as executor:
+            yield from executor.map(fetch, ranges)
 
     def _write_ranges(self, path: str, data: bytes, mode: str | None) -> None:
         def push(rng: tuple[int, int]) -> None:
             offset, length = rng
-            params: dict[str, Any] = {"path": path, "offset": offset}
-            if mode is not None:
-                params["mode"] = mode
-            self._sandbox._request("PUT", "/files/write", params=params, content=data[offset : offset + length])
+            self._put_range(path, offset, data[offset : offset + length], mode, total=len(data))
 
         self._parallel(self._ranges(len(data)), push)
 
+    def _write_ranges_from_file(self, path: str, source: BinaryIO, size: int, mode: str | None) -> None:
+        """Upload `source` in parallel ranges, reading each range on demand.
+
+        `pread` rather than `seek`+`read`, because the workers share one
+        descriptor and a seeking read would race with its siblings.
+        """
+        fileno = source.fileno()
+
+        def push(rng: tuple[int, int]) -> None:
+            offset, length = rng
+            self._put_range(path, offset, os.pread(fileno, length, offset), mode, total=size)
+
+        self._parallel(self._ranges(size), push)
+
+    def _put_range(self, path: str, offset: int, chunk: bytes, mode: str | None, *, total: int) -> None:
+        params: dict[str, Any] = {"path": path, "offset": offset}
+        if mode is not None:
+            params["mode"] = mode
+        # Tell the server the intended final size. A ranged write does not
+        # truncate, so overwriting a larger file with a smaller one used to leave
+        # the old bytes past the new end.
+        if offset + len(chunk) >= total:
+            params["truncate_to"] = total
+        self._sandbox._request("PUT", "/files/write", params=params, content=chunk)
+
     def list(self, path: str) -> List[FileEntry]:
-        """List a directory in the sandbox."""
-        response = self._sandbox._request("GET", "/files/list", params={"path": path})
-        return [FileEntry(**entry) for entry in response.json()["entries"]]
+        """List a directory in the sandbox.
+
+        The server paginates, so this follows the cursor and returns everything --
+        a caller sees one list regardless of how the directory is chunked on the
+        wire.
+        """
+        entries: List[FileEntry] = []
+        after: str | None = None
+        while True:
+            params: dict[str, Any] = {"path": path}
+            if after is not None:
+                params["after"] = after
+            payload = self._sandbox._request("GET", "/files/list", params=params).json()
+            entries.extend(FileEntry(**entry) for entry in payload["entries"])
+            after = payload.get("next")
+            if not after:
+                return entries
 
     def stat(self, path: str) -> FileEntry:
         """Get metadata of a file or directory in the sandbox."""
@@ -860,6 +940,7 @@ def run(
         on_stdout: Callable[[str], None] | None = ...,
         on_stderr: Callable[[str], None] | None = ...,
         check: bool = ...,
+        capture_output: bool = ...,
         background: Literal[False] = ...,
     ) -> SandboxCommandResult: ...
 
@@ -886,6 +967,7 @@ def run(
         on_stdout: Callable[[str], None] | None = None,
         on_stderr: Callable[[str], None] | None = None,
         check: bool = True,
+        capture_output: bool = True,
         background: bool = False,
     ) -> SandboxCommandResult | SandboxProcess:
         """Run a command in the sandbox and wait for it, streaming output live.
@@ -919,6 +1001,11 @@ def run(
                 Callback invoked with stderr chunks as they arrive.
             check (`bool`, *optional*, defaults to `True`):
                 If True, raise [`SandboxCommandError`] on non-zero exit.
+            capture_output (`bool`, *optional*, defaults to `True`):
+                If True, accumulate stdout/stderr into the returned result. Pass `False`
+                when you only want `on_stdout`/`on_stderr`: output is then handed to the
+                callbacks and dropped, so a command producing gigabytes does not have to fit
+                in memory. `result.stdout`/`result.stderr` are empty in that case.
             background (`bool`, *optional*, defaults to `False`):
                 If True, start the command detached and return a [`SandboxProcess`] right
                 away instead of waiting for it and returning a [`SandboxCommandResult`].
@@ -939,17 +1026,35 @@ def run(
         if stdin is not None:
             payload["stdin"] = stdin
 
+        # Output was accumulated unconditionally, even when a callback was
+        # consuming it -- so a streaming consumer still paid full memory for the
+        # command's entire output. `capture_output=False` opts out.
         stdout_parts: list[str] = []
         stderr_parts: list[str] = []
+        captured = 0
         result: SandboxCommandResult | None = None
+
+        def keep(parts: list[str], data: str) -> None:
+            nonlocal captured
+            if not capture_output:
+                return
+            if captured + len(data) > MAX_CAPTURED_OUTPUT_CHARS:
+                raise SandboxError(
+                    f"command produced more than {MAX_CAPTURED_OUTPUT_CHARS} characters of output. "
+                    "Pass `capture_output=False` with `on_stdout`/`on_stderr` to stream it instead, "
+                    "or redirect it to a file in the sandbox and download that."
+                )
+            captured += len(data)
+            parts.append(data)
+
         with self._stream("POST", "/exec", json=payload) as response:
             for event in _iter_events(response):
                 if event["event"] == "stdout":
-                    stdout_parts.append(event["data"])
+                    keep(stdout_parts, event["data"])
                     if on_stdout is not None:
                         on_stdout(event["data"])
                 elif event["event"] == "stderr":
-                    stderr_parts.append(event["data"])
+                    keep(stderr_parts, event["data"])
                     if on_stderr is not None:
                         on_stderr(event["data"])
                 elif event["event"] == "exit":

tests/test_sandbox.py

@@ -219,6 +219,18 @@ def do_GET(self) -> None:
             self.send_header("Content-Length", "5")
             self.end_headers()
             self.wfile.write(b"hello")
+        elif "/files/list" in self.path:
+            # Paginated, like the real server: `list_pages` lets a test hand out a
+            # cursor and check the client follows it.
+            pages = getattr(cls, "list_pages", None)
+            if pages:
+                after = ""
+                if "after=" in self.path:
+                    after = self.path.split("after=")[1].split("&")[0]
+                page = pages[0] if not after else next((p for p in pages if p.get("_after") == after), pages[-1])
+                self._json({"entries": page["entries"], "next": page.get("next")})
+            else:
+                self._json({"entries": [], "next": None})
         elif self.path.endswith("/processes"):  # list background processes
             self._json(cls.processes)
         elif self.path == "/v1/sandboxes":
@@ -378,6 +390,54 @@ def test_context_manager_closes_when_reattached(self, fake_server: str) -> None:
         assert sandbox._server._client.is_closed
 
 
+class TestResourceBounds:
+    """Client-side memory bounds. Each of these was unbounded: the caller's
+    process grew with whatever the sandbox produced, whether or not it was
+    already consuming it."""
+
+    def test_output_is_not_accumulated_when_the_caller_opts_out(self, fake_server: str) -> None:
+        sandbox = _make_sandbox(fake_server)
+        chunks: list = []
+        result = sandbox.run("echo", on_stdout=chunks.append, capture_output=False)
+
+        # The callback still sees everything; the result deliberately holds nothing.
+        assert chunks == ["out1"]
+        assert result.stdout == ""
+        assert result.stderr == ""
+        assert result.exit_code == 0
+
+    def test_output_is_accumulated_by_default(self, fake_server: str) -> None:
+        sandbox = _make_sandbox(fake_server)
+        assert sandbox.run("echo").stdout == "out1"
+
+    def test_runaway_output_raises_instead_of_growing_without_bound(self, fake_server: str, monkeypatch) -> None:
+        # Lowered so the test need not actually produce 64 MB.
+        monkeypatch.setattr(sandbox_mod, "MAX_CAPTURED_OUTPUT_CHARS", 2)
+        sandbox = _make_sandbox(fake_server)
+        with pytest.raises(SandboxError, match="capture_output=False"):
+            sandbox.run("echo")
+        # ...and the escape hatch the error names actually works.
+        assert sandbox.run("echo", capture_output=False).stdout == ""
+
+    def test_read_refuses_a_file_too_large_to_hold_in_memory(self, fake_server: str, monkeypatch) -> None:
+        monkeypatch.setattr(sandbox_mod.SandboxFiles, "MAX_READ_BYTES", 1)
+        sandbox = _make_sandbox(fake_server)
+        with pytest.raises(SandboxError, match="files.download"):
+            sandbox.files.read("big.bin")
+
+    def test_list_follows_the_servers_pagination(self, fake_server: str) -> None:
+        # The server pages; a caller should still see one list.
+        sandbox = _make_sandbox(fake_server)
+        _FakeServer.list_pages = [
+            {"entries": [{"name": "a", "path": "/a", "type": "file", "size": 1}], "next": "a"},
+            {"_after": "a", "entries": [{"name": "b", "path": "/b", "type": "file", "size": 1}]},
+        ]
+        try:
+            assert [entry.name for entry in sandbox.files.list("/dir")] == ["a", "b"]
+        finally:
+            _FakeServer.list_pages = None
+
+
 class TestSharedSandbox:
     """A shared sandbox routes operations under /v1/sandboxes/<local_id>/ and is
     terminated with a DELETE on the host (the host job keeps running)."""