feat: add ws CLI and support optional timeout/cache (#1066)

* feat: add ws CLI and support optional timeout/cache

* lint

* fix bugs

* convert extra_volumes to dict for multiprocess

* lint

Author: you-n-g

rdagent/app/utils/ws.py

@@ -0,0 +1,51 @@
+from typing import Optional
+
+import typer
+
+from rdagent.app.data_science.conf import DS_RD_SETTING
+from rdagent.components.coder.data_science.conf import get_ds_env
+from rdagent.utils.agent.tpl import T
+
+app = typer.Typer(help="Run data-science environment commands.")
+
+
+@app.command()
+def run(competition: str, cmd: str, local_path: str = "./"):
+    """
+    Launch the data-science environment for a specific competition and run the
+    provided command.
+
+    Example:
+        1) start the container:
+        dotenv run -- python -m rdagent.app.utils.ws nomad2018-predict-transparent-conductors "sleep 3600" --local-path your_workspace
+
+        2) then run the following command to enter the latest container:
+        - docker exec -it `docker ps --filter 'status=running' -l --format '{{.Names}}'` bash
+        Or you can attach to the container by specifying the container name (find it in the run info)
+        - docker exec -it sweet_robinson bash
+
+    Arguments:
+        competition: The competition slug/folder name.
+        cmd: The shell command or script entry point to execute inside
+             the environment.
+    """
+    data_path = DS_RD_SETTING.local_data_path
+
+    data_path = (
+        f"{data_path}/{competition}" if DS_RD_SETTING.sample_data_by_LLM else f"{data_path}/sample/{competition}"
+    )
+    target_path = T("scenarios.data_science.share:scen.input_path").r()
+    extra_volumes = {data_path: target_path}
+
+    # Don't set time limitation and always disable cache
+    env = get_ds_env(
+        extra_volumes=extra_volumes,
+        running_timeout_period=None,
+        enable_cache=False,
+    )
+
+    env.run(entry=cmd, local_path=local_path)
+
+
+if __name__ == "__main__":  # pragma: no cover
+    app()

rdagent/components/coder/data_science/conf.py

@@ -27,7 +27,8 @@ class Config:
 def get_ds_env(
     conf_type: Literal["kaggle", "mlebench"] = "kaggle",
     extra_volumes: dict = {},
-    running_timeout_period: int = DS_RD_SETTING.debug_timeout,
+    running_timeout_period: int | None = DS_RD_SETTING.debug_timeout,
+    enable_cache: bool | None = None,
 ) -> Env:
     """
     Retrieve the appropriate environment configuration based on the env_type setting.
@@ -52,8 +53,10 @@ def get_ds_env(
         )
     else:
         raise ValueError(f"Unknown env type: {conf.env_type}")
-    env.conf.extra_volumes = extra_volumes
+    env.conf.extra_volumes = extra_volumes.copy()
     env.conf.running_timeout_period = running_timeout_period
+    if enable_cache is not None:
+        env.conf.enable_cache = enable_cache
     env.prepare()
     return env
 

rdagent/utils/env.py

@@ -117,15 +117,19 @@ def pull_image_with_progress(image: str) -> None:
 
 
 class EnvConf(ExtendedBaseSettings):
-    # TODO: add prefix ....
     default_entry: str
     extra_volumes: dict = {}
-    running_timeout_period: int = 3600  # 10 minutes
+    running_timeout_period: int | None = 3600  # 10 minutes
     # helper settings to support transparent;
     enable_cache: bool = True
     retry_count: int = 5  # retry count for the docker run
     retry_wait_seconds: int = 10  # retry wait seconds for the docker run
 
+    model_config = SettingsConfigDict(
+        # TODO: add prefix ....
+        env_parse_none_str="None",  # Nthis is the key to accept `RUNNING_TIMEOUT_PERIOD=None`
+    )
+
 
 ASpecificEnvConf = TypeVar("ASpecificEnvConf", bound=EnvConf)
 
@@ -225,7 +229,7 @@ def __run_with_retry(
                 )
                 end = time.time()
                 logger.info(f"Running time: {end - start} seconds")
-                if end - start + 1 >= self.conf.running_timeout_period:
+                if self.conf.running_timeout_period is not None and end - start + 1 >= self.conf.running_timeout_period:
                     logger.warning(
                         f"The running time exceeds {self.conf.running_timeout_period} seconds, so the process is killed."
                     )
@@ -299,9 +303,13 @@ def _get_path_stem(path: str) -> str | None:
             chmod_cmd += ")"
             return chmod_cmd
 
+        if self.conf.running_timeout_period is None:
+            timeout_cmd = entry
+        else:
+            timeout_cmd = f"timeout --kill-after=10 {self.conf.running_timeout_period} {entry}"
         entry_add_timeout = (
-            f"/bin/sh -c 'timeout --kill-after=10 {self.conf.running_timeout_period} {entry}; "
-            + "entry_exit_code=$?; "
+            f"/bin/sh -c '"  # start of the sh command
+            + f"{timeout_cmd}; entry_exit_code=$?; "
             + (
                 f"{_get_chmod_cmd(self.conf.mount_path)}; "
                 # We don't have to change the permission of the cache and input folder to remove it
@@ -310,7 +318,8 @@ def _get_path_stem(path: str) -> str | None:
                 if isinstance(self.conf, DockerConf)
                 else ""
             )
-            + "exit $entry_exit_code'"
+            + "exit $entry_exit_code"
+            + "'"  # end of the sh command
         )
 
         if self.conf.enable_cache:
@@ -635,7 +644,7 @@ class DockerConf(EnvConf):
     mem_limit: str | None = "48g"  # Add memory limit attribute
     cpu_count: int | None = None  # Add CPU limit attribute
 
-    running_timeout_period: int = 3600  # 1 hour
+    running_timeout_period: int | None = 3600  # 1 hour
 
     enable_cache: bool = True  # enable the cache mechanism
 
@@ -678,7 +687,10 @@ def prepare(self) -> None:
 
 
 class QlibDockerConf(DockerConf):
-    model_config = SettingsConfigDict(env_prefix="QLIB_DOCKER_")
+    model_config = SettingsConfigDict(
+        env_prefix="QLIB_DOCKER_",
+        env_parse_none_str="None",  # Nthis is the key to accept `RUNNING_TIMEOUT_PERIOD=None`
+    )
 
     build_from_dockerfile: bool = True
     dockerfile_folder_path: Path = Path(__file__).parent.parent / "scenarios" / "qlib" / "docker"
@@ -707,7 +719,7 @@ class KGDockerConf(DockerConf):
     #     Path("git_ignore_folder/data").resolve(): "/root/.data/"
     # }
 
-    running_timeout_period: int = 600
+    running_timeout_period: int | None = 600
     mem_limit: str | None = (
         "48g"  # Add memory limit attribute # new-york-city-taxi-fare-prediction may need more memory
     )
@@ -722,7 +734,7 @@ class DSDockerConf(DockerConf):
     mount_path: str = "/kaggle/workspace"
     default_entry: str = "python main.py"
 
-    running_timeout_period: int = 600
+    running_timeout_period: int | None = 600
     mem_limit: str | None = (
         "48g"  # Add memory limit attribute # new-york-city-taxi-fare-prediction may need more memory
     )

test/utils/test_conf.py

@@ -9,11 +9,13 @@ def test_conf(self):
         from rdagent.utils.env import EnvConf, QlibDockerConf
 
         os.environ["MEM_LIMIT"] = "200g"
+        os.environ["RUNNING_TIMEOUT_PERIOD"] = "None"
         assert QlibDockerConf().mem_limit == "200g"  # base class will affect subclasses
         os.environ["QLIB_DOCKER_MEM_LIMIT"] = "300g"
         assert QlibDockerConf().mem_limit == "300g"  # more accurate subclass will override the base class
+        assert QlibDockerConf().running_timeout_period is None
 
-        os.environ["default_entry"] = "which python"
+        os.environ["DEFAULT_ENTRY"] = "which python"
         os.environ["ENABLE_CACHE"] = "False"
 
         assert EnvConf().enable_cache is False