From f03e6a9812992eda802341c8861128fd7f70f556 Mon Sep 17 00:00:00 2001 From: Copilot <198982749+Copilot@users.noreply.github.com> Date: Thu, 19 Jun 2025 18:32:50 +0800 Subject: [PATCH] fix: docker container cleanup to prevent accumulation and system slowdown (#975) * Initial plan for issue * Fix Docker container cleanup issue by using try-finally block Co-authored-by: peteryang1 <25981102+peteryang1@users.noreply.github.com> * Fix additional Docker container leaks in health_check and GPU test functions Co-authored-by: peteryang1 <25981102+peteryang1@users.noreply.github.com> * Remove temporary test files and finalize Docker container cleanup fix Co-authored-by: peteryang1 <25981102+peteryang1@users.noreply.github.com> * Refactor container cleanup code to reduce duplication as requested in review feedback Co-authored-by: peteryang1 <25981102+peteryang1@users.noreply.github.com> * Refactor container cleanup to use shared function and always stop before remove Co-authored-by: peteryang1 <25981102+peteryang1@users.noreply.github.com> * fix CI * Fix mypy type checking errors for Docker container cleanup Co-authored-by: peteryang1 <25981102+peteryang1@users.noreply.github.com> * fix CI * Remove unnecessary _cleanup_container wrapper method in DockerEnv class Co-authored-by: peteryang1 <25981102+peteryang1@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: peteryang1 <25981102+peteryang1@users.noreply.github.com> Co-authored-by: Xu Yang --- rdagent/app/utils/health_check.py | 5 +++- rdagent/core/conf.py | 3 +-- rdagent/core/proposal.py | 26 ++++++++++----------- rdagent/utils/env.py | 38 +++++++++++++++++++++++++++---- test/utils/test_env.py | 13 +++++++++++ 5 files changed, 64 insertions(+), 21 deletions(-) diff --git a/rdagent/app/utils/health_check.py b/rdagent/app/utils/health_check.py index 8e4a1576..4e235221 100644 --- a/rdagent/app/utils/health_check.py +++ b/rdagent/app/utils/health_check.py @@ -3,22 +3,25 @@ import socket import docker from rdagent.log import rdagent_logger as logger +from rdagent.utils.env import cleanup_container def check_docker() -> None: + container = None try: client = docker.from_env() client.images.pull("hello-world") container = client.containers.run("hello-world", detach=True) logs = container.logs().decode("utf-8") print(logs) - container.remove() logger.info(f"The docker status is normal") except docker.errors.DockerException as e: logger.error(f"An error occurred: {e}") logger.warning( f"Docker status is exception, please check the docker configuration or reinstall it. Refs: https://docs.docker.com/engine/install/ubuntu/." ) + finally: + cleanup_container(container, "health check") def is_port_in_use(port): diff --git a/rdagent/core/conf.py b/rdagent/core/conf.py index 74d3a19d..c006672a 100644 --- a/rdagent/core/conf.py +++ b/rdagent/core/conf.py @@ -87,8 +87,7 @@ class RDAgentSettings(ExtendedBaseSettings): """Based on the setting of semaphore, return the maximum number of parallel loops""" if isinstance(self.step_semaphore, int): return self.step_semaphore - else: - return max(self.step_semaphore.values()) + return max(self.step_semaphore.values()) # NOTE: for debug # the following function only serves as debugging and is necessary in main logic. diff --git a/rdagent/core/proposal.py b/rdagent/core/proposal.py index 137fc20c..f88512c4 100644 --- a/rdagent/core/proposal.py +++ b/rdagent/core/proposal.py @@ -4,7 +4,7 @@ from __future__ import annotations import asyncio from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Generic, List, Tuple, TypeVar +from typing import TYPE_CHECKING, Generic, TypeVar from rdagent.core.conf import RD_AGENT_SETTINGS from rdagent.core.evaluation import Feedback @@ -108,7 +108,7 @@ ASpecificKB = TypeVar("ASpecificKB", bound=KnowledgeBase) class Trace(Generic[ASpecificScen, ASpecificKB]): NodeType = tuple[Experiment, ExperimentFeedback] # Define NodeType as a new type representing the tuple - NEW_ROOT: Tuple = () + NEW_ROOT: tuple = () def __init__(self, scen: ASpecificScen, knowledge_base: ASpecificKB | None = None) -> None: self.scen: ASpecificScen = scen @@ -163,35 +163,33 @@ class Trace(Generic[ASpecificScen, ASpecificKB]): return [self.hist[i] for i in self.get_parents(selection[0])] - def exp2idx(self, exp: Experiment | List[Experiment]) -> int | List[int] | None: + def exp2idx(self, exp: Experiment | list[Experiment]) -> int | list[int] | None: if isinstance(exp, list): - exps: List[Experiment] = exp + exps: list[Experiment] = exp # keep the order exp_to_index: dict[Experiment, int] = {_exp: i for i, (_exp, _) in enumerate(self.hist)} return [exp_to_index[_exp] for _exp in exps] - else: - for i, (_exp, _) in enumerate(self.hist): - if _exp == exp: - return i + for i, (_exp, _) in enumerate(self.hist): + if _exp == exp: + return i return None - def idx2exp(self, idx: int | List[int]) -> Experiment | List[Experiment]: + def idx2exp(self, idx: int | list[int]) -> Experiment | list[Experiment]: if isinstance(idx, list): - idxs: List[int] = idx + idxs: list[int] = idx return [self.hist[_idx][0] for _idx in idxs] - else: - return self.hist[idx][0] + return self.hist[idx][0] def is_parent(self, parent_idx: int, child_idx: int) -> bool: ancestors = self.get_parents(child_idx) return parent_idx in ancestors - def get_parents(self, child_idx: int) -> List[int]: + def get_parents(self, child_idx: int) -> list[int]: if self.is_selection_new_tree((child_idx,)): return [] - ancestors: List[int] = [] + ancestors: list[int] = [] curr = child_idx while True: ancestors.insert(0, curr) diff --git a/rdagent/utils/env.py b/rdagent/utils/env.py index 7c668442..26f7e96f 100644 --- a/rdagent/utils/env.py +++ b/rdagent/utils/env.py @@ -44,6 +44,29 @@ from rdagent.utils.agent.tpl import T from rdagent.utils.workflow import wait_retry +def cleanup_container(container: docker.models.containers.Container | None, context: str = "") -> None: # type: ignore[no-any-unimported] + """ + Shared helper function to clean up a Docker container. + Always stops the container before removing it. + + Parameters + ---------- + container : docker container object or None + The container to clean up, or None if no container to clean up + context : str + Additional context for logging (e.g., "health check", "GPU test") + """ + if container is not None: + try: + # Always stop first - stop() doesn't raise error if already stopped + container.stop() + container.remove() + except Exception as cleanup_error: + # Log cleanup error but don't mask the original exception + context_str = f" {context}" if context else "" + logger.warning(f"Failed to cleanup{context_str} container {container.id}: {cleanup_error}") + + # Normalize all bind paths in volumes to absolute paths using the workspace (working_dir). def normalize_volumes(vols: dict[str, str | dict[str, str]], working_dir: str) -> dict: abs_vols: dict[str, str | dict[str, str]] = {} @@ -785,12 +808,17 @@ class DockerEnv(Env[DockerConf]): @wait_retry(5, 10) def _f() -> dict: + container = None try: get_image(self.conf.image) - client.containers.run(self.conf.image, "nvidia-smi", **gpu_kwargs) + container = client.containers.run(self.conf.image, "nvidia-smi", detach=True, **gpu_kwargs) + # Wait for container to complete + container.wait() logger.info("GPU Devices are available.") except docker.errors.APIError: return {} + finally: + cleanup_container(container, context="GPU test") return gpu_kwargs return _f() @@ -835,9 +863,10 @@ class DockerEnv(Env[DockerConf]): volumes = normalize_volumes(cast(dict[str, str | dict[str, str]], volumes), self.conf.mount_path) log_output = "" + container: docker.models.containers.Container | None = None # type: ignore[no-any-unimported] try: - container: docker.models.containers.Container = client.containers.run( # type: ignore[no-any-unimported] + container = client.containers.run( image=self.conf.image, command=entry, volumes=volumes, @@ -851,6 +880,7 @@ class DockerEnv(Env[DockerConf]): cpu_count=self.conf.cpu_count, # Set CPU limit **self._gpu_kwargs(client), ) + assert container is not None # Ensure container was created successfully logs = container.logs(stream=True) print(Rule("[bold green]Docker Logs Begin[/bold green]", style="dark_orange")) table = Table(title="Run Info", show_header=False) @@ -869,8 +899,6 @@ class DockerEnv(Env[DockerConf]): Console().print(decoded_log, markup=False) log_output += decoded_log + "\n" exit_status = container.wait()["StatusCode"] - container.stop() - container.remove() print(Rule("[bold green]Docker Logs End[/bold green]", style="dark_orange")) return log_output, exit_status except docker.errors.ContainerError as e: @@ -879,6 +907,8 @@ class DockerEnv(Env[DockerConf]): raise RuntimeError("Docker image not found.") except docker.errors.APIError as e: raise RuntimeError(f"Error while running the container: {e}") + finally: + cleanup_container(container) class QTDockerEnv(DockerEnv): diff --git a/test/utils/test_env.py b/test/utils/test_env.py index 92321ec3..a8ee6723 100644 --- a/test/utils/test_env.py +++ b/test/utils/test_env.py @@ -13,6 +13,7 @@ from rdagent.utils.env import ( LocalEnv, QlibDockerConf, QTDockerEnv, + cleanup_container, ) DIRNAME = Path(__file__).absolute().resolve().parent @@ -142,6 +143,18 @@ class EnvUtils(unittest.TestCase): # docker run --memory=10m -it --rm local_qlib:latest python -c 'import numpy as np; print(123); size_mb = 1; size = size_mb * 1024 * 1024 // 8; array = np.random.randn(size).astype(np.float64); array[0], array[-1] = 1.0, 1.0; print(321)' # docker run --memory=10g -it --rm local_qlib:latest python -c 'import numpy as np; print(123); size_mb = 1; size = size_mb * 1024 * 1024 // 8; array = np.random.randn(size).astype(np.float64); array[0], array[-1] = 1.0, 1.0; print(321)' + def test_cleanup_container_import(self): + """Test that cleanup_container function can be imported and has correct interface.""" + # Test that the function exists and can be called + self.assertTrue(callable(cleanup_container)) + + # Test with None (should not raise an exception) + cleanup_container(None, "test context") + + # The function should accept positional and keyword arguments + cleanup_container(None) + cleanup_container(None, context="test") + if __name__ == "__main__": unittest.main()