Explorar el Código

feat(sandbox): add container process timeout kill machenism (#316)

* add timeout

* set timeout
Yufan Song hace 2 años
padre
commit
d993162801
Se han modificado 1 ficheros con 16 adiciones y 3 borrados
  1. 16 3
      opendevin/sandbox/sandbox.py

+ 16 - 3
opendevin/sandbox/sandbox.py

@@ -8,6 +8,7 @@ from collections import namedtuple
 from typing import Dict, List, Tuple
 
 import docker
+import concurrent.futures
 
 InputType = namedtuple("InputType", ["content"])
 OutputType = namedtuple("OutputType", ["content"])
@@ -148,9 +149,21 @@ class DockerInteractive:
 
     def execute(self, cmd: str) -> Tuple[int, str]:
         # TODO: each execute is not stateful! We need to keep track of the current working directory
-        exit_code, logs = self.container.exec_run(
-            self.get_exec_cmd(cmd), workdir="/workspace"
-        )
+        def run_command(container, command):
+            return container.exec_run(command,workdir="/workspace")
+        # Use ThreadPoolExecutor to control command and set timeout
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            future = executor.submit(run_command, self.container, self.get_exec_cmd(cmd))
+            try:
+                exit_code, logs = future.result(timeout=self.timeout)
+            except concurrent.futures.TimeoutError:
+                print("Command timed out, killing process...")
+                pid = self.get_pid(cmd)
+                if pid is not None:
+                    self.container.exec_run(
+                        f"kill -9 {pid}", workdir="/workspace"
+                    )
+                return -1, f"Command: \"{cmd}\" timed out"
         return exit_code, logs.decode("utf-8")
 
     def execute_in_background(self, cmd: str) -> BackgroundCommand: