mirror of
https://github.com/NicolasBohn/NexQuant.git
synced 2026-07-27 15:37:44 +00:00
refactor: replace run_ret_code with run and adjust helper methods (#1013)
This commit is contained in:
@@ -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),
|
||||
|
||||
@@ -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 = ".",
|
||||
|
||||
@@ -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"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user