refactor: replace run_ret_code with run and adjust helper methods (#1013)

This commit is contained in:
you-n-g
2025-07-03 11:24:05 +08:00
committed by GitHub
parent fb8f2b561e
commit 856fbde807
3 changed files with 18 additions and 18 deletions
+1 -1
View File
@@ -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),
+9 -9
View File
@@ -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 = ".",
+8 -8
View File
@@ -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"