Files
NexQuant/rdagent/components/coder/data_science/conf.py
T
you-n-g 918b0b0d43 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 <peteryang@vip.qq.com>
2025-03-12 11:36:28 +08:00

51 lines
1.4 KiB
Python

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):
"""Data Science CoSTEER settings"""
class Config:
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