Fix late Docker creation cleanup and preserve query artifacts

deploy/benchmark_docker/README.md

@@ -109,6 +109,8 @@ are reported in `CaseResult.error` and make `success` false for cases that used
 that worker; failures for an unused worker apply to the whole run. Cancellation
 waits for owned-container cleanup before propagating. A failed `close()` retains
 unfinished Docker handles so another call can retry removal.
+If container creation outlasts the cleanup timeout, a tracked cleanup task
+remains in the running host process and removes the container once creation settles.
 The build log also records configuration, construction, build job, checkpoint,
 and Application lifecycle failures, including available tracebacks.
 

deploy/benchmark_docker/direct_docker_workspace.py

@@ -8,6 +8,7 @@ would add an unrelated network-dependent image layer and process.
 from __future__ import annotations
 
 import asyncio
+import logging
 import os
 import sys
 from typing import Any
@@ -16,6 +17,7 @@ from aiodocker.exceptions import DockerError
 from agentscope.workspace import DockerWorkspace
 
 DOCKER_CLEANUP_TIMEOUT = 30.0
+_PENDING_CLEANUPS: set[asyncio.Task[None]] = set()
 
 
 class DirectDockerWorkspace(DockerWorkspace):
@@ -31,11 +33,22 @@ class DirectDockerWorkspace(DockerWorkspace):
         super().__init__(**kwargs)
         self._lifecycle_lock = asyncio.Lock()
         self._stopped = False
+        # Retain the task until its creation result has been reconciled.
+        self._creation_task: asyncio.Task[None] | None = None
+        self._deferred_cleanup_task: asyncio.Task[None] | None = None
 
     async def _build_or_reuse_image(self) -> None:
         """Select the caller-provided image instead of deriving a gateway image."""
         self._image_tag = self.base_image
 
+    async def _create_and_start_container(self) -> None:
+        """Keep creation owned until its result can be reconciled on cleanup."""
+        self._creation_task = asyncio.create_task(super()._create_and_start_container())
+        # Let the request settle even when initialize is cancelled. Otherwise a
+        # name lookup could return 404 before Docker finishes creating it.
+        await asyncio.shield(self._creation_task)
+        self._creation_task = None
+
     async def initialize(self) -> None:
         """Start one Docker container and expose its filesystem/exec backend."""
         async with self._lifecycle_lock:
@@ -81,6 +94,31 @@ class DirectDockerWorkspace(DockerWorkspace):
     async def _teardown_backend(self) -> None:
         """Remove the container, preserving handles until cleanup is confirmed."""
         self.is_alive = False
+        if self._creation_task is not None:
+            done, _ = await asyncio.wait({self._creation_task}, timeout=DOCKER_CLEANUP_TIMEOUT)
+            if not done:
+                if self._deferred_cleanup_task is None or self._deferred_cleanup_task.done():
+                    self._deferred_cleanup_task = asyncio.create_task(
+                        self._finish_creation_cleanup(self._creation_task),
+                    )
+                    # Keep cleanup owned even if factory creation returns no workspace.
+                    _PENDING_CLEANUPS.add(self._deferred_cleanup_task)
+                    self._deferred_cleanup_task.add_done_callback(_PENDING_CLEANUPS.discard)
+                raise TimeoutError("Docker container creation has not settled; deferred cleanup is pending")
+            # initialize already reports the provisioning failure. A cancelled
+            # initialize may instead leave a failed creation to reconcile here.
+            if not self._creation_task.cancelled():
+                self._creation_task.exception()
+            if self._container is None:
+                try:
+                    self._container = await asyncio.wait_for(
+                        self._client.containers.get(f"as_ws_{self.workspace_id}"),
+                        timeout=DOCKER_CLEANUP_TIMEOUT,
+                    )
+                except DockerError as exc:
+                    if exc.status != 404:
+                        raise
+            self._creation_task = None
         if self._container is not None:
             # Preserve AgentScope's bind-mount ownership repair on native Linux.
             if (
@@ -110,6 +148,18 @@ class DirectDockerWorkspace(DockerWorkspace):
             self._client = None
         self._backend = None
 
+    async def _finish_creation_cleanup(self, creation: asyncio.Task[None]) -> None:
+        """Remove a late creation result after bounded foreground cleanup expires."""
+        await asyncio.wait({creation})
+        async with self._lifecycle_lock:
+            # A caller may have already closed and even reinitialized this workspace.
+            if self._creation_task is not creation:
+                return
+            try:
+                await self._teardown_backend()
+            except Exception:
+                logging.getLogger(__name__).exception("Deferred Docker cleanup failed for %s", self.workspace_id)
+
     async def close(self) -> None:
         """Report cleanup errors and allow another call to retry unfinished work."""
         # The inherited close clears backend state in a finally block. This

deploy/benchmark_docker/docker_sandbox.py

@@ -665,10 +665,10 @@ class DockerReMeSandbox:
 
     async def run_queries(self, queries: list[EvaluationQuery | dict[str, Any]]) -> dict[str, Any]:
         """Answer multiple queries in one candidate Application with isolated logs."""
-        serialized, query_ids = self._serialize_queries(queries)
+        serialized = self._serialize_queries(queries)
         if not serialized:
             raise ValueError("run_queries requires at least one query")
-        if len(query_ids) != len(set(query_ids)):
+        if len(serialized) != len({query["query_id"] for query in serialized}):
             raise ValueError("query IDs must be unique")
         async with self._operation_lock:
             return await self._run_worker_request("queries", {"queries": serialized})
@@ -676,7 +676,7 @@ class DockerReMeSandbox:
     async def run_query(self, query: EvaluationQuery | dict[str, Any]) -> dict[str, Any]:
         """Answer one append-only query for a host-side lease."""
 
-        serialized, _ = self._serialize_queries([query])
+        serialized = self._serialize_queries([query])
         async with self._operation_lock:
             response = await self._run_worker_request("query", {"queries": serialized})
         results = response.get("queries")
@@ -687,11 +687,10 @@ class DockerReMeSandbox:
     def _serialize_queries(
         self,
         queries: list[EvaluationQuery | dict[str, Any]],
-    ) -> tuple[list[dict[str, Any]], list[str]]:
+    ) -> list[dict[str, Any]]:
         """Validate and serialize query payloads shared by batch and lease APIs."""
 
         serialized: list[dict[str, Any]] = []
-        query_ids: list[str] = []
         for query in queries:
             value = (
                 {
@@ -716,7 +715,6 @@ class DockerReMeSandbox:
             if leaked:
                 names = ", ".join(sorted(leaked))
                 raise ValueError(f"query {query_id!r} answer_arguments contain private judge fields: {names}")
-            query_ids.append(query_id)
             # This allowlist is a security boundary.  Golden answers, rubrics,
             # question types and scoring configuration remain on the host and
             # are not written to the candidate container's inbox.
@@ -728,7 +726,7 @@ class DockerReMeSandbox:
                     "answer_job": str(value.get("answer_job") or "agentic_answer"),
                 },
             )
-        return serialized, query_ids
+        return serialized
 
     async def ingest_session(
         self,
@@ -834,46 +832,58 @@ class DockerReMeSandbox:
         except FileNotFoundError as exc:
             raise SandboxCommandError(f"{description} is missing required artifact: {path}") from exc
 
+    async def _artifact_directory_exists(self, path: str) -> bool:
+        """Check for artifacts independently of completion markers."""
+        result = await self._exec("check_artifact_directory", ["test", "-d", path])
+        if result.exit_code not in (0, 1):
+            raise SandboxCommandError(f"cannot check artifact directory {path!r}: {_decode(result.stderr)}")
+        return result.exit_code == 0
+
+    async def _require_artifact_directory(self, path: str, description: str) -> None:
+        """Require an artifact directory while allowing incomplete diagnostics."""
+        if not await self._artifact_directory_exists(path):
+            raise SandboxCommandError(f"{description} is missing required artifact: {path}")
+
     async def export_memory_construction(self, destination: str | Path) -> Path:
         """Download the workspace and build log; reusable with upload_workspace."""
         async with self._operation_lock:
+            await self._require_artifact(f"{CASE_ROOT}/build_log/build.log", "memory construction export")
             return await self._export_archive(
                 destination,
                 "export_memory_construction",
                 CASE_ROOT,
                 ["reme_workspace", "build_log"],
-                required=("build_log/build.log", "memory construction export"),
             )
 
     async def export_queries(self, destination: str | Path) -> Path:
-        """Download per-query logs, results, and their aggregate summary."""
+        """Download per-query logs, results, and any aggregate summary."""
         async with self._operation_lock:
+            await self._require_artifact_directory(f"{CASE_ROOT}/queries", "query export")
             return await self._export_archive(
                 destination,
                 "export_queries",
                 CASE_ROOT,
                 ["queries"],
-                required=("queries/summary.json", "query export"),
             )
 
     async def export_query(self, query_id: str, destination: str | Path) -> Path:
-        """Download one leased query's logs and result without a case summary."""
+        """Download one leased query's available logs and result without a case summary."""
 
         self._validate_query_artifact_id(query_id)
         async with self._operation_lock:
+            await self._require_artifact_directory(f"{CASE_ROOT}/queries/{query_id}", "query export")
             return await self._export_archive(
                 destination,
                 "export_query",
                 CASE_ROOT,
                 [f"queries/{query_id}"],
-                required=(f"queries/{query_id}/result.json", "query export"),
             )
 
     async def export_evaluation(self, destination: str | Path) -> Path:
         """Download reusable workspace and query artifacts, plus any build log."""
         async with self._operation_lock:
             assert self.backend is not None, "sandbox is not initialized"
-            await self._require_artifact(f"{CASE_ROOT}/queries/summary.json", "evaluation export")
+            await self._require_artifact_directory(f"{CASE_ROOT}/queries", "evaluation export")
             members = ["reme_workspace", "queries"]
             try:
                 await self.backend.read_file(f"{CASE_ROOT}/build_log/build.log")
@@ -898,20 +908,28 @@ class DockerReMeSandbox:
         members: list[str],
         *,
         archive_path: str = EXPORT_ARCHIVE,
-        required: tuple[str, str] | None = None,
     ) -> Path:
         """Archive selected sandbox paths and download the archive while locked."""
         assert self.backend is not None, "sandbox is not initialized"
-        if required is not None:
-            path, description = required
-            await self._require_artifact(f"{root}/{path}", description)
         if root == CASE_ROOT:
             await self._write_export_manifest(
                 name.removeprefix("export_"),
                 workspace_root="reme_workspace" if "reme_workspace" in members else None,
             )
             members = ["manifest.json", *members]
-        await self._exec_checked(name, ["tar", "-czf", archive_path, "-C", root, *members])
+        return await self._pack_and_download(destination, name, ["-C", root, *members], archive_path=archive_path)
+
+    async def _pack_and_download(
+        self,
+        destination: str | Path,
+        name: str,
+        arguments: list[str],
+        *,
+        archive_path: str = EXPORT_ARCHIVE,
+    ) -> Path:
+        """Pack selected paths, download the archive, and save it on the host."""
+        assert self.backend is not None, "sandbox is not initialized"
+        await self._exec_checked(name, ["tar", "-czf", archive_path, *arguments])
         return self._save_host_file(destination, await self.backend.read_file(archive_path))
 
     async def upload_workspace(self, source: str | Path, *, clear: bool = True) -> None:
@@ -1043,10 +1061,7 @@ class DockerReMeSandbox:
         )
         if profile == "analysis":
             workspace_root, analysis_excludes = await self._analysis_archive_spec()
-            command = [
-                "tar",
-                "-czf",
-                EXPORT_ARCHIVE,
+            arguments = [
                 "-C",
                 CASE_ROOT,
                 *(f"--exclude={path}" for path in analysis_excludes),
@@ -1057,11 +1072,10 @@ class DockerReMeSandbox:
                 workspace_root,
             ]
         else:
-            command = ["tar", "-czf", EXPORT_ARCHIVE, "-C", CASE_ROOT, "."]
+            arguments = ["-C", CASE_ROOT, "."]
         if include_candidate:
-            command.extend(["-C", "/workspace", "candidate"])
-        await self._exec_checked("export_case", command)
-        self._save_host_file(target, await self.backend.read_file(EXPORT_ARCHIVE))
+            arguments.extend(["-C", "/workspace", "candidate"])
+        await self._pack_and_download(target, "export_case", arguments)
         if profile == "analysis":
             await self._export_analysis_logs(artifacts)
         return target
@@ -1078,9 +1092,7 @@ class DockerReMeSandbox:
             log_path.parent.mkdir(parents=True, exist_ok=True)
             log_path.write_bytes(build_log)
 
-        try:
-            await self.backend.read_file(f"{CASE_ROOT}/queries/summary.json")
-        except FileNotFoundError:
+        if not await self._artifact_directory_exists(f"{CASE_ROOT}/queries"):
             return
         await self._exec_checked(
             "export_analysis_queries",

deploy/benchmark_docker/pool.py

@@ -53,6 +53,18 @@ class CaseResult:
         )
 
 
+@dataclass(slots=True)
+class _WorkerSlots:
+    """Mutable worker availability owned by one pool run."""
+
+    workers: list[DockerReMeSandbox]
+    available: asyncio.Queue[DockerReMeSandbox | None]
+    unclaimed: set[str]
+    live_slots: int
+    next_worker_index: int
+    exhaustion_error: str = "no available workers"
+
+
 class BenchmarkDockerPool:
     """Reuse at most ``concurrency`` containers across case plans.
 
@@ -95,10 +107,7 @@ class BenchmarkDockerPool:
         available: asyncio.Queue[DockerReMeSandbox | None] = asyncio.Queue()
         for worker in workers:
             available.put_nowait(worker)
-        unclaimed = set(by_id)
-        live_slots = len(workers)
-        next_worker_index = len(workers)
-        exhaustion_error = "no available workers"
+        slots = _WorkerSlots(workers, available, set(by_id), len(workers), len(workers))
         worker_results: dict[int, list[CaseResult]] = {}
         results: list[CaseResult] = []
         run_error: BaseException | None = None
@@ -111,27 +120,21 @@ class BenchmarkDockerPool:
                 tasks: dict[str, asyncio.Task[CaseResult]] = {}
 
                 async def execute(plan: CasePlan) -> CaseResult:
-                    nonlocal live_slots, next_worker_index, exhaustion_error
                     result = CaseResult(plan.case_id)
                     source = plan.workspace_source
                     if plan.reuse_workspace_from is not None:
                         parent = await tasks[plan.reuse_workspace_from]
                         if parent.workspace_snapshot is None:
-                            unclaimed.discard(plan.case_id)
+                            slots.unclaimed.discard(plan.case_id)
                             return CaseResult(plan.case_id, error=f"source case {parent.case_id!r} failed")
                         source = parent.workspace_snapshot
 
                     worker = await available.get()
-                    unclaimed.discard(plan.case_id)
+                    slots.unclaimed.discard(plan.case_id)
                     if worker is None:
-                        return CaseResult(plan.case_id, error=exhaustion_error)
+                        return CaseResult(plan.case_id, error=slots.exhaustion_error)
                     worker_results.setdefault(id(worker), []).append(result)
 
-                    def record_lifecycle_error(phase: str, exc: Exception) -> str:
-                        detail = f"{phase} failed: {format_error(exc)}"
-                        result.error = f"{result.error}; {detail}" if result.error else detail
-                        return detail
-
                     try:
                         await worker.reset_case(plan.case_id)
                         if source is not None:
@@ -165,42 +168,7 @@ class BenchmarkDockerPool:
                         result.error = format_error(exc)
                         return result
                     finally:
-                        if worker.reusable:
-                            available.put_nowait(worker)
-                        else:
-                            # Confirm cleanup (or retry a failed removal) before
-                            # allowing another case into this slot.
-                            closed = False
-                            try:
-                                await worker.close()
-                            except Exception as exc:
-                                exhaustion_error = (
-                                    f"no available workers: {record_lifecycle_error('worker close', exc)}"
-                                )
-                            else:
-                                closed = True
-
-                            replaced = False
-                            if closed and unclaimed and not asyncio.current_task().cancelling():
-                                worker_index = next_worker_index
-                                next_worker_index += 1
-                                try:
-                                    replacement = await self.factory.create_case(f"pool-{worker_index}")
-                                except Exception as exc:
-                                    exhaustion_error = (
-                                        f"no available workers: {record_lifecycle_error('replacement', exc)}"
-                                    )
-                                else:
-                                    workers.append(replacement)
-                                    available.put_nowait(replacement)
-                                    replaced = True
-                            if not replaced:
-                                live_slots -= 1
-                                if live_slots == 0:
-                                    # Every future case, including one waiting on a
-                                    # parent, must be able to leave available.get().
-                                    for _ in unclaimed:
-                                        available.put_nowait(None)
+                        await self._release_worker(worker, result, slots)
 
                 for plan in plans:
                     tasks[plan.case_id] = asyncio.create_task(execute(plan))
@@ -221,6 +189,41 @@ class BenchmarkDockerPool:
         finally:
             await self._close_workers(workers, worker_results, results, run_error)
 
+    async def _release_worker(self, worker: DockerReMeSandbox, result: CaseResult, slots: _WorkerSlots) -> None:
+        """Return a reusable worker or retire its slot after confirmed cleanup."""
+        if worker.reusable:
+            slots.available.put_nowait(worker)
+            return
+
+        def record_lifecycle_error(phase: str, exc: Exception) -> None:
+            detail = f"{phase} failed: {format_error(exc)}"
+            result.error = f"{result.error}; {detail}" if result.error else detail
+            slots.exhaustion_error = f"no available workers: {detail}"
+
+        # Do not reuse the slot while its old container may still be running.
+        try:
+            await worker.close()
+        except Exception as exc:
+            record_lifecycle_error("worker close", exc)
+        else:
+            if slots.unclaimed and not asyncio.current_task().cancelling():
+                worker_index = slots.next_worker_index
+                slots.next_worker_index += 1
+                try:
+                    replacement = await self.factory.create_case(f"pool-{worker_index}")
+                except Exception as exc:
+                    record_lifecycle_error("replacement", exc)
+                else:
+                    slots.workers.append(replacement)
+                    slots.available.put_nowait(replacement)
+                    return
+
+        slots.live_slots -= 1
+        if slots.live_slots == 0:
+            # Also wake cases that are still waiting for a parent to complete.
+            for _ in slots.unclaimed:
+                slots.available.put_nowait(None)
+
     @staticmethod
     async def _retain_failed_build_log(worker: DockerReMeSandbox, plan: CasePlan, result: CaseResult) -> None:
         """Keep diagnostics without replacing the original build failure."""

tests/unit/test_benchmark_docker.py

@@ -768,6 +768,8 @@ class FakeBackend:
         self.commands.append(command)
         output = b""
         exit_code = 0
+        if command[:2] == ["test", "-d"]:
+            exit_code = 0 if any(path.startswith(command[2] + "/") for path in self.files) else 1
         if len(command) > 2 and command[1] == "-c":
             output = json.dumps(
                 {
@@ -1148,11 +1150,11 @@ async def test_pool_replaces_worker_after_backend_timeout(tmp_path):
             "memory construction export",
             ["manifest.json", "reme_workspace", "build_log"],
         ),
-        ("export_queries", "queries/summary.json", "query export", ["manifest.json", "queries"]),
-        ("export_query", "queries/q1/result.json", "query export", ["manifest.json", "queries/q1"]),
+        ("export_queries", "queries", "query export", ["manifest.json", "queries"]),
+        ("export_query", "queries/q1", "query export", ["manifest.json", "queries/q1"]),
         (
             "export_evaluation",
-            "queries/summary.json",
+            "queries",
             "evaluation export",
             ["manifest.json", "reme_workspace", "queries"],
         ),
@@ -1168,7 +1170,8 @@ async def test_archive_exports_check_required_artifacts(tmp_path, method, requir
         with pytest.raises(SandboxCommandError) as failure:
             await export(*args)
         assert str(failure.value) == f"{description} is missing required artifact: /workspace/case/{required}"
-        backend.files[f"/workspace/case/{required}"] = b"{}"
+        artifact = "queries/q1/answer.log" if required.startswith("queries") else required
+        backend.files[f"/workspace/case/{artifact}"] = b"{}"
         assert await export(*args) == destination
         assert ["tar", "-czf", "/tmp/reme-case-export.tar.gz", "-C", "/workspace/case", *members] in backend.commands
     finally:
@@ -1264,17 +1267,18 @@ def test_workspace_import_checks_case_archive_format(tmp_path, manifest, workspa
 
 
 @pytest.mark.asyncio
-async def test_default_analysis_export_saves_logs_outside_archive(tmp_path):
+@pytest.mark.parametrize("query_state", ["batch", "single", "partial", "absent"])
+async def test_default_analysis_export_saves_logs_outside_archive(tmp_path, query_state):
     class AnalysisBackend(FakeBackend):
         async def exec_shell(self, command, *, cwd=None, timeout=None):
             result = await super().exec_shell(command, cwd=cwd, timeout=timeout)
             if command[:2] == ["tar", "-czf"] and command[-1] == "queries":
                 stream = io.BytesIO()
                 with tarfile.open(fileobj=stream, mode="w:gz") as archive:
-                    for name, content in (
-                        ("queries/summary.json", b"{}"),
-                        ("queries/q1/answer.log", b"answer log"),
-                    ):
+                    for path, content in self.files.items():
+                        if not path.startswith("/workspace/case/queries/"):
+                            continue
+                        name = path.removeprefix("/workspace/case/")
                         info = tarfile.TarInfo(name)
                         info.size = len(content)
                         archive.addfile(info, io.BytesIO(content))
@@ -1290,14 +1294,32 @@ async def test_default_analysis_export_saves_logs_outside_archive(tmp_path):
             b'{"workspace_root":"reme_workspace","analysis_excludes":[]}'
         )
         backend.files["/workspace/case/build_log/build.log"] = b"build log"
-        backend.files["/workspace/case/queries/summary.json"] = b"{}"
+        if query_state == "single":
+            backend.worker_responses["query"] = {"success": True, "queries": [{"answer": "ok"}]}
+            assert await case.run_query({"query_id": "q1", "question": "Question?"}) == {"answer": "ok"}
+        if query_state != "absent":
+            backend.files["/workspace/case/queries/q1/answer.log"] = b"answer log"
+        if query_state in {"batch", "single"}:
+            backend.files["/workspace/case/queries/q1/result.json"] = b'{"answer":"ok"}'
+        if query_state == "batch":
+            backend.files["/workspace/case/queries/summary.json"] = b"{}"
         archive = tmp_path / "analysis.tar.gz"
         assert await case.export(archive) == archive
         assert archive.read_bytes() == b"workspace archive"
         artifacts = tmp_path / "analysis.tar.gz.artifacts"
         assert (artifacts / "build_log" / "build.log").read_bytes() == b"build log"
-        assert (artifacts / "queries" / "summary.json").read_bytes() == b"{}"
-        assert (artifacts / "queries" / "q1" / "answer.log").read_bytes() == b"answer log"
+        query_files = {
+            path.removeprefix("/workspace/case/"): content
+            for path, content in backend.files.items()
+            if path.startswith("/workspace/case/queries/")
+        }
+        assert {
+            str(path.relative_to(artifacts)): path.read_bytes()
+            for path in (artifacts / "queries").rglob("*")
+            if path.is_file()
+        } == query_files
+        query_commands = [command for command in backend.commands if command[:2] == ["tar", "-czf"]]
+        assert len(query_commands) == (1 if query_state == "absent" else 2)
         archive_command = next(
             command for command in backend.commands if command[:3] == ["tar", "-czf", "/tmp/reme-case-export.tar.gz"]
         )

tests/unit/test_benchmark_docker_queries.py

@@ -13,8 +13,7 @@ from deploy.benchmark_docker.docker_sandbox import DockerReMeSandbox
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("method", ["run_query", "run_queries"])
-@pytest.mark.parametrize("as_mapping", [False, True])
+@pytest.mark.parametrize(("method", "as_mapping"), [("run_query", False), ("run_queries", True)])
 async def test_query_upload_ignores_uncopyable_golden_answer(monkeypatch, method, as_mapping):
     case = DockerReMeSandbox(
         case_id="query",

tests/unit/test_benchmark_docker_workspace.py

@@ -10,6 +10,7 @@ import pytest
 
 DockerError = pytest.importorskip("aiodocker.exceptions").DockerError
 from deploy.benchmark_docker.direct_docker_workspace import DirectDockerWorkspace  # noqa: E402
+from deploy.benchmark_docker import DockerReMeSandboxFactory, ImageCandidate  # noqa: E402
 
 
 def make_workspace():
@@ -21,6 +22,139 @@ def make_workspace():
     return workspace
 
 
+def make_creating_workspace(monkeypatch):
+    """Exercise AgentScope's real provisioning with only the Docker API mocked."""
+    workspace = DirectDockerWorkspace(workspace_id="creation-test")
+    container = SimpleNamespace(start=AsyncMock(), delete=AsyncMock())
+    client = SimpleNamespace(
+        containers=SimpleNamespace(
+            create_or_replace=AsyncMock(return_value=container),
+            get=AsyncMock(return_value=container),
+        ),
+        close=AsyncMock(),
+    )
+    monkeypatch.setattr("aiodocker.Docker", lambda: client)
+    return workspace, container, client
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("cleanup_times_out", [False, True])
+async def test_cancelled_creation_is_cleaned_up_after_response(monkeypatch, cleanup_times_out):
+    from deploy.benchmark_docker import direct_docker_workspace
+
+    monkeypatch.setattr(direct_docker_workspace, "DOCKER_CLEANUP_TIMEOUT", 0.01)
+    workspace, container, client = make_creating_workspace(monkeypatch)
+    created, respond = asyncio.Event(), asyncio.Event()
+
+    async def create(**_kwargs):
+        created.set()  # Docker has created the container, but not returned it.
+        await respond.wait()
+        return container
+
+    client.containers.create_or_replace.side_effect = create
+    task = asyncio.create_task(workspace.initialize())
+    await asyncio.wait_for(created.wait(), timeout=2)
+    task.cancel()
+    await asyncio.sleep(0)
+    assert not workspace._creation_task.cancelled()
+    client.containers.get.assert_not_awaited()
+    client.close.assert_not_awaited()
+    if cleanup_times_out:
+        with pytest.raises(asyncio.CancelledError) as error:
+            await task
+        assert any("creation has not settled" in note for note in error.value.__notes__)
+        assert not workspace._creation_task.done()
+        assert workspace._client is client
+        client.close.assert_not_awaited()
+
+    respond.set()
+    if not cleanup_times_out:
+        with pytest.raises(asyncio.CancelledError):
+            await task
+    await workspace.close()
+    container.delete.assert_awaited_once_with(force=True)
+    client.close.assert_awaited_once()
+    assert workspace._creation_task is workspace._container is workspace._client is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("response_lost", [False, True])
+async def test_factory_cleans_up_creation_after_both_close_attempts_time_out(monkeypatch, response_lost):
+    from deploy.benchmark_docker import direct_docker_workspace
+
+    monkeypatch.setattr(direct_docker_workspace, "DOCKER_CLEANUP_TIMEOUT", 0.01)
+    workspace, container, client = make_creating_workspace(monkeypatch)
+    created, respond = asyncio.Event(), asyncio.Event()
+
+    async def create(**_kwargs):
+        created.set()
+        await respond.wait()
+        if response_lost:
+            raise ConnectionError("creation response lost")
+        return container
+
+    client.containers.create_or_replace.side_effect = create
+    factory = DockerReMeSandboxFactory(ImageCandidate("test"), workspace_builder=lambda *_args: workspace)
+    task = asyncio.create_task(factory.create_case("late-creation"))
+    await asyncio.wait_for(created.wait(), timeout=2)
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError) as failure:
+        await task
+    assert sum("creation has not settled" in note for note in failure.value.__notes__) == 2
+    cleanup = workspace._deferred_cleanup_task
+    assert cleanup in direct_docker_workspace._PENDING_CLEANUPS
+    container.delete.assert_not_awaited()
+
+    respond.set()
+    await asyncio.wait_for(asyncio.shield(cleanup), timeout=2)
+    container.delete.assert_awaited_once_with(force=True)
+    client.close.assert_awaited_once()
+    assert cleanup not in direct_docker_workspace._PENDING_CLEANUPS
+    assert workspace._creation_task is workspace._container is workspace._client is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("failure", "missing"),
+    [(ConnectionError(), False), (asyncio.CancelledError(), False), (ConnectionError(), True)],
+)
+async def test_creation_without_handle_is_reconciled_by_name(monkeypatch, failure, missing):
+    workspace, container, client = make_creating_workspace(monkeypatch)
+    client.containers.create_or_replace.side_effect = failure
+    if missing:
+        client.containers.get.side_effect = DockerError(404, {"message": "not found"})
+
+    with pytest.raises(type(failure)):
+        await workspace.initialize()
+    client.containers.get.assert_awaited_once_with("as_ws_creation-test")
+    if missing:
+        container.delete.assert_not_awaited()
+    else:
+        container.delete.assert_awaited_once_with(force=True)
+    client.close.assert_awaited_once()
+    assert workspace._client is None
+
+
+@pytest.mark.asyncio
+async def test_uncertain_creation_lookup_failure_retains_state_for_retry(monkeypatch):
+    workspace, container, client = make_creating_workspace(monkeypatch)
+    client.containers.create_or_replace.side_effect = ConnectionError("response lost")
+    client.containers.get.side_effect = [DockerError(500, {"message": "lookup failed"}), container]
+
+    with pytest.raises(ConnectionError, match="response lost") as error:
+        await workspace.initialize()
+    assert any("workspace close failed" in note for note in error.value.__notes__)
+    assert workspace._creation_task is not None
+    assert workspace._client is client
+    client.close.assert_not_awaited()
+    with pytest.raises(RuntimeError, match="cleanup must complete"):
+        await workspace.initialize()
+
+    await workspace.close()
+    container.delete.assert_awaited_once_with(force=True)
+    client.close.assert_awaited_once()
+
+
 @pytest.mark.asyncio
 async def test_stop_kills_and_confirms_exit_before_removal():
     workspace = make_workspace()
@@ -58,44 +192,28 @@ async def test_docker_cleanup_timeout_retains_handles(monkeypatch, operation):
 
 
 @pytest.mark.asyncio
-async def test_delete_failure_retains_handles_and_can_be_retried():
+@pytest.mark.parametrize("status", [404, 500])
+async def test_delete_retains_handles_until_removal_is_confirmed(status):
     workspace = make_workspace()
     container, client, backend = workspace._container, workspace._client, workspace._backend
-    container.delete.side_effect = [RuntimeError("delete failed"), None]
-
-    with pytest.raises(RuntimeError, match="delete failed"):
-        await workspace.close()
-    assert workspace._container is container
-    assert workspace._client is client
-    assert workspace._backend is backend
-    assert not workspace.is_alive
-    client.close.assert_not_awaited()
-    with pytest.raises(RuntimeError, match="cleanup must complete"):
-        await workspace.initialize()
-
-    await workspace.close()
-    await workspace.close()
-    assert container.delete.await_count == 2
-    container.delete.assert_awaited_with(force=True)
-    client.close.assert_awaited_once()
-    assert workspace._container is workspace._client is workspace._backend is None
-
-
-@pytest.mark.asyncio
-@pytest.mark.parametrize("status", [404, 500])
-async def test_only_missing_container_is_a_successful_delete(status):
-    workspace = make_workspace()
-    container, client = workspace._container, workspace._client
-    container.delete.side_effect = DockerError(status, {"message": "delete failed"})
-    if status == 404:
-        await workspace.close()
-        assert workspace._container is None
-        client.close.assert_awaited_once()
-    else:
+    container.delete.side_effect = [DockerError(status, {"message": "delete failed"}), None]
+    if status == 500:
         with pytest.raises(DockerError):
             await workspace.close()
         assert workspace._container is container
+        assert workspace._client is client
+        assert workspace._backend is backend
+        assert not workspace.is_alive
         client.close.assert_not_awaited()
+        with pytest.raises(RuntimeError, match="cleanup must complete"):
+            await workspace.initialize()
+
+    await workspace.close()
+    await workspace.close()
+    assert container.delete.await_count == (2 if status == 500 else 1)
+    container.delete.assert_awaited_with(force=True)
+    client.close.assert_awaited_once()
+    assert workspace._container is workspace._client is workspace._backend is None
 
 
 @pytest.mark.asyncio