From 667af3e1eaaa7f735be2de18830fc26676c93d58 Mon Sep 17 00:00:00 2001 From: Yuante Li <104308117+WinstonLiyt@users.noreply.github.com> Date: Wed, 2 Jul 2025 15:11:18 +0800 Subject: [PATCH] feat: added running time statistics for the DS scenario experiment (#1007) * added running time statistics for the DS scenario experiment * update execute_ret_code to return running_time * fix * fix * update describe * add EnvResult * update corresponding calls * add RunningInfo class * fix * fix * fix * fix ci * rename function name * fix ci * fix * refine running_time logic * fix ci --- .../coder/data_science/ensemble/eval.py | 4 +- .../coder/data_science/feature/eval.py | 8 +-- .../coder/data_science/model/eval.py | 6 ++- .../coder/data_science/pipeline/eval.py | 12 +++-- .../data_science/raw_data_loader/eval.py | 4 +- .../coder/data_science/workflow/eval.py | 6 +-- rdagent/core/experiment.py | 47 +++++++++++------ .../scenarios/data_science/dev/prompts.yaml | 5 +- .../data_science/dev/runner/__init__.py | 1 + .../scenarios/data_science/dev/runner/eval.py | 6 ++- rdagent/scenarios/data_science/share.yaml | 3 ++ rdagent/scenarios/data_science/test_eval.py | 14 +++--- .../scenarios/kaggle/experiment/workspace.py | 2 +- rdagent/scenarios/kaggle/kaggle_crawler.py | 14 +++--- rdagent/scenarios/qlib/experiment/utils.py | 2 +- .../scenarios/qlib/experiment/workspace.py | 4 +- rdagent/utils/env.py | 50 ++++++++++++------- test/utils/test_env.py | 46 ++++++++--------- 18 files changed, 143 insertions(+), 91 deletions(-) diff --git a/rdagent/components/coder/data_science/ensemble/eval.py b/rdagent/components/coder/data_science/ensemble/eval.py index 04b9fc65..d844e17d 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, ret_code = implementation.execute_ret_code(env=env, entry=f"python {fname}") + result = implementation.run(env=env, entry=f"python {fname}") + stdout = result.stdout + ret_code = result.ret_code stdout += f"\nNOTE: the above scripts run with return code {ret_code}" diff --git a/rdagent/components/coder/data_science/feature/eval.py b/rdagent/components/coder/data_science/feature/eval.py index 01fa857e..18b79a49 100644 --- a/rdagent/components/coder/data_science/feature/eval.py +++ b/rdagent/components/coder/data_science/feature/eval.py @@ -50,9 +50,9 @@ class FeatureCoSTEEREvaluator(CoSTEEREvaluator): test_code = (DIRNAME / "eval_tests" / "feature_test.txt").read_text() implementation.inject_files(**{fname: test_code}) - stdout, ret_code = implementation.execute_ret_code(env=env, entry=f"python {fname}") + result = implementation.run(env=env, entry=f"python {fname}") - if "main.py" in implementation.file_dict and ret_code == 0: + if "main.py" in implementation.file_dict and result.ret_code == 0: workflow_stdout = implementation.execute(env=env, entry="python main.py") workflow_stdout = remove_eda_part(workflow_stdout) else: @@ -66,7 +66,7 @@ class FeatureCoSTEEREvaluator(CoSTEEREvaluator): workflow_code=implementation.all_codes, ) user_prompt = T(".prompts:feature_eval.user").r( - stdout=shrink_text(stdout), + stdout=shrink_text(result.stdout), workflow_stdout=workflow_stdout, ) @@ -76,6 +76,6 @@ class FeatureCoSTEEREvaluator(CoSTEEREvaluator): user_prompt=user_prompt, init_kwargs_update_func=FeatureEvalFeedback.val_and_update_init_dict, ) - fb.final_decision = fb.final_decision and ret_code == 0 + fb.final_decision = fb.final_decision and result.ret_code == 0 return fb diff --git a/rdagent/components/coder/data_science/model/eval.py b/rdagent/components/coder/data_science/model/eval.py index 229512bb..75e21133 100644 --- a/rdagent/components/coder/data_science/model/eval.py +++ b/rdagent/components/coder/data_science/model/eval.py @@ -67,7 +67,9 @@ class ModelGeneralCaseSpecEvaluator(CoSTEEREvaluator): (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, ret_code = implementation.execute_ret_code(env=env, entry=f"python {fname}") + result = implementation.run(env=env, entry=f"python {fname}") + stdout = result.stdout + ret_code = result.ret_code if stdout is None: raise CoderError( @@ -113,6 +115,6 @@ class ModelGeneralCaseSpecEvaluator(CoSTEEREvaluator): user_prompt=user_prompt, init_kwargs_update_func=ModelSingleFeedback.val_and_update_init_dict, ) - fb.final_decision = fb.final_decision and ret_code == 0 + fb.final_decision = fb.final_decision and result.ret_code == 0 return fb diff --git a/rdagent/components/coder/data_science/pipeline/eval.py b/rdagent/components/coder/data_science/pipeline/eval.py index 9f8859df..e7e3e8df 100644 --- a/rdagent/components/coder/data_science/pipeline/eval.py +++ b/rdagent/components/coder/data_science/pipeline/eval.py @@ -57,8 +57,10 @@ class PipelineCoSTEEREvaluator(CoSTEEREvaluator): # Clean the scores.csv & submission.csv. implementation.execute(env=env, entry=get_clear_ws_cmd()) - stdout, execute_ret_code = implementation.execute_ret_code(env=env, entry=f"python -m coverage run main.py") - stdout = remove_eda_part(stdout) + result = implementation.run(env=env, entry=f"python -m coverage run main.py") + implementation.running_info.running_time = result.running_time + execute_ret_code = result.ret_code + stdout = remove_eda_part(result.stdout) stdout += f"The code executed {'successfully' if execute_ret_code == 0 else 'failed'}." score_fp = implementation.workspace_path / "scores.csv" @@ -105,9 +107,9 @@ class PipelineCoSTEEREvaluator(CoSTEEREvaluator): base_check_code = T(".eval_tests.submission_format_test", ftype="txt").r() 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=env, entry="python test/submission_format_test.py" - ) + submission_result = implementation.run(env=env, entry="python test/submission_format_test.py") + submission_check_out = submission_result.stdout + submission_ret_code = submission_result.ret_code if DS_RD_SETTING.rule_base_eval: if execute_ret_code == 0 and score_ret_code == 0 and submission_ret_code == 0: return PipelineSingleFeedback( diff --git a/rdagent/components/coder/data_science/raw_data_loader/eval.py b/rdagent/components/coder/data_science/raw_data_loader/eval.py index 26dc6758..69154e3a 100644 --- a/rdagent/components/coder/data_science/raw_data_loader/eval.py +++ b/rdagent/components/coder/data_science/raw_data_loader/eval.py @@ -52,7 +52,9 @@ class DataLoaderCoSTEEREvaluator(CoSTEEREvaluator): fname = "test/data_loader_test.py" test_code = (DIRNAME / "eval_tests" / "data_loader_test.txt").read_text() implementation.inject_files(**{fname: test_code}) - stdout, ret_code = implementation.execute_ret_code(env=env, entry=f"python {fname}") + result = implementation.run(env=env, entry=f"python {fname}") + stdout = result.stdout + ret_code = result.ret_code 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 diff --git a/rdagent/components/coder/data_science/workflow/eval.py b/rdagent/components/coder/data_science/workflow/eval.py index 927317ba..72798c40 100644 --- a/rdagent/components/coder/data_science/workflow/eval.py +++ b/rdagent/components/coder/data_science/workflow/eval.py @@ -121,9 +121,9 @@ class WorkflowGeneralCaseSpecEvaluator(CoSTEEREvaluator): base_check_code = T(".eval_tests.submission_format_test", ftype="txt").r() 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=env, entry="python test/submission_format_test.py" - ) + submission_result = implementation.run(env=env, entry="python test/submission_format_test.py") + submission_check_out = submission_result.stdout + submission_ret_code = submission_result.ret_code stdout += "\n" + submission_check_out system_prompt = T(".prompts:workflow_eval.system").r( diff --git a/rdagent/core/experiment.py b/rdagent/core/experiment.py index 776306b2..6a56e722 100644 --- a/rdagent/core/experiment.py +++ b/rdagent/core/experiment.py @@ -9,12 +9,17 @@ import uuid from abc import ABC, abstractmethod from collections.abc import Sequence from copy import deepcopy +from dataclasses import dataclass from pathlib import Path -from typing import Any, Generic, TypeVar +from typing import TYPE_CHECKING, Any, Generic, TypeVar from rdagent.core.conf import RD_AGENT_SETTINGS from rdagent.core.evaluation import Feedback from rdagent.utils import filter_redundant_text + +if TYPE_CHECKING: + from rdagent.utils.env import EnvResult + from rdagent.utils.fmt import shrink_text if typing.TYPE_CHECKING: @@ -59,6 +64,12 @@ ASpecificTask = TypeVar("ASpecificTask", bound=Task) ASpecificFeedback = TypeVar("ASpecificFeedback", bound=Feedback) +@dataclass +class RunningInfo: + result: object = None # The result of the experiment, can be different types in different scenarios. + running_time: float | None = None + + class Workspace(ABC, Generic[ASpecificTask, ASpecificFeedback]): """ A workspace is a place to store the task implementation. It evolves as the developer implements the task. @@ -68,6 +79,7 @@ class Workspace(ABC, Generic[ASpecificTask, ASpecificFeedback]): def __init__(self, target_task: ASpecificTask | None = None) -> None: self.target_task: ASpecificTask | None = target_task self.feedback: ASpecificFeedback | None = None + self.running_info: RunningInfo = RunningInfo() @abstractmethod def execute(self, *args: Any, **kwargs: Any) -> object | None: @@ -250,26 +262,25 @@ class FBWorkspace(Workspace): """ Before each execution, make sure to prepare and inject code. """ - stdout, _ = self.execute_ret_code(env, entry) - return stdout + result = self.run(env, entry) + return result.stdout - def execute_ret_code(self, env: Env, entry: str) -> tuple[str, int]: + def run(self, env: Env, entry: str) -> EnvResult: """ - Execute the code in the environment and return both the stdout and the exit code. + Execute the code in the environment and return an EnvResult object (stdout, exit_code, running_time). Before each execution, make sure to prepare and inject code. """ self.prepare() self.inject_files(**self.file_dict) - stdout, return_code = env.run_ret_code(entry, str(self.workspace_path), env={"PYTHONPATH": "./"}) - return ( - shrink_text( - filter_redundant_text(stdout), - context_lines=RD_AGENT_SETTINGS.stdout_context_len, - line_len=RD_AGENT_SETTINGS.stdout_line_len, - ), - return_code, + result = env.run_ret_code(entry, str(self.workspace_path), env={"PYTHONPATH": "./"}) + # result is EnvResult + result.stdout = shrink_text( + filter_redundant_text(result.stdout), + context_lines=RD_AGENT_SETTINGS.stdout_context_len, + line_len=RD_AGENT_SETTINGS.stdout_line_len, ) + return result def __str__(self) -> str: return f"Workspace[{self.workspace_path=}" + ( @@ -319,7 +330,7 @@ class Experiment( # NOTE: Assumption # - only runner will assign this variable # - We will always create a new Experiment without copying previous results when we goto the next new loop. - self.result: object = None # The result of the experiment, can be different types in different scenarios. + self.running_info = RunningInfo() self.sub_results: dict[str, float] = ( {} ) # TODO: in Kaggle, now sub results are all saved in self.result, remove this in the future. @@ -327,6 +338,14 @@ class Experiment( # For parallel multi-trace support self.local_selection: tuple[int, ...] | None = None + @property + def result(self) -> object: + return self.running_info.result + + @result.setter + def result(self, value: object) -> None: + self.running_info.result = value + ASpecificExp = TypeVar("ASpecificExp", bound=Experiment) diff --git a/rdagent/scenarios/data_science/dev/prompts.yaml b/rdagent/scenarios/data_science/dev/prompts.yaml index 912d948a..18922803 100644 --- a/rdagent/scenarios/data_science/dev/prompts.yaml +++ b/rdagent/scenarios/data_science/dev/prompts.yaml @@ -98,8 +98,9 @@ exp_feedback: 1. Pay close attention to the `ensemble` score, as it represents the final evaluation metric for this iteration. 2. If any individual model significantly outperforms the ensemble, this may indicate an issue in the ensemble method. But if the final `ensemble` score surpasses the current SOTA, you should update the SOTA record. However, it seems that there are noticeable issues in the ensemble component, be sure to highlight them explicitly. - Below are the results for this experiment: - {{ cur_exp.result }} + Below are the results and running time for this experiment: + Running time: {{ cur_exp.running_info.running_time }} seconds. + Results: {{ cur_exp.result }} {% if cur_vs_sota_score is not none %} Below is the comparison of the current `ensemble` performance with the SOTA results: diff --git a/rdagent/scenarios/data_science/dev/runner/__init__.py b/rdagent/scenarios/data_science/dev/runner/__init__.py index aaa1c455..51583da4 100644 --- a/rdagent/scenarios/data_science/dev/runner/__init__.py +++ b/rdagent/scenarios/data_science/dev/runner/__init__.py @@ -130,6 +130,7 @@ class DSCoSTEERRunner(CoSTEER): logger.error("Metrics file (scores.csv) is not generated.") raise RunnerError(f"Metrics file (scores.csv) is not generated") exp.result = pd.read_csv(score_fp, index_col=0) + exp.running_info.running_time = exp.experiment_workspace.running_info.running_time # 2) if mle-bench, then the submission format checking will be used. # DockerEnv for MLEBench submission validation diff --git a/rdagent/scenarios/data_science/dev/runner/eval.py b/rdagent/scenarios/data_science/dev/runner/eval.py index dc94eb18..3f85269a 100644 --- a/rdagent/scenarios/data_science/dev/runner/eval.py +++ b/rdagent/scenarios/data_science/dev/runner/eval.py @@ -53,7 +53,11 @@ class DSCoSTEERCoSTEEREvaluator(CoSTEEREvaluator): ) # Remove previous submission and scores files generated by worklfow. # execute workflow - stdout, execute_ret_code = implementation.execute_ret_code(env=env, entry="python -m coverage run main.py") + result = implementation.run(env=env, entry="python -m coverage run main.py") + stdout = result.stdout + execute_ret_code = result.ret_code + implementation.running_info.running_time = result.running_time + match = re.search(r"(.*?)=== Start of EDA part ===(.*)=== End of EDA part ===", stdout, re.DOTALL) eda_output = match.groups()[1] if match else None if eda_output is None: diff --git a/rdagent/scenarios/data_science/share.yaml b/rdagent/scenarios/data_science/share.yaml index 14898533..4d490690 100644 --- a/rdagent/scenarios/data_science/share.yaml +++ b/rdagent/scenarios/data_science/share.yaml @@ -22,6 +22,9 @@ describe: # some template to describe some object Submission format check result is: {{ exp.format_check_result }} {% endif %} + {% if exp.running_info.running_time is not none %} + Running time: {{ exp.running_info.running_time }} seconds + {% endif %} {% endif %} {% else %} diff --git a/rdagent/scenarios/data_science/test_eval.py b/rdagent/scenarios/data_science/test_eval.py index aa88a853..9db0b8ec 100644 --- a/rdagent/scenarios/data_science/test_eval.py +++ b/rdagent/scenarios/data_science/test_eval.py @@ -74,15 +74,15 @@ class TestEval(TestEvalBase): raise NoTestEvalError(err_msg) workspace.inject_files(**{"submission_format_valid.py": (eval_path / "valid.py").read_text()}) workspace.inject_files(**{"submission_test.csv": (eval_path / "submission_test.csv").read_text()}) - submission_check_out, submission_ret_code = workspace.execute_ret_code( + submission_result = workspace.run( env=self.env, entry=f"python submission_format_valid.py {competition}", ) workspace.inject_files( **{file: workspace.DEL_KEY for file in ["submission_format_valid.py", "submission_test.csv"]} ) - workspace.inject_files(**{"test/mle_submission_format_test.output": submission_check_out}) - return submission_check_out, submission_ret_code + workspace.inject_files(**{"test/mle_submission_format_test.output": submission_result.stdout}) + return submission_result.stdout, submission_result.ret_code def enabled(self, competition) -> bool: return Path( @@ -116,12 +116,10 @@ class MLETestEval(TestEvalBase): .replace("", competition) ) workspace.inject_files(**{"test/mle_submission_format_test.py": mle_check_code}) - submission_check_out, submission_ret_code = workspace.execute_ret_code( - env=self.env, entry="python test/mle_submission_format_test.py" - ) + submission_result = workspace.run(env=self.env, entry="python test/mle_submission_format_test.py") - workspace.inject_files(**{"test/mle_submission_format_test.output": submission_check_out}) - return submission_check_out, submission_ret_code + workspace.inject_files(**{"test/mle_submission_format_test.output": submission_result.stdout}) + return submission_result.stdout, submission_result.ret_code def enabled(self, competition) -> bool: return True diff --git a/rdagent/scenarios/kaggle/experiment/workspace.py b/rdagent/scenarios/kaggle/experiment/workspace.py index 5d4b0497..a9814d63 100644 --- a/rdagent/scenarios/kaggle/experiment/workspace.py +++ b/rdagent/scenarios/kaggle/experiment/workspace.py @@ -83,7 +83,7 @@ class KGFBWorkspace(FBWorkspace): else: running_extra_volume = {} - execute_log = kgde.run( + execute_log = kgde.check_output( local_path=str(self.workspace_path), env=run_env, running_extra_volume=running_extra_volume, diff --git a/rdagent/scenarios/kaggle/kaggle_crawler.py b/rdagent/scenarios/kaggle/kaggle_crawler.py index 1bc8825d..94989577 100644 --- a/rdagent/scenarios/kaggle/kaggle_crawler.py +++ b/rdagent/scenarios/kaggle/kaggle_crawler.py @@ -118,7 +118,7 @@ def download_data(competition: str, settings: ExtendedBaseSettings, enable_creat mleb_env = MLEBDockerEnv() mleb_env.prepare() (Path(zipfile_path)).mkdir(parents=True, exist_ok=True) - mleb_env.run( + mleb_env.check_output( f"mlebench prepare -c {competition} --data-dir ./zip_files", local_path=local_path, running_extra_volume={str(Path("~/.kaggle").expanduser().absolute()): "/root/.kaggle"}, @@ -129,17 +129,19 @@ def download_data(competition: str, settings: ExtendedBaseSettings, enable_creat mleb_env = MLEBDockerEnv() mleb_env.prepare() - mleb_env.run(f"cp -r ./zip_files/{competition}/prepared/public/* ./{competition}", local_path=local_path) + mleb_env.check_output( + f"cp -r ./zip_files/{competition}/prepared/public/* ./{competition}", local_path=local_path + ) for zip_path in competition_local_path.rglob("*.zip"): with zipfile.ZipFile(zip_path, "r") as zip_ref: if len(zip_ref.namelist()) == 1: - mleb_env.run( + mleb_env.check_output( f"unzip -o ./{zip_path.relative_to(competition_local_path)} -d {zip_path.parent.relative_to(competition_local_path)}", local_path=competition_local_path, ) else: - mleb_env.run( + mleb_env.check_output( f"mkdir -p ./{zip_path.parent.relative_to(competition_local_path)}/{zip_path.stem}; unzip -o ./{zip_path.relative_to(competition_local_path)} -d ./{zip_path.parent.relative_to(competition_local_path)}/{zip_path.stem}", local_path=competition_local_path, ) @@ -150,13 +152,13 @@ def download_data(competition: str, settings: ExtendedBaseSettings, enable_creat is_gzip_file = open(tar_path, "rb").read(2) == b"\x1f\x8b" with tarfile.open(tar_path, "r:gz") if is_gzip_file else tarfile.open(tar_path, "r") as tar_ref: if len(tar_ref.getmembers()) == 1: - mleb_env.run( + mleb_env.check_output( f"tar -{'xzf' if is_gzip_file else 'xf'} ./{tar_path.relative_to(competition_local_path)} -C {tar_path.parent.relative_to(competition_local_path)}", local_path=competition_local_path, ) else: folder_name = tar_path.name.replace(".tar", "").replace(".gz", "") - mleb_env.run( + mleb_env.check_output( f"mkdir -p ./{tar_path.parent.relative_to(competition_local_path)}/{folder_name}; tar -{'xzf' if is_gzip_file else 'xf'} ./{tar_path.relative_to(competition_local_path)} -C ./{tar_path.parent.relative_to(competition_local_path)}/{folder_name}", local_path=competition_local_path, ) diff --git a/rdagent/scenarios/qlib/experiment/utils.py b/rdagent/scenarios/qlib/experiment/utils.py index a1a1fa44..d58a6af1 100644 --- a/rdagent/scenarios/qlib/experiment/utils.py +++ b/rdagent/scenarios/qlib/experiment/utils.py @@ -16,7 +16,7 @@ def generate_data_folder_from_qlib(): qtde.prepare() # Run the Qlib backtest - execute_log = qtde.run( + execute_log = qtde.check_output( local_path=str(template_path), entry=f"python generate.py", ) diff --git a/rdagent/scenarios/qlib/experiment/workspace.py b/rdagent/scenarios/qlib/experiment/workspace.py index ec245014..37797ecc 100644 --- a/rdagent/scenarios/qlib/experiment/workspace.py +++ b/rdagent/scenarios/qlib/experiment/workspace.py @@ -26,7 +26,7 @@ class QlibFBWorkspace(FBWorkspace): qtde.prepare() # Run the Qlib backtest - execute_qlib_log = qtde.run( + execute_qlib_log = qtde.check_output( local_path=str(self.workspace_path), entry=f"qrun {qlib_config_name}", env=run_env, @@ -34,7 +34,7 @@ class QlibFBWorkspace(FBWorkspace): logger.log_object(execute_qlib_log, tag="Qlib_execute_log") # TODO: We should handle the case when Docker times out. - execute_log = qtde.run( + execute_log = qtde.check_output( local_path=str(self.workspace_path), entry="python read_exp_res.py", env=run_env, diff --git a/rdagent/utils/env.py b/rdagent/utils/env.py index 709a6631..930cecd0 100644 --- a/rdagent/utils/env.py +++ b/rdagent/utils/env.py @@ -19,6 +19,7 @@ import time import uuid import zipfile from abc import abstractmethod +from dataclasses import dataclass from pathlib import Path from types import MappingProxyType from typing import Any, Generator, Generic, Mapping, Optional, TypeVar, cast @@ -129,6 +130,18 @@ class EnvConf(ExtendedBaseSettings): ASpecificEnvConf = TypeVar("ASpecificEnvConf", bound=EnvConf) +@dataclass +class EnvResult: + """ + The result of running the environment. + It contains the stdout, the exit code, and the running time in seconds. + """ + + stdout: str + exit_code: int + running_time: float + + class Env(Generic[ASpecificEnvConf]): """ We use BaseModel as the setting due to the features it provides @@ -168,7 +181,9 @@ class Env(Generic[ASpecificEnvConf]): Prepare for the environment based on it's configure """ - def run(self, entry: str | None = None, local_path: str = ".", env: dict | None = None, **kwargs: dict) -> str: + def check_output( + self, entry: str | None = None, local_path: str = ".", env: dict | None = None, **kwargs: dict + ) -> str: """ Run the folder under the environment. @@ -189,8 +204,8 @@ class Env(Generic[ASpecificEnvConf]): ------- the stdout """ - stdout, _ = self.run_ret_code(entry=entry, local_path=local_path, env=env, **kwargs) - return stdout + result = self.run_ret_code(entry=entry, local_path=local_path, env=env, **kwargs) + return result.stdout def __run_ret_code_with_retry( self, @@ -199,7 +214,7 @@ class Env(Generic[ASpecificEnvConf]): env: dict | None = None, running_extra_volume: Mapping = MappingProxyType({}), remove_timestamp: bool = True, - ) -> tuple[str, int]: + ) -> EnvResult: # TODO: remove_timestamp can be implemented in a shallower way... for retry_index in range(self.conf.retry_count + 1): try: @@ -214,7 +229,8 @@ class Env(Generic[ASpecificEnvConf]): f"The running time exceeds {self.conf.running_timeout_period} seconds, so the process is killed." ) log_output += f"\n\nThe running time exceeds {self.conf.running_timeout_period} seconds, so the process is killed." - return log_output, return_code + log_output += f"\nTotal running time: {end - start:.3f} seconds." + return EnvResult(log_output, return_code, end - start) except Exception as e: if retry_index == self.conf.retry_count: raise @@ -230,9 +246,9 @@ class Env(Generic[ASpecificEnvConf]): local_path: str = ".", env: dict | None = None, **kwargs: dict, - ) -> tuple[str, int]: + ) -> EnvResult: """ - Run the folder under the environment and return both the stdout and the exit code. + Run the folder under the environment and return the stdout, exit code, and running time. Parameters ---------- @@ -249,7 +265,7 @@ class Env(Generic[ASpecificEnvConf]): Returns ------- - A tuple containing the stdout and the exit code + EnvResult: An object containing the stdout, the exit code, and the running time in seconds. """ running_extra_volume = kwargs.get("running_extra_volume", {}) if entry is None: @@ -297,13 +313,13 @@ class Env(Generic[ASpecificEnvConf]): ) if self.conf.enable_cache: - stdout, return_code = self.cached_run(entry_add_timeout, local_path, env, running_extra_volume) + result = self.cached_run(entry_add_timeout, local_path, env, running_extra_volume) else: - stdout, return_code = self.__run_ret_code_with_retry( + result = self.__run_ret_code_with_retry( entry_add_timeout, local_path, env, running_extra_volume, remove_timestamp=False ) - return stdout, return_code + return result def cached_run( self, @@ -312,7 +328,7 @@ class Env(Generic[ASpecificEnvConf]): env: dict | None = None, running_extra_volume: Mapping = MappingProxyType({}), remove_timestamp: bool = True, - ) -> tuple[str, int]: + ) -> EnvResult: """ Run the folder under the environment. Will cache the output and the folder diff for next round of running. @@ -345,14 +361,14 @@ class Env(Generic[ASpecificEnvConf]): ) 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) + ret = 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 + return cast(EnvResult, ret) @abstractmethod def _run_ret_code( @@ -380,7 +396,7 @@ class Env(Generic[ASpecificEnvConf]): Returns ------- tuple[str, int] - A tuple containing the standard output and the exit code of the execution. + A tuple containing the standard output and the exit code. """ pass @@ -400,7 +416,7 @@ class Env(Generic[ASpecificEnvConf]): 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=dict(running_extra_volume)) + log_output = self.check_output(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: @@ -933,7 +949,7 @@ class QTDockerEnv(DockerEnv): if not (Path(qlib_data_path) / "qlib_data" / "cn_data").exists(): logger.info("We are downloading!") cmd = "python -m qlib.run.get_data qlib_data --target_dir ~/.qlib/qlib_data/cn_data --region cn --interval 1d --delete_old False" - self.run(entry=cmd) + self.check_output(entry=cmd) else: logger.info("Data already exists. Download skipped.") diff --git a/test/utils/test_env.py b/test/utils/test_env.py index a8ee6723..84f36e5b 100644 --- a/test/utils/test_env.py +++ b/test/utils/test_env.py @@ -22,7 +22,7 @@ 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( + self.check_output( entry="python -m qlib.run.get_data qlib_data --target_dir ~/.qlib/qlib_data/cn_data --region cn", ) else: @@ -48,7 +48,7 @@ class EnvUtils(unittest.TestCase): qle = QlibLocalEnv(conf=local_conf) qle.prepare() conf_path = str(DIRNAME / "env_tpl" / "conf.yaml") - qle.run(entry="qrun " + conf_path) + qle.check_output(entry="qrun " + conf_path) mlrun_p = DIRNAME / "env_tpl" / "mlruns" self.assertTrue(mlrun_p.exists(), f"Expected output file {mlrun_p} not found") @@ -63,8 +63,8 @@ class EnvUtils(unittest.TestCase): print(local_conf) le = LocalEnv(conf=local_conf) le.prepare() - res, code = le.run_ret_code(local_path=str(code_path)) - print(res, code) + result = le.run_ret_code(local_path=str(code_path)) + print(result.stdout, result.exit_code, result.running_time) def test_conda_simple(self): conda_conf = CondaConf(default_entry="which python", conda_env_name="MLE") @@ -72,8 +72,8 @@ class EnvUtils(unittest.TestCase): 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) + result = le.run_ret_code(local_path=str(code_path)) + print(result.stdout, result.exit_code, result.running_time) def test_conda_error(self): conda_conf = CondaConf(conda_env_name="MLE") @@ -82,9 +82,9 @@ class EnvUtils(unittest.TestCase): file_name = f"{time.time()}.py" with open(self.test_workspace / file_name, "w") as f: f.write('import json \njson.loads(b\'{"name": "\xa1"}\')') - res, code = le.run_ret_code(local_path=str(self.test_workspace), entry=f"python {file_name}") - assert code == 1 - assert "bytes can only contain ASCII literal characters" in res + result = le.run_ret_code(local_path=str(self.test_workspace), entry=f"python {file_name}") + assert result.exit_code == 1 + assert "bytes can only contain ASCII literal characters" in result.stdout def test_docker(self): """We will mount `env_tpl` into the docker image. @@ -94,13 +94,13 @@ class EnvUtils(unittest.TestCase): qtde.prepare() # you can prepare for multiple times. It is expected to handle it correctly # qtde.run("nvidia-smi") # NOTE: you can check your GPU with this command # the stdout are returned as result - result = qtde.run(local_path=str(DIRNAME / "env_tpl"), entry="qrun conf.yaml") + result = qtde.check_output(local_path=str(DIRNAME / "env_tpl"), entry="qrun conf.yaml") mlrun_p = DIRNAME / "env_tpl" / "mlruns" self.assertTrue(mlrun_p.exists(), f"Expected output file {mlrun_p} not found") # read experiment - result = qtde.run(local_path=str(DIRNAME / "env_tpl"), entry="python read_exp_res.py") + result = qtde.check_output(local_path=str(DIRNAME / "env_tpl"), entry="python read_exp_res.py") print(result) def test_run_ret_code(self): @@ -109,34 +109,34 @@ class EnvUtils(unittest.TestCase): 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" + result = qtde.run_ret_code(entry='echo "Hello, World!"', local_path=str(self.test_workspace)) + print(result.exit_code) + assert result.exit_code == 0, f"Expected return code 0, but got {result.exit_code}" + assert "Hello, World!" in result.stdout, "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" + result = qtde.run_ret_code(entry="invalid_command", local_path=str(self.test_workspace)) + print(result.exit_code) + assert result.exit_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" + result = qtde.run_ret_code(entry="sleep 2", local_path=str(self.test_workspace)) + print(result.exit_code) + assert result.exit_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")\'' qtde = QTDockerEnv(QlibDockerConf(mem_limit="10m")) qtde.prepare() - result = qtde.run(local_path=str(DIRNAME / "env_tpl"), entry=cmd) + result = qtde.check_output(local_path=str(DIRNAME / "env_tpl"), entry=cmd) self.assertTrue(not result.strip().endswith("success")) qtde = QTDockerEnv(QlibDockerConf(mem_limit="1g")) qtde.prepare() - result = qtde.run(local_path=str(DIRNAME / "env_tpl"), entry=cmd) + result = qtde.check_output(local_path=str(DIRNAME / "env_tpl"), entry=cmd) self.assertTrue(result.strip().endswith("success")) # The above command equals to the follow commands with dockr cli.sh