Support specify timeout for long running tasks (All-Hands-AI#1756)

* support specify timeout for long running tasks

* add timeout for all existing sandbox impl

* Update opendevin/runtime/docker/local_box.py

Co-authored-by: Yufan Song <[email protected]>

* Update opendevin/runtime/docker/exec_box.py

Co-authored-by: Yufan Song <[email protected]>

* Update opendevin/runtime/docker/ssh_box.py

Co-authored-by: Yufan Song <[email protected]>

* Update opendevin/runtime/e2b/sandbox.py

Co-authored-by: Yufan Song <[email protected]>

* Update opendevin/runtime/sandbox.py

Co-authored-by: Yufan Song <[email protected]>

---------

Co-authored-by: Yufan Song <[email protected]>

opendevin/runtime/docker/exec_box.py

@@ -106,7 +106,9 @@ def read_logs(self, id) -> str:
         bg_cmd = self.background_commands[id]
         return bg_cmd.read_logs()
 
-    def execute(self, cmd: str) -> tuple[int, str]:
+    def execute(self, cmd: str, timeout: int | None = None) -> tuple[int, str]:
+        timeout = timeout if timeout is not None else self.timeout
+
         # TODO: each execute is not stateful! We need to keep track of the current working directory
         def run_command(container, command):
             return container.exec_run(
@@ -119,7 +121,7 @@ def run_command(container, command):
                 run_command, self.container, self.get_exec_cmd(cmd)
             )
             try:
-                exit_code, logs = future.result(timeout=self.timeout)
+                exit_code, logs = future.result(timeout=timeout)
             except concurrent.futures.TimeoutError:
                 logger.exception(
                     'Command timed out, killing process...', exc_info=False

opendevin/runtime/docker/local_box.py

@@ -33,14 +33,15 @@ def __init__(self, timeout: int = 120):
         atexit.register(self.cleanup)
         super().__init__()
 
-    def execute(self, cmd: str) -> tuple[int, str]:
+    def execute(self, cmd: str, timeout: int | None = None) -> tuple[int, str]:
+        timeout = timeout if timeout is not None else self.timeout
         try:
             completed_process = subprocess.run(
                 cmd,
                 shell=True,
                 text=True,
                 capture_output=True,
-                timeout=self.timeout,
+                timeout=timeout,
                 cwd=config.workspace_base,
                 env=self._env,
             )

opendevin/runtime/docker/ssh_box.py

@@ -8,7 +8,6 @@
 import uuid
 from collections import namedtuple
 from glob import glob
-from typing import Tuple
 
 import docker
 from pexpect import pxssh
@@ -343,7 +342,9 @@ def _send_interrupt(
             f'Command: "{cmd}" timed out. Sending SIGINT to the process: {command_output}',
         )
 
-    def execute(self, cmd: str) -> Tuple[int, str]:
+    def execute(self, cmd: str, timeout: int | None = None) -> tuple[int, str]:
+        timeout = timeout if timeout is not None else self.timeout
+
         commands = split_bash_commands(cmd)
         if len(commands) > 1:
             all_output = ''
@@ -356,7 +357,7 @@ def execute(self, cmd: str) -> Tuple[int, str]:
                     return exit_code, all_output
             return 0, all_output
         self.ssh.sendline(cmd)
-        success = self.ssh.prompt(timeout=self.timeout)
+        success = self.ssh.prompt(timeout=timeout)
         if not success:
             logger.exception('Command timed out, killing process...', exc_info=False)
             return self._send_interrupt(cmd)
@@ -389,7 +390,7 @@ def execute(self, cmd: str) -> Tuple[int, str]:
             self.ssh.prompt()
             exit_code_str = self.ssh.before.decode('utf-8')
             logger.debug(f'WAITING FOR exit code: {exit_code_str}')
-            if time.time() - _start_time > self.timeout:
+            if time.time() - _start_time > timeout:
                 return self._send_interrupt(
                     cmd, command_output, ignore_last_output=True
                 )

opendevin/runtime/e2b/sandbox.py

@@ -72,10 +72,11 @@ def read_logs(self, process_id: int) -> str:
         assert isinstance(proc, E2BProcess)
         return '\n'.join([m.line for m in proc.output_messages])
 
-    def execute(self, cmd: str) -> tuple[int, str]:
+    def execute(self, cmd: str, timeout: int | None = None) -> tuple[int, str]:
+        timeout = timeout if timeout is not None else self.timeout
         process = self.sandbox.process.start(cmd, env_vars=self._env)
         try:
-            process_output = process.wait(timeout=self.timeout)
+            process_output = process.wait(timeout=timeout)
         except TimeoutException:
             logger.info('Command timed out, killing process...')
             process.kill()

opendevin/runtime/sandbox.py

@@ -19,7 +19,7 @@ def add_to_env(self, key: str, value: str):
         self._env[key] = value
 
     @abstractmethod
-    def execute(self, cmd: str) -> tuple[int, str]:
+    def execute(self, cmd: str, timeout: int | None = None) -> tuple[int, str]:
         pass
 
     @abstractmethod