From 13cfd1b3a86f32a81f698e1cc1303c28ebc448a0 Mon Sep 17 00:00:00 2001
From: XianBW <36835909+XianBW@users.noreply.github.com>
Date: Thu, 13 Mar 2025 16:57:10 +0800
Subject: [PATCH] timeout message adjustment for Run Env (#677)

---
 rdagent/utils/env.py | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/rdagent/utils/env.py b/rdagent/utils/env.py
index 41de3ce8..825caf79 100644
--- a/rdagent/utils/env.py
+++ b/rdagent/utils/env.py
@@ -128,9 +128,17 @@ class Env(Generic[ASpecificEnvConf]):
         # TODO: remove_timestamp can be implemented in a shallower way...
         for retry_index in range(self.conf.retry_count + 1):
             try:
-                return self._run_ret_code(
+                start = time.time()
+                log_output, return_code = self._run_ret_code(
                     entry, local_path, env, running_extra_volume=running_extra_volume, remove_timestamp=remove_timestamp
                 )
+                end = time.time()
+                if end - start >= self.conf.running_timeout_period:
+                    print(
+                        f"[red]The running time exceeds {self.conf.running_timeout_period} seconds, so the process is killed.[/red]"
+                    )
+                    log_output += f"\n\nThe running time exceeds {self.conf.running_timeout_period} seconds, so the process is killed."
+                return log_output, return_code
             except Exception as e:
                 if retry_index == self.conf.retry_count:
                     raise
@@ -331,6 +339,7 @@ class LocalEnv(Env[ASpecificLocalConf]):
             "env": env,
             "volumes": volumns,
         }
+        print(Rule("[bold green]LocalEnv Logs Begin[/bold green]", style="dark_orange"))
         print(Pretty(summary))
 
         cwd = None
@@ -340,6 +349,7 @@ class LocalEnv(Env[ASpecificLocalConf]):
         result = subprocess.run(entry, cwd=cwd, env={**os.environ, **env}, capture_output=True, text=True, shell=True)
         combined_output = result.stderr + result.stdout  # Combine stdout and stderr
         print(combined_output)  # Display the combined output in the console
+        print(Rule("[bold green]LocalEnv Logs End[/bold green]", style="dark_orange"))
 
         return combined_output, result.returncode
 
@@ -604,7 +614,6 @@ class DockerEnv(Env[DockerConf]):
         log_output = ""
 
         try:
-            start = time.time()
             container: docker.models.containers.Container = client.containers.run(  # type: ignore[no-any-unimported]
                 image=self.conf.image,
                 command=entry,
@@ -638,12 +647,6 @@ class DockerEnv(Env[DockerConf]):
             exit_status = container.wait()["StatusCode"]
             container.stop()
             container.remove()
-            end = time.time()
-            if end - start >= self.conf.running_timeout_period:
-                print(
-                    f"[red]The running time exceeds {self.conf.running_timeout_period} seconds, so the process is killed.[/red]"
-                )
-                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, exit_status
         except docker.errors.ContainerError as e: