Execution inside container (microsoft#217)

- add option for running code inside container
- add stop function at chat end

---------

Co-authored-by: Shilin HE <[email protected]>

README.md

@@ -23,12 +23,13 @@ Unlike many agent frameworks that only track the chat history with LLMs in text,
 
 
 ## 🆕 News
+- 📅2024-03-04: TaskWeaver now supports a [container](https://microsoft.github.io/TaskWeaver/docs/code_execution) mode, which provides a more secure environment for code execution.🐳
 - 📅2024-02-28: TaskWeaver now offers a [CLI-only](https://microsoft.github.io/TaskWeaver/docs/cli_only) mode, enabling users to interact seamlessly with the Command Line Interface (CLI) using natural language.📟
 - 📅2024-02-01: TaskWeaver now has a plugin [document_retriever](https://github.com/microsoft/TaskWeaver/blob/main/project/plugins/README.md#document_retriever) for RAG based on a knowledge base.📚
 - 📅2024-01-30: TaskWeaver introduces a new plugin-only mode that securely generates calls to specified plugins without producing extraneous code.🪡 
 - 📅2024-01-23: TaskWeaver can now be personalized by transforming your chat histories into enduring [experiences](https://microsoft.github.io/TaskWeaver/docs/experience) 🎉
 - 📅2024-01-17: TaskWeaver now has a plugin [vision_web_explorer](https://github.com/microsoft/TaskWeaver/blob/main/project/plugins/README.md#vision_web_explorer) that can open a web browser and explore websites.🌐
-- 📅2024-01-15: TaskWeaver now supports Streaming♒ in both UI and command line.✌️
+<!-- - 📅2024-01-15: TaskWeaver now supports Streaming♒ in both UI and command line.✌️ -->
 <!-- - 📅2024-01-01: Welcome join TaskWeaver [Discord](https://discord.gg/Z56MXmZgMb). -->
 <!-- - 📅2023-12-21: TaskWeaver now supports a number of LLMs, such as LiteLLM, Ollama, Gemini, and QWen🎈.) -->
 <!-- - 📅2023-12-21: TaskWeaver Website is now [available]&#40;https://microsoft.github.io/TaskWeaver/&#41; with more documentations.) -->

ces_container/Dockerfile

@@ -0,0 +1,27 @@
+# Use the official Python 3.10 image as the base image  
+FROM python:3.10
+
+# Set the working directory to /app
+WORKDIR /app
+
+# Copy the requrements file
+COPY requirements.txt /app/requirements.txt
+
+# Install any necessary dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Install additional dependencies below for the plugins
+# e.g., RUN pip install your-package-name
+
+# Copy the project code
+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
+COPY ces_container/docker_entry.py /app/docker_entry.py
+
+ENV PYTHONPATH "${PYTHONPATH}:/app"
+
+CMD ["python", "docker_entry.py"]
+
+

ces_container/docker_entry.py

@@ -0,0 +1,62 @@
+import os
+import signal
+import time
+
+from taskweaver.ces import Environment, EnvMode
+
+env_id = os.getenv(
+    "TASKWEAVER_ENV_ID",
+    "local",
+)
+env_dir = os.getenv(
+    "TASKWEAVER_ENV_DIR",
+    "/app",
+)
+session_id = os.getenv(
+    "TASKWEAVER_SESSION_ID",
+    "session_id",
+)
+port_start = int(
+    os.getenv(
+        "TASKWEAVER_PORT_START",
+        "12345",
+    ),
+)
+kernel_id = os.getenv(
+    "TASKWEAVER_KERNEL_ID",
+    "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.start_session(
+        session_id=session_id,
+        port_start_inside_container=port_start,
+        kernel_id_inside_container=kernel_id,
+    )
+
+    print(f"Session {session_id} is running at {env_dir} inside a container.")
+
+    # Keep the script running until it receives a termination signal
+    try:
+        # Keep the script running indefinitely
+        while True:
+            time.sleep(10)  # Sleep for 10 seconds
+    except KeyboardInterrupt:
+        # Handle Ctrl-C or other interruption signals
+        pass
+    finally:
+        # Clean up and shut down the kernel
+        env.stop_session(session_id)

playground/UI/app.py

@@ -365,6 +365,27 @@ def format_message(self, message: str, is_end: bool) -> str:
 async def start():
     user_session_id = cl.user_session.get("id")
     app_session_dict[user_session_id] = app.get_session()
+    exec_kernel_mode = app_session_dict[user_session_id].code_executor.get_execution_mode()
+    print(f"Starting session in `{exec_kernel_mode}` mode")
+    if exec_kernel_mode == "local":
+        print(
+            "Code running in local mode "
+            "may incur security risks, such as file system access. "
+            "Please be cautious when executing code. "
+            "For higher security, consider using the `container` mode by setting "
+            "the `execution_service.kernel_mode` to `container`. "
+            "For more information, please refer to the documentation ("
+            "https://microsoft.github.io/TaskWeaver/docs/code_execution).",
+        )
+
+
+@cl.on_chat_end
+async def end():
+    user_session_id = cl.user_session.get("id")
+    app_session = app_session_dict[user_session_id]
+    print(f"Stopping session {app_session.session_id}")
+    app_session.stop()
+    app_session_dict.pop(user_session_id)
 
 
 @cl.on_message

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
 
 
+
+

scripts/build.ps1

@@ -0,0 +1,20 @@
+$scriptDirectory = $PSScriptRoot
+Write-Host "The script directory is: $scriptDirectory"
+
+$imageName = "taskweaver/executor"
+$taskweaverPath = Join-Path -Path $scriptDirectory -ChildPath "..\taskweaver"
+$dockerfilePath = Join-Path -Path $scriptDirectory -ChildPath "..\ces_container\Dockerfile"
+$contextPath = Join-Path -Path $scriptDirectory -ChildPath "..\"
+
+if (Test-Path $taskweaverPath) {
+    Write-Host "Found module files from $taskweaverPath"
+    Write-Host "Dockerfile path: $dockerfilePath"
+    Write-Host "Context path: $contextPath"
+} else {
+    Write-Host "Local files not found."
+    exit 1
+}
+
+# Build the Docker image
+docker build -t $imageName -f $dockerfilePath $contextPath
+

scripts/build.sh

@@ -0,0 +1,22 @@
+#!/bin/bash
+
+# Get the directory containing the script file
+scriptDirectory="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
+echo "The script directory is: $scriptDirectory"
+
+imageName="taskweaver/executor"
+taskweaverPath="$scriptDirectory/../taskweaver"
+dockerfilePath="$scriptDirectory/../ces_container/Dockerfile"
+contextPath="$scriptDirectory/../"
+
+if [ -d "$taskweaverPath" ]; then
+    echo "Found module files from $taskweaverPath"
+    echo "Dockerfile path: $dockerfilePath"
+    echo "Context path: $contextPath"
+else
+    echo "Local files not found."
+    exit 1
+fi
+
+# Build the Docker image
+docker build -t "$imageName" -f "$dockerfilePath" "$contextPath"

taskweaver/app/app.py

@@ -49,6 +49,12 @@ def get_session(
     ) -> Session:
         return self.session_manager.get_session(session_id, prev_round_id)
 
+    def stop_session(self, session_id: str) -> None:
+        self.session_manager.stop_session(session_id)
+
+    def stop_all_sessions(self) -> None:
+        self.session_manager.stop_all_sessions()
+
     @staticmethod
     def discover_app_dir(
         app_dir: Optional[str] = None,

taskweaver/app/session_manager.py

@@ -47,6 +47,17 @@ def update_session(self, session: Session) -> None:
         """update session in session store"""
         self.session_store.set_session(session.session_id, session)
 
+    def stop_session(self, session_id: str) -> None:
+        """stop session in session store"""
+        session = self._get_session_from_store(session_id, False)
+        if session is not None:
+            session.stop()
+            self.session_store.remove_session(session_id)
+
+    def stop_all_sessions(self) -> None:
+        for session_id in self.session_store.get_all_session_ids():
+            self.stop_session(session_id)
+
     @overload
     def _get_session_from_store(
         self,

taskweaver/app/session_store.py

@@ -1,5 +1,5 @@
 import abc
-from typing import Dict, Optional
+from typing import Dict, Iterable, Optional
 
 from ..session.session import Session
 
@@ -21,6 +21,10 @@ def remove_session(self, session_id: str) -> None:
     def has_session(self, session_id: str) -> bool:
         pass
 
+    @abc.abstractmethod
+    def get_all_session_ids(self) -> Iterable[str]:
+        pass
+
 
 class InMemorySessionStore(SessionStore):
     def __init__(self) -> None:
@@ -37,3 +41,6 @@ def remove_session(self, session_id: str) -> None:
 
     def has_session(self, session_id: str) -> bool:
         return session_id in self.sessions
+
+    def get_all_session_ids(self) -> Iterable[str]:
+        return self.sessions.keys()

taskweaver/ces/__init__.py

@@ -1,6 +1,15 @@
+from typing import Literal
+
 from taskweaver.ces.common import Manager
+from taskweaver.ces.environment import Environment, EnvMode
 from taskweaver.ces.manager.sub_proc import SubProcessManager
 
 
-def code_execution_service_factory(env_dir: str) -> Manager:
-    return SubProcessManager(env_dir=env_dir)
+def code_execution_service_factory(
+    env_dir: str,
+    kernel_mode: Literal["local", "container"] = "local",
+) -> Manager:
+    return SubProcessManager(
+        env_dir=env_dir,
+        kernel_mode=kernel_mode,
+    )

taskweaver/ces/common.py

@@ -126,3 +126,7 @@ def get_session_client(
         cwd: Optional[str] = None,
     ) -> Client:
         ...
+
+    @abstractmethod
+    def get_kernel_mode(self) -> Literal["local", "container"] | None:
+        ...