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

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}"
 

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

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

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(

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

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(

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)
 

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:

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

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:

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 %}

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_id>", 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

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,

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,
                         )

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",
     )

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,

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.")
 

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