Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 21 additions & 8 deletions tests/tools/test_docker_environment.py
Original file line number Diff line number Diff line change
Expand Up @@ -207,21 +207,34 @@ def test_non_persistent_cleanup_removes_container(monkeypatch):
monkeypatch.setattr(docker_env, "find_docker", lambda: "/usr/bin/docker")
calls = _mock_subprocess_run(monkeypatch)

popen_cmds = []
monkeypatch.setattr(
docker_env.subprocess, "Popen",
lambda cmd, **kw: (popen_cmds.append(cmd), type("P", (), {"poll": lambda s: 0, "wait": lambda s, **k: None, "returncode": 0, "stdout": iter([]), "stdin": None})())[1],
)
class ImmediateThread:
def __init__(self, target, daemon=False):
self.target = target
self.daemon = daemon

def start(self):
self.target()

monkeypatch.setattr(docker_env, "threading", types.SimpleNamespace(Thread=ImmediateThread))
monkeypatch.setattr(docker_env, "time", types.SimpleNamespace(sleep=lambda _seconds: None))

env = _make_dummy_env(persistent_filesystem=False, task_id="ephemeral-task")
assert env._container_id
container_id = env._container_id
env._docker_exe = r"C:\Program Files\Docker\Docker\resources\bin\docker.exe"

env.cleanup()

# Should have stop and rm calls via Popen
stop_cmds = [c for c in popen_cmds if container_id in str(c) and "stop" in str(c)]
assert len(stop_cmds) >= 1, f"cleanup() should schedule docker stop for {container_id}"
stop_cmds = [
cmd for cmd, _kwargs in calls
if cmd == [env._docker_exe, "stop", container_id]
]
rm_cmds = [
cmd for cmd, _kwargs in calls
if cmd == [env._docker_exe, "rm", "-f", container_id]
]
assert len(stop_cmds) == 1, f"cleanup() should stop {container_id} with argv"
assert len(rm_cmds) == 1, f"cleanup() should remove {container_id} with argv"


class _FakePopen:
Expand Down
87 changes: 75 additions & 12 deletions tools/environments/docker.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@
import shutil
import subprocess
import sys
import threading
import time
import uuid
from typing import Optional

Expand All @@ -33,6 +35,46 @@
_ENV_VAR_NAME_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")


def _run_docker_cleanup_cmd(
cmd: list[str],
*,
timeout: int,
action: str,
container_id: str,
) -> bool:
"""Run a Docker cleanup command without invoking a shell."""
try:
result = subprocess.run(
cmd,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
timeout=timeout,
)
if result.returncode == 0:
return True
logger.debug(
"Docker cleanup %s failed for container %s with exit code %s",
action,
container_id,
result.returncode,
)
except subprocess.TimeoutExpired:
logger.debug(
"Docker cleanup %s timed out for container %s after %ss",
action,
container_id,
timeout,
)
except Exception as e:
logger.debug(
"Docker cleanup %s failed for container %s: %s",
action,
container_id,
e,
)
return False


def _normalize_forward_env_names(forward_env: list[str] | None) -> list[str]:
"""Return a deduplicated list of valid environment variable names."""
normalized: list[str] = []
Expand Down Expand Up @@ -629,23 +671,44 @@ def _storage_opt_supported() -> bool:
def cleanup(self):
"""Stop and remove the container. Bind-mount dirs persist if persistent=True."""
if self._container_id:
try:
# Stop in background so cleanup doesn't block
stop_cmd = (
f"(timeout 60 {self._docker_exe} stop {self._container_id} || "
f"{self._docker_exe} rm -f {self._container_id}) >/dev/null 2>&1 &"
container_id = self._container_id
docker_exe = self._docker_exe
remove_after_stop = not self._persistent

def _stop_container():
stopped = _run_docker_cleanup_cmd(
[docker_exe, "stop", container_id],
timeout=60,
action="stop",
container_id=container_id,
)
if not stopped:
_run_docker_cleanup_cmd(
[docker_exe, "rm", "-f", container_id],
timeout=10,
action="force remove",
container_id=container_id,
)

def _remove_container_after_delay():
time.sleep(3)
_run_docker_cleanup_cmd(
[docker_exe, "rm", "-f", container_id],
timeout=10,
action="remove",
container_id=container_id,
)
subprocess.Popen(stop_cmd, shell=True)

try:
# Stop in background so cleanup doesn't block.
threading.Thread(target=_stop_container, daemon=True).start()
except Exception as e:
logger.warning("Failed to stop container %s: %s", self._container_id, e)
logger.warning("Failed to stop container %s: %s", container_id, e)

if not self._persistent:
if remove_after_stop:
# Also schedule removal (stop only leaves it as stopped)
try:
subprocess.Popen(
f"sleep 3 && {self._docker_exe} rm -f {self._container_id} >/dev/null 2>&1 &",
shell=True,
)
threading.Thread(target=_remove_container_after_delay, daemon=True).start()
except Exception:
pass
self._container_id = None
Expand Down
Loading