Merge pull request #14 from questcollector/feat/jupyter

add Jupyter server code executor

Author: questcollector

python/packages/autogen-kubernetes-mcp/README.md

@@ -23,34 +23,52 @@ uv pip install autogen-kubernetes-mcp
 3. Run via `uvx`
 
 ```bash
-uvx autogen-kubernetes-mcp --namespace my-namespace --image python:3.11-slim
+uvx autogen-kubernetes-mcp commandline --namespace my-namespace --image python:3.11-slim
 ```
-> Note: When using `uvx`, arguments must be passed after `--`.
 
 4. Run via Python module
 
 ```bash
-python -m autogen_kubernetes_mcp --namespace my-namespace --image python:3.11-slim
+python -m autogen_kubernetes_mcp commandline --namespace my-namespace --image python:3.11-slim
 ```
 
 ## Command-line Arguments
 
-Command-line arguments are used when creating PodCommandLineCodeExecutor and MCP server.
+Command-line arguments are used when creating `PodCommandLineCodeExecutor`/`PodJupyterCodeExecutor` and MCP server.
+
+`type` arguments: jupyter type creates `PodJupyterCodeExecutor` for python code executor tool.
+
+MCP server have to run in same kubernetes cluster because `PodJupyterCodeExecutor` only uses service FQDN to connect jupyter server.
+
+Not validated for non-text outputs yet.
+
+```mermaid
+flowchart LR
+  subgraph K8sCluster["Kubernetes Cluster"]
+    E[PodJupyterCodeExecutor]
+    J[JupyterServerPod]
+  end
+
+  E --generate--> J
+```
 
 All the arguments are optional
 
 |Argument|Description|Default|
 |--|--|--|
+|`type`|Code Executor type, commandline(stateless), jupyter(stateful) supported||
 |`--host`|MCP server host address|`0.0.0.0`|
 |`--port`|MCP server port|`8000`|
-|`--kubeconfig`|Path to the kubeconfig file|(auto-detected)|
-|`--image`|Pod container image name|`python:3-slim`|
+|`--kubeconfig`|Path to the kubeconfig file|`None`(auto-detected)|
+|`--image`|Pod container image name|`None`|
 |`--pod-name`|Pod name|(auto-generated)|
 |`--timeout`|Code execution timeout(seconds)|`60`|
 |`--workspace-path`|Path inside the container where scripts are stored|`/workspace`|
 |`--namespace`, `-n`|Kubernetes namespace for Pod creation|`default`|
 |`--volume`|Kubernetes volume to mount into the Pod/container, accepts YAML format string, YAML file path|`None`|
 |`--pod-spec`|Custom Pod spec definition(YAML format string, YAML file path)|`None`|
+|`--command`|Custom executor container command|`None`|
+|`--args`|Custom executor container arguments|`None`|
 
 ## License
 

python/packages/autogen-kubernetes-mcp/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "autogen-kubernetes-mcp"
-version = "0.5.2"
+version = "0.6.0"
 description = "MPC server provides code interpreter with kubernetes workload"
 readme = "README.md"
 requires-python = ">=3.10"

python/packages/autogen-kubernetes-mcp/src/autogen_kubernetes_mcp/_executor.py

@@ -1,32 +1,73 @@
 from typing import Any
 
 from autogen_core import CancellationToken, ComponentModel
-from autogen_core.code_executor import CodeBlock
+from autogen_core.code_executor import CodeBlock, CodeExecutor
 from autogen_kubernetes.code_executors import (
     PodCommandLineCodeExecutor,
     PodCommandLineCodeExecutorConfig,
+    PodJupyterCodeExecutor,
+    PodJupyterCodeExecutorConfig,
+    PodJupyterServer,
+    PodJupyterServerConfig,
 )
 
 
-def make_executor(args: dict[str, Any]) -> PodCommandLineCodeExecutor:
-    pod_commandline_executor_config = PodCommandLineCodeExecutorConfig(**args)
-    component_model = ComponentModel(
-        provider="autogen_kubernetes.code_executors.PodCommandLineCodeExecutor",
-        component_type=PodCommandLineCodeExecutor.component_type,
-        version=PodCommandLineCodeExecutor.component_version,
-        component_version=PodCommandLineCodeExecutor.component_version,
-        description=PodCommandLineCodeExecutor.component_description,
-        label=PodCommandLineCodeExecutor.__name__,
-        config=pod_commandline_executor_config.model_dump(exclude_none=True),
-    )
-    return PodCommandLineCodeExecutor.load_component(component_model)
+async def make_executor(args: dict[str, Any]) -> list[Any]:
+    closable_instances = []
+    if args["type"] == "jupyter":
+        jupyter_server_config = PodJupyterServerConfig(**args)
+        jupyter_server = PodJupyterServer.load_component(
+            ComponentModel(
+                provider="autogen_kubernetes.code_executors.PodJupyterServer",
+                component_type=PodJupyterServer.component_type,
+                version=PodJupyterServer.component_version,
+                component_version=PodJupyterServer.component_version,
+                description=PodJupyterServer.component_description,
+                label=PodJupyterServer.__name__,
+                config=jupyter_server_config.model_dump(),
+            )
+        )
+        await jupyter_server.start()
+        jupyter_executor_config = PodJupyterCodeExecutorConfig(jupyter_server=jupyter_server, **args)
+        jupyter_executor = PodJupyterCodeExecutor.load_component(
+            ComponentModel(
+                provider="autogen_kubernetes.code_executors.PodJupyterCodeExecutor",
+                component_type=PodJupyterCodeExecutor.component_type,
+                version=PodJupyterCodeExecutor.component_version,
+                component_version=PodJupyterCodeExecutor.component_version,
+                description=PodJupyterCodeExecutor.component_description,
+                label=PodJupyterCodeExecutor.__name__,
+                config=jupyter_executor_config.model_dump(),
+            )
+        )
+        await jupyter_executor.start()
+        closable_instances.extend([jupyter_executor, jupyter_server])
+    else:
+        pod_commandline_executor_config = PodCommandLineCodeExecutorConfig(**args)
+        component_model = ComponentModel(
+            provider="autogen_kubernetes.code_executors.PodCommandLineCodeExecutor",
+            component_type=PodCommandLineCodeExecutor.component_type,
+            version=PodCommandLineCodeExecutor.component_version,
+            component_version=PodCommandLineCodeExecutor.component_version,
+            description=PodCommandLineCodeExecutor.component_description,
+            label=PodCommandLineCodeExecutor.__name__,
+            config=pod_commandline_executor_config.model_dump(exclude_none=True),
+        )
+        cmd_executor = PodCommandLineCodeExecutor.load_component(component_model)
+        await cmd_executor.start()
+        closable_instances.append(cmd_executor)
+    return closable_instances
 
 
-async def run_code(executor: PodCommandLineCodeExecutor, code: str) -> str:
+async def run_code(
+    executor: CodeExecutor,
+    code: str,
+    cancellation_token: CancellationToken,
+) -> str:
     code_result = await executor.execute_code_blocks(
         code_blocks=[
             CodeBlock(language="python", code=code),
         ],
-        cancellation_token=CancellationToken(),
+        cancellation_token=cancellation_token,
     )
     return code_result.output

python/packages/autogen-kubernetes-mcp/src/autogen_kubernetes_mcp/server.py

@@ -1,51 +1,92 @@
 import argparse
-from typing import Any
+import asyncio
+import uuid
+from contextlib import asynccontextmanager
+from typing import Any, AsyncIterator, TypedDict
 
-from mcp.server.fastmcp import FastMCP
+from autogen_core import CancellationToken
+from mcp.server.fastmcp import Context, FastMCP
 from mcp.types import ToolAnnotations
 
 from autogen_kubernetes_mcp._executor import make_executor, run_code
 
+STATEFUL_DESCRIPTION = r"""Use this tool to execute Python code in your chain of thought. The code will not be shown to the user. This tool should be used for internal reasoning, but not for code that is intended to be visible to the user (e.g. when creating plots, tables, or files).
+When you send a message containing Python code to python, it will be executed in a stateful Jupyter notebook environment. python will respond with the output of the execution or time out after 120.0 seconds. The drive at '/mnt/data' can be used to save and persist user files. Internet access for this session is UNKNOWN. Depends on the cluster"""
+STATELESS_DESCRIPTION = r"""Use this tool to execute Python code in your chain of thought. The code will not be shown to the user. This tool should be used for internal reasoning, but not for code that is intended to be visible to the user (e.g. when creating plots, tables, or files).
+When you send a message containing python code to python, it will be executed in a stateless docker container, and the stdout of that process will be returned to you."""
+
+sessions: dict[str, list[Any]] = {}
+
 
 def build_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser(description="autogen-kubernetes arguments")
-    parser.add_argument("--host", default="0.0.0.0")
-    parser.add_argument("--port", type=int, default=8000)
+    parser = argparse.ArgumentParser(description="autogen-kubernetes-mcp arguments")
+    parser.add_argument(
+        dest="type",
+        choices=["commandline", "jupyter"],
+        help="choose a code executor type (commandline, jupyter)",
+    )
+    parser.add_argument("--host", default="0.0.0.0", help="MCP server host")
+    parser.add_argument("--port", type=int, default=8000, help="MCP server port")
     parser.add_argument("--kubeconfig", dest="kube_config_file", default=None)
-    parser.add_argument("--image", default="python:3-slim")
+    parser.add_argument("--image", default=None, help="image for code execution pod")
     parser.add_argument("--pod-name", default=None)
-    parser.add_argument("--timeout", type=int, default=60)
+    parser.add_argument("--timeout", type=int, default=argparse.SUPPRESS)
     parser.add_argument("--workspace-path", default="/workspace")
     parser.add_argument("-n", "--namespace", default="default")
     parser.add_argument("--volume", default=None)
     parser.add_argument("--pod-spec", default=None)
+    parser.add_argument("--command", nargs="+", help="jupyter server pod commands", default=None)
+    parser.add_argument("--args", nargs="+", help="jupyter server pod arguments", default=None)
 
     return parser
 
 
+class State(TypedDict):
+    sid: str
+
+
 def build_server(args: dict[str, Any]) -> FastMCP:
+    tool_description = STATELESS_DESCRIPTION
+    if args["type"] == "jupyter":
+        tool_description = STATEFUL_DESCRIPTION
+        if "timeout" not in args:
+            args["timeout"] = 120
+
+    @asynccontextmanager
+    async def session_lifespan(app: FastMCP) -> AsyncIterator[State]:
+        sid = str(uuid.uuid4())
+        executor = await make_executor(args)
+        sessions[sid] = executor
+
+        yield {"sid": sid}
+
+        for instance in executor:
+            await instance.stop()
+        sessions.pop(sid, None)
+
     mcp = FastMCP(
         name="python",
-        instructions=r"""
-    Use this tool to execute Python code in your chain of thought. The code will not be shown to the user. This tool should be used for internal reasoning, but not for code that is intended to be visible to the user (e.g. when creating plots, tables, or files).
-    When you send a message containing python code to python, it will be executed in a stateless docker container, and the stdout of that process will be returned to you.
-    """.strip(),
+        instructions=tool_description.strip(),
         host=args["host"],
         port=args["port"],
+        lifespan=session_lifespan,
     )
 
     @mcp.tool(
         name="python",
         title="Execute Python code",
-        description="""
-    Use this tool to execute Python code in your chain of thought. The code will not be shown to the user. This tool should be used for internal reasoning, but not for code that is intended to be visible to the user (e.g. when creating plots, tables, or files).
-    When you send a message containing python code to python, it will be executed in a stateless docker container, and the stdout of that process will be returned to you.
-        """,
+        description=tool_description,
     )
-    async def python(code: str) -> str:
-        async with make_executor(args) as executor:
-            result: str = await run_code(executor, code)
+    async def python(code: str, ctx: Context) -> str:  # type: ignore
+        sid = ctx.request_context.lifespan_context["sid"]
+        executor = sessions[sid][0]
+        cancellation_token = CancellationToken()
+        try:
+            result: str = await run_code(executor, code, cancellation_token)
             return result
+        except asyncio.CancelledError:
+            cancellation_token.cancel()
+            return ""
 
     return mcp
 

python/packages/autogen-kubernetes-mcp/tests/test_executor.py

@@ -1,13 +1,42 @@
+import socket
+
 import pytest
-from autogen_kubernetes.code_executors import PodCommandLineCodeExecutor
+from autogen_core import CancellationToken
+from autogen_kubernetes.code_executors import (
+    PodCommandLineCodeExecutor,
+    PodJupyterCodeExecutor,
+)
 from autogen_kubernetes_mcp._executor import make_executor, run_code
 from conftest import kubernetes_enabled, state_kubernetes_enabled
 
 
+def can_resolve_svc_fqdn() -> bool:
+    try:
+        socket.gethostbyname("kubernetes.default")
+        return True
+    except socket.error:
+        return False
+
+
 @pytest.mark.skipif(not state_kubernetes_enabled, reason="kubernetes not accessible")
 @pytest.mark.asyncio
-async def test_vanilla_executor() -> None:
-    async with make_executor({}) as executor:
-        assert isinstance(executor, PodCommandLineCodeExecutor)
-        result = await run_code(executor, 'print("Hello")')
-        assert "Hello" in result
+async def test_vanilla_commandline_executor() -> None:
+    instances = await make_executor({"type": "commandline"})
+    executor = instances[0]
+    assert isinstance(executor, PodCommandLineCodeExecutor)
+    result = await run_code(executor, 'print("Hello")', CancellationToken())
+    assert "Hello" in result
+    for instance in instances:
+        await instance.stop()
+
+
+@pytest.mark.skipif(not state_kubernetes_enabled or not can_resolve_svc_fqdn(), reason="kubernetes not accessible")
+@pytest.mark.asyncio
+async def test_vanilla_jupyter_executor() -> None:
+    instances = await make_executor({"type": "jupyter"})
+    executor = instances[0]
+    assert isinstance(executor, PodJupyterCodeExecutor)
+    result = await run_code(executor, 'print("Hello")', CancellationToken())
+    assert "Hello" in result
+    for instance in instances:
+        await instance.stop()

python/packages/autogen-kubernetes/README.md

@@ -420,6 +420,63 @@ async with PodCommandLineCodeExecutor(functions=[test_function]) as executor:
 CommandLineCodeResult(exit_code=0, output='kube_config_path not provided and default location (~/.kube/config) does not exist. Using inCluster Config. This might not work.\nTrue\n', code_file='/workspace/tmp_code_c61a3c1e421357bd54041ad195e242d8205b86a3a4a0778b8e2684bc373aac22.py')
 ```
 
+### PodJupyterServer
+
+- Creates a Jupyter Server Pod using `jupyter-kernel-gateway`, along with the required token secret and service resources.
+- Similar to `PodCommandLineCodeExecutor`, the `service_spec`, and `secret_spec` can be customized in multiple formats: Python dictionary, YAML/json string, YAML/json file path, or Kubernetes client model.
+- When used with default arguments, it runs on the `quay.io/jupyter/docker-stacks-foundation` image with `jupyter-kernel-gateway` and `ipykernal` installed, and executes via `jupyter-kernel-gateway`. For efficiency in production, building a custom image is recommended. Below is the sample Dockerfile
+
+```Dockerfile
+FROM quay.io/jupyter/docker-stacks-foundation
+
+RUN mamba install --yes jupyter_kernel_gateway ipykernel && \
+    mamba clean --all -f -y && \
+    fix-permissions "${CONDA_DIR}" && \
+    fix-permissions "/home/${NB_USER}"
+CMD python -m jupyter kernelgateway --KernelGatewayApp.ip=0.0.0.0 \
+        --JupyterApp.answer_yes=true \
+        --JupyterWebsocketPersonality.list_kernels=true
+```
+
+### PodJupyterCodeExecutor
+
+- A `CodeExecutor` that leverages PodJupyterServer and the jupyter-kernel-gateway server to executor code statefully and retrieve results
+- When used together with PodJupyterServer,note that the server generates PodJupyterConnectionInfo based on the service FQDN. This means it cannot be directly used in a non-incluster environment.
+  - After creating the PodJupyterServer, you must construct and provide PodJupyterConnectionInfo so that PodJupyterCodeExecutor can access it properly.
+
+Using with PodJupyterServer
+```python
+async with PodJupyeterServer() as jupyter_server:
+    async with PodJupyterCodeExecutor(jupyter_server) as executor:
+        code_result = await executor.execute_code_blocks(
+            code_blocks=[
+                CodeBlock(language="python", code="print('Hello, World!')"),
+            ],
+            cancellation_token=CancellationToken(),
+        )
+        print(code_result)
+```
+
+Using custom PodJupyterConnectionInfo
+```python
+async with PodJupyeterServer() as jupyter_server:
+    # connection info for created jupyter server pod
+    connection_info = PodJupyetrConnectionInfo(
+        host="https://jupyter-server/access/path",
+        port="443",
+        token=SecretStr("token-string")
+    )
+    async with PodJupyterCodeExecutor(connection_info) as executor:
+        code_result = await executor.execute_code_blocks(
+            code_blocks=[
+                CodeBlock(language="python", code="print('Hello, World!')"),
+            ],
+            cancellation_token=CancellationToken(),
+        )
+        print(code_result)
+```
+
+- Even without using PodJupyterServer, you can configure a custom Jupyter server that meets the required conditions and provide PodJupyterConnectionInfo for PodJupyterCodeExecutor to work with
 
 ## Contribute
 

python/packages/autogen-kubernetes/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "autogen-kubernetes"
-version = "0.5.2"
+version = "0.6.0"
 license = {file = "LICENSE-CODE"}
 description = "autogen kubernetes extension"
 keywords = ["autogen", "agent", "AI", "kubernetes"]