|
|
@@ -4,7 +4,7 @@ import uuid
|
|
|
import time
|
|
|
import select
|
|
|
import docker
|
|
|
-from typing import Tuple
|
|
|
+from typing import Tuple, Dict, List
|
|
|
from collections import namedtuple
|
|
|
import atexit
|
|
|
|
|
|
@@ -12,6 +12,35 @@ InputType = namedtuple("InputType", ["content"])
|
|
|
OutputType = namedtuple("OutputType", ["content"])
|
|
|
|
|
|
CONTAINER_IMAGE = os.getenv("SANDBOX_CONTAINER_IMAGE", "opendevin/sandbox:v0.1")
|
|
|
+# FIXME: On some containers, the devin user doesn't have enough permission, e.g. to install packages
|
|
|
+# How do we make this more flexible?
|
|
|
+RUN_AS_DEVIN = os.getenv("RUN_AS_DEVIN", "true").lower() != "false"
|
|
|
+
|
|
|
+class BackgroundCommand:
|
|
|
+ def __init__(self, id: int, command: str, result):
|
|
|
+ self.id = id
|
|
|
+ self.command = command
|
|
|
+ self.result = result
|
|
|
+
|
|
|
+ def read_logs(self) -> str:
|
|
|
+ # TODO: get an exit code if process is exited
|
|
|
+ logs = ""
|
|
|
+ while True:
|
|
|
+ ready_to_read, _, _ = select.select([self.result.output], [], [], .1) # type: ignore[has-type]
|
|
|
+ if ready_to_read:
|
|
|
+ data = self.result.output.read(4096) # type: ignore[has-type]
|
|
|
+ if not data:
|
|
|
+ break
|
|
|
+ # FIXME: we're occasionally seeing some escape characters like `\x02` and `\x00` in the logs...
|
|
|
+ chunk = data.decode('utf-8')
|
|
|
+ logs += chunk
|
|
|
+ else:
|
|
|
+ break
|
|
|
+ return logs
|
|
|
+
|
|
|
+ def kill(self):
|
|
|
+ # FIXME: this doesn't actually kill the process!
|
|
|
+ self.result.output.close()
|
|
|
|
|
|
|
|
|
USER_ID = 1000
|
|
|
@@ -21,6 +50,9 @@ elif hasattr(os, "getuid"):
|
|
|
USER_ID = os.getuid()
|
|
|
|
|
|
class DockerInteractive:
|
|
|
+ closed = False
|
|
|
+ cur_background_id = 0
|
|
|
+ background_commands : Dict[int, BackgroundCommand] = {}
|
|
|
|
|
|
def __init__(
|
|
|
self,
|
|
|
@@ -55,45 +87,55 @@ class DockerInteractive:
|
|
|
self.container_name = f"sandbox-{self.instance_id}"
|
|
|
|
|
|
self.restart_docker_container()
|
|
|
+ if RUN_AS_DEVIN:
|
|
|
+ self.setup_devin_user()
|
|
|
+ atexit.register(self.cleanup)
|
|
|
+
|
|
|
+ def setup_devin_user(self):
|
|
|
+ uid = os.getuid()
|
|
|
exit_code, logs = self.container.exec_run([
|
|
|
'/bin/bash', '-c',
|
|
|
- f'useradd --shell /bin/bash -u {USER_ID} -o -c \"\" -m devin'
|
|
|
+ f'useradd --shell /bin/bash -u {uid} -o -c \"\" -m devin'
|
|
|
],
|
|
|
workdir="/workspace"
|
|
|
)
|
|
|
- # regester container cleanup function
|
|
|
- atexit.register(self.cleanup)
|
|
|
|
|
|
- def read_logs(self) -> str:
|
|
|
- if not hasattr(self, "log_generator"):
|
|
|
- return ""
|
|
|
- logs = ""
|
|
|
- while True:
|
|
|
- ready_to_read, _, _ = select.select([self.log_generator], [], [], .1) # type: ignore[has-type]
|
|
|
- if ready_to_read:
|
|
|
- data = self.log_generator.read(4096) # type: ignore[has-type]
|
|
|
- if not data:
|
|
|
- break
|
|
|
- # FIXME: we're occasionally seeing some escape characters like `\x02` and `\x00` in the logs...
|
|
|
- chunk = data.decode('utf-8')
|
|
|
- logs += chunk
|
|
|
- else:
|
|
|
- break
|
|
|
- return logs
|
|
|
+ def get_exec_cmd(self, cmd: str) -> List[str]:
|
|
|
+ if RUN_AS_DEVIN:
|
|
|
+ return ['su', 'devin', '-c', cmd]
|
|
|
+ else:
|
|
|
+ return ['/bin/bash', '-c', cmd]
|
|
|
+
|
|
|
+ def read_logs(self, id) -> str:
|
|
|
+ if id not in self.background_commands:
|
|
|
+ raise ValueError("Invalid background command id")
|
|
|
+ bg_cmd = self.background_commands[id]
|
|
|
+ return bg_cmd.read_logs()
|
|
|
|
|
|
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(['su', 'devin', '-c', cmd], workdir="/workspace")
|
|
|
+ exit_code, logs = self.container.exec_run(self.get_exec_cmd(cmd), workdir="/workspace")
|
|
|
return exit_code, logs.decode('utf-8')
|
|
|
|
|
|
- def execute_in_background(self, cmd: str) -> None:
|
|
|
- self.log_time = time.time()
|
|
|
- result = self.container.exec_run(['su', 'devin', '-c', cmd], socket=True, workdir="/workspace")
|
|
|
- self.log_generator = result.output # socket.SocketIO
|
|
|
- self.log_generator._sock.setblocking(0)
|
|
|
+ def execute_in_background(self, cmd: str) -> BackgroundCommand:
|
|
|
+ result = self.container.exec_run(self.get_exec_cmd(cmd), socket=True, workdir="/workspace")
|
|
|
+ result.output._sock.setblocking(0)
|
|
|
+ bg_cmd = BackgroundCommand(self.cur_background_id, cmd, result)
|
|
|
+ self.background_commands[bg_cmd.id] = bg_cmd
|
|
|
+ self.cur_background_id += 1
|
|
|
+ return bg_cmd
|
|
|
+
|
|
|
+ def kill_background(self, id: int) -> BackgroundCommand:
|
|
|
+ if id not in self.background_commands:
|
|
|
+ raise ValueError("Invalid background command id")
|
|
|
+ bg_cmd = self.background_commands[id]
|
|
|
+ bg_cmd.kill()
|
|
|
+ self.background_commands.pop(id)
|
|
|
+ return bg_cmd
|
|
|
|
|
|
def close(self):
|
|
|
self.stop_docker_container()
|
|
|
+ self.closed = True
|
|
|
|
|
|
def stop_docker_container(self):
|
|
|
docker_client = docker.from_env()
|
|
|
@@ -145,8 +187,9 @@ class DockerInteractive:
|
|
|
|
|
|
# clean up the container, cannot do it in __del__ because the python interpreter is already shutting down
|
|
|
def cleanup(self):
|
|
|
+ if self.closed:
|
|
|
+ return
|
|
|
self.container.remove(force=True)
|
|
|
- print("Finish cleaning up Docker container")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
import argparse
|
|
|
@@ -165,10 +208,7 @@ if __name__ == "__main__":
|
|
|
)
|
|
|
print("Interactive Docker container started. Type 'exit' or use Ctrl+C to exit.")
|
|
|
|
|
|
- bg = DockerInteractive(
|
|
|
- workspace_dir=args.directory,
|
|
|
- )
|
|
|
- bg.execute_in_background("while true; do echo 'dot ' && sleep 1; done")
|
|
|
+ bg_cmd = docker_interactive.execute_in_background("while true; do echo 'dot ' && sleep 1; done")
|
|
|
|
|
|
sys.stdout.flush()
|
|
|
try:
|
|
|
@@ -181,11 +221,16 @@ if __name__ == "__main__":
|
|
|
if user_input.lower() == "exit":
|
|
|
print("Exiting...")
|
|
|
break
|
|
|
+ if user_input.lower() == "kill":
|
|
|
+ docker_interactive.kill_background(bg_cmd.id)
|
|
|
+ print("Background process killed")
|
|
|
+ continue
|
|
|
exit_code, output = docker_interactive.execute(user_input)
|
|
|
print("exit code:", exit_code)
|
|
|
print(output + "\n", end="")
|
|
|
- logs = bg.read_logs()
|
|
|
- print("background logs:", logs, "\n")
|
|
|
+ if bg_cmd.id in docker_interactive.background_commands:
|
|
|
+ logs = docker_interactive.read_logs(bg_cmd.id)
|
|
|
+ print("background logs:", logs, "\n")
|
|
|
sys.stdout.flush()
|
|
|
except KeyboardInterrupt:
|
|
|
print("\nExiting...")
|