From 891fce9efaf3ce79b962af9d319c487cc020dce6 Mon Sep 17 00:00:00 2001
From: Liqun Li <liqli@microsoft.com>
Date: Fri, 23 Feb 2024 17:40:59 +0800
Subject: [PATCH] fix bug

---
 ces_container/Dockerfile                          |  3 +--
 ces_container/build.ps1                           |  7 ++-----
 ces_container/build.sh                            |  7 ++-----
 {taskweaver/ces => ces_container}/docker_entry.py |  3 ---
 taskweaver/ces/environment.py                     | 13 +++++--------
 5 files changed, 10 insertions(+), 23 deletions(-)
 rename {taskweaver/ces => ces_container}/docker_entry.py (96%)

diff --git a/ces_container/Dockerfile b/ces_container/Dockerfile
index 864b1f49..0094432e 100644
--- a/ces_container/Dockerfile
+++ b/ces_container/Dockerfile
@@ -15,8 +15,7 @@ COPY taskweaver/ces /app/taskweaver/ces
 COPY taskweaver/plugin /app/taskweaver/plugin
 COPY taskweaver/module /app/taskweaver/module
 COPY taskweaver/__init__.py /app/taskweaver/__init__.py
-
-RUN mv /app/taskweaver/ces/docker_entry.py /app/docker_entry.py
+COPY docker_entry.py /app/docker_entry.py
 
 ENV PYTHONPATH "${PYTHONPATH}:/app"
 
diff --git a/ces_container/build.ps1 b/ces_container/build.ps1
index 3a685acb..d908d3d6 100644
--- a/ces_container/build.ps1
+++ b/ces_container/build.ps1
@@ -15,11 +15,8 @@ if (Test-Path $taskweaverPath) {
     Copy-Item -Path $taskweaverPath -Destination ".\taskweaver" -Recurse
     Copy-Item -Path $requirementsPath -Destination ".\requirements.txt"
 } else {
-    Write-Host "Cloning repository and using repository files"
-    git clone https://github.com/microsoft/TaskWeaver.git
-    Move-Item -Path ".\TaskWeaver\taskweaver" -Destination ".\taskweaver"
-    Move-Item -Path ".\TaskWeaver\requirements.txt" -Destination ".\requirements.txt"
-    Remove-Item -Path ".\TaskWeaver" -Recurse -Force
+    Write-Host "Local files not found."
+    exit 1
 }
 
 # Build the Docker image
diff --git a/ces_container/build.sh b/ces_container/build.sh
index 9386c7e0..b867bb9a 100644
--- a/ces_container/build.sh
+++ b/ces_container/build.sh
@@ -14,11 +14,8 @@ if [ -d "$taskweaverPath" ]; then
     cp -r "$taskweaverPath" .
     cp "$requirementsPath" .
 else
-    echo "Cloning repository and using repository files"
-    git clone https://github.com/microsoft/TaskWeaver.git
-    mv ".TaskWeaver/taskweaver" .
-    mv ".TaskWeaver/requirements.txt" .
-    rm -rf ".TaskWeaver"
+    echo "Local files not found."
+    exit 1
 fi
 
 # Build the Docker image
diff --git a/taskweaver/ces/docker_entry.py b/ces_container/docker_entry.py
similarity index 96%
rename from taskweaver/ces/docker_entry.py
rename to ces_container/docker_entry.py
index 3baf9a1d..a9aef638 100644
--- a/taskweaver/ces/docker_entry.py
+++ b/ces_container/docker_entry.py
@@ -1,9 +1,6 @@
 import os
-import sys
 import time
 
-sys.path.append("/app")
-
 from taskweaver.ces import Environment, EnvMode
 
 # Flag to control the main loop
diff --git a/taskweaver/ces/environment.py b/taskweaver/ces/environment.py
index b226fe60..4f78b4a3 100644
--- a/taskweaver/ces/environment.py
+++ b/taskweaver/ces/environment.py
@@ -138,7 +138,10 @@ def __init__(
         self.env_dir = env_dir if env_dir is not None else os.getcwd()
         self.mode = env_mode
         if self.mode == EnvMode.SubProcess or self.mode == EnvMode.InsideContainer:
-            self.multi_kernel_manager = self._init_kernel_manager()
+            self.multi_kernel_manager = TaskWeaverMultiKernelManager(
+                default_kernel_name="taskweaver",
+                kernel_spec_manager=KernelSpecProvider(),
+            )
         elif self.mode == EnvMode.OutsideContainer:
             try:
                 import docker
@@ -150,6 +153,7 @@ def __init__(
             self.port_start = 12300
         else:
             raise ValueError(f"Unsupported environment mode {env_mode}")
+        atexit.register(self.clean_up)
 
     def clean_up(self) -> None:
         for session in self.session_dict.values():
@@ -158,13 +162,6 @@ def clean_up(self) -> None:
             except Exception as e:
                 logger.error(e)
 
-    def _init_kernel_manager(self) -> MultiKernelManager:
-        atexit.register(self.clean_up)
-        return TaskWeaverMultiKernelManager(
-            default_kernel_name="taskweaver",
-            kernel_spec_manager=KernelSpecProvider(),
-        )
-
     def _get_connection_file(self, session_id: str, kernel_id: str) -> str:
         return os.path.join(
             self._get_session(session_id).session_dir,