From 7a7f86573322fe8916ca9e77d7c1c4eae160807d Mon Sep 17 00:00:00 2001 From: you-n-g Date: Thu, 3 Jul 2025 11:24:05 +0800 Subject: [PATCH] refactor: replace run_ret_code with run and adjust helper methods (#1013) --- rdagent/core/experiment.py | 2 +- rdagent/utils/env.py | 18 +++++++++--------- test/utils/test_env.py | 16 ++++++++-------- 3 files changed, 18 insertions(+), 18 deletions(-) diff --git a/rdagent/core/experiment.py b/rdagent/core/experiment.py index 6a56e722..168638ea 100644 --- a/rdagent/core/experiment.py +++ b/rdagent/core/experiment.py @@ -273,7 +273,7 @@ class FBWorkspace(Workspace): """ self.prepare() self.inject_files(**self.file_dict) - result = env.run_ret_code(entry, str(self.workspace_path), env={"PYTHONPATH": "./"}) + result = env.run(entry, str(self.workspace_path), env={"PYTHONPATH": "./"}) # result is EnvResult result.stdout = shrink_text( filter_redundant_text(result.stdout), diff --git a/rdagent/utils/env.py b/rdagent/utils/env.py index 930cecd0..c265e0bb 100644 --- a/rdagent/utils/env.py +++ b/rdagent/utils/env.py @@ -204,10 +204,10 @@ class Env(Generic[ASpecificEnvConf]): ------- the stdout """ - result = self.run_ret_code(entry=entry, local_path=local_path, env=env, **kwargs) + result = self.run(entry=entry, local_path=local_path, env=env, **kwargs) return result.stdout - def __run_ret_code_with_retry( + def __run_with_retry( self, entry: str | None = None, local_path: str = ".", @@ -219,7 +219,7 @@ class Env(Generic[ASpecificEnvConf]): for retry_index in range(self.conf.retry_count + 1): try: start = time.time() - log_output, return_code = self._run_ret_code( + log_output, return_code = self._run( entry, local_path, env, running_extra_volume=running_extra_volume, remove_timestamp=remove_timestamp ) end = time.time() @@ -240,7 +240,7 @@ class Env(Generic[ASpecificEnvConf]): time.sleep(self.conf.retry_wait_seconds) raise RuntimeError # for passing CI - def run_ret_code( + def run( self, entry: str | None = None, local_path: str = ".", @@ -315,7 +315,7 @@ class Env(Generic[ASpecificEnvConf]): if self.conf.enable_cache: result = self.cached_run(entry_add_timeout, local_path, env, running_extra_volume) else: - result = self.__run_ret_code_with_retry( + result = self.__run_with_retry( entry_add_timeout, local_path, env, running_extra_volume, remove_timestamp=False ) @@ -364,14 +364,14 @@ class Env(Generic[ASpecificEnvConf]): 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) + ret = self.__run_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 cast(EnvResult, ret) @abstractmethod - def _run_ret_code( + def _run( self, entry: str | None, local_path: str = ".", @@ -452,7 +452,7 @@ class LocalEnv(Env[ASpecificLocalConf]): def prepare(self) -> None: ... - def _run_ret_code( + def _run( self, entry: str | None = None, local_path: str | None = None, @@ -853,7 +853,7 @@ class DockerEnv(Env[DockerConf]): output_string = re.sub(datetime_pattern, "[DATETIME]", input_string) return output_string - def _run_ret_code( + def _run( self, entry: str | None = None, local_path: str = ".", diff --git a/test/utils/test_env.py b/test/utils/test_env.py index 84f36e5b..a550b4d4 100644 --- a/test/utils/test_env.py +++ b/test/utils/test_env.py @@ -63,7 +63,7 @@ class EnvUtils(unittest.TestCase): print(local_conf) le = LocalEnv(conf=local_conf) le.prepare() - result = le.run_ret_code(local_path=str(code_path)) + result = le.run(local_path=str(code_path)) print(result.stdout, result.exit_code, result.running_time) def test_conda_simple(self): @@ -72,7 +72,7 @@ class EnvUtils(unittest.TestCase): le.prepare() code_path = DIRNAME / "tmp_code" code_path.mkdir(exist_ok=True) - result = le.run_ret_code(local_path=str(code_path)) + result = le.run(local_path=str(code_path)) print(result.stdout, result.exit_code, result.running_time) def test_conda_error(self): @@ -82,7 +82,7 @@ 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"}\')') - result = le.run_ret_code(local_path=str(self.test_workspace), entry=f"python {file_name}") + result = le.run(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 @@ -103,26 +103,26 @@ class EnvUtils(unittest.TestCase): result = qtde.check_output(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.""" + def test_run(self): + """Test the run method of QTDockerEnv with both valid and invalid commands.""" qtde = QTDockerEnv() qtde.prepare() # Test with a valid command - result = qtde.run_ret_code(entry='echo "Hello, World!"', local_path=str(self.test_workspace)) + result = qtde.run(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 - result = qtde.run_ret_code(entry="invalid_command", local_path=str(self.test_workspace)) + result = qtde.run(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 = qtde.run_ret_code(entry="sleep 2", local_path=str(self.test_workspace)) + result = qtde.run(entry="sleep 2", local_path=str(self.test_workspace)) print(result.exit_code) assert result.exit_code == 124, "Expected return code 124 for timeout"