Files
NexQuant/rdagent/components/coder/data_science/conf.py
T
TPTBusiness 64ed6b0cce fix(security): resolve path-injection, B701, B101, B112 Bandit alerts
- Path injection (B614): centralized safe_resolve_path in core/utils.py,
  refactored 6 UI modules to use it with safe_root validation
- B701: added explicit autoescape=select_autoescape() to Jinja2
  Environment() calls in 3 files
- B101: replaced assert statements with proper if/raise patterns in
  12+ files (partial)
- B112: added logger.warning() to bare except:continue blocks in
  5 files
2026-05-01 13:42:59 +02:00

90 lines
2.8 KiB
Python

from typing import Literal
from rdagent.app.data_science.conf import DS_RD_SETTING
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_multiplier: int = 4
env_type: str = "docker"
# TODO: extract a function for env and conf.
extra_evaluator: list[str] = []
"""Extra evaluators to use"""
extra_eval: list[str] = []
"""
Extra evaluators
The evaluator follows the following assumptions:
- It runs after previous evaluator (So the running results are already there)
It is not a complete feature due to it is only implemented in DS Pipeline & Coder.
TODO: The complete version should be implemented in the CoSTEERSettings.
"""
def get_ds_env(
conf_type: Literal["kaggle", "mlebench"] = "kaggle",
extra_volumes: dict = {},
running_timeout_period: int | None = DS_RD_SETTING.debug_timeout,
enable_cache: bool | None = None,
) -> Env:
"""
Retrieve the appropriate environment configuration based on the env_type setting.
Returns:
Env: An instance of the environment configured either as DockerEnv or LocalEnv.
Raises:
ValueError: If the env_type is not recognized.
"""
conf = DSCoderCoSTEERSettings()
if conf_type not in ["kaggle", "mlebench"]:
raise ValueError(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}")
env.conf.extra_volumes = extra_volumes.copy()
env.conf.running_timeout_period = running_timeout_period
if enable_cache is not None:
env.conf.enable_cache = enable_cache
env.prepare()
return env
def get_clear_ws_cmd(stage: Literal["before_training", "before_inference"] = "before_training") -> str:
"""
Clean the files in workspace to a specific stage
"""
if stage not in ["before_training", "before_inference"]:
raise ValueError(f"Unknown stage: {stage}")
if DS_RD_SETTING.enable_model_dump and stage == "before_training":
cmd = "rm -r submission.csv scores.csv models trace.log"
else:
cmd = "rm submission.csv scores.csv trace.log"
return cmd