diff --git a/docs/source/en/concepts/sandbox.md b/docs/source/en/concepts/sandbox.md index e15793f33a..3d5b8c235d 100644 --- a/docs/source/en/concepts/sandbox.md +++ b/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//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 diff --git a/src/huggingface_hub/_sandbox.py b/src/huggingface_hub/_sandbox.py index d1c9d8140b..e91518f0ec 100644 --- a/src/huggingface_hub/_sandbox.py +++ b/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": diff --git a/tests/test_sandbox.py b/tests/test_sandbox.py index b02c98dad1..6ec1a4d951 100644 --- a/tests/test_sandbox.py +++ b/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// and is terminated with a DELETE on the host (the host job keeps running)."""