diff --git a/src/agent_env/artifact/artifacts/docker_image.py b/src/agent_env/artifact/artifacts/docker_image.py index af595cb7..19af2c0e 100644 --- a/src/agent_env/artifact/artifacts/docker_image.py +++ b/src/agent_env/artifact/artifacts/docker_image.py @@ -10,6 +10,7 @@ import shlex import subprocess import tempfile +import threading import uuid from dataclasses import dataclass from pathlib import Path @@ -145,16 +146,7 @@ def put( stderr=subprocess.PIPE, ) with gzip.open(tmp_path, "wb") as gz_file: - while chunk := save_proc.stdout.read(8192): - gz_file.write(chunk) - try: - save_proc.wait(timeout=cls.DOCKER_SAVE_TIMEOUT_SECONDS) - except subprocess.TimeoutExpired: - save_proc.kill() - raise RuntimeError(f"docker save {image_ref} timed out after {cls.DOCKER_SAVE_TIMEOUT_SECONDS} seconds") - if save_proc.returncode != 0: - stderr = save_proc.stderr.read().decode() if save_proc.stderr else "" - raise RuntimeError(f"docker save {image_ref} failed: {stderr}") + _stream_save_output(save_proc, gz_file, image_ref, cls.DOCKER_SAVE_TIMEOUT_SECONDS) tar_gz_object_url = objects.put_file_at( f"{prefix}{fs_safe(id)}-v{version}.tar.gz", str(tmp_path), "application/gzip" @@ -615,6 +607,53 @@ def docker_login(auth) -> None: raise RuntimeError(f"docker login failed: {login.stderr}") +# How much of a failed save's stderr its error message carries. +_SAVE_STDERR_TAIL_BYTES = 65536 + + +def _stream_save_output(save_proc: subprocess.Popen, gz_file, image_ref: str, timeout_seconds: int) -> None: + """Stream a ``docker save`` process's stdout into ``gz_file`` until it exits, bounded by ``timeout_seconds``. + + stderr drains on a thread into a bounded tail, and a watchdog kills the process at the + deadline: reading stdout to EOF first deadlocks once the process fills its stderr pipe, and + a stalled process blocks the read itself, so timing out only the wait after EOF never fires. + The process is reaped on every path. + """ + stderr_tail = bytearray() + + def drain_stderr() -> None: + while chunk := save_proc.stderr.read(8192): + stderr_tail.extend(chunk) + del stderr_tail[:-_SAVE_STDERR_TAIL_BYTES] + + drainer = threading.Thread(target=drain_stderr, daemon=True) + drainer.start() + timed_out = threading.Event() + + def watchdog() -> None: + timed_out.set() + save_proc.kill() + + timer = threading.Timer(timeout_seconds, watchdog) + timer.start() + try: + while chunk := save_proc.stdout.read(8192): + gz_file.write(chunk) + except BaseException: + save_proc.kill() + raise + finally: + # The process is reaped on every path: it either exited, or the watchdog killed it + # (which is also what unblocks a read stuck on a stalled process's stdout). + save_proc.wait() + timer.cancel() + drainer.join(timeout=5) + if timed_out.is_set(): + raise RuntimeError(f"docker save {image_ref} timed out after {timeout_seconds} seconds") + if save_proc.returncode != 0: + raise RuntimeError(f"docker save {image_ref} failed: {stderr_tail.decode(errors='replace')}") + + def _parse_copy_sources(dockerfile_text: str, dockerfile_rel_path: str | None = None) -> list[str]: """Parse COPY/ADD source paths from Dockerfile text. Returns top-level directories to include in context tar.""" paths: list[str] = [] diff --git a/src/agent_env/artifact/artifacts/file.py b/src/agent_env/artifact/artifacts/file.py index 708f7fb3..58b34267 100644 --- a/src/agent_env/artifact/artifacts/file.py +++ b/src/agent_env/artifact/artifacts/file.py @@ -10,12 +10,18 @@ from agent_env.artifact.artifact import Artifact, _write_twin from agent_env.config import get_config +from agent_env.store.base import NotFoundError, ObjectAlreadyExistsError if TYPE_CHECKING: + from agent_env.artifact.store import ArtifactStore from agent_env.bundle.authoring import AuthoringContext from agent_env.entity.export import ExportContext +# How many reserved versions one upload tries before giving up on the version path. +_UPLOAD_VERSION_ATTEMPTS = 3 + + class FileArtifact(Artifact): """A generic file artifact that stores any file in the object store. @@ -69,19 +75,10 @@ def put( if content_type is None: content_type = "application/octet-stream" - version = store.next_version(id) - # put_object_file, not put_object: a hosted store uploads a file in parts, where a # single put is capped (5 GB on S3). It streams from disk, so large files don't get # loaded into RAM either. - stored_url = store.put_object_file( - artifact_type="file", - id=id, - version=version, - object_name=filename, - file_path=file_path, - content_type=content_type, - ) + stored_url = cls._upload(store, id, filename, content_type, file_path=file_path) # The object is keyed by the reserved version; the document takes the next free one, the same unless a # concurrent write took it first. @@ -108,16 +105,7 @@ def put_bytes( from agent_env.artifact.store import get_artifact_store store = get_artifact_store() - version = store.next_version(id) - - stored_url = store.put_object( - artifact_type="file", - id=id, - version=version, - object_name=filename, - data=content, - content_type=content_type, - ) + stored_url = cls._upload(store, id, filename, content_type, content=content) # The object is keyed by the reserved version; the document takes the next free one, the same unless a # concurrent write took it first. @@ -130,6 +118,61 @@ def put_bytes( ) return store.put_document(instance) + @classmethod + def _upload( + cls, + store: "ArtifactStore", + id: str, + filename: str, + content_type: str, + *, + file_path: str | None = None, + content: bytes | None = None, + ) -> str: + """Upload to the reserved version's write-once path, working around what is already there. + + A collision with a published version means a concurrent writer won the reservation: + reserve the next version and retry. A collision with no document behind it is an + orphan an earlier attempt left when its document insert failed; write-once objects + are never deleted or overwritten, so this attempt uploads under a prefix of its own + (``ArtifactStore.attempt_prefix``) instead of staying blocked by the orphan. + """ + version = store.next_version(id) + for _ in range(_UPLOAD_VERSION_ATTEMPTS): + try: + if file_path is not None: + return store.put_object_file( + artifact_type="file", + id=id, + version=version, + object_name=filename, + file_path=file_path, + content_type=content_type, + ) + return store.put_object( + artifact_type="file", + id=id, + version=version, + object_name=filename, + data=content, + content_type=content_type, + ) + except ObjectAlreadyExistsError: + try: + store.get(id, version) + except NotFoundError: + break + version = store.next_version(id) + else: + raise ObjectAlreadyExistsError( + f"every reserved version of artifact {id!r} was published by a concurrent writer" + ) + prefix = store.attempt_prefix("file", id) + objects = get_config().get_object_store_to_write(prefix, id) + if file_path is not None: + return objects.put_file_at(prefix + filename, file_path, content_type) + return objects.put(objects.get_object_key(prefix + filename), content, content_type) + @classmethod def from_toml(cls, data: dict, ctx: AuthoringContext) -> Self: """Write the folder's one file (``put_attempt``); ``description`` defaults to the file's name.""" diff --git a/tst/unit/artifact/test_docker_image_save.py b/tst/unit/artifact/test_docker_image_save.py new file mode 100644 index 00000000..1b70f5c9 --- /dev/null +++ b/tst/unit/artifact/test_docker_image_save.py @@ -0,0 +1,80 @@ +"""DockerImageArtifact.put's ``docker save`` streaming: bounded, drained, and reaped. + +Regression guard: put used to read the save process's stdout to EOF before waiting +on it, with stderr never drained. A save that filled its stderr pipe deadlocked, and +a stalled save blocked the read itself, so the timeout in the wait never fired. +The docker CLI is a PATH stub; the image store and the push are stubbed at the same +boundaries the other docker image tests use. No Docker daemon is involved. +""" + +import gzip +import os +import time +from types import SimpleNamespace + +import pytest + +import agent_env.artifact.artifacts.docker_image as docker_image +from agent_env.artifact.artifacts.docker_image import DockerImageArtifact +from agent_env.config import get_config + +_STUB = """#!/bin/sh +if [ "$1" = "save" ]; then + case "$FAKE_DOCKER_MODE" in + flood) python3 -c 'import sys; sys.stderr.write("x" * 300000); sys.stdout.write("FAKETAR")' ;; + stall) exec sleep 60 ;; + fail) echo "boom: no such image" >&2; exit 3 ;; + *) printf 'FAKETAR' ;; + esac + exit 0 +fi +exit 0 +""" + + +@pytest.fixture +def stubbed_put(local_stores, tmp_path, monkeypatch): + bindir = tmp_path / "bin" + bindir.mkdir() + stub = bindir / "docker" + stub.write_text(_STUB) + stub.chmod(0o755) + monkeypatch.setenv("PATH", f"{bindir}:{os.environ['PATH']}") + monkeypatch.setattr(docker_image, "push_local_image", lambda *a, **k: None) + image_store = SimpleNamespace( + image_ref=lambda repo, tag: f"localhost:5000/{repo}:{tag}", + ensure_repository=lambda repo: None, + ) + monkeypatch.setattr(get_config(), "get_image_store_for", lambda id: image_store) + monkeypatch.setattr(DockerImageArtifact, "DOCKER_SAVE_TIMEOUT_SECONDS", 2) + return monkeypatch + + +def _put(monkeypatch, mode, artifact_id): + monkeypatch.setenv("FAKE_DOCKER_MODE", mode) + return DockerImageArtifact.put(artifact_id, description="d", image_name="img:latest") + + +def test_save_with_saturated_stderr_completes(stubbed_put): + art = _put(stubbed_put, "flood", "img-flood") + + assert gzip.decompress(art.load()) == b"FAKETAR" + + +def test_stalled_save_is_killed_at_the_timeout(stubbed_put): + start = time.monotonic() + with pytest.raises(RuntimeError, match="timed out after 2 seconds"): + _put(stubbed_put, "stall", "img-stall") + assert time.monotonic() - start < 20 + + +def test_failed_save_reports_its_stderr(stubbed_put): + with pytest.raises(RuntimeError, match="boom: no such image"): + _put(stubbed_put, "fail", "img-fail") + + +def test_save_streams_the_archive_to_the_object_store(stubbed_put): + art = _put(stubbed_put, "ok", "img-ok") + + assert art.version == 1 + assert gzip.decompress(art.load()) == b"FAKETAR" diff --git a/tst/unit/artifact/test_file_artifact_attempts.py b/tst/unit/artifact/test_file_artifact_attempts.py new file mode 100644 index 00000000..e4415524 --- /dev/null +++ b/tst/unit/artifact/test_file_artifact_attempts.py @@ -0,0 +1,137 @@ +"""FileArtifact.put / put_bytes work around objects already on the version path. + +Regression guard: both upload to the write-once path of the version they reserve, +before the document is inserted. When the insert then failed, the orphaned object +stayed on that path and every retry collided with it (ObjectAlreadyExistsError); +two writers reserving the same version collided the same way. The upload now +retries past a published version and falls back to a prefix of the attempt's own +past an orphan. Network-free on the local stores. +""" + +import threading + +import pytest + +from agent_env.artifact.artifacts.file import FileArtifact +from agent_env.artifact.store import ArtifactStore +from agent_env.store.object_store.local.store import LocalFilesystemObjectStore + + +@pytest.fixture +def fail_first_insert(monkeypatch): + """Make the first ArtifactStore.put_document call raise, as a failed insert does.""" + real = ArtifactStore.put_document + calls = {"n": 0} + + def flaky(self, artifact): + calls["n"] += 1 + if calls["n"] == 1: + raise RuntimeError("simulated document insert failure") + return real(self, artifact) + + monkeypatch.setattr(ArtifactStore, "put_document", flaky) + return calls + + +def test_put_bytes_retry_after_document_insert_failure(local_stores, fail_first_insert): + with pytest.raises(RuntimeError, match="simulated document insert failure"): + FileArtifact.put_bytes("f1", description="d", filename="a.txt", content=b"hello") + + art = FileArtifact.put_bytes("f1", description="d", filename="a.txt", content=b"hello") + + assert art.version == 1 + assert art.load() == b"hello" + + +def test_put_retry_after_document_insert_failure(local_stores, tmp_path, fail_first_insert): + path = tmp_path / "a.txt" + path.write_bytes(b"from file") + with pytest.raises(RuntimeError, match="simulated document insert failure"): + FileArtifact.put("f2", description="d", file_path=str(path)) + + art = FileArtifact.put("f2", description="d", file_path=str(path)) + + assert art.version == 1 + assert art.load() == b"from file" + + +def _barrier_uploads(monkeypatch, method_name): + """Hold the first two uploads in ``method_name`` until both are in flight at once.""" + real = getattr(LocalFilesystemObjectStore, method_name) + barrier = threading.Barrier(2) + calls = {"n": 0} + lock = threading.Lock() + + def gated(self, *args, **kwargs): + with lock: + calls["n"] += 1 + gated_call = calls["n"] <= 2 + if gated_call: + barrier.wait(timeout=10) + return real(self, *args, **kwargs) + + monkeypatch.setattr(LocalFilesystemObjectStore, method_name, gated) + + +def _put_two_concurrently(put_one): + results, errors = {}, {} + + def run(slot): + try: + results[slot] = put_one(slot) + except Exception as e: # collected and asserted below, not swallowed + errors[slot] = e + + threads = [threading.Thread(target=run, args=(slot,)) for slot in (0, 1)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=30) + assert not errors, errors + return results + + +def test_concurrent_put_bytes_uploads_do_not_collide(local_stores, monkeypatch): + _barrier_uploads(monkeypatch, "put") + + results = _put_two_concurrently( + lambda slot: FileArtifact.put_bytes( + "conc", description="d", filename="a.txt", content=f"content-{slot}".encode() + ) + ) + + assert sorted(art.version for art in results.values()) == [1, 2] + assert len({art.object_url for art in results.values()}) == 2 + for slot, art in results.items(): + assert art.load() == f"content-{slot}".encode() + + +def test_concurrent_put_uploads_do_not_collide(local_stores, tmp_path, monkeypatch): + _barrier_uploads(monkeypatch, "put_file_at") + paths = [] + for slot in (0, 1): + path = tmp_path / f"a{slot}.txt" + path.write_bytes(f"file-{slot}".encode()) + paths.append(path) + + results = _put_two_concurrently( + lambda slot: FileArtifact.put("concf", description="d", file_path=str(paths[slot])) + ) + + assert sorted(art.version for art in results.values()) == [1, 2] + assert len({art.object_url for art in results.values()}) == 2 + for slot, art in results.items(): + assert art.load() == f"file-{slot}".encode() + + +def test_a_second_put_does_not_overwrite_the_first_object(local_stores, tmp_path): + path = tmp_path / "a.txt" + path.write_bytes(b"first") + first = FileArtifact.put("seq", description="d", file_path=str(path)) + path.write_bytes(b"second") + second = FileArtifact.put("seq", description="d", file_path=str(path)) + + assert (first.version, second.version) == (1, 2) + assert first.object_url != second.object_url + assert first.load() == b"first" + assert second.load() == b"second"