From a77b5a62a4a8a362173d77f73e45ebefd699a851 Mon Sep 17 00:00:00 2001
From: Liqun Li <liqli@microsoft.com>
Date: Wed, 28 Feb 2024 11:06:29 +0800
Subject: [PATCH] fix issue on deleting the conn file inside docker

---
 ces_container/docker_entry.py | 18 ++++++++++++++----
 requirements.txt              |  3 +++
 taskweaver/ces/environment.py | 10 ++++++++--
 3 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/ces_container/docker_entry.py b/ces_container/docker_entry.py
index cc12afd7..ba399f14 100644
--- a/ces_container/docker_entry.py
+++ b/ces_container/docker_entry.py
@@ -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",
@@ -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,
diff --git a/requirements.txt b/requirements.txt
index cc07fac1..367c0def 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -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
@@ -24,3 +25,5 @@ vcrpy>=5.0.0
 colorama>=0.4.6
 
 
+
+
diff --git a/taskweaver/ces/environment.py b/taskweaver/ces/environment.py
index 44996e37..aac0c87b 100644
--- a/taskweaver/ces/environment.py
+++ b/taskweaver/ces/environment.py
@@ -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:
@@ -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(