Bound memory during native Docker artifact collection

benchmark/runtime/sforge_backend.py

@@ -2,11 +2,11 @@
 
 from __future__ import annotations
 
-import io
 import inspect
 import ipaddress
 import json
 import tarfile
+import tempfile
 from pathlib import Path, PurePosixPath
 
 from docker.errors import ImageNotFound
@@ -191,6 +191,24 @@ def start_feedback(self, handle):
             raise RuntimeError("Best-only publisher is unavailable")
         self.feedback.start(self, handle)
 
+    def _extract_artifact_archive(self, handle, remote, target):
+        """Collect one native archive without materializing it as Python bytes."""
+        chunks, _ = self._raw(handle).get_archive(str(remote))
+        try:
+            # A seekable spool preserves the existing tar extraction/filter
+            # behavior. Its lifetime is independent of TarFile reference cycles.
+            with tempfile.TemporaryFile() as archive_file:
+                for chunk in chunks:
+                    archive_file.write(chunk)
+                archive_file.seek(0)
+                target.mkdir(exist_ok=True)
+                with tarfile.open(fileobj=archive_file) as archive:
+                    archive.extractall(target, filter="data")
+        finally:
+            close = getattr(chunks, "close", None)
+            if close is not None:
+                close()
+
     def cleanup_container(self, handle, logger=None):
         native_failure = getattr(self, "native_failure_evidence", None) if handle is not None else None
         try:
@@ -223,11 +241,8 @@ def cleanup_container(self, handle, logger=None):
                     sources.append((RESEARCH_DIRECTORY, "research"))
                 for remote, name in sources:
                     try:
-                        data = self.copy_from_container(handle, PurePosixPath(remote))
-                        target = self.log_dir / name
-                        target.mkdir(exist_ok=True)
-                        with tarfile.open(fileobj=io.BytesIO(data)) as archive:
-                            archive.extractall(target, filter="data")
+                        self._extract_artifact_archive(
+                            handle, PurePosixPath(remote), self.log_dir / name)
                         observations.append({"artifact": name, "collected": True})
                     except Exception as error:
                         observations.append({"artifact": name, "collected": False,

benchmark/tests/test_sforge_artifact_spooling.py

@@ -0,0 +1,165 @@
+"""Collection keeps the native extraction, privacy and cleanup contracts."""
+import errno
+import io
+import json
+import tarfile
+from types import SimpleNamespace
+
+import pytest
+import requests
+
+pytest.importorskip("sforge")
+pytest.importorskip("harbor")
+from benchmark.runtime import sforge_backend as transport
+
+
+def archive_bytes(*, unsafe=False):
+    with io.BytesIO() as payload:
+        with tarfile.open(fileobj=payload, mode="w") as archive:
+            file = tarfile.TarInfo("artifact/nested/data")
+            file.mode = 0o640
+            file.size = 7
+            archive.addfile(file, io.BytesIO(b"payload"))
+            link = tarfile.TarInfo("artifact/link")
+            link.type = tarfile.SYMTYPE
+            link.linkname = "/outside" if unsafe else "nested/data"
+            archive.addfile(link)
+        return payload.getvalue()
+
+
+def collector(tmp_path):
+    backend = object.__new__(transport.RecordingDockerBackend)
+    backend.checkpoint = backend.feedback = None
+    backend.research_sidecar = False
+    backend.log_dir = tmp_path / "collected"
+    return backend
+
+
+@pytest.mark.parametrize("failure", [None, "http", "space", "unsafe-link"])
+def test_collection_spool_closes_with_original_error_record_and_native_cleanup(
+        tmp_path, monkeypatch, failure):
+    backend = collector(tmp_path)
+    payload = archive_bytes(unsafe=failure == "unsafe-link")
+    streams, spools, cleaned = [], [], []
+
+    def get_archive(remote):
+        def chunks():
+            try:
+                yield payload[:512]
+                if failure == "http":
+                    raise requests.ConnectionError("synthetic interrupted download")
+                yield payload[512:]
+            finally:
+                streams.append(remote)
+        return chunks(), {}
+
+    monkeypatch.setattr(backend, "_raw", lambda handle: SimpleNamespace(get_archive=get_archive))
+    monkeypatch.setattr(backend, "copy_from_container", lambda *args: pytest.fail("bytes API used"))
+    monkeypatch.setattr(transport.DockerBackend, "cleanup_container",
+                        lambda self, handle, logger: cleaned.append(handle))
+    real = transport.tempfile.TemporaryFile
+
+    def temporary():
+        spool = real(dir=tmp_path)
+        spools.append(spool)
+        if failure == "space":
+            def full(*args):
+                raise OSError(errno.ENOSPC, "synthetic spool full")
+            monkeypatch.setattr(spool, "write", full)
+        return spool
+
+    monkeypatch.setattr(transport.tempfile, "TemporaryFile", temporary)
+    backend.cleanup_container("owned-container")
+    assert cleaned == ["owned-container"]
+    assert len(streams) == len(spools) == 5
+    assert all(spool.closed for spool in spools)
+    observed = json.loads((backend.log_dir / "artifact-collection.json").read_text())
+    assert all(x["collected"] is (failure is None) for x in observed)
+    if failure:
+        expected = {"http": "ConnectionError", "space": "OSError", "unsafe-link": "AbsoluteLinkError"}
+        assert {x["error_kind"] for x in observed} == {expected[failure]}
+        assert not (tmp_path / "outside").exists()
+        if failure in ("http", "space"):
+            assert not (backend.log_dir / "runtime").exists()
+    else:
+        file = backend.log_dir / "runtime/artifact/nested/data"
+        assert file.read_bytes() == b"payload"
+        assert file.stat().st_mode & 0o777 == 0o640
+        assert (file.parents[1] / "link").is_symlink()
+
+
+def test_native_retention_still_avoids_normal_stop_and_removal(tmp_path, monkeypatch):
+    backend = collector(tmp_path)
+    calls = []
+    preservation = SimpleNamespace(retain_source=True, stop_source=lambda handle: calls.append("retain"))
+    backend.checkpoint = SimpleNamespace(preservation=preservation, close=lambda: calls.append("close"))
+    monkeypatch.setattr(backend, "_raw", lambda handle: SimpleNamespace(
+        get_archive=lambda remote: (iter([archive_bytes()]), {})))
+    monkeypatch.setattr(transport.DockerBackend, "cleanup_container",
+                        lambda *args: pytest.fail("retained source removed"))
+    backend.cleanup_container(SimpleNamespace(raw=object()))
+    assert calls == ["retain", "close", "retain"]
+
+
+def test_collection_failure_does_not_release_online_capacity_while_container_remains(
+        tmp_path, monkeypatch):
+    backend = collector(tmp_path)
+    calls = []
+    backend.feedback = SimpleNamespace(close=lambda: calls.append("observer-close"),
+        release_registration=lambda: calls.append("released"))
+    monkeypatch.setattr(backend, "_raw", lambda handle: SimpleNamespace(
+        get_archive=lambda remote: (_ for _ in ()).throw(requests.ConnectionError("synthetic"))))
+    monkeypatch.setattr(transport.DockerBackend, "cleanup_container", lambda *args: calls.append("native-cleanup"))
+    monkeypatch.setattr(backend, "container_exists", lambda name: True)
+    with pytest.raises(RuntimeError, match="capacity was not released"):
+        backend.cleanup_container(SimpleNamespace(name="owned-container"))
+    assert calls == ["observer-close", "native-cleanup"]
+
+
+def test_interrupted_stream_closes_spool_and_still_calls_native_cleanup(tmp_path, monkeypatch):
+    backend = collector(tmp_path)
+    closed, cleaned, spools = [], [], []
+
+    def chunks():
+        try:
+            yield b"partial archive"
+            raise KeyboardInterrupt
+        finally:
+            closed.append(True)
+
+    monkeypatch.setattr(backend, "_raw", lambda handle: SimpleNamespace(
+        get_archive=lambda remote: (chunks(), {})))
+    real = transport.tempfile.TemporaryFile
+
+    def temporary():
+        file = real(dir=tmp_path)
+        spools.append(file)
+        return file
+
+    monkeypatch.setattr(transport.tempfile, "TemporaryFile", temporary)
+    monkeypatch.setattr(transport.DockerBackend, "cleanup_container",
+                        lambda self, handle, logger: cleaned.append(handle))
+    with pytest.raises(KeyboardInterrupt):
+        backend.cleanup_container("owned-container")
+    assert closed == [True] and cleaned == ["owned-container"]
+    assert len(spools) == 1 and spools[0].closed
+
+
+@pytest.mark.parametrize("remains", [False, True])
+def test_native_cleanup_error_keeps_original_capacity_release_rule(tmp_path, monkeypatch, remains):
+    backend = collector(tmp_path)
+    calls = []
+    backend.feedback = SimpleNamespace(close=lambda: calls.append("observer-close"),
+        release_registration=lambda: calls.append("released"))
+    monkeypatch.setattr(backend, "_raw", lambda handle: SimpleNamespace(
+        get_archive=lambda remote: (_ for _ in ()).throw(requests.ConnectionError("synthetic"))))
+
+    def native_cleanup(*args):
+        calls.append("native-cleanup")
+        raise OSError("synthetic native removal failure")
+
+    monkeypatch.setattr(transport.DockerBackend, "cleanup_container", native_cleanup)
+    monkeypatch.setattr(backend, "container_exists", lambda name: remains)
+    with pytest.raises(RuntimeError if remains else OSError):
+        backend.cleanup_container(SimpleNamespace(name="owned-container"))
+    assert calls == ["observer-close", "native-cleanup"] + ([] if remains else ["released"])

benchmark/tests/test_sforge_runtime.py

@@ -441,7 +441,7 @@ def test_artifact_collection_preserves_both_session_homes_without_auth(tmp_path,
     backend.log_dir = tmp_path / "artifacts"
     backend.feedback = None
     copied, cleaned = [], []
-    def archive(handle, remote):
+    def archive(remote):
         copied.append(str(remote))
         payload = io.BytesIO()
         event = json.dumps({"type": "event_msg", "payload": {"type": "task_complete",
@@ -451,8 +451,9 @@ def archive(handle, remote):
             member = tarfile.TarInfo("sessions/fixture/session.jsonl")
             member.size = len(event)
             bundle.addfile(member, io.BytesIO(event))
-        return payload.getvalue()
-    monkeypatch.setattr(backend, "copy_from_container", archive)
+        return iter([payload.getvalue()]), {}
+    from types import SimpleNamespace
+    monkeypatch.setattr(backend, "_raw", lambda handle: SimpleNamespace(get_archive=archive))
     monkeypatch.setattr(DockerBackend, "cleanup_container",
                         lambda self, handle, logger: cleaned.append(handle))
     backend.cleanup_container("fixture-container")
@@ -486,7 +487,9 @@ def test_provider_observation_failure_cannot_skip_native_cleanup(tmp_path, monke
     payload = io.BytesIO()
     with tarfile.open(fileobj=payload, mode="w"):
         pass
-    monkeypatch.setattr(backend, "copy_from_container", lambda *args: payload.getvalue())
+    from types import SimpleNamespace
+    monkeypatch.setattr(backend, "_raw", lambda handle: SimpleNamespace(
+        get_archive=lambda remote: (iter([payload.getvalue()]), {})))
     monkeypatch.setattr("benchmark.runtime.sforge_backend.observe_terminal_provider_failure",
                         lambda *args: (_ for _ in ()).throw(TypeError("unknown event shape")))
     cleaned = []