gavindiaz/NexQuant
1
0
Fork 0
mirror of https://github.com/NicolasBohn/NexQuant.git synced 2026-07-27 23:47:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: condaenv & full docker env (#668)

Browse Source 
* 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 <peteryang@vip.qq.com>
This commit is contained in:
you-n-g
2025-03-12 11:36:28 +08:00
committed by GitHub
parent 2b9763c4a7
commit 20778a9b87
20 changed files with 835 additions and 247 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1 -1
View File
@@ -170,6 +170,6 @@ mlruns/
# shell script
*.out
*.sh
/*.sh
.aider*
rdagent/app/benchmark/factor/example.json
rdagent/components/coder/data_science/conf.py
+40
View File
@@ -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
rdagent/components/coder/data_science/ensemble/eval.py
+5 -9
View File
@@ -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
rdagent/components/coder/data_science/feature/eval.py
+5 -10
View File
@@ -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
rdagent/components/coder/data_science/model/eval.py
+5 -8
View File
@@ -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
rdagent/components/coder/data_science/model/exp.py
-8
View File
@@ -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
rdagent/components/coder/data_science/raw_data_loader/__init__.py
+8 -6
View File
@@ -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
rdagent/components/coder/data_science/raw_data_loader/eval.py
+5 -9
View File
@@ -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
rdagent/components/coder/data_science/raw_data_loader/exp.py
-11
View File
@@ -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
rdagent/components/coder/data_science/workflow/eval.py
+6 -11
View File
@@ -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
rdagent/log/mle_summary.py
+3 -5
View File
@@ -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()
rdagent/scenarios/data_science/dev/runner/__init__.py
+3 -5
View File
@@ -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 = (
rdagent/scenarios/data_science/dev/runner/eval.py
+12 -11
View File
@@ -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())
rdagent/scenarios/data_science/scen/__init__.py
+3 -4
View File
@@ -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:
rdagent/scenarios/data_science/sing_docker/Dockerfile
+47
View File
@@ -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/
rdagent/scenarios/data_science/sing_docker/entrypoint.sh
Executable
+47
View File
@@ -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
rdagent/scenarios/data_science/sing_docker/kaggle_environment.yaml
+367
View File
@@ -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
rdagent/scenarios/kaggle/kaggle_crawler.py
+4 -2
View File
@@ -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"):
rdagent/utils/env.py
+236 -144
View File
@@ -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 <path1>:<path2>:<path3>, 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,
test/utils/test_env.py
+38 -3
View File
@@ -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`