refactor: add run_ret_code method and update run method to use it (#623)

* refactor: Add run_ret_code method and update run method to use it

* feat: Add kwargs support to run methods and test for run_ret_code

* fix: preserve exit code after chmod in DockerEnv entry command

* chore: Change file permissions from 755 to 644 in env_tpl directory

* refactor: Return execution code and update evaluator logic

* lint

* refactor: Use MappingProxyType for running_extra_volume in DockerEnv methods

* lint

rdagent/components/coder/data_science/ensemble/eval.py

@@ -63,7 +63,9 @@ class EnsembleCoSTEEREvaluator(CoSTEEREvaluator):
         )
 
         implementation.inject_files(**{fname: test_code})
-        stdout = implementation.execute(env=de, entry=f"python {fname}")
+        stdout, ret_code = implementation.execute_ret_code(env=de, entry=f"python {fname}")
+
+        stdout += f"\nNOTE: the above scripts run with return code {ret_code}"
 
         if "main.py" in implementation.file_dict:
             workflow_stdout = implementation.execute(env=de, entry="python main.py")
@@ -81,6 +83,6 @@ class EnsembleCoSTEEREvaluator(CoSTEEREvaluator):
             stdout=stdout,
             workflow_stdout=workflow_stdout,
         )
-        return build_cls_from_json_with_retry(
-            EnsembleEvalFeedback, system_prompt=system_prompt, user_prompt=user_prompt
-        )
+        efb = build_cls_from_json_with_retry(EnsembleEvalFeedback, system_prompt=system_prompt, user_prompt=user_prompt)
+        efb.final_decision = efb.final_decision and ret_code == 0
+        return efb

rdagent/core/experiment.py

@@ -232,21 +232,29 @@ class FBWorkspace(Workspace):
         shutil.rmtree(self.workspace_path, ignore_errors=True)
         self.file_dict = {}
 
-    def execute(self, env: Env | None = None, entry: str | None = None) -> object | None:
+    def execute(self, env: Env, entry: str) -> str:
         """
-        Before each execution, make sure to prepare and inject code
+        Before each execution, make sure to prepare and inject code.
+        """
+        stdout, _ = self.execute_ret_code(env, entry)
+        return stdout
+
+    def execute_ret_code(self, env: Env, entry: str) -> tuple[str, int]:
+        """
+        Execute the code in the environment and return both the stdout and the exit code.
+
+        Before each execution, make sure to prepare and inject code.
         """
         self.prepare()
         self.inject_files(**self.file_dict)
-        # TODO: env should be not None in new design (no code can run without environment)
-        if env is not None and entry is not None:
-            # We believe that removing the progress bar might resolve the issue.
-            # If needed, cutting off the information in the middle should be a backup plan.
-            return shrink_text(
-                filter_progress_bar(env.run(entry, str(self.workspace_path))),
+        stdout, return_code = env.run_ret_code(entry, str(self.workspace_path))
+        return (
+            shrink_text(
+                filter_progress_bar(stdout),
                 context_lines=RD_AGENT_SETTINGS.stdout_context_len,
-            )
-        return None
+            ),
+            return_code,
+        )
 
     def __str__(self) -> str:
         return f"Workspace[{self.workspace_path=}" + (

rdagent/utils/env.py

@@ -18,7 +18,8 @@ import uuid
 import zipfile
 from abc import abstractmethod
 from pathlib import Path
-from typing import Generic, Optional, TypeVar
+from types import MappingProxyType
+from typing import Generic, Mapping, Optional, TypeVar
 
 import docker  # type: ignore[import-untyped]
 import docker.models  # type: ignore[import-untyped]
@@ -48,6 +49,7 @@ class Env(Generic[ASpecificBaseModel]):
     """
 
     conf: ASpecificBaseModel  # different env have different conf.
+
     # last_exit_code:  # TODO: get the more concrete information about the exit code.
 
     def __init__(self, conf: ASpecificBaseModel):
@@ -59,8 +61,7 @@ class Env(Generic[ASpecificBaseModel]):
         Prepare for the environment based on it's configure
         """
 
-    @abstractmethod
-    def run(self, entry: str | None, local_path: str = ".", env: dict | None = None) -> str:
+    def run(self, entry: str | None, local_path: str = ".", env: dict | None = None, **kwargs: dict) -> str:
         """
         Run the folder under the environment.
 
@@ -81,6 +82,33 @@ class Env(Generic[ASpecificBaseModel]):
         -------
             the stdout
         """
+        stdout, _ = self.run_ret_code(entry=entry, local_path=local_path, env=env, **kwargs)
+        return stdout
+
+    @abstractmethod
+    def run_ret_code(
+        self, entry: str | None, local_path: str = ".", env: dict | None = None, **kwargs: dict
+    ) -> tuple[str, int]:
+        """
+        Run the folder under the environment and return both the stdout and the exit code.
+
+        Parameters
+        ----------
+        entry : str | None
+            We may we the entry point when we run it.
+            For example, we may have different entries when we run and summarize the project.
+        local_path : str | None
+            the local path (to project, mainly for code) will be mounted into the docker
+            Here are some examples for a None local path
+            - for example, run docker for updating the data in the extra_volumes.
+            - simply run the image. The results are produced by output or network
+        env : dict | None
+            Run the code with your specific environment.
+
+        Returns
+        -------
+            A tuple containing the stdout and the exit code
+        """
 
 
 ## Local Environment -----
@@ -104,7 +132,13 @@ class LocalEnv(Env[LocalConf]):
         else:
             print("Data already exists. Download skipped.")
 
-    def run(self, entry: str | None = None, local_path: Optional[str] = None, env: dict | None = None) -> str:
+    def run_ret_code(
+        self,
+        entry: str | None = None,
+        local_path: str | None = None,
+        env: dict | None = None,
+        **kwargs: dict,
+    ) -> tuple[str, int]:
         if env is None:
             env = {}
 
@@ -118,10 +152,7 @@ class LocalEnv(Env[LocalConf]):
             cwd = Path(local_path).resolve()
         result = subprocess.run(command, cwd=cwd, env={**os.environ, **env}, capture_output=True, text=True)
 
-        if result.returncode != 0:
-            raise RuntimeError(f"Error while running the command: {result.stderr}")
-
-        return result.stdout
+        return result.stdout, result.returncode
 
 
 ## Docker Environment -----
@@ -337,14 +368,14 @@ class DockerEnv(Env[DockerConf]):
         output_string = re.sub(datetime_pattern, "[DATETIME]", input_string)
         return output_string
 
-    def __run(
+    def __run_ret_code(
         self,
         entry: str | None = None,
         local_path: str = ".",
         env: dict | None = None,
-        running_extra_volume: dict | None = None,
+        running_extra_volume: Mapping = MappingProxyType({}),
         remove_timestamp: bool = True,
-    ) -> str:
+    ) -> tuple[str, int]:
         if env is None:
             env = {}
         env["PYTHONWARNINGS"] = "ignore"
@@ -359,9 +390,8 @@ class DockerEnv(Env[DockerConf]):
         if self.conf.extra_volumes is not None:
             for lp, rp in self.conf.extra_volumes.items():
                 volumns[lp] = {"bind": rp, "mode": self.conf.extra_volume_mode}
-        if running_extra_volume is not None:
-            for lp, rp in running_extra_volume.items():
-                volumns[lp] = {"bind": rp, "mode": self.conf.extra_volume_mode}
+        for lp, rp in running_extra_volume.items():
+            volumns[lp] = {"bind": rp, "mode": self.conf.extra_volume_mode}
 
         log_output = ""
 
@@ -397,7 +427,7 @@ class DockerEnv(Env[DockerConf]):
                 decoded_log = self.replace_time_info(decoded_log) if remove_timestamp else decoded_log
                 Console().print(decoded_log, markup=False)
                 log_output += decoded_log + "\n"
-            container.wait()
+            exit_status = container.wait()["StatusCode"]
             container.stop()
             container.remove()
             end = time.time()
@@ -407,7 +437,7 @@ class DockerEnv(Env[DockerConf]):
                 )
                 log_output += f"\n\nThe running time exceeds {self.conf.running_timeout_period} seconds, so the process is killed."
             print(Rule("[bold green]Docker Logs End[/bold green]", style="dark_orange"))
-            return log_output
+            return log_output, exit_status
         except docker.errors.ContainerError as e:
             raise RuntimeError(f"Error while running the container: {e}")
         except docker.errors.ImageNotFound:
@@ -415,17 +445,17 @@ class DockerEnv(Env[DockerConf]):
         except docker.errors.APIError as e:
             raise RuntimeError(f"Error while running the container: {e}")
 
-    def __run_with_retry(
+    def __run_ret_code_with_retry(
         self,
         entry: str | None = None,
         local_path: str = ".",
         env: dict | None = None,
-        running_extra_volume: dict | None = None,
+        running_extra_volume: Mapping = MappingProxyType({}),
         remove_timestamp: bool = True,
-    ) -> str:
+    ) -> tuple[str, int]:
         for retry_index in range(self.conf.retry_count):
             try:
-                return self.__run(entry, local_path, env, running_extra_volume, remove_timestamp)
+                return self.__run_ret_code(entry, local_path, env, running_extra_volume, remove_timestamp)
             except Exception as e:
                 logger.warning(
                     f"Error while running the container: {e}, current try index: {retry_index + 1}, {self.conf.retry_count - retry_index - 1} retries left."
@@ -459,9 +489,9 @@ class DockerEnv(Env[DockerConf]):
         entry: str | None = None,
         local_path: str = ".",
         env: dict | None = None,
-        running_extra_volume: dict | None = None,
+        running_extra_volume: Mapping = MappingProxyType({}),
         remove_timestamp: bool = True,
-    ) -> str:
+    ) -> tuple[str, int]:
         """
         Run the folder under the environment.
         Will cache the output and the folder diff for next round of running.
@@ -488,42 +518,47 @@ class DockerEnv(Env[DockerConf]):
                     for path in sorted(Path(local_path).rglob("*.py"))
                 ]
             )
-            + json.dumps({"entry": entry, "running_extra_volume": running_extra_volume})
+            + json.dumps({"entry": entry, "running_extra_volume": dict(running_extra_volume)})
             + json.dumps({"extra_volumes": self.conf.extra_volumes})
             + json.dumps(data_key)
         )
         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: str = pickle.load(f)
+                ret: tuple[str, int] = pickle.load(f)
             self.unzip_a_file_into_a_folder(str(target_folder / f"{key}.zip"), local_path)
         else:
-            ret = self.__run_with_retry(entry, local_path, env, running_extra_volume, remove_timestamp)
+            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
 
-    def run(
+    def run_ret_code(
         self,
         entry: str | None = None,
         local_path: str = ".",
         env: dict | None = None,
-        running_extra_volume: dict | None = None,
-    ) -> str:
+        **kwargs: dict,
+    ) -> tuple[str, int]:
+        running_extra_volume = kwargs.get("running_extra_volume", {})
         if entry is None:
             entry = self.conf.default_entry
+
         entry_add_timeout = (
-            f"/bin/sh -c 'timeout {self.conf.running_timeout_period} {entry}; chmod -R 777 {self.conf.mount_path}'"
+            f"/bin/sh -c 'timeout {self.conf.running_timeout_period} {entry}; "
+            f"entry_exit_code=$?; "
+            f"chmod -R 777 {self.conf.mount_path}; "
+            f"exit $entry_exit_code'"
         )
 
         if self.conf.enable_cache:
-            out = self.cached_run(entry_add_timeout, local_path, env, running_extra_volume)
+            stdout, return_code = self.cached_run(entry_add_timeout, local_path, env, running_extra_volume)
         else:
-            out = self.__run_with_retry(
+            stdout, return_code = self.__run_ret_code_with_retry(
                 entry_add_timeout, local_path, env, running_extra_volume, remove_timestamp=False
             )
 
-        return out
+        return stdout, return_code
 
     def dump_python_code_run_and_get_results(
         self,
@@ -531,7 +566,7 @@ class DockerEnv(Env[DockerConf]):
         dump_file_names: list[str],
         local_path: str,
         env: dict | None = None,
-        running_extra_volume: dict | None = None,
+        running_extra_volume: Mapping = MappingProxyType({}),
         code_dump_file_py_name: Optional[str] = None,
     ) -> tuple[str, list]:
         """
@@ -541,7 +576,7 @@ class DockerEnv(Env[DockerConf]):
         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=running_extra_volume)
+        log_output = self.run(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:

test/utils/test_env.py

@@ -13,14 +13,12 @@ DIRNAME = Path(__file__).absolute().resolve().parent
 
 class EnvUtils(unittest.TestCase):
     def setUp(self):
-        pass
+        self.test_workspace = DIRNAME / "test_workspace"
+        self.test_workspace.mkdir(exist_ok=True)
 
     def tearDown(self):
-        # NOTE: For a docker file, the output are generated with root permission.
-        # mlrun_p = DIRNAME / "env_tpl" / "mlruns"
-        # if mlrun_p.exists():
-        #     shutil.rmtree(mlrun_p)
-        ...
+        if self.test_workspace.exists():
+            shutil.rmtree(self.test_workspace)
 
     # NOTE: Since I don't know the exact environment in which it will be used, here's just an example.
     # NOTE: Because you need to download the data during the prepare process. So you need to have pyqlib in your environment.
@@ -53,6 +51,29 @@ class EnvUtils(unittest.TestCase):
         result = qtde.run(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."""
+        qtde = QTDockerEnv()
+        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"
+
+        # 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"
+
+        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"
+
     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")\''