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 <xuyang1@microsoft.com>

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):

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.

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)

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):

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()