Skip to content

Commit

Permalink
fix issue on deleting the conn file inside docker
Browse files Browse the repository at this point in the history
  • Loading branch information
liqul committed Feb 28, 2024
1 parent 6ce6e90 commit a77b5a6
Show file tree
Hide file tree
Showing 3 changed files with 25 additions and 6 deletions.
18 changes: 14 additions & 4 deletions ces_container/docker_entry.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,9 @@
import os
import signal
import time

from taskweaver.ces import Environment, EnvMode

# Flag to control the main loop
keep_running = True

env_id = os.getenv(
"TASKWEAVER_ENV_ID",
"local",
Expand All @@ -29,8 +27,20 @@
"kernel_id",
)

env = Environment(env_id, env_dir, env_mode=EnvMode.InsideContainer)


def signal_handler(sig, frame):
print("Received termination signal. Shutting down the environment.")
env.stop_session(session_id)
exit(0)


# Register the signal handler
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)

if __name__ == "__main__":
env = Environment(env_id, env_dir, env_mode=EnvMode.InsideContainer)
env.start_session(
session_id=session_id,
port_start_inside_container=port_start,
Expand Down
3 changes: 3 additions & 0 deletions requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ jsonschema==4.20.0
injector>=0.21.0
ijson>=3.2.3
requests>=2.31.0
docker>=7.0.0

# Code Execution related
ipykernel==6.26.0
Expand All @@ -24,3 +25,5 @@ vcrpy>=5.0.0
colorama>=0.4.6




10 changes: 8 additions & 2 deletions taskweaver/ces/environment.py
Original file line number Diff line number Diff line change
Expand Up @@ -153,9 +153,15 @@ def __init__(
elif self.mode == EnvMode.OutsideContainer:
try:
import docker
import docker.errors
except ImportError:
raise ImportError("docker package is required for container-based kernel.")
self.docker_client = docker.from_env()

try:
self.docker_client = docker.from_env()
except docker.errors.DockerException as e:
raise docker.errors.DockerException(f"Failed to connect to Docker daemon: {e}. ")

self.session_container_dict: Dict[str, str] = {}
self.port_start_inside_container = port_start_inside_container
else:
Expand Down Expand Up @@ -310,7 +316,7 @@ def start_session(
os.environ["TASKWEAVER_SESSION_DIR"] = session.session_dir
ces_session_dir = os.path.join(session.session_dir, "ces")
connection_file = self._get_connection_file(session_id, kernel_id_inside_container)
cwd = cwd if cwd is not None else os.path.join(session.session_dir, "cwd")
cwd = os.path.join(session.session_dir, "cwd")

kernel_env = os.environ.copy()
kernel_env.update(
Expand Down

0 comments on commit a77b5a6

Please sign in to comment.