diff --git a/rdagent/app/utils/ws.py b/rdagent/app/utils/ws.py new file mode 100644 index 00000000..ab1221a5 --- /dev/null +++ b/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() diff --git a/rdagent/components/coder/data_science/conf.py b/rdagent/components/coder/data_science/conf.py index 08ebb168..3719b86c 100644 --- a/rdagent/components/coder/data_science/conf.py +++ b/rdagent/components/coder/data_science/conf.py @@ -27,7 +27,8 @@ class DSCoderCoSTEERSettings(CoSTEERSettings): 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 diff --git a/rdagent/utils/env.py b/rdagent/utils/env.py index d2c0607e..01c620f8 100644 --- a/rdagent/utils/env.py +++ b/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 @@ class Env(Generic[ASpecificEnvConf]): ) 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 @@ class Env(Generic[ASpecificEnvConf]): 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 @@ class Env(Generic[ASpecificEnvConf]): 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 @@ class QlibCondaEnv(LocalEnv[QlibCondaConf]): 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 ) diff --git a/test/utils/test_conf.py b/test/utils/test_conf.py index 79b0e8ff..3a5a38a7 100644 --- a/test/utils/test_conf.py +++ b/test/utils/test_conf.py @@ -9,11 +9,13 @@ class ConfUtils(unittest.TestCase): 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