From 0631de51036d9105030a895dbb848bd325db5cb0 Mon Sep 17 00:00:00 2001 From: you-n-g Date: Thu, 20 Feb 2025 00:42:10 +0800 Subject: [PATCH] refactor: add run_ret_code method and update run method to use it (#623) * refactor: Add run_ret_code method and update run method to use it * feat: Add kwargs support to run methods and test for run_ret_code * fix: preserve exit code after chmod in DockerEnv entry command * chore: Change file permissions from 755 to 644 in env_tpl directory * refactor: Return execution code and update evaluator logic * lint * refactor: Use MappingProxyType for running_extra_volume in DockerEnv methods * lint --- .../coder/data_science/ensemble/eval.py | 10 +- rdagent/core/experiment.py | 28 +++-- rdagent/utils/env.py | 103 ++++++++++++------ test/utils/test_env.py | 33 +++++- 4 files changed, 120 insertions(+), 54 deletions(-) diff --git a/rdagent/components/coder/data_science/ensemble/eval.py b/rdagent/components/coder/data_science/ensemble/eval.py index 6e46c34f..6abe6e72 100644 --- a/rdagent/components/coder/data_science/ensemble/eval.py +++ b/rdagent/components/coder/data_science/ensemble/eval.py @@ -63,7 +63,9 @@ class EnsembleCoSTEEREvaluator(CoSTEEREvaluator): ) implementation.inject_files(**{fname: test_code}) - stdout = implementation.execute(env=de, entry=f"python {fname}") + stdout, ret_code = implementation.execute_ret_code(env=de, 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") @@ -81,6 +83,6 @@ class EnsembleCoSTEEREvaluator(CoSTEEREvaluator): stdout=stdout, workflow_stdout=workflow_stdout, ) - return build_cls_from_json_with_retry( - EnsembleEvalFeedback, system_prompt=system_prompt, user_prompt=user_prompt - ) + efb = build_cls_from_json_with_retry(EnsembleEvalFeedback, system_prompt=system_prompt, user_prompt=user_prompt) + efb.final_decision = efb.final_decision and ret_code == 0 + return efb diff --git a/rdagent/core/experiment.py b/rdagent/core/experiment.py index 96f08352..7d94cf8b 100644 --- a/rdagent/core/experiment.py +++ b/rdagent/core/experiment.py @@ -232,21 +232,29 @@ class FBWorkspace(Workspace): shutil.rmtree(self.workspace_path, ignore_errors=True) self.file_dict = {} - def execute(self, env: Env | None = None, entry: str | None = None) -> object | None: + def execute(self, env: Env, entry: str) -> str: """ - Before each execution, make sure to prepare and inject code + Before each execution, make sure to prepare and inject code. + """ + stdout, _ = self.execute_ret_code(env, entry) + return stdout + + def execute_ret_code(self, env: Env, entry: str) -> tuple[str, int]: + """ + Execute the code in the environment and return both the stdout and the exit code. + + Before each execution, make sure to prepare and inject code. """ self.prepare() self.inject_files(**self.file_dict) - # TODO: env should be not None in new design (no code can run without environment) - if env is not None and entry is not None: - # We believe that removing the progress bar might resolve the issue. - # If needed, cutting off the information in the middle should be a backup plan. - return shrink_text( - filter_progress_bar(env.run(entry, str(self.workspace_path))), + stdout, return_code = env.run_ret_code(entry, str(self.workspace_path)) + return ( + shrink_text( + filter_progress_bar(stdout), context_lines=RD_AGENT_SETTINGS.stdout_context_len, - ) - return None + ), + return_code, + ) def __str__(self) -> str: return f"Workspace[{self.workspace_path=}" + ( diff --git a/rdagent/utils/env.py b/rdagent/utils/env.py index 11caae5d..8017e1c7 100644 --- a/rdagent/utils/env.py +++ b/rdagent/utils/env.py @@ -18,7 +18,8 @@ import uuid import zipfile from abc import abstractmethod from pathlib import Path -from typing import Generic, Optional, TypeVar +from types import MappingProxyType +from typing import Generic, Mapping, Optional, TypeVar import docker # type: ignore[import-untyped] import docker.models # type: ignore[import-untyped] @@ -48,6 +49,7 @@ class Env(Generic[ASpecificBaseModel]): """ conf: ASpecificBaseModel # different env have different conf. + # last_exit_code: # TODO: get the more concrete information about the exit code. def __init__(self, conf: ASpecificBaseModel): @@ -59,8 +61,7 @@ class Env(Generic[ASpecificBaseModel]): Prepare for the environment based on it's configure """ - @abstractmethod - def run(self, entry: str | None, local_path: str = ".", env: dict | None = None) -> str: + def run(self, entry: str | None, local_path: str = ".", env: dict | None = None, **kwargs: dict) -> str: """ Run the folder under the environment. @@ -81,6 +82,33 @@ class Env(Generic[ASpecificBaseModel]): ------- the stdout """ + stdout, _ = self.run_ret_code(entry=entry, local_path=local_path, env=env, **kwargs) + return stdout + + @abstractmethod + def run_ret_code( + self, entry: str | 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. + + Parameters + ---------- + entry : str | None + We may we the entry point when we run it. + For example, we may have different entries when we run and summarize the project. + local_path : str | None + the local path (to project, mainly for code) will be mounted into the docker + Here are some examples for a None local path + - for example, run docker for updating the data in the extra_volumes. + - simply run the image. The results are produced by output or network + env : dict | None + Run the code with your specific environment. + + Returns + ------- + A tuple containing the stdout and the exit code + """ ## Local Environment ----- @@ -104,7 +132,13 @@ class LocalEnv(Env[LocalConf]): else: print("Data already exists. Download skipped.") - def run(self, entry: str | None = None, local_path: Optional[str] = None, env: dict | None = None) -> str: + def run_ret_code( + self, + entry: str | None = None, + local_path: str | None = None, + env: dict | None = None, + **kwargs: dict, + ) -> tuple[str, int]: if env is None: env = {} @@ -118,10 +152,7 @@ class LocalEnv(Env[LocalConf]): cwd = Path(local_path).resolve() result = subprocess.run(command, cwd=cwd, env={**os.environ, **env}, capture_output=True, text=True) - if result.returncode != 0: - raise RuntimeError(f"Error while running the command: {result.stderr}") - - return result.stdout + return result.stdout, result.returncode ## Docker Environment ----- @@ -337,14 +368,14 @@ class DockerEnv(Env[DockerConf]): output_string = re.sub(datetime_pattern, "[DATETIME]", input_string) return output_string - def __run( + def __run_ret_code( self, entry: str | None = None, local_path: str = ".", env: dict | None = None, - running_extra_volume: dict | None = None, + running_extra_volume: Mapping = MappingProxyType({}), remove_timestamp: bool = True, - ) -> str: + ) -> tuple[str, int]: if env is None: env = {} env["PYTHONWARNINGS"] = "ignore" @@ -359,9 +390,8 @@ class DockerEnv(Env[DockerConf]): if self.conf.extra_volumes is not None: for lp, rp in self.conf.extra_volumes.items(): volumns[lp] = {"bind": rp, "mode": self.conf.extra_volume_mode} - if running_extra_volume is not None: - for lp, rp in running_extra_volume.items(): - volumns[lp] = {"bind": rp, "mode": self.conf.extra_volume_mode} + for lp, rp in running_extra_volume.items(): + volumns[lp] = {"bind": rp, "mode": self.conf.extra_volume_mode} log_output = "" @@ -397,7 +427,7 @@ class DockerEnv(Env[DockerConf]): decoded_log = self.replace_time_info(decoded_log) if remove_timestamp else decoded_log Console().print(decoded_log, markup=False) log_output += decoded_log + "\n" - container.wait() + exit_status = container.wait()["StatusCode"] container.stop() container.remove() end = time.time() @@ -407,7 +437,7 @@ class DockerEnv(Env[DockerConf]): ) log_output += f"\n\nThe running time exceeds {self.conf.running_timeout_period} seconds, so the process is killed." print(Rule("[bold green]Docker Logs End[/bold green]", style="dark_orange")) - return log_output + return log_output, exit_status except docker.errors.ContainerError as e: raise RuntimeError(f"Error while running the container: {e}") except docker.errors.ImageNotFound: @@ -415,17 +445,17 @@ class DockerEnv(Env[DockerConf]): except docker.errors.APIError as e: raise RuntimeError(f"Error while running the container: {e}") - def __run_with_retry( + def __run_ret_code_with_retry( self, entry: str | None = None, local_path: str = ".", env: dict | None = None, - running_extra_volume: dict | None = None, + running_extra_volume: Mapping = MappingProxyType({}), remove_timestamp: bool = True, - ) -> str: + ) -> tuple[str, int]: for retry_index in range(self.conf.retry_count): try: - return self.__run(entry, local_path, env, running_extra_volume, remove_timestamp) + 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." @@ -459,9 +489,9 @@ class DockerEnv(Env[DockerConf]): entry: str | None = None, local_path: str = ".", env: dict | None = None, - running_extra_volume: dict | None = None, + running_extra_volume: Mapping = MappingProxyType({}), remove_timestamp: bool = True, - ) -> str: + ) -> tuple[str, int]: """ Run the folder under the environment. Will cache the output and the folder diff for next round of running. @@ -488,42 +518,47 @@ class DockerEnv(Env[DockerConf]): for path in sorted(Path(local_path).rglob("*.py")) ] ) - + json.dumps({"entry": entry, "running_extra_volume": running_extra_volume}) + + 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: str = pickle.load(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_with_retry(entry, local_path, env, running_extra_volume, remove_timestamp) + 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( + def run_ret_code( self, entry: str | None = None, local_path: str = ".", env: dict | None = None, - running_extra_volume: dict | None = None, - ) -> str: + **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}; chmod -R 777 {self.conf.mount_path}'" + 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: - out = self.cached_run(entry_add_timeout, local_path, env, running_extra_volume) + stdout, return_code = self.cached_run(entry_add_timeout, local_path, env, running_extra_volume) else: - out = self.__run_with_retry( + stdout, return_code = self.__run_ret_code_with_retry( entry_add_timeout, local_path, env, running_extra_volume, remove_timestamp=False ) - return out + return stdout, return_code def dump_python_code_run_and_get_results( self, @@ -531,7 +566,7 @@ class DockerEnv(Env[DockerConf]): dump_file_names: list[str], local_path: str, env: dict | None = None, - running_extra_volume: dict | None = None, + running_extra_volume: Mapping = MappingProxyType({}), code_dump_file_py_name: Optional[str] = None, ) -> tuple[str, list]: """ @@ -541,7 +576,7 @@ class DockerEnv(Env[DockerConf]): with open(os.path.join(local_path, random_file_name), "w") as f: f.write(code) entry = f"python {random_file_name}" - log_output = self.run(entry, local_path, env, running_extra_volume=running_extra_volume) + log_output = self.run(entry, local_path, env, running_extra_volume=dict(running_extra_volume)) results = [] os.remove(os.path.join(local_path, random_file_name)) for name in dump_file_names: diff --git a/test/utils/test_env.py b/test/utils/test_env.py index 54d005e5..976814d9 100644 --- a/test/utils/test_env.py +++ b/test/utils/test_env.py @@ -13,14 +13,12 @@ DIRNAME = Path(__file__).absolute().resolve().parent class EnvUtils(unittest.TestCase): def setUp(self): - pass + self.test_workspace = DIRNAME / "test_workspace" + self.test_workspace.mkdir(exist_ok=True) def tearDown(self): - # NOTE: For a docker file, the output are generated with root permission. - # mlrun_p = DIRNAME / "env_tpl" / "mlruns" - # if mlrun_p.exists(): - # shutil.rmtree(mlrun_p) - ... + if self.test_workspace.exists(): + shutil.rmtree(self.test_workspace) # NOTE: Since I don't know the exact environment in which it will be used, here's just an example. # NOTE: Because you need to download the data during the prepare process. So you need to have pyqlib in your environment. @@ -53,6 +51,29 @@ class EnvUtils(unittest.TestCase): result = qtde.run(local_path=str(DIRNAME / "env_tpl"), entry="python read_exp_res.py") print(result) + def test_run_ret_code(self): + """Test the run_ret_code method of QTDockerEnv with both valid and invalid commands.""" + qtde = QTDockerEnv() + qtde.prepare() + + # Test with a valid command + result, return_code = qtde.run_ret_code(entry='echo "Hello, World!"', local_path=str(self.test_workspace)) + print(return_code) + assert return_code == 0, f"Expected return code 0, but got {return_code}" + assert "Hello, World!" in result, "Expected output not found in result" + + # Test with an invalid command + _, return_code = qtde.run_ret_code(entry="invalid_command", local_path=str(self.test_workspace)) + print(return_code) + assert return_code != 0, "Expected non-zero return code for invalid command" + + dc = QlibDockerConf() + dc.running_timeout_period = 1 + qtde = QTDockerEnv(dc) + result, return_code = qtde.run_ret_code(entry="sleep 2", local_path=str(self.test_workspace)) + print(result) + assert return_code == 124, "Expected return code 124 for timeout" + def test_docker_mem(self): cmd = 'python -c \'print("start"); import numpy as np; size_mb = 500; size = size_mb * 1024 * 1024 // 8; array = np.random.randn(size).astype(np.float64); print("success")\''