diff --git a/rdagent/utils/env.py b/rdagent/utils/env.py index 347c3717..41de3ce8 100644 --- a/rdagent/utils/env.py +++ b/rdagent/utils/env.py @@ -41,7 +41,7 @@ from rdagent.oai.llm_utils import md5_hash from rdagent.utils.workflow import wait_retry -class EnvConf(BaseModel): +class EnvConf(ExtendedBaseSettings): default_entry: str extra_volumes: dict = {} running_timeout_period: int = 600 # 10 minutes @@ -365,7 +365,7 @@ class MLECondaConf(CondaConf): ## Docker Environment ----- -class DockerConf(EnvConf, ExtendedBaseSettings): +class DockerConf(EnvConf): 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 diff --git a/test/utils/test_conf.py b/test/utils/test_conf.py index 029976d4..79b0e8ff 100644 --- a/test/utils/test_conf.py +++ b/test/utils/test_conf.py @@ -6,13 +6,23 @@ class ConfUtils(unittest.TestCase): def test_conf(self): import os - from rdagent.utils.env import QlibDockerConf + from rdagent.utils.env import EnvConf, QlibDockerConf os.environ["MEM_LIMIT"] = "200g" assert QlibDockerConf().mem_limit == "200g" # base class will affect subclasses os.environ["QLIB_DOCKER_MEM_LIMIT"] = "300g" assert QlibDockerConf().mem_limit == "300g" # more accurate subclass will override the base class + os.environ["default_entry"] = "which python" + os.environ["ENABLE_CACHE"] = "False" + + assert EnvConf().enable_cache is False + assert QlibDockerConf().enable_cache is False + + os.environ["ENABLE_CACHE"] = "True" + assert EnvConf().enable_cache is True + assert QlibDockerConf().enable_cache is True + if __name__ == "__main__": unittest.main()