From 20778a9b8745dc7640b0f69b856813e701e74402 Mon Sep 17 00:00:00 2001 From: you-n-g Date: Wed, 12 Mar 2025 11:36:28 +0800 Subject: [PATCH] feat: condaenv & full docker env (#668) * use conda to run kaggle and mlebench code * refactor: Simplify environment configuration and execution logic * add setting to use local env in ds * refine dockerfile * fix: Move MLEBDockerEnv initialization inside conditionals & fix condaenv * refactor: reformat code for better readability and consistency * feat: add conda env to all envs. * fix: fix bugs when run loop * refactor: Simplify DockerEnv configuration in mle_summary.py * fix image bug * style: reformat code for better readability and consistency * change commit * feat: Add entrypoint script for sing_docker scenario in rdagent * refactor: add Any type hints and comments for clarity in env.py * feat: Create log directory if it doesn't exist in entrypoint script * feat: Add debug mode and list root directory in entrypoint script * fix: Remove specific branch checkout in Dockerfile for RD-Agent * fix: Add competition argument to loop.py script execution * fix: Correct directory navigation and dependency installation in entrypoint.sh * fix: Correct user ownership assignment in entrypoint script * refactor: Comment out redundant log copying to RD_OUTPUT_DIR * fix: Unset LOG_TRACE_PATH to prevent log contamination in entrypoint.sh --------- Co-authored-by: Xu Yang --- .gitignore | 2 +- rdagent/components/coder/data_science/conf.py | 40 ++ .../coder/data_science/ensemble/eval.py | 14 +- .../coder/data_science/feature/eval.py | 15 +- .../coder/data_science/model/eval.py | 13 +- .../coder/data_science/model/exp.py | 8 - .../data_science/raw_data_loader/__init__.py | 14 +- .../data_science/raw_data_loader/eval.py | 14 +- .../coder/data_science/raw_data_loader/exp.py | 11 - .../coder/data_science/workflow/eval.py | 17 +- rdagent/log/mle_summary.py | 8 +- .../data_science/dev/runner/__init__.py | 8 +- .../scenarios/data_science/dev/runner/eval.py | 23 +- .../scenarios/data_science/scen/__init__.py | 7 +- .../data_science/sing_docker/Dockerfile | 47 +++ .../data_science/sing_docker/entrypoint.sh | 47 +++ .../sing_docker/kaggle_environment.yaml | 367 +++++++++++++++++ rdagent/scenarios/kaggle/kaggle_crawler.py | 6 +- rdagent/utils/env.py | 380 +++++++++++------- test/utils/test_env.py | 41 +- 20 files changed, 835 insertions(+), 247 deletions(-) create mode 100644 rdagent/scenarios/data_science/sing_docker/Dockerfile create mode 100755 rdagent/scenarios/data_science/sing_docker/entrypoint.sh create mode 100644 rdagent/scenarios/data_science/sing_docker/kaggle_environment.yaml diff --git a/.gitignore b/.gitignore index 2ebc0030..d3e11fea 100644 --- a/.gitignore +++ b/.gitignore @@ -170,6 +170,6 @@ mlruns/ # shell script *.out -*.sh +/*.sh .aider* rdagent/app/benchmark/factor/example.json diff --git a/rdagent/components/coder/data_science/conf.py b/rdagent/components/coder/data_science/conf.py index 8b46985c..a1941b87 100644 --- a/rdagent/components/coder/data_science/conf.py +++ b/rdagent/components/coder/data_science/conf.py @@ -1,4 +1,15 @@ +from typing import Literal + from rdagent.components.coder.CoSTEER.config import CoSTEERSettings +from rdagent.utils.env import ( + CondaConf, + DockerEnv, + DSDockerConf, + Env, + LocalEnv, + MLEBDockerConf, + MLECondaConf, +) class DSCoderCoSTEERSettings(CoSTEERSettings): @@ -8,3 +19,32 @@ class DSCoderCoSTEERSettings(CoSTEERSettings): env_prefix = "DS_Coder_CoSTEER_" max_seconds: int = 2400 + env_type: str = "docker" + # TODO: extract a function for env and conf. + + +def get_ds_env(conf_type: Literal["kaggle", "mlebench"] = "kaggle") -> Env: + """ + Retrieve the appropriate environment configuration based on the env_type setting. + + Returns: + Env: An instance of the environment configured either as DockerEnv or LocalEnv. + + Raises: + ValueError: If the env_type is not recognized. + """ + conf = DSCoderCoSTEERSettings() + assert conf_type in ["kaggle", "mlebench"], f"Unknown conf_type: {conf_type}" + + if conf.env_type == "docker": + env_conf = DSDockerConf() if conf_type == "kaggle" else MLEBDockerConf() + env = DockerEnv(conf=env_conf) + elif conf.env_type == "conda": + env = LocalEnv( + conf=( + CondaConf(conda_env_name=conf_type) if conf_type == "kaggle" else MLECondaConf(conda_env_name=conf_type) + ) + ) + else: + raise ValueError(f"Unknown env type: {conf.env_type}") + return env diff --git a/rdagent/components/coder/data_science/ensemble/eval.py b/rdagent/components/coder/data_science/ensemble/eval.py index 96830882..e6a4b72f 100644 --- a/rdagent/components/coder/data_science/ensemble/eval.py +++ b/rdagent/components/coder/data_science/ensemble/eval.py @@ -9,12 +9,11 @@ from rdagent.components.coder.CoSTEER.evaluators import ( CoSTEEREvaluator, CoSTEERSingleFeedback, ) +from rdagent.components.coder.data_science.conf import get_ds_env from rdagent.core.evolving_framework import QueriedKnowledge from rdagent.core.experiment import FBWorkspace, Task -from rdagent.oai.llm_utils import APIBackend from rdagent.utils.agent.tpl import T from rdagent.utils.agent.workflow import build_cls_from_json_with_retry -from rdagent.utils.env import DockerEnv, DSDockerConf DIRNAME = Path(__file__).absolute().resolve().parent @@ -45,11 +44,8 @@ class EnsembleCoSTEEREvaluator(CoSTEEREvaluator): final_decision=False, ) - ds_docker_conf = DSDockerConf() - ds_docker_conf.extra_volumes = { - f"{DS_RD_SETTING.local_data_path}/sample/{self.scen.competition}": "/kaggle/input" - } - de = DockerEnv(conf=ds_docker_conf) + env = get_ds_env() + env.conf.extra_volumes = {f"{DS_RD_SETTING.local_data_path}/sample/{self.scen.competition}": "/kaggle/input"} fname = "test/ensemble_test.txt" test_code = (DIRNAME / "eval_tests" / "ensemble_test.txt").read_text() @@ -64,12 +60,12 @@ class EnsembleCoSTEEREvaluator(CoSTEEREvaluator): ) implementation.inject_files(**{fname: test_code}) - stdout, ret_code = implementation.execute_ret_code(env=de, entry=f"python {fname}") + stdout, ret_code = implementation.execute_ret_code(env=env, entry=f"python {fname}") stdout += f"\nNOTE: the above scripts run with return code {ret_code}" if "main.py" in implementation.file_dict: - workflow_stdout = implementation.execute(env=de, entry="python main.py") + workflow_stdout = implementation.execute(env=env, entry="python main.py") workflow_stdout = re.sub(r"=== Start of EDA part ===(.*)=== End of EDA part ===", "", workflow_stdout) else: workflow_stdout = None diff --git a/rdagent/components/coder/data_science/feature/eval.py b/rdagent/components/coder/data_science/feature/eval.py index 818da0d9..f09f8eda 100644 --- a/rdagent/components/coder/data_science/feature/eval.py +++ b/rdagent/components/coder/data_science/feature/eval.py @@ -7,12 +7,11 @@ from rdagent.components.coder.CoSTEER.evaluators import ( CoSTEEREvaluator, CoSTEERSingleFeedback, ) +from rdagent.components.coder.data_science.conf import get_ds_env from rdagent.core.evolving_framework import QueriedKnowledge from rdagent.core.experiment import FBWorkspace, Task -from rdagent.oai.llm_utils import APIBackend from rdagent.utils.agent.tpl import T from rdagent.utils.agent.workflow import build_cls_from_json_with_retry -from rdagent.utils.env import DockerEnv, DSDockerConf from rdagent.utils.fmt import shrink_text DIRNAME = Path(__file__).absolute().resolve().parent @@ -45,22 +44,18 @@ class FeatureCoSTEEREvaluator(CoSTEEREvaluator): final_decision=False, ) - ds_docker_conf = DSDockerConf() - # TODO: we should /= 20 for the timeout period on debug component - ds_docker_conf.extra_volumes = { - f"{DS_RD_SETTING.local_data_path}/sample/{self.scen.competition}": "/kaggle/input" - } - de = DockerEnv(conf=ds_docker_conf) + env = get_ds_env() + env.conf.extra_volumes = {f"{DS_RD_SETTING.local_data_path}/sample/{self.scen.competition}": "/kaggle/input"} # TODO: do we need to clean the generated temporary content? fname = "test/feature_test.py" test_code = (DIRNAME / "eval_tests" / "feature_test.txt").read_text() implementation.inject_files(**{fname: test_code}) - stdout = implementation.execute(env=de, entry=f"python {fname}") + stdout = implementation.execute(env=env, entry=f"python {fname}") if "main.py" in implementation.file_dict: - workflow_stdout = implementation.execute(env=de, entry="python main.py") + workflow_stdout = implementation.execute(env=env, entry="python main.py") workflow_stdout = re.sub(r"=== Start of EDA part ===(.*)=== End of EDA part ===", "", workflow_stdout) else: workflow_stdout = None diff --git a/rdagent/components/coder/data_science/model/eval.py b/rdagent/components/coder/data_science/model/eval.py index df03341d..a1b03fb2 100644 --- a/rdagent/components/coder/data_science/model/eval.py +++ b/rdagent/components/coder/data_science/model/eval.py @@ -12,13 +12,13 @@ from rdagent.components.coder.CoSTEER.evaluators import ( CoSTEEREvaluator, CoSTEERSingleFeedback, ) +from rdagent.components.coder.data_science.conf import get_ds_env from rdagent.core.evolving_framework import QueriedKnowledge from rdagent.core.exception import CoderError from rdagent.core.experiment import FBWorkspace, Task from rdagent.oai.llm_utils import APIBackend from rdagent.utils.agent.tpl import T from rdagent.utils.agent.workflow import build_cls_from_json_with_retry -from rdagent.utils.env import DockerEnv, DSDockerConf DIRNAME = Path(__file__).absolute().resolve().parent ModelSingleFeedback = CoSTEERSingleFeedback @@ -56,18 +56,15 @@ class ModelGeneralCaseSpecEvaluator(CoSTEEREvaluator): final_decision=False, ) - ds_docker_conf = DSDockerConf() - ds_docker_conf.extra_volumes = { - f"{DS_RD_SETTING.local_data_path}/sample/{self.scen.competition}": "/kaggle/input" - } - de = DockerEnv(conf=ds_docker_conf) + env = get_ds_env() + env.conf.extra_volumes = {f"{DS_RD_SETTING.local_data_path}/sample/{self.scen.competition}": "/kaggle/input"} fname = "test/model_test.py" test_code = ( (DIRNAME / "eval_tests" / "model_test.txt").read_text().replace("model01", target_task.name) ) # only check the model changed this time implementation.inject_files(**{fname: test_code}) - stdout = implementation.execute(env=de, entry=f"python {fname}") + stdout = implementation.execute(env=env, entry=f"python {fname}") if stdout is None: raise CoderError( @@ -75,7 +72,7 @@ class ModelGeneralCaseSpecEvaluator(CoSTEEREvaluator): ) if "main.py" in implementation.file_dict: - workflow_stdout = implementation.execute(env=de, entry="python main.py") + workflow_stdout = implementation.execute(env=env, entry="python main.py") workflow_stdout = re.sub(r"=== Start of EDA part ===(.*)=== End of EDA part ===", "", workflow_stdout) else: workflow_stdout = None diff --git a/rdagent/components/coder/data_science/model/exp.py b/rdagent/components/coder/data_science/model/exp.py index 43879fa8..19e8afef 100644 --- a/rdagent/components/coder/data_science/model/exp.py +++ b/rdagent/components/coder/data_science/model/exp.py @@ -1,14 +1,6 @@ -import pickle -import site -import traceback -from pathlib import Path from typing import Dict, Optional from rdagent.components.coder.CoSTEER.task import CoSTEERTask -from rdagent.core.experiment import Experiment, FBWorkspace -from rdagent.core.utils import cache_with_pickle -from rdagent.oai.llm_utils import md5_hash -from rdagent.utils.env import DockerEnv, DSDockerConf # Because we use isinstance to distinguish between different types of tasks, we need to use sub classes to represent different types of tasks diff --git a/rdagent/components/coder/data_science/raw_data_loader/__init__.py b/rdagent/components/coder/data_science/raw_data_loader/__init__.py index e1b84044..7450d400 100644 --- a/rdagent/components/coder/data_science/raw_data_loader/__init__.py +++ b/rdagent/components/coder/data_science/raw_data_loader/__init__.py @@ -38,7 +38,10 @@ from rdagent.components.coder.CoSTEER.evolving_strategy import ( from rdagent.components.coder.CoSTEER.knowledge_management import ( CoSTEERQueriedKnowledge, ) -from rdagent.components.coder.data_science.conf import DSCoderCoSTEERSettings +from rdagent.components.coder.data_science.conf import ( + DSCoderCoSTEERSettings, + get_ds_env, +) from rdagent.components.coder.data_science.raw_data_loader.eval import ( DataLoaderCoSTEEREvaluator, ) @@ -48,7 +51,6 @@ from rdagent.core.experiment import FBWorkspace from rdagent.core.scenario import Scenario from rdagent.oai.llm_utils import APIBackend from rdagent.utils.agent.tpl import T -from rdagent.utils.env import DockerEnv, DSDockerConf class DataLoaderMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy): @@ -214,10 +216,10 @@ class DataLoaderCoSTEER(CoSTEER): def develop(self, exp): new_exp = super().develop(exp) - ds_docker_conf = DSDockerConf() - ds_docker_conf.extra_volumes = {f"{DS_RD_SETTING.local_data_path}/{self.scen.competition}": "/kaggle/input"} - de = DockerEnv(conf=ds_docker_conf) - stdout = new_exp.experiment_workspace.execute(env=de, entry=f"python test/data_loader_test.py") + env = get_ds_env() + env.conf.extra_volumes = {f"{DS_RD_SETTING.local_data_path}/{self.scen.competition}": "/kaggle/input"} + + stdout = new_exp.experiment_workspace.execute(env=env, entry=f"python test/data_loader_test.py") match = re.search(r"(.*?)=== Start of EDA part ===(.*)=== End of EDA part ===", stdout, re.DOTALL) eda_output = match.groups()[1] if match else None self.scen.eda_output = eda_output diff --git a/rdagent/components/coder/data_science/raw_data_loader/eval.py b/rdagent/components/coder/data_science/raw_data_loader/eval.py index 335022f4..58d84327 100644 --- a/rdagent/components/coder/data_science/raw_data_loader/eval.py +++ b/rdagent/components/coder/data_science/raw_data_loader/eval.py @@ -12,11 +12,10 @@ from rdagent.components.coder.CoSTEER.evaluators import ( from rdagent.components.coder.CoSTEER.knowledge_management import ( CoSTEERQueriedKnowledgeV2, ) +from rdagent.components.coder.data_science.conf import get_ds_env from rdagent.core.experiment import FBWorkspace, Task -from rdagent.oai.llm_utils import APIBackend from rdagent.utils.agent.tpl import T from rdagent.utils.agent.workflow import build_cls_from_json_with_retry -from rdagent.utils.env import DockerEnv, DSDockerConf DIRNAME = Path(__file__).absolute().resolve().parent @@ -48,17 +47,14 @@ class DataLoaderCoSTEEREvaluator(CoSTEEREvaluator): final_decision=False, ) - ds_docker_conf = DSDockerConf() - ds_docker_conf.extra_volumes = { - f"{DS_RD_SETTING.local_data_path}/sample/{self.scen.competition}": "/kaggle/input" - } - de = DockerEnv(conf=ds_docker_conf) + env = get_ds_env() + env.conf.extra_volumes = {f"{DS_RD_SETTING.local_data_path}/sample/{self.scen.competition}": "/kaggle/input"} # TODO: do we need to clean the generated temporary content? fname = "test/data_loader_test.py" test_code = (DIRNAME / "eval_tests" / "data_loader_test.txt").read_text() implementation.inject_files(**{fname: test_code}) - stdout = implementation.execute(env=de, entry=f"python {fname}") + stdout = implementation.execute(env=env, entry=f"python {fname}") match = re.search(r"(.*?)=== Start of EDA part ===(.*)=== End of EDA part ===(.*)", stdout, re.DOTALL) stdout_part_1, eda_output, stdout_part_2 = match.groups() if match else (stdout, None, "") stdout = stdout_part_1 + stdout_part_2 @@ -66,7 +62,7 @@ class DataLoaderCoSTEEREvaluator(CoSTEEREvaluator): eda_output += "Length of EDA output is too long, truncated. Please reject this implementation and motivate it to reduce the length of EDA output." if "main.py" in implementation.file_dict: - workflow_stdout = implementation.execute(env=de, entry="python main.py") + workflow_stdout = implementation.execute(env=env, entry="python main.py") workflow_stdout = re.sub(r"=== Start of EDA part ===(.*)=== End of EDA part ===", "", workflow_stdout) else: workflow_stdout = None diff --git a/rdagent/components/coder/data_science/raw_data_loader/exp.py b/rdagent/components/coder/data_science/raw_data_loader/exp.py index 9263e209..d17f27f8 100644 --- a/rdagent/components/coder/data_science/raw_data_loader/exp.py +++ b/rdagent/components/coder/data_science/raw_data_loader/exp.py @@ -1,15 +1,4 @@ -import pickle -import site -import traceback -from pathlib import Path -from typing import Dict, Optional - from rdagent.components.coder.CoSTEER.task import CoSTEERTask -from rdagent.core.experiment import Experiment, FBWorkspace -from rdagent.core.utils import cache_with_pickle -from rdagent.oai.llm_utils import md5_hash -from rdagent.utils.agent.tpl import T -from rdagent.utils.env import DockerEnv, DSDockerConf # Because we use isinstance to distinguish between different types of tasks, we need to use sub classes to represent different types of tasks diff --git a/rdagent/components/coder/data_science/workflow/eval.py b/rdagent/components/coder/data_science/workflow/eval.py index 50efbf31..eb830b7f 100644 --- a/rdagent/components/coder/data_science/workflow/eval.py +++ b/rdagent/components/coder/data_science/workflow/eval.py @@ -10,12 +10,11 @@ from rdagent.components.coder.CoSTEER.evaluators import ( CoSTEERMultiFeedback, CoSTEERSingleFeedback, ) +from rdagent.components.coder.data_science.conf import get_ds_env from rdagent.core.evolving_framework import QueriedKnowledge from rdagent.core.experiment import FBWorkspace, Task -from rdagent.oai.llm_utils import APIBackend from rdagent.utils.agent.tpl import T from rdagent.utils.agent.workflow import build_cls_from_json_with_retry -from rdagent.utils.env import DockerEnv, DSDockerConf, MLEBDockerConf DIRNAME = Path(__file__).absolute().resolve().parent @@ -54,12 +53,8 @@ class WorkflowGeneralCaseSpecEvaluator(CoSTEEREvaluator): final_decision=False, ) - # DockerEnv for Kaggle Competition - ds_docker_conf = DSDockerConf() - ds_docker_conf.extra_volumes = { - f"{DS_RD_SETTING.local_data_path}/sample/{self.scen.competition}": "/kaggle/input" - } - de = DockerEnv(conf=ds_docker_conf) + env = get_ds_env() + env.conf.extra_volumes = {f"{DS_RD_SETTING.local_data_path}/sample/{self.scen.competition}": "/kaggle/input"} # # DockerEnv for MLEBench submission validation # mle_de_conf = MLEBDockerConf() @@ -70,9 +65,9 @@ class WorkflowGeneralCaseSpecEvaluator(CoSTEEREvaluator): # mde.prepare() # Clean the scores.csv & submission.csv. - implementation.execute(env=de, entry=f"rm submission.csv scores.csv") + implementation.execute(env=env, entry=f"rm submission.csv scores.csv") - stdout = implementation.execute(env=de, entry=f"python main.py") + stdout = implementation.execute(env=env, entry=f"python main.py") stdout = re.sub(r"=== Start of EDA part ===(.*)=== End of EDA part ===", "", stdout) # Check score file @@ -102,7 +97,7 @@ class WorkflowGeneralCaseSpecEvaluator(CoSTEEREvaluator): implementation.inject_files(**{"test/submission_format_test.py": base_check_code}) # stdout += "----Submission Check 1-----\n" submission_check_out, submission_ret_code = implementation.execute_ret_code( - env=de, entry="python test/submission_format_test.py" + env=env, entry="python test/submission_format_test.py" ) stdout += "\n" + submission_check_out diff --git a/rdagent/log/mle_summary.py b/rdagent/log/mle_summary.py index 1d1d5a0e..be56e955 100644 --- a/rdagent/log/mle_summary.py +++ b/rdagent/log/mle_summary.py @@ -7,6 +7,7 @@ import fire import pandas as pd from rdagent.app.data_science.conf import DS_RD_SETTING +from rdagent.components.coder.data_science.conf import get_ds_env from rdagent.core.experiment import FBWorkspace from rdagent.core.proposal import ExperimentFeedback from rdagent.log.storage import FileStorage @@ -14,11 +15,8 @@ from rdagent.scenarios.data_science.experiment.experiment import DSExperiment from rdagent.scenarios.kaggle.kaggle_crawler import score_rank from rdagent.utils.env import DockerEnv, MLEBDockerConf -mle_de_conf = MLEBDockerConf() -mle_de_conf.extra_volumes = { - f"{DS_RD_SETTING.local_data_path}/zip_files": "/mle/data", -} -de = DockerEnv(conf=mle_de_conf) +de = get_ds_env("mlebench") +de.conf.extra_volumes = {f"{DS_RD_SETTING.local_data_path}/zip_files": "/mle/data"} de.prepare() diff --git a/rdagent/scenarios/data_science/dev/runner/__init__.py b/rdagent/scenarios/data_science/dev/runner/__init__.py index a51cca4e..51229361 100644 --- a/rdagent/scenarios/data_science/dev/runner/__init__.py +++ b/rdagent/scenarios/data_science/dev/runner/__init__.py @@ -17,6 +17,7 @@ from rdagent.components.coder.CoSTEER.evolving_strategy import ( MultiProcessEvolvingStrategy, ) from rdagent.components.coder.CoSTEER.task import CoSTEERTask +from rdagent.components.coder.data_science.conf import get_ds_env from rdagent.core.exception import RunnerError from rdagent.core.scenario import Scenario from rdagent.log import rdagent_logger as logger @@ -115,11 +116,8 @@ class DSCoSTEERRunner(CoSTEER): exp.result = pd.read_csv(score_fp, index_col=0) # DockerEnv for MLEBench submission validation - mle_de_conf = MLEBDockerConf() - mle_de_conf.extra_volumes = { - f"{DS_RD_SETTING.local_data_path}/zip_files": "/mle/data", - } - mde = DockerEnv(conf=mle_de_conf) + mde = get_ds_env(conf_type="mlebench") + mde.conf.extra_volumes = {f"{DS_RD_SETTING.local_data_path}/zip_files": "/mle/data"} mde.prepare() # MLEBench Check mle_check_code = ( diff --git a/rdagent/scenarios/data_science/dev/runner/eval.py b/rdagent/scenarios/data_science/dev/runner/eval.py index 9aceed34..527d2bca 100644 --- a/rdagent/scenarios/data_science/dev/runner/eval.py +++ b/rdagent/scenarios/data_science/dev/runner/eval.py @@ -8,13 +8,17 @@ from rdagent.components.coder.CoSTEER.evaluators import ( CoSTEEREvaluator, CoSTEERSingleFeedback, ) +from rdagent.components.coder.data_science.conf import ( + DSCoderCoSTEERSettings, + get_ds_env, +) from rdagent.core.evolving_framework import QueriedKnowledge from rdagent.core.experiment import FBWorkspace, Task from rdagent.log import rdagent_logger as logger from rdagent.oai.llm_utils import APIBackend from rdagent.utils.agent.tpl import T from rdagent.utils.agent.workflow import build_cls_from_json_with_retry -from rdagent.utils.env import DockerEnv, DSDockerConf, MLEBDockerConf +from rdagent.utils.env import DockerEnv, MLEBDockerConf from rdagent.utils.fmt import shrink_text DIRNAME = Path(__file__).absolute().resolve().parent @@ -32,18 +36,16 @@ class DSCoSTEERCoSTEEREvaluator(CoSTEEREvaluator): queried_knowledge: QueriedKnowledge = None, **kwargs, ) -> DSCoSTEEREvalFeedback: - ds_docker_conf = DSDockerConf() - ds_docker_conf.extra_volumes = {f"{DS_RD_SETTING.local_data_path}/{self.scen.competition}": "/kaggle/input"} - ds_docker_conf.running_timeout_period = DS_RD_SETTING.full_timeout - de = DockerEnv(conf=ds_docker_conf) + env = get_ds_env() + env.conf.extra_volumes = {f"{DS_RD_SETTING.local_data_path}/{self.scen.competition}": "/kaggle/input"} stdout = implementation.execute( - env=de, entry=f"rm submission.csv scores.csv" + env=env, entry=f"rm submission.csv scores.csv" ) # Remove previous submission and scores files generated by worklfow. # execute workflow - stdout = implementation.execute(env=de, entry="coverage run main.py") + stdout = implementation.execute(env=env, entry="coverage run main.py") stdout = re.sub(r"=== Start of EDA part ===(.*)=== End of EDA part ===", "", stdout) score_fp = implementation.workspace_path / "scores.csv" @@ -57,11 +59,10 @@ class DSCoSTEERCoSTEEREvaluator(CoSTEEREvaluator): stdout += "\n Submission file (submission.csv) is not generated!" else: # DockerEnv for MLEBench submission validation - mle_de_conf = MLEBDockerConf() - mle_de_conf.extra_volumes = { + mde = get_ds_env("mlebench") + mde.conf.extra_volumes = { f"{DS_RD_SETTING.local_data_path}/zip_files": "/mle/data", } - mde = DockerEnv(conf=mle_de_conf) mde.prepare() # MLEBench Check mle_check_code = ( @@ -88,7 +89,7 @@ class DSCoSTEERCoSTEEREvaluator(CoSTEEREvaluator): if feedback: # remove unused files - implementation.execute(env=de, entry="coverage json -o coverage.json") + implementation.execute(env=env, entry="coverage json -o coverage.json") if Path(implementation.workspace_path / "coverage.json").exists(): with open(implementation.workspace_path / "coverage.json") as f: used_files = set(json.load(f)["files"].keys()) diff --git a/rdagent/scenarios/data_science/scen/__init__.py b/rdagent/scenarios/data_science/scen/__init__.py index b64b82ee..93c46b9e 100644 --- a/rdagent/scenarios/data_science/scen/__init__.py +++ b/rdagent/scenarios/data_science/scen/__init__.py @@ -7,6 +7,7 @@ import pandas as pd from PIL import Image, TiffTags from rdagent.app.data_science.conf import DS_RD_SETTING +from rdagent.components.coder.data_science.conf import get_ds_env from rdagent.core.experiment import FBWorkspace from rdagent.core.scenario import Scenario from rdagent.log import rdagent_logger as logger @@ -16,7 +17,6 @@ from rdagent.scenarios.kaggle.kaggle_crawler import ( leaderboard_scores, ) from rdagent.utils.agent.tpl import T -from rdagent.utils.env import DockerEnv, DSDockerConf def read_csv_head(file_path, indent=0, lines=5, max_col_width=100): @@ -314,14 +314,13 @@ class DataScienceScen(Scenario): def get_runtime_environment(self) -> str: # TODO: add it into base class. Environment should(i.e. `DSDockerConf`) should be part of the scenario class. - ds_docker_conf = DSDockerConf() - de = DockerEnv(conf=ds_docker_conf) + env = get_ds_env() implementation = FBWorkspace() fname = "temp.py" implementation.inject_files( **{fname: (Path(__file__).absolute().resolve().parent / "runtime_info.py").read_text()} ) - stdout = implementation.execute(env=de, entry=f"python {fname}") + stdout = implementation.execute(env=env, entry=f"python {fname}") return stdout def _get_data_folder_description(self) -> str: diff --git a/rdagent/scenarios/data_science/sing_docker/Dockerfile b/rdagent/scenarios/data_science/sing_docker/Dockerfile new file mode 100644 index 00000000..32cd8116 --- /dev/null +++ b/rdagent/scenarios/data_science/sing_docker/Dockerfile @@ -0,0 +1,47 @@ +# Use the official PyTorch image as the base image +FROM pytorch/pytorch:latest +# FROM pytorch/pytorch:2.4.1-cuda12.1-cudnn9-runtime +# torch.__version__ == '2.4.1+cu121' in "gcr.io/kaggle-gpu-images/python" + +# Install additional tools +RUN apt-get update && apt-get install -y \ + curl \ + vim \ + git \ + build-essential \ + git-lfs \ + unzip && \ + rm -rf /var/lib/apt/lists/* + +# Default command to keep the container running +ENTRYPOINT ["/workspace/run/entrypoint.sh"] + +RUN conda init bash + +# MLE-Bench +RUN conda create -n mlebench python==3.11 pip -y +RUN cd /workspace && git clone https://github.com/openai/mle-bench.git +RUN cd /workspace/mle-bench && git lfs fetch --all +RUN cd /workspace/mle-bench && git lfs pull +RUN cd /workspace/mle-bench && conda run -n mlebench pip install -e . + +# Kaggle Environment +COPY ./kaggle_environment.yaml /workspace +RUN cd /workspace && conda env create -f /workspace/kaggle_environment.yaml + +# RD-Agent +RUN cd /workspace && git clone https://github.com/microsoft/RD-Agent +RUN cd RD-Agent && git fetch && make dev + + +# litellm +RUN cd /workspace && mkdir -p litellm-srv +RUN cd /workspace/litellm-srv && curl https://raw.githubusercontent.com/you-n-g/deploy/refs/heads/master/configs/python/litellm.trapi.yaml -o litellm.trapi.yaml +RUN pip install 'litellm[proxy]' +RUN pip install git+https://github.com/you-n-g/litellm@add_mi_cred_pr + +run cd /workspace && mkdir -p run +COPY ./entrypoint.sh /workspace/run + + +WORKDIR /workspace/RD-Agent/ diff --git a/rdagent/scenarios/data_science/sing_docker/entrypoint.sh b/rdagent/scenarios/data_science/sing_docker/entrypoint.sh new file mode 100755 index 00000000..3dde42ea --- /dev/null +++ b/rdagent/scenarios/data_science/sing_docker/entrypoint.sh @@ -0,0 +1,47 @@ +#!/bin/sh +set -x + +DIR="$( cd "$(dirname "$(readlink -f "$0")")" || exit ; pwd -P )" + +sudo mkdir -p /mle/ /kaggle/ + +CURRENT_USER=$(id -un) +sudo chown -R $CURRENT_USER:$CURRENT_USER /workspace/ /mle/ /kaggle/ + +ls -lat / + +cd $DIR/../RD-Agent +mkdir -p log/ +git fetch +git checkout ${RD_COMMIT:-ee8d97c52062607cac778b8aeb10769b075a8d11} +make dev +pip install 'litellm[proxy]' +pip install git+https://github.com/you-n-g/litellm@add_mi_cred_pr + + +cd $DIR/../litellm-srv/ +export AZURE_CLIENT_ID +export AZURE_SCOPE=api://trapi/.default +export AZURE_CREDENTIAL=ManagedIdentityCredential +sed -i '/proxy_handler_instance/d' litellm.trapi.yaml # remove useless handler in production +nohup litellm --config litellm.trapi.yaml & + +sleep 10 # wait for litellm to start + + +cd $DIR/../RD-Agent +script -c "timeout ${RD_TIMEOUT:-24h} python rdagent/app/data_science/loop.py --competition $DS_COMPETITION" log/stdout.${DS_COMPETITION}.log + +unset LOG_TRACE_PATH # avoid make the original log dirty. +python rdagent/log/mle_summary.py grade_summary --log_folder=./log/ + +tar cf log.tar log + +# NOTE: when we have $AMLT_OUTPUT_DIR, maybe we don't have to copy file actively to azure blob now. +# RD_OUTPUT_DIR=${RD_OUTPUT_DIR:-/data/rdagent}/ +# mkdir -p $RD_OUTPUT_DIR +# cp -r log.tar $RD_OUTPUT_DIR/${RD_RES_NAME:-log.tar} + +cp -r log.tar $AMLT_OUTPUT_DIR/${RD_RES_NAME:-log.tar} + +set > $AMLT_OUTPUT_DIR/env diff --git a/rdagent/scenarios/data_science/sing_docker/kaggle_environment.yaml b/rdagent/scenarios/data_science/sing_docker/kaggle_environment.yaml new file mode 100644 index 00000000..5cf48e09 --- /dev/null +++ b/rdagent/scenarios/data_science/sing_docker/kaggle_environment.yaml @@ -0,0 +1,367 @@ +name: kaggle +channels: + - defaults + - https://repo.anaconda.com/pkgs/main + - https://repo.anaconda.com/pkgs/r +dependencies: + - _libgcc_mutex=0.1=main + - _openmp_mutex=5.1=1_gnu + - bzip2=1.0.8=h5eee18b_6 + - ca-certificates=2025.2.25=h06a4308_0 + - ld_impl_linux-64=2.40=h12ee557_0 + - libffi=3.4.4=h6a678d5_1 + - libgcc-ng=11.2.0=h1234567_1 + - libgomp=11.2.0=h1234567_1 + - libstdcxx-ng=11.2.0=h1234567_1 + - libuuid=1.41.5=h5eee18b_0 + - ncurses=6.4=h6a678d5_0 + - openssl=3.0.15=h5eee18b_0 + - pip=25.0=py311h06a4308_0 + - python=3.11.11=he870216_0 + - readline=8.2=h5eee18b_0 + - setuptools=75.8.0=py311h06a4308_0 + - sqlite=3.45.3=h5eee18b_0 + - tk=8.6.14=h39e8969_0 + - wheel=0.45.1=py311h06a4308_0 + - xz=5.6.4=h5eee18b_1 + - zlib=1.2.13=h5eee18b_1 + - pip: + - absl-py==2.1.0 + - accelerate==0.33.0 + - aideml==0.1.4 + - aiohappyeyeballs==2.4.6 + - aiohttp==3.11.13 + - aiosignal==1.3.2 + - albucore==0.0.23 + - albumentations==1.4.14 + - alembic==1.14.1 + - annotated-types==0.7.0 + - anthropic==0.34.1 + - antlr4-python3-runtime==4.9.3 + - anyio==4.8.0 + - arrow==1.3.0 + - asttokens==3.0.0 + - astunparse==1.6.3 + - attrs==25.1.0 + - audioread==3.0.1 + - azure-ai-formrecognizer==3.3.3 + - azure-common==1.1.28 + - azure-core==1.32.0 + - azure-identity==1.20.0 + - azure-storage-blob==12.24.1 + - backoff==2.2.1 + - bayesian-optimization==1.5.1 + - bayespy==0.5.1 + - biopython==1.84 + - black==24.3.0 + - bleach==6.2.0 + - blis==0.7.11 + - brotli==1.1.0 + - bson==0.5.10 + - cachetools==5.5.2 + - catalogue==2.0.10 + - catboost==1.2.5 + - certifi==2025.1.31 + - cffi==1.17.1 + - charset-normalizer==3.4.1 + - click==8.1.8 + - cloudpathlib==0.20.0 + - cloudpickle==3.1.1 + - colorama==0.4.6 + - colorlog==6.9.0 + - comm==0.2.2 + - confection==0.1.5 + - contourpy==1.3.1 + - coolname==2.2.0 + - cryptography==44.0.2 + - cycler==0.12.1 + - cymem==2.0.11 + - cython==3.0.11 + - dacite==1.8.1 + - dataclasses-json==0.6.7 + - datasets==2.1.0 + - debugpy==1.8.12 + - decorator==5.2.1 + - defusedxml==0.7.1 + - dill==0.3.9 + - distro==1.9.0 + - efficientnet-pytorch==0.7.1 + - eval-type-backport==0.2.2 + - evaluate==0.4.2 + - executing==2.2.0 + - fastai==2.7.17 + - fastcore==1.7.29 + - fastdownload==0.0.7 + - fastdtw==0.3.4 + - fastjsonschema==2.21.1 + - fastprogress==1.0.3 + - faust-cchardet==2.1.19 + - filelock==3.17.0 + - flatbuffers==25.2.10 + - fonttools==4.56.0 + - frozenlist==1.5.0 + - fsspec==2025.2.0 + - funcy==2.0 + - future==1.0.0 + - gast==0.6.0 + - gdcm==1.1 + - gensim==4.3.3 + - genson==1.3.0 + - geographiclib==2.0 + - geopy==2.4.1 + - graphviz==0.20.3 + - greenlet==3.1.1 + - grpcio==1.71.0rc2 + - gym==0.26.2 + - gym-notices==0.0.8 + - h11==0.14.0 + - h5py==3.11.0 + - hmmlearn==0.3.2 + - httpcore==1.0.7 + - httplib2==0.22.0 + - httpx==0.27.2 + - huggingface-hub==0.29.1 + - humanize==4.8.0 + - hyperopt==0.2.7 + - idna==3.10 + - igraph==0.11.6 + - imagecodecs==2024.6.1 + - imageio==2.37.0 + - imbalanced-learn==0.12.3 + - imgaug==0.4.0 + - implicit==0.7.2 + - inflate64==1.0.1 + - iniconfig==2.0.0 + - ipykernel==6.29.5 + - ipython==8.27.0 + - isodate==0.7.2 + - jedi==0.19.2 + - jinja2==3.1.5 + - jiter==0.8.2 + - joblib==1.4.2 + - jsonlines==4.0.0 + - jsonpatch==1.33 + - jsonpointer==3.0.0 + - jsonschema==4.19.2 + - jsonschema-specifications==2024.10.1 + - jupyter-client==8.6.3 + - jupyter-core==5.7.2 + - kaggle==1.6.17 + - keras==3.5.0 + - kiwisolver==1.4.8 + - kornia==0.6.10 + - kornia-rs==0.1.8 + - langchain==0.2.15 + - langchain-anthropic==0.1.23 + - langchain-core==0.2.43 + - langchain-text-splitters==0.2.4 + - langcodes==3.5.0 + - langsmith==0.1.147 + - language-data==1.3.0 + - lazy-loader==0.4 + - levenshtein==0.25.1 + - libclang==18.1.1 + - librosa==0.10.2.post1 + - lightgbm==4.5.0 + - lightning-utilities==0.12.0 + - littleutils==0.2.4 + - llvmlite==0.43.0 + - loguru==0.7.2 + - lxml==5.3.1 + - mako==1.3.9 + - marisa-trie==1.2.1 + - markdown==3.7 + - markdown-it-py==3.0.0 + - markovify==0.9.4 + - markupsafe==3.0.2 + - marshmallow==3.26.1 + - matplotlib==3.9.2 + - matplotlib-inline==0.1.7 + - mdurl==0.1.2 + - ml-dtypes==0.4.1 + - mpmath==1.3.0 + - msal==1.31.1 + - msal-extensions==1.2.0 + - msgpack==1.1.0 + - msgpack-numpy==0.4.8 + - msrest==0.7.1 + - multidict==6.1.0 + - multiprocess==0.70.17 + - multivolumefile==0.2.3 + - munch==4.0.0 + - murmurhash==1.0.12 + - mypy-extensions==1.0.0 + - namex==0.0.8 + - nbformat==5.10.4 + - nest-asyncio==1.6.0 + - networkx==3.3 + - nltk==3.9.1 + - numba==0.60.0 + - numpy==1.26.2 + - nvidia-cublas-cu12==12.1.3.1 + - nvidia-cuda-cupti-cu12==12.1.105 + - nvidia-cuda-nvcc-cu12==12.3.107 + - nvidia-cuda-nvrtc-cu12==12.1.105 + - nvidia-cuda-runtime-cu12==12.1.105 + - nvidia-cudnn-cu12==8.9.2.26 + - nvidia-cufft-cu12==11.0.2.54 + - nvidia-curand-cu12==10.3.2.106 + - nvidia-cusolver-cu12==11.4.5.107 + - nvidia-cusparse-cu12==12.1.0.106 + - nvidia-nccl-cu12==2.19.3 + - nvidia-nvjitlink-cu12==12.3.101 + - nvidia-nvtx-cu12==12.1.105 + - oauthlib==3.2.2 + - ogb==1.3.6 + - omegaconf==2.3.0 + - openai==1.48.0 + - opencv-python==4.10.0.84 + - opencv-python-headless==4.11.0.86 + - opt-einsum==3.4.0 + - optree==0.14.1 + - optuna==4.0.0 + - orjson==3.10.15 + - outdated==0.2.2 + - packaging==24.2 + - pandas==2.1.4 + - parso==0.8.4 + - pathspec==0.12.1 + - pdf2image==1.17.0 + - peft==0.12.0 + - pexpect==4.9.0 + - pillow==10.4.0 + - platformdirs==4.3.6 + - plotly==5.24.0 + - pluggy==1.5.0 + - pooch==1.8.2 + - portalocker==2.10.1 + - preshed==3.0.9 + - pretrainedmodels==0.7.4 + - prompt-toolkit==3.0.50 + - propcache==0.3.0 + - proto-plus==1.26.0 + - psutil==7.0.0 + - ptyprocess==0.7.0 + - pure-eval==0.2.3 + - py4j==0.10.9.9 + - py7zr==0.22.0 + - pyaml==25.1.0 + - pyarrow==17.0.0 + - pyasn1==0.6.1 + - pyasn1-modules==0.4.1 + - pybcj==1.0.3 + - pycparser==2.22 + - pycryptodomex==3.21.0 + - pydantic==2.9.2 + - pydantic-core==2.23.4 + - pydantic-settings==2.6.1 + - pydicom==2.4.4 + - pygments==2.19.1 + - pyjwt==2.10.1 + - pylibjpeg==2.0.1 + - pyocr==0.8.5 + - pyparsing==3.1.4 + - pypdf==4.3.1 + - pyppmd==1.1.1 + - pytest==7.4.3 + - python-dateutil==2.9.0.post0 + - python-dotenv==1.0.1 + - python-slugify==8.0.4 + - pytorch-lightning==2.4.0 + - pytz==2024.1 + - pyyaml==6.0.2 + - pyzmq==26.2.1 + - pyzstd==0.16.2 + - ranger21==0.1.0 + - rapidfuzz==3.12.2 + - referencing==0.36.2 + - regex==2024.11.6 + - requests==2.31.0 + - requests-oauthlib==2.0.0 + - requests-toolbelt==1.0.0 + - resampy==0.4.3 + - responses==0.18.0 + - rich==13.7.0 + - rouge-score==0.1.2 + - rpds-py==0.23.1 + - rsa==4.9 + - sacrebleu==2.4.3 + - safetensors==0.5.3 + - scikit-image==0.24.0 + - scikit-learn==1.2.2 + - scikit-optimize==0.10.2 + - scikit-surprise==1.1.4 + - scipy==1.11.4 + - seaborn==0.13.2 + - segmentation-models-pytorch==0.3.4 + - sentence-transformers==3.0.1 + - sentencepiece==0.2.0 + - shapely==2.0.7 + - shellingham==1.5.4 + - shutup==0.2.0 + - simsimd==6.2.1 + - six==1.17.0 + - sklearn-pandas==2.2.0 + - smart-open==7.1.0 + - sniffio==1.3.1 + - soundfile==0.13.1 + - soxr==0.5.0.post1 + - spacy==3.7.6 + - spacy-legacy==3.0.12 + - spacy-loggers==1.0.5 + - sqlalchemy==2.0.38 + - srsly==2.5.1 + - stack-data==0.6.3 + - stringzilla==3.12.2 + - sympy==1.13.2 + - tabulate==0.9.0 + - tenacity==8.5.0 + - tensorboard==2.17.1 + - tensorboard-data-server==0.7.2 + - tensorflow==2.17.0 + - tensorflow-hub==0.16.1 + - tensorflow-io-gcs-filesystem==0.37.1 + - tensorpack==0.11 + - termcolor==2.5.0 + - text-unidecode==1.3 + - textblob==0.18.0.post0 + - texttable==1.7.0 + - tf-keras==2.17.0 + - thinc==8.2.5 + - threadpoolctl==3.5.0 + - tifffile==2025.2.18 + - tiktoken==0.7.0 + - timm==0.9.7 + - tokenizers==0.19.1 + - torch==2.2.0 + - torch-geometric==2.3.1 + - torchaudio==2.2.0 + - torchdata==0.7.1 + - torchinfo==1.8.0 + - torchmetrics==1.3.1 + - torchtext==0.17.0 + - torchvision==0.17.0 + - tornado==6.4.2 + - tqdm==4.66.2 + - traitlets==5.14.3 + - transformers==4.44.2 + - triton==2.2.0 + - typer==0.15.2 + - types-python-dateutil==2.9.0.20241206 + - typing-extensions==4.12.2 + - typing-inspect==0.9.0 + - tzdata==2025.1 + - unidecode==1.3.8 + - uritemplate==4.1.1 + - urllib3==2.3.0 + - wasabi==1.1.3 + - wcwidth==0.2.13 + - weasel==0.4.1 + - webencodings==0.5.1 + - werkzeug==3.1.3 + - wrapt==1.17.2 + - xgboost==2.1.1 + - xlrd==2.0.1 + - xxhash==3.5.0 + - yarl==1.18.3 +prefix: /opt/conda/envs/kaggle diff --git a/rdagent/scenarios/kaggle/kaggle_crawler.py b/rdagent/scenarios/kaggle/kaggle_crawler.py index 0fac2c5a..c1f37a80 100644 --- a/rdagent/scenarios/kaggle/kaggle_crawler.py +++ b/rdagent/scenarios/kaggle/kaggle_crawler.py @@ -114,9 +114,9 @@ def download_data(competition: str, settings: ExtendedBaseSettings = KAGGLE_IMPL zipfile_path = f"{local_path}/zip_files" zip_competition_path = Path(zipfile_path) / competition - mleb_env = MLEBDockerEnv() - mleb_env.prepare() if not zip_competition_path.exists(): + mleb_env = MLEBDockerEnv() + mleb_env.prepare() (Path(zipfile_path)).mkdir(parents=True, exist_ok=True) mleb_env.run( f"mlebench prepare -c {competition} --data-dir ./zip_files", @@ -127,6 +127,8 @@ def download_data(competition: str, settings: ExtendedBaseSettings = KAGGLE_IMPL if not (Path(local_path) / competition).exists() or list((Path(local_path) / competition).iterdir()) == []: (Path(local_path) / competition).mkdir(parents=True, exist_ok=True) + mleb_env = MLEBDockerEnv() + mleb_env.prepare() mleb_env.run(f"cp -r ./zip_files/{competition}/prepared/public/* ./{competition}", local_path=local_path) for zip_path in (Path(local_path) / competition).rglob("*.zip"): diff --git a/rdagent/utils/env.py b/rdagent/utils/env.py index dc3336c2..347c3717 100644 --- a/rdagent/utils/env.py +++ b/rdagent/utils/env.py @@ -19,16 +19,17 @@ import zipfile from abc import abstractmethod from pathlib import Path from types import MappingProxyType -from typing import Generic, Mapping, Optional, TypeVar +from typing import Any, Generic, Mapping, Optional, TypeVar import docker # type: ignore[import-untyped] import docker.models # type: ignore[import-untyped] import docker.models.containers # type: ignore[import-untyped] import docker.types # type: ignore[import-untyped] -from pydantic import BaseModel +from pydantic import BaseModel, model_validator from pydantic_settings import SettingsConfigDict from rich import print from rich.console import Console +from rich.pretty import Pretty from rich.progress import Progress, SpinnerColumn, TextColumn from rich.rule import Rule from rich.table import Table @@ -39,23 +40,53 @@ from rdagent.log import rdagent_logger as logger from rdagent.oai.llm_utils import md5_hash from rdagent.utils.workflow import wait_retry -ASpecificBaseModel = TypeVar("ASpecificBaseModel", bound=BaseModel) + +class EnvConf(BaseModel): + default_entry: str + extra_volumes: dict = {} + running_timeout_period: int = 600 # 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 -class Env(Generic[ASpecificBaseModel]): +ASpecificEnvConf = TypeVar("ASpecificEnvConf", bound=EnvConf) + + +class Env(Generic[ASpecificEnvConf]): """ We use BaseModel as the setting due to the features it provides - It provides base typing and checking features. - loading and dumping the information will be easier: for example, we can use package like `pydantic-yaml` """ - conf: ASpecificBaseModel # different env have different conf. + conf: ASpecificEnvConf # different env have different conf. - # last_exit_code: # TODO: get the more concrete information about the exit code. - - def __init__(self, conf: ASpecificBaseModel): + def __init__(self, conf: ASpecificEnvConf): self.conf = conf + def zip_a_folder_into_a_file(self, folder_path: str, zip_file_path: str) -> None: + """ + Zip a folder into a file, use zipfile instead of subprocess + """ + with zipfile.ZipFile(zip_file_path, "w") as z: + for root, _, files in os.walk(folder_path): + for file in files: + z.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), folder_path)) + + def unzip_a_file_into_a_folder(self, zip_file_path: str, folder_path: str) -> None: + """ + Unzip a file into a folder, use zipfile instead of subprocess + """ + # Clear folder_path before extracting + if os.path.exists(folder_path): + shutil.rmtree(folder_path) + os.makedirs(folder_path) + + with zipfile.ZipFile(zip_file_path, "r") as z: + z.extractall(folder_path) + @abstractmethod def prepare(self, *args, **kwargs) -> None: # type: ignore[no-untyped-def] """ @@ -86,9 +117,35 @@ class Env(Generic[ASpecificBaseModel]): stdout, _ = self.run_ret_code(entry=entry, local_path=local_path, env=env, **kwargs) return stdout - @abstractmethod + def __run_ret_code_with_retry( + self, + entry: str | None = None, + local_path: str = ".", + env: dict | None = None, + running_extra_volume: Mapping = MappingProxyType({}), + remove_timestamp: bool = True, + ) -> tuple[str, int]: + # TODO: remove_timestamp can be implemented in a shallower way... + for retry_index in range(self.conf.retry_count + 1): + try: + return self._run_ret_code( + entry, local_path, env, running_extra_volume=running_extra_volume, remove_timestamp=remove_timestamp + ) + except Exception as e: + if retry_index == self.conf.retry_count: + raise + logger.warning( + f"Error while running the container: {e}, current try index: {retry_index + 1}, {self.conf.retry_count - retry_index - 1} retries left." + ) + time.sleep(self.conf.retry_wait_seconds) + raise RuntimeError # for passing CI + def run_ret_code( - self, entry: str | None, local_path: str = ".", env: dict | None = None, **kwargs: dict + self, + entry: str | None = None, + local_path: str = ".", + env: dict | None = None, + **kwargs: dict, ) -> tuple[str, int]: """ Run the folder under the environment and return both the stdout and the exit code. @@ -110,56 +167,205 @@ class Env(Generic[ASpecificBaseModel]): ------- A tuple containing the stdout and the exit code """ + running_extra_volume = kwargs.get("running_extra_volume", {}) + if entry is None: + entry = self.conf.default_entry + entry_add_timeout = ( + f"/bin/sh -c 'timeout {self.conf.running_timeout_period} {entry}; " + + "entry_exit_code=$?; " + + (f"chmod -R 777 {self.conf.mount_path}; " if hasattr(self.conf, "mount_path") else "") + + "exit $entry_exit_code'" + ) + + if self.conf.enable_cache: + stdout, return_code = self.cached_run(entry_add_timeout, local_path, env, running_extra_volume) + else: + stdout, return_code = self.__run_ret_code_with_retry( + entry_add_timeout, local_path, env, running_extra_volume, remove_timestamp=False + ) + + return stdout, return_code + + def cached_run( + self, + entry: str | None = None, + local_path: str = ".", + env: dict | None = None, + running_extra_volume: Mapping = MappingProxyType({}), + remove_timestamp: bool = True, + ) -> tuple[str, int]: + """ + Run the folder under the environment. + Will cache the output and the folder diff for next round of running. + Use the python codes and the parameters(entry, running_extra_volume) as key to hash the input. + """ + target_folder = Path(RD_AGENT_SETTINGS.pickle_cache_folder_path_str) / f"utils.env.run" + target_folder.mkdir(parents=True, exist_ok=True) + + # we must add the information of data (beyound code) into the key. + # Otherwise, all commands operating on data will become invalue (e.g. rm -r submission.csv) + # So we recursively walk in the folder and add the sorted relative filename list as part of the key. + data_key = [] + for path in Path(local_path).rglob("*"): + p = str(path.relative_to(Path(local_path))) + if p.startswith("__pycache__"): + continue + data_key.append(p) + data_key = sorted(data_key) + + key = md5_hash( + json.dumps( + [ + [str(path.relative_to(Path(local_path))), path.read_text()] + for path in sorted(Path(local_path).rglob("*.py")) + ] + ) + + json.dumps({"entry": entry, "running_extra_volume": dict(running_extra_volume)}) + + json.dumps({"extra_volumes": self.conf.extra_volumes}) + + json.dumps(data_key) + ) + if Path(target_folder / f"{key}.pkl").exists() and Path(target_folder / f"{key}.zip").exists(): + with open(target_folder / f"{key}.pkl", "rb") as f: + ret: tuple[str, int] = pickle.load(f) + self.unzip_a_file_into_a_folder(str(target_folder / f"{key}.zip"), local_path) + else: + ret = self.__run_ret_code_with_retry(entry, local_path, env, running_extra_volume, remove_timestamp) + with open(target_folder / f"{key}.pkl", "wb") as f: + pickle.dump(ret, f) + self.zip_a_folder_into_a_file(local_path, str(target_folder / f"{key}.zip")) + return ret + + @abstractmethod + def _run_ret_code( + self, + entry: str | None, + local_path: str = ".", + env: dict | None = None, + running_extra_volume: Mapping = MappingProxyType({}), + **kwargs: Any, + ) -> tuple[str, int]: + """ + Execute the specified entry point within the given environment and local path. + + Parameters + ---------- + entry : str | None + The entry point to execute. If None, defaults to the configured entry. + local_path : str + The local directory path where the execution should occur. + env : dict | None + Environment variables to set during execution. + kwargs : dict + Additional keyword arguments for execution customization. + + Returns + ------- + tuple[str, int] + A tuple containing the standard output and the exit code of the execution. + """ + pass + + +# class EnvWithCache +# ## Local Environment ----- -class LocalConf(BaseModel): - py_bin: str - default_entry: str +class LocalConf(EnvConf): + bin_path: str = "" + """path like ::, which will be prepend to bin path.""" + + retry_count: int = 0 # retry count for; run `retry_count + 1` times -class LocalEnv(Env[LocalConf]): +ASpecificLocalConf = TypeVar("ASpecificLocalConf", bound=LocalConf) + + +class LocalEnv(Env[ASpecificLocalConf]): """ Sometimes local environment may be more convinient for testing """ - def prepare(self) -> None: - if not (Path("~/.qlib/qlib_data/cn_data").expanduser().resolve().exists()): - self.run( - entry="python -m qlib.run.get_data qlib_data --target_dir ~/.qlib/qlib_data/cn_data --region cn", - ) - else: - print("Data already exists. Download skipped.") + def prepare(self) -> None: ... - def run_ret_code( + def _run_ret_code( self, entry: str | None = None, local_path: str | None = None, env: dict | None = None, + running_extra_volume: Mapping = MappingProxyType({}), **kwargs: dict, ) -> tuple[str, int]: + + # mocking the volumns + volumns = {} + if self.conf.extra_volumes is not None: + for lp, rp in self.conf.extra_volumes.items(): + volumns[lp] = rp + for lp, rp in running_extra_volume.items(): + volumns[lp] = rp + + for rp, lp in volumns.items(): + link_path = Path(lp) + real_path = Path(rp) + if not link_path.parent.exists(): + link_path.parent.mkdir(parents=True, exist_ok=True) + if link_path.exists() or link_path.is_symlink(): + link_path.unlink() + link_path.symlink_to(real_path) + if env is None: env = {} + path = [*self.conf.bin_path.split(":"), "/bin/", "/usr/bin/", *env.get("PATH", "").split(":")] + env["PATH"] = ":".join(path) + if entry is None: entry = self.conf.default_entry - command = str(Path(self.conf.py_bin).joinpath(entry)).split(" ") + summary = { + "entry": entry, + "local_path": local_path, + "env": env, + "volumes": volumns, + } + print(Pretty(summary)) cwd = None if local_path: cwd = Path(local_path).resolve() - result = subprocess.run(command, cwd=cwd, env={**os.environ, **env}, capture_output=True, text=True) - return result.stdout, result.returncode + result = subprocess.run(entry, cwd=cwd, env={**os.environ, **env}, capture_output=True, text=True, shell=True) + combined_output = result.stderr + result.stdout # Combine stdout and stderr + print(combined_output) # Display the combined output in the console + + return combined_output, result.returncode + + +class CondaConf(LocalConf): + conda_env_name: str + default_entry: str = "python main.py" + + @model_validator(mode="after") + def change_bin_path(self, **data: Any) -> "CondaConf": + conda_path_result = subprocess.run( + f"conda run -n {self.conda_env_name} --no-capture-output env | grep '^PATH='", + capture_output=True, + text=True, + shell=True, + ) + self.bin_path = conda_path_result.stdout.strip().split("=")[1] if conda_path_result.returncode == 0 else "" + return self + + +class MLECondaConf(CondaConf): + enable_cache: bool = False # aligning with the docker settings. ## Docker Environment ----- - - -class DockerConf(ExtendedBaseSettings): +class DockerConf(EnvConf, ExtendedBaseSettings): build_from_dockerfile: bool = False dockerfile_folder_path: Optional[Path] = ( None # the path to the dockerfile optional path provided when build_from_dockerfile is False @@ -369,13 +575,14 @@ class DockerEnv(Env[DockerConf]): output_string = re.sub(datetime_pattern, "[DATETIME]", input_string) return output_string - def __run_ret_code( + def _run_ret_code( self, entry: str | None = None, local_path: str = ".", env: dict | None = None, running_extra_volume: Mapping = MappingProxyType({}), remove_timestamp: bool = True, + **kwargs: Any, ) -> tuple[str, int]: if env is None: env = {} @@ -446,121 +653,6 @@ class DockerEnv(Env[DockerConf]): except docker.errors.APIError as e: raise RuntimeError(f"Error while running the container: {e}") - def __run_ret_code_with_retry( - self, - entry: str | None = None, - local_path: str = ".", - env: dict | None = None, - running_extra_volume: Mapping = MappingProxyType({}), - remove_timestamp: bool = True, - ) -> tuple[str, int]: - for retry_index in range(self.conf.retry_count): - try: - return self.__run_ret_code(entry, local_path, env, running_extra_volume, remove_timestamp) - except Exception as e: - logger.warning( - f"Error while running the container: {e}, current try index: {retry_index + 1}, {self.conf.retry_count - retry_index - 1} retries left." - ) - time.sleep(self.conf.retry_wait_seconds) - raise RuntimeError("Error while running the container. Retry count exceeded.") - - def zip_a_folder_into_a_file(self, folder_path: str, zip_file_path: str) -> None: - """ - Zip a folder into a file, use zipfile instead of subprocess - """ - with zipfile.ZipFile(zip_file_path, "w") as z: - for root, _, files in os.walk(folder_path): - for file in files: - z.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), folder_path)) - - def unzip_a_file_into_a_folder(self, zip_file_path: str, folder_path: str) -> None: - """ - Unzip a file into a folder, use zipfile instead of subprocess - """ - # Clear folder_path before extracting - if os.path.exists(folder_path): - shutil.rmtree(folder_path) - os.makedirs(folder_path) - - with zipfile.ZipFile(zip_file_path, "r") as z: - z.extractall(folder_path) - - def cached_run( - self, - entry: str | None = None, - local_path: str = ".", - env: dict | None = None, - running_extra_volume: Mapping = MappingProxyType({}), - remove_timestamp: bool = True, - ) -> tuple[str, int]: - """ - Run the folder under the environment. - Will cache the output and the folder diff for next round of running. - Use the python codes and the parameters(entry, running_extra_volume) as key to hash the input. - """ - target_folder = Path(RD_AGENT_SETTINGS.pickle_cache_folder_path_str) / f"utils.env.run" - target_folder.mkdir(parents=True, exist_ok=True) - - # we must add the information of data (beyound code) into the key. - # Otherwise, all commands operating on data will become invalue (e.g. rm -r submission.csv) - # So we recursively walk in the folder and add the sorted relative filename list as part of the key. - data_key = [] - for path in Path(local_path).rglob("*"): - p = str(path.relative_to(Path(local_path))) - if p.startswith("__pycache__"): - continue - data_key.append(p) - data_key = sorted(data_key) - - key = md5_hash( - json.dumps( - [ - [str(path.relative_to(Path(local_path))), path.read_text()] - for path in sorted(Path(local_path).rglob("*.py")) - ] - ) - + json.dumps({"entry": entry, "running_extra_volume": dict(running_extra_volume)}) - + json.dumps({"extra_volumes": self.conf.extra_volumes}) - + json.dumps(data_key) - ) - if Path(target_folder / f"{key}.pkl").exists() and Path(target_folder / f"{key}.zip").exists(): - with open(target_folder / f"{key}.pkl", "rb") as f: - ret: tuple[str, int] = pickle.load(f) - self.unzip_a_file_into_a_folder(str(target_folder / f"{key}.zip"), local_path) - else: - ret = self.__run_ret_code_with_retry(entry, local_path, env, running_extra_volume, remove_timestamp) - with open(target_folder / f"{key}.pkl", "wb") as f: - pickle.dump(ret, f) - self.zip_a_folder_into_a_file(local_path, str(target_folder / f"{key}.zip")) - return ret - - def run_ret_code( - self, - entry: str | None = None, - local_path: str = ".", - env: dict | None = None, - **kwargs: dict, - ) -> tuple[str, int]: - running_extra_volume = kwargs.get("running_extra_volume", {}) - if entry is None: - entry = self.conf.default_entry - - entry_add_timeout = ( - f"/bin/sh -c 'timeout {self.conf.running_timeout_period} {entry}; " - f"entry_exit_code=$?; " - f"chmod -R 777 {self.conf.mount_path}; " - f"exit $entry_exit_code'" - ) - - if self.conf.enable_cache: - stdout, return_code = self.cached_run(entry_add_timeout, local_path, env, running_extra_volume) - else: - stdout, return_code = self.__run_ret_code_with_retry( - entry_add_timeout, local_path, env, running_extra_volume, remove_timestamp=False - ) - - return stdout, return_code - def dump_python_code_run_and_get_results( self, code: str, diff --git a/test/utils/test_env.py b/test/utils/test_env.py index 976814d9..66a1c7c0 100644 --- a/test/utils/test_env.py +++ b/test/utils/test_env.py @@ -6,11 +6,27 @@ from pathlib import Path sys.path.append(str(Path(__file__).resolve().parent.parent)) import shutil -from rdagent.utils.env import LocalConf, LocalEnv, QlibDockerConf, QTDockerEnv +from rdagent.utils.env import ( + CondaConf, + LocalConf, + LocalEnv, + QlibDockerConf, + QTDockerEnv, +) DIRNAME = Path(__file__).absolute().resolve().parent +class QlibLocalEnv(LocalEnv): + def prepare(self) -> None: + if not (Path("~/.qlib/qlib_data/cn_data").expanduser().resolve().exists()): + self.run( + entry="python -m qlib.run.get_data qlib_data --target_dir ~/.qlib/qlib_data/cn_data --region cn", + ) + else: + print("Data already exists. Download skipped.") + + class EnvUtils(unittest.TestCase): def setUp(self): self.test_workspace = DIRNAME / "test_workspace" @@ -24,16 +40,35 @@ class EnvUtils(unittest.TestCase): # NOTE: Because you need to download the data during the prepare process. So you need to have pyqlib in your environment. def test_local(self): local_conf = LocalConf( - py_bin="/home/v-linlanglv/miniconda3/envs/RD-Agent-310/bin", + bin_path="/home/v-linlanglv/miniconda3/envs/RD-Agent-310/bin", default_entry="qrun conf.yaml", ) - qle = LocalEnv(conf=local_conf) + qle = QlibLocalEnv(conf=local_conf) qle.prepare() conf_path = str(DIRNAME / "env_tpl" / "conf.yaml") qle.run(entry="qrun " + conf_path) mlrun_p = DIRNAME / "env_tpl" / "mlruns" self.assertTrue(mlrun_p.exists(), f"Expected output file {mlrun_p} not found") + def test_local_simple(self): + local_conf = LocalConf(bin_path="/home/xiaoyang/miniconda3/bin/", default_entry="which python") + le = LocalEnv(conf=local_conf) + print(local_conf) + le.prepare() + code_path = DIRNAME / "tmp_code" + code_path.mkdir(exist_ok=True) + res, code = le.run_ret_code(local_path=str(code_path)) + print(res, code) + + def test_conda_simple(self): + conda_conf = CondaConf(default_entry="which python", conda_env_name="MLE") + le = LocalEnv(conf=conda_conf) + le.prepare() + code_path = DIRNAME / "tmp_code" + code_path.mkdir(exist_ok=True) + res, code = le.run_ret_code(local_path=str(code_path)) + print(res, code) + def test_docker(self): """We will mount `env_tpl` into the docker image. And run the docker image with `qrun conf.yaml`