Fix container related issues (#426)

1. support custom image 
2. fix the issue of import failure

.gitattributes

@@ -0,0 +1 @@
+*.sh text eol=lf

docker/ces_container/Dockerfile

@@ -17,7 +17,7 @@ COPY taskweaver/__init__.py /app/taskweaver/__init__.py
 COPY docker/ces_container/entrypoint.sh /app/entrypoint.sh
 RUN chmod +x /app/entrypoint.sh
 
-ENV PYTHONPATH "${PYTHONPATH}:/app"
+ENV PYTHONPATH="/app"
 
 ENTRYPOINT ["/app/entrypoint.sh"]
 

scripts/build_executor.ps1

@@ -1,7 +1,7 @@
 $scriptDirectory = $PSScriptRoot
 Write-Host "The script directory is: $scriptDirectory"
 
-$version = "0.2"
+$version = "0.3"
 $imageName = "taskweavercontainers/taskweaver-executor"
 $imageFullName = "${imageName}:${version}"
 

taskweaver/ces/__init__.py

@@ -1,4 +1,4 @@
-from typing import Literal
+from typing import Literal, Optional
 
 from taskweaver.ces.common import Manager
 from taskweaver.ces.manager.defer import DeferredManager
@@ -8,11 +8,13 @@
 def code_execution_service_factory(
     env_dir: str,
     kernel_mode: Literal["local", "container"] = "local",
+    custom_image: Optional[str] = None,
 ) -> Manager:
     def sub_proc_manager_factory() -> SubProcessManager:
         return SubProcessManager(
             env_dir=env_dir,
             kernel_mode=kernel_mode,
+            custom_image=custom_image,
         )
 
     return DeferredManager(

taskweaver/ces/environment.py

@@ -111,12 +111,15 @@ class EnvMode(enum.Enum):
 
 
 class Environment:
+    DEFAULT_IMAGE = "taskweavercontainers/taskweaver-executor:latest"
+
     def __init__(
         self,
         env_id: Optional[str] = None,
         env_dir: Optional[str] = None,
         env_mode: Optional[EnvMode] = EnvMode.Local,
         port_start_inside_container: Optional[int] = 12345,
+        custom_image: Optional[str] = None,
     ) -> None:
         self.session_dict: Dict[str, EnvSession] = {}
         self.id = get_id(prefix="env") if env_id is None else env_id
@@ -145,19 +148,27 @@ def __init__(
             except docker.errors.DockerException as e:
                 raise docker.errors.DockerException(f"Failed to connect to Docker daemon: {e}. ")
 
-            self.image_name = "taskweavercontainers/taskweaver-executor:latest"
-            try:
-                local_image = self.docker_client.images.get(self.image_name)
-                registry_image = self.docker_client.images.get_registry_data(self.image_name)
-                if local_image.id != registry_image.id:
-                    logger.info(f"Local image {local_image.id} does not match registry image {registry_image.id}.")
-                    raise docker.errors.ImageNotFound("Local image is outdated.")
-            except docker.errors.ImageNotFound:
-                logger.info("Pulling image from docker.io.")
+            if custom_image:
+                logger.info(f"Using custom image {custom_image}.")
+                self.image_name = custom_image
+                try:
+                    self.docker_client.images.get(self.image_name)
+                except docker.errors.ImageNotFound:
+                    raise docker.errors.ImageNotFound(f"Custom image {self.image_name} not found.")
+            else:
+                self.image_name = self.DEFAULT_IMAGE
                 try:
-                    self.docker_client.images.pull(self.image_name)
-                except docker.errors.DockerException as e:
-                    raise docker.errors.DockerException(f"Failed to pull image: {e}. ")
+                    local_image = self.docker_client.images.get(self.image_name)
+                    registry_image = self.docker_client.images.get_registry_data(self.image_name)
+                    if local_image.id != registry_image.id:
+                        logger.info(f"Local image {local_image.id} does not match registry image {registry_image.id}.")
+                        raise docker.errors.ImageNotFound("Local image is outdated.")
+                except docker.errors.ImageNotFound:
+                    logger.info("Pulling image from docker.io.")
+                    try:
+                        self.docker_client.images.pull(self.image_name)
+                    except docker.errors.DockerException as e:
+                        raise docker.errors.DockerException(f"Failed to pull image: {e}. ")
 
             self.session_container_dict: Dict[str, str] = {}
             self.port_start_inside_container = port_start_inside_container

taskweaver/ces/kernel/launcher.py

@@ -1,7 +1,6 @@
 import os
 import sys
 
-from taskweaver.ces.kernel.ext import TaskWeaverZMQShellDisplayHook
 from taskweaver.ces.kernel.kernel_logging import logger
 
 kernel_mode = os.getenv("TASKWEAVER_KERNEL_MODE", "local")
@@ -56,6 +55,8 @@ def start_app():
     from ipykernel.kernelapp import IPKernelApp
     from ipykernel.zmqshell import ZMQInteractiveShell
 
+    from taskweaver.ces.kernel.ext import TaskWeaverZMQShellDisplayHook
+
     # override displayhook_class for skipping output suppress token issue
     ZMQInteractiveShell.displayhook_class = TaskWeaverZMQShellDisplayHook
 
@@ -82,6 +83,11 @@ def start_app():
 if __name__ == "__main__":
     if sys.path[0] == "":
         del sys.path[0]
+    import site
+
+    user_site_packages = site.getusersitepackages()
+    if user_site_packages not in sys.path:
+        sys.path.append(site.getusersitepackages())
     logger.info("Starting process...")
     logger.info("sys.path: %s", sys.path)
     logger.info("os.getcwd(): %s", os.getcwd())

taskweaver/ces/manager/sub_proc.py

@@ -57,6 +57,7 @@ def __init__(
         env_id: Optional[str] = None,
         env_dir: Optional[str] = None,
         kernel_mode: KernelModeType = "local",
+        custom_image: Optional[str] = None,
     ) -> None:
         from taskweaver.ces.environment import Environment, EnvMode
 
@@ -76,6 +77,7 @@ def __init__(
             env_id,
             env_dir,
             env_mode=env_mode,
+            custom_image=custom_image,
         )
 
     def initialize(self) -> None:

taskweaver/module/execution_service.py

@@ -19,6 +19,7 @@ def _configure(self) -> None:
             "kernel_mode",
             "container",
         )
+        assert self.kernel_mode in ["local", "container"], f"Invalid kernel mode: {self.kernel_mode}"
         if self.kernel_mode == "local":
             print(
                 "TaskWeaver is running in the `local` mode. This implies that "
@@ -27,6 +28,11 @@ def _configure(self) -> None:
                 "More information can be found in the documentation "
                 "(https://microsoft.github.io/TaskWeaver/docs/code_execution/).",
             )
+        self.custom_image = self._get_str(
+            "custom_image",
+            default=None,
+            required=False,
+        )
 
 
 class ExecutionServiceModule(Module):
@@ -39,5 +45,6 @@ def provide_executor_manager(self, config: ExecutionServiceConfig) -> Manager:
             self.manager = code_execution_service_factory(
                 env_dir=config.env_dir,
                 kernel_mode=config.kernel_mode,
+                custom_image=config.custom_image,
             )
         return self.manager