mcp-colab-gpu 0.2.0__tar.gz

mcp_colab_gpu-0.2.0/.github/workflows/publish.yml

@@ -0,0 +1,31 @@
+name: Publish to PyPI
+
+on:
+  push:
+    tags:
+      - "v*"
+
+permissions:
+  contents: read
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Install build tools
+        run: pip install build
+
+      - name: Build package
+        run: python -m build
+
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

mcp_colab_gpu-0.2.0/.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+*.pyc
+dist/
+build/
+*.egg-info/
+.venv/
+.eggs/

mcp_colab_gpu-0.2.0/LICENSE

@@ -0,0 +1,22 @@
+MIT License
+
+Copyright (c) 2026 Paritosh Dwivedi (original mcp-server-colab-exec)
+Copyright (c) 2026 Masaya Hirano (mcp-colab-gpu — extended fork)
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

mcp_colab_gpu-0.2.0/PKG-INFO

@@ -0,0 +1,21 @@
+Metadata-Version: 2.4
+Name: mcp-colab-gpu
+Version: 0.2.0
+Summary: MCP server for executing Python on Google Colab GPU/TPU runtimes — supports all Colab Pro accelerators (T4/L4/A100/H100/G4) with high-memory option and security hardening. Fork of mcp-server-colab-exec by Paritosh Dwivedi.
+Project-URL: Homepage, https://github.com/mio-github/mcp-colab-gpu
+Project-URL: Repository, https://github.com/mio-github/mcp-colab-gpu
+Project-URL: Original Project, https://github.com/pdwi2020/mcp-server-colab-exec
+Author: Masaya Hirano
+Maintainer: Masaya Hirano
+License-Expression: MIT
+License-File: LICENSE
+Keywords: colab,colab-pro,cuda,gpu,machine-learning,mcp,tpu
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Requires-Dist: google-auth-oauthlib>=1.0
+Requires-Dist: google-auth>=2.0
+Requires-Dist: mcp[cli]>=1.6.0
+Requires-Dist: requests>=2.28
+Requires-Dist: websocket-client>=1.5

mcp_colab_gpu-0.2.0/README.md

@@ -0,0 +1,248 @@
+# mcp-colab-gpu
+
+Extended MCP server for Google Colab GPU/TPU runtimes
+
+> Based on [mcp-server-colab-exec](https://github.com/pdwi2020/mcp-server-colab-exec) by [Paritosh Dwivedi](https://github.com/pdwi2020) (MIT License). Thank you for the original idea and implementation! / 素晴らしいアイデアと実装に感謝します！
+
+MCP server that allocates Google Colab GPU/TPU runtimes and executes Python code on them. Lets any MCP-compatible AI assistant -- Claude Code, Claude Desktop, Gemini CLI, Cline, and others -- run GPU/TPU-accelerated code (CUDA, PyTorch, TensorFlow, JAX) without local GPU hardware.
+
+## What's different from the original
+
+| Feature | mcp-server-colab-exec | mcp-colab-gpu |
+|---|---|---|
+| GPU support | T4, L4 | **T4, L4, A100, H100, G4** |
+| TPU support | -- | **V5E1, V6E1** |
+| High-memory runtime | -- | **Supported** |
+| Input validation | -- | **Accelerator + timeout validation** |
+| Path traversal protection | -- | **.py-only + resolved symlinks** |
+| Zip slip protection | -- | **Member path validation** |
+| Token file permissions | Default | **0600 (owner-only)** |
+| Token refresh error logging | Silent | **Logged with re-auth fallback** |
+
+## Supported accelerators
+
+| Accelerator | VRAM / Memory | Tier |
+|---|---|---|
+| `T4` | 16 GB | Free |
+| `L4` | 22 GB | Colab Pro |
+| `A100` | 40 GB | Colab Pro / Pro+ |
+| `H100` | 80 GB | Colab Pro+ |
+| `G4` | 95 GB | Colab Pro+ |
+| `V5E1` | TPU v5e-1 | Colab Pro+ |
+| `V6E1` | TPU v6e-1 | Colab Pro+ |
+
+## Prerequisites
+
+- Python 3.10+
+- A Google account with access to [Google Colab](https://colab.research.google.com)
+- On first run, a browser window opens for OAuth2 consent. The token is cached at `~/.config/colab-exec/token.json` for subsequent runs.
+
+## Installation
+
+### With uvx (recommended)
+
+```bash
+uvx mcp-colab-gpu
+```
+
+### With pip
+
+```bash
+pip install mcp-colab-gpu
+```
+
+### Claude Code configuration
+
+Add to your project's `.mcp.json` or `~/.claude/.mcp.json`:
+
+```json
+{
+  "mcpServers": {
+    "colab-gpu": {
+      "command": "uvx",
+      "args": ["mcp-colab-gpu"]
+    }
+  }
+}
+```
+
+Or via the CLI:
+
+```bash
+claude mcp add colab-gpu -- uvx mcp-colab-gpu
+```
+
+### Claude Desktop configuration
+
+Add to `claude_desktop_config.json`:
+
+```json
+{
+  "mcpServers": {
+    "colab-gpu": {
+      "command": "uvx",
+      "args": ["mcp-colab-gpu"]
+    }
+  }
+}
+```
+
+## Tools
+
+### `colab_execute`
+
+Execute inline Python code on a Colab GPU/TPU runtime.
+
+| Parameter | Type | Default | Description |
+|---|---|---|---|
+| `code` | string | -- | Python code to execute (required) |
+| `accelerator` | string | `"T4"` | GPU/TPU type: `T4`, `L4`, `A100`, `H100`, `G4`, `V5E1`, `V6E1` |
+| `high_memory` | bool | `false` | Enable high-memory runtime (more RAM) |
+| `timeout` | int | `300` | Max execution time in seconds (10--3600) |
+
+Returns JSON with per-cell output, errors, and stderr.
+
+### `colab_execute_file`
+
+Execute a local `.py` file on a Colab GPU/TPU runtime.
+
+| Parameter | Type | Default | Description |
+|---|---|---|---|
+| `file_path` | string | -- | Path to a local `.py` file (required) |
+| `accelerator` | string | `"T4"` | GPU/TPU type: `T4`, `L4`, `A100`, `H100`, `G4`, `V5E1`, `V6E1` |
+| `high_memory` | bool | `false` | Enable high-memory runtime (more RAM) |
+| `timeout` | int | `300` | Max execution time in seconds (10--3600) |
+
+Returns JSON with per-cell output, errors, and stderr.
+
+### `colab_execute_notebook`
+
+Execute code and collect all generated artifacts (images, CSVs, models, etc.).
+
+| Parameter | Type | Default | Description |
+|---|---|---|---|
+| `code` | string | -- | Python code to execute (required) |
+| `output_dir` | string | -- | Local directory for downloaded artifacts (required) |
+| `accelerator` | string | `"T4"` | GPU/TPU type: `T4`, `L4`, `A100`, `H100`, `G4`, `V5E1`, `V6E1` |
+| `high_memory` | bool | `false` | Enable high-memory runtime (more RAM) |
+| `timeout` | int | `300` | Max execution time in seconds (10--3600) |
+
+Artifacts are downloaded as a zip and extracted into `output_dir`.
+
+## Examples
+
+### Check GPU availability
+
+```
+colab_execute(code="import torch; print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0))")
+```
+
+### Matrix multiplication benchmark on A100
+
+```
+colab_execute(
+    code="""
+import torch
+import time
+
+device = torch.device('cuda')
+print(f"GPU: {torch.cuda.get_device_name(0)}")
+print(f"Memory: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB")
+
+# Benchmark: large matrix multiplication
+a = torch.randn(8192, 8192, device=device)
+b = torch.randn(8192, 8192, device=device)
+
+torch.cuda.synchronize()
+start = time.time()
+c = torch.mm(a, b)
+torch.cuda.synchronize()
+elapsed = time.time() - start
+
+tflops = 2 * 8192**3 / elapsed / 1e12
+print(f"8192x8192 matmul: {elapsed:.3f}s ({tflops:.1f} TFLOPS)")
+""",
+    accelerator="A100",
+    high_memory=True,
+)
+```
+
+### LLM inference on H100
+
+```
+colab_execute(
+    code="""
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+model_name = "microsoft/phi-2"
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True
+)
+
+prompt = "Explain quantum computing in one paragraph:"
+inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+outputs = model.generate(**inputs, max_new_tokens=200)
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+""",
+    accelerator="H100",
+    timeout=600,
+)
+```
+
+### Train and download model weights
+
+```
+colab_execute_notebook(
+    code="""
+import torch
+import torch.nn as nn
+
+model = nn.Sequential(nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10))
+model = model.cuda()
+# ... training code ...
+torch.save(model.state_dict(), '/tmp/model.pt')
+print("Model saved!")
+""",
+    output_dir="./outputs",
+    accelerator="T4",
+)
+```
+
+## Authentication
+
+On first use, the server opens a browser window for Google OAuth2 consent. The access token and refresh token are cached at `~/.config/colab-exec/token.json`. Subsequent runs use the cached token and refresh it automatically.
+
+The OAuth2 client credentials are the same ones used by the official Google Colab VS Code extension (`google.colab@0.3.0`). They are intentionally public.
+
+## Security improvements
+
+This fork includes the following security hardening over the original:
+
+- **Path validation in `colab_execute_file`**: Only `.py` files are accepted. Paths are resolved through `pathlib.Path.resolve()` to prevent symlink-based traversal attacks.
+- **Zip slip protection in `colab_execute_notebook`**: Every member in a downloaded artifact zip is validated to ensure its resolved path stays within the target `output_dir`, preventing directory traversal via crafted zip entries.
+- **Token file permissions**: The OAuth token cache file (`~/.config/colab-exec/token.json`) is created with `0o600` permissions (owner read/write only) using `os.open` instead of plain `open`.
+- **Input validation**: The `accelerator` parameter is validated against the known set of supported accelerators, and `timeout` is bounded to 10--3600 seconds, before any network calls are made.
+- **Token refresh error logging**: When automatic token refresh fails, the error is logged to stderr with a warning message before falling back to re-authentication, instead of silently discarding the error.
+
+## Troubleshooting
+
+**"GPU quota exceeded"** -- Colab has usage limits. Wait and retry, or use a different Google account.
+
+**"Timed out creating kernel session"** -- The runtime took too long to start. Retry -- Colab sometimes has delays during peak usage.
+
+**"Authentication failed"** -- Delete `~/.config/colab-exec/token.json` and re-authenticate.
+
+**OAuth browser window doesn't open** -- Ensure you're running in an environment with a browser. For headless servers, authenticate on a machine with a browser first and copy the token file.
+
+## About the maintainer
+
+Masaya Hirano -- CEO of [Mio System Co., Ltd.](https://miosystem.co.jp), CRO of [TrustedAI Corporation](https://trustedai.co.jp)
+
+## License
+
+[MIT](LICENSE)
+
+Original work: Copyright (c) 2026 Paritosh Dwivedi
+Extended fork: Copyright (c) 2026 Masaya Hirano

mcp_colab_gpu-0.2.0/mcp_colab_gpu/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.2.0"

mcp_colab_gpu-0.2.0/mcp_colab_gpu/colab_runtime.py

@@ -0,0 +1,448 @@
+"""Colab GPU runtime engine --- auth, allocate, execute, cleanup.
+
+Provides a clean importable API for allocating Google Colab GPU runtimes
+and executing Python code on them via the Jupyter kernel WebSocket protocol.
+
+Original: mcp-server-colab-exec v0.1.0 by Paritosh Dwivedi (MIT License)
+Extended: mcp-colab-gpu by Masaya Hirano --- all Colab Pro GPUs, high-memory, security fixes.
+"""
+
+import io
+import json
+import os
+import re
+import sys
+import threading
+import time
+import uuid
+
+import requests
+import websocket
+from google.auth.transport.requests import Request as GoogleRequest
+from google.oauth2.credentials import Credentials
+from google_auth_oauthlib.flow import InstalledAppFlow
+
+os.environ.setdefault("OAUTHLIB_RELAX_TOKEN_SCOPE", "1")
+
+COLAB_API = "https://colab.research.google.com"
+SCOPES = [
+    "https://www.googleapis.com/auth/colaboratory",
+    "profile",
+    "email",
+]
+
+# OAuth2 client credentials from the Colab VS Code extension (google.colab@0.3.0)
+# These are intentionally public (the extension names them "ClientNotSoSecret")
+CLIENT_CONFIG = {
+    "installed": {
+        "client_id": "1014160490159-cvot3bea7tgkp72a4m29h20d9ddo6bne.apps.googleusercontent.com",
+        "client_secret": "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",
+        "auth_uri": "https://accounts.google.com/o/oauth2/v2/auth",
+        "token_uri": "https://oauth2.googleapis.com/token",
+        "redirect_uris": ["http://localhost"],
+    }
+}
+
+TOKEN_CACHE_DIR = os.path.expanduser("~/.config/colab-exec")
+TOKEN_CACHE_PATH = os.path.join(TOKEN_CACHE_DIR, "token.json")
+
+HIGHMEM_REQUIRED_ACCELERATORS = {"V5E1", "V6E1"}
+VALID_ACCELERATORS = {"T4", "L4", "A100", "H100", "G4", "V5E1", "V6E1"}
+MAX_TIMEOUT = 3600
+MIN_TIMEOUT = 10
+EPHEMERAL_AUTH_TYPES = {"dfs_ephemeral", "auth_user_ephemeral"}
+
+
+def validate_params(accelerator: str, timeout: int) -> None:
+    """Validate accelerator and timeout parameters."""
+    if accelerator not in VALID_ACCELERATORS:
+        raise ValueError(
+            f"Invalid accelerator '{accelerator}'. "
+            f"Must be one of: {sorted(VALID_ACCELERATORS)}"
+        )
+    if not (MIN_TIMEOUT <= timeout <= MAX_TIMEOUT):
+        raise ValueError(
+            f"timeout must be between {MIN_TIMEOUT} and {MAX_TIMEOUT}, got {timeout}"
+        )
+
+
+def get_credentials() -> Credentials:
+    """Load cached credentials or run the browser OAuth2 flow."""
+    creds = None
+
+    if os.path.exists(TOKEN_CACHE_PATH):
+        creds = Credentials.from_authorized_user_file(TOKEN_CACHE_PATH, SCOPES)
+
+    if creds and creds.expired and creds.refresh_token:
+        try:
+            creds.refresh(GoogleRequest())
+            _save_credentials(creds)
+        except Exception as e:
+            print(f"[colab-gpu] Warning: token refresh failed ({e}), re-authenticating...", file=sys.stderr)
+            creds = None
+
+    if not creds or not creds.valid:
+        flow = InstalledAppFlow.from_client_config(CLIENT_CONFIG, SCOPES)
+        creds = flow.run_local_server(
+            port=0,
+            access_type="offline",
+            prompt="consent",
+            success_message="Authentication successful! You can close this tab.",
+        )
+        _save_credentials(creds)
+
+    return creds
+
+
+def _save_credentials(creds: Credentials):
+    os.makedirs(TOKEN_CACHE_DIR, exist_ok=True)
+    fd = os.open(TOKEN_CACHE_PATH, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
+    with os.fdopen(fd, "w") as f:
+        f.write(creds.to_json())
+
+
+def _strip_xssi(text: str) -> dict:
+    """Strip the )]}' XSSI prefix and parse JSON."""
+    if text.startswith(")]}'"):
+        text = text[text.index("\n") + 1:]
+    return json.loads(text)
+
+
+def _colab_headers(token: str, extra: dict = None) -> dict:
+    headers = {
+        "Authorization": f"Bearer {token}",
+        "Accept": "application/json",
+        "X-Colab-Client-Agent": "vscode",
+    }
+    if extra:
+        headers.update(extra)
+    return headers
+
+
+def generate_notebook_hash() -> str:
+    """Generate a Colab-valid notebook hash."""
+    raw_uuid = str(uuid.uuid4())
+    return raw_uuid.replace("-", "_") + "." * (44 - len(raw_uuid))
+
+
+def _build_assign_params(
+    nbh: str,
+    accelerator: str,
+    high_memory: bool = False,
+) -> dict:
+    params = {
+        "nbh": nbh,
+        "authuser": "0",
+    }
+    if accelerator:
+        params["variant"] = "GPU"
+        params["accelerator"] = accelerator
+        if high_memory or accelerator in HIGHMEM_REQUIRED_ACCELERATORS:
+            params["shape"] = "hm"
+    return params
+
+
+def _parse_assignment(assignment: dict) -> dict:
+    endpoint = assignment.get("endpoint")
+    proxy_info = assignment.get("runtimeProxyInfo", {})
+    proxy_url = (proxy_info.get("url") or "").rstrip("/")
+    proxy_token = proxy_info.get("token")
+    return {
+        "endpoint": endpoint,
+        "proxy_url": proxy_url,
+        "proxy_token": proxy_token,
+    }
+
+
+def allocate_runtime(
+    token: str,
+    accelerator: str = "T4",
+    high_memory: bool = False,
+) -> dict:
+    """Allocate a Colab GPU runtime. Returns dict with endpoint + proxy info."""
+    nbh = generate_notebook_hash()
+    params = _build_assign_params(nbh, accelerator, high_memory)
+    headers = _colab_headers(token)
+
+    hm_label = " (high-memory)" if high_memory else ""
+    print(f"[colab-gpu] Requesting {accelerator}{hm_label} runtime...", file=sys.stderr)
+    r = requests.get(f"{COLAB_API}/tun/m/assign", params=params, headers=headers, timeout=30)
+    r.raise_for_status()
+    data = _strip_xssi(r.text)
+
+    parsed = _parse_assignment(data)
+    if parsed["endpoint"] and parsed["proxy_url"] and parsed["proxy_token"]:
+        print(f"[colab-gpu] Reusing existing runtime: endpoint={parsed['endpoint']}", file=sys.stderr)
+        return {
+            **parsed,
+            "xsrf_token": None,
+            "nbh": nbh,
+            "reused": True,
+        }
+
+    xsrf_token = data.get("token") or data.get("xsrfToken")
+    if not xsrf_token:
+        raise RuntimeError("No XSRF token in Colab assign response")
+
+    post_headers = _colab_headers(token, {"X-Goog-Colab-Token": xsrf_token})
+    r = requests.post(
+        f"{COLAB_API}/tun/m/assign",
+        params=params,
+        headers=post_headers,
+        timeout=30,
+    )
+    r.raise_for_status()
+    assignment = _strip_xssi(r.text)
+
+    parsed = _parse_assignment(assignment)
+    if not parsed["endpoint"] or not parsed["proxy_url"] or not parsed["proxy_token"]:
+        raise RuntimeError("Incomplete assignment response from Colab")
+
+    print(f"[colab-gpu] Runtime allocated: endpoint={parsed['endpoint']}", file=sys.stderr)
+    return {
+        **parsed,
+        "xsrf_token": xsrf_token,
+        "nbh": nbh,
+        "reused": False,
+    }
+
+
+def unassign_runtime(token: str, endpoint: str) -> bool:
+    """Release the runtime."""
+    headers = _colab_headers(token)
+    url = f"{COLAB_API}/tun/m/unassign/{endpoint}"
+    params = {"authuser": "0"}
+    try:
+        r = requests.get(url, headers=headers, params=params, timeout=20)
+        r.raise_for_status()
+        data = _strip_xssi(r.text)
+        xsrf = data.get("token", "")
+        post_headers = _colab_headers(token, {"X-Goog-Colab-Token": xsrf})
+        r = requests.post(url, headers=post_headers, params=params, timeout=20)
+        r.raise_for_status()
+        print(f"[colab-gpu] Runtime {endpoint} released.", file=sys.stderr)
+        return True
+    except Exception as e:
+        print(f"[colab-gpu] Warning: failed to unassign runtime: {e}", file=sys.stderr)
+        return False
+
+
+def propagate_credentials(token: str, endpoint: str, auth_type: str, dry_run: bool) -> dict:
+    """Call Colab credentials propagation API for ephemeral auth challenges."""
+    url = f"{COLAB_API}/tun/m/credentials-propagation/{endpoint}"
+    params = {
+        "authuser": "0",
+        "authtype": auth_type,
+        "version": "2",
+        "dryrun": str(bool(dry_run)).lower(),
+        "propagate": "true",
+        "record": "false",
+    }
+    headers = _colab_headers(token)
+    r = requests.get(url, headers=headers, params=params, timeout=30)
+    r.raise_for_status()
+    data = _strip_xssi(r.text)
+    xsrf = data.get("token") or data.get("xsrfToken")
+    if not xsrf:
+        raise RuntimeError("No XSRF token from credentials propagation")
+    post_headers = _colab_headers(token, {"X-Goog-Colab-Token": xsrf})
+    r = requests.post(url, headers=post_headers, params=params, timeout=30)
+    r.raise_for_status()
+    return _strip_xssi(r.text)
+
+
+def start_keepalive(token: str, endpoint: str) -> threading.Event:
+    """Start a background thread that pings keep-alive every 60s."""
+    stop_event = threading.Event()
+    def loop():
+        headers = _colab_headers(token, {"X-Colab-Tunnel": "Google"})
+        url = f"{COLAB_API}/tun/m/{endpoint}/keep-alive/"
+        params = {"authuser": "0"}
+        while not stop_event.is_set():
+            try:
+                requests.get(url, headers=headers, params=params, timeout=10)
+            except Exception:
+                pass
+            stop_event.wait(60)
+    t = threading.Thread(target=loop, daemon=True)
+    t.start()
+    return stop_event
+
+
+def create_session(proxy_url: str, proxy_token: str, startup_timeout: int = 180) -> str:
+    """Create a Jupyter session and return the kernel ID."""
+    headers = {
+        "X-Colab-Runtime-Proxy-Token": proxy_token,
+        "X-Colab-Client-Agent": "vscode",
+        "Content-Type": "application/json",
+    }
+    body = {
+        "kernel": {"name": "python3"},
+        "name": "colab-exec",
+        "path": "colab-exec",
+        "type": "notebook",
+    }
+    last_error = None
+    deadline = time.time() + startup_timeout
+    attempt = 0
+    while time.time() < deadline:
+        attempt += 1
+        try:
+            r = requests.post(f"{proxy_url}/api/sessions", headers=headers, json=body, timeout=30)
+            r.raise_for_status()
+            data = r.json()
+            kernel_id = data["kernel"]["id"]
+            print(f"[colab-gpu] Kernel ready: {kernel_id}", file=sys.stderr)
+            return kernel_id
+        except Exception as e:
+            last_error = e
+            remaining = int(deadline - time.time())
+            if remaining <= 0:
+                break
+            print(f"[colab-gpu] Waiting for runtime readiness (attempt {attempt}, {remaining}s left)...", file=sys.stderr)
+            time.sleep(3)
+    raise RuntimeError(f"Timed out creating kernel session: {last_error}")
+
+
+def _make_colab_input_reply(client_session_id: str, colab_msg_id, err: str = None) -> dict:
+    value = {"type": "colab_reply", "colab_msg_id": colab_msg_id}
+    if err:
+        value["error"] = err
+    return {
+        "header": {
+            "msg_id": uuid.uuid4().hex,
+            "msg_type": "input_reply",
+            "session": client_session_id,
+            "date": time.strftime("%Y-%m-%dT%H:%M:%S.000Z", time.gmtime()),
+            "username": "username",
+            "version": "5.0",
+        },
+        "content": {"value": value},
+        "channel": "stdin",
+        "metadata": {},
+        "parent_header": {},
+    }
+
+
+def execute_code(
+    proxy_url: str,
+    proxy_token: str,
+    kernel_id: str,
+    code: str,
+    timeout: int = 300,
+    access_token: str = None,
+    endpoint: str = None,
+) -> tuple[str, str, int]:
+    """Execute code on a Colab kernel via WebSocket.
+    Returns (stdout, stderr, exit_code).
+    """
+    stdout_buf = io.StringIO()
+    stderr_buf = io.StringIO()
+
+    execute_session_id = uuid.uuid4().hex
+    ws_url = proxy_url.replace("https://", "wss://").replace("http://", "ws://")
+    ws_url = f"{ws_url}/api/kernels/{kernel_id}/channels?session_id={execute_session_id}"
+
+    ws = websocket.create_connection(
+        ws_url,
+        header=[f"X-Colab-Runtime-Proxy-Token: {proxy_token}", "X-Colab-Client-Agent: vscode"],
+        timeout=timeout,
+    )
+
+    msg_id = uuid.uuid4().hex
+    execute_msg = {
+        "header": {
+            "msg_id": msg_id, "msg_type": "execute_request",
+            "username": "colab-exec", "session": execute_session_id, "version": "5.3",
+        },
+        "parent_header": {},
+        "metadata": {},
+        "content": {
+            "code": code, "silent": False, "store_history": True,
+            "user_expressions": {}, "allow_stdin": False, "stop_on_error": True,
+        },
+        "channel": "shell",
+    }
+    ws.send(json.dumps(execute_msg))
+
+    had_error = False
+    saw_idle = False
+    deadline = time.time() + timeout
+
+    while time.time() < deadline:
+        try:
+            raw = ws.recv()
+        except websocket.WebSocketTimeoutException:
+            continue
+        if not raw:
+            continue
+
+        msg = json.loads(raw)
+        msg_type = msg.get("msg_type") or msg.get("header", {}).get("msg_type", "")
+        content = msg.get("content", {})
+
+        if msg_type == "colab_request":
+            metadata = msg.get("metadata", {})
+            request_type = metadata.get("colab_request_type")
+            colab_msg_id = metadata.get("colab_msg_id")
+            auth_type = (content.get("request", {}).get("authType", "") if isinstance(content, dict) else "")
+            auth_type = str(auth_type).lower()
+            if request_type == "request_auth" and colab_msg_id is not None:
+                error_text = None
+                if not access_token or not endpoint:
+                    error_text = "missing auth context for credentials propagation"
+                elif auth_type not in EPHEMERAL_AUTH_TYPES:
+                    error_text = f"unsupported auth type: {auth_type}"
+                else:
+                    try:
+                        dry = propagate_credentials(access_token, endpoint, auth_type, dry_run=True)
+                        if dry.get("success"):
+                            propagate_credentials(access_token, endpoint, auth_type, dry_run=False)
+                        elif dry.get("unauthorizedRedirectUri"):
+                            error_text = f"{auth_type} requires interactive browser consent: {dry['unauthorizedRedirectUri']}"
+                        else:
+                            error_text = f"{auth_type} dry-run failed: {dry}"
+                    except Exception as e:
+                        error_text = f"{auth_type} propagation failed: {e}"
+                reply = _make_colab_input_reply(execute_session_id, colab_msg_id, error_text)
+                ws.send(json.dumps(reply))
+                if error_text:
+                    stderr_buf.write(f"[colab-gpu] Warning: {error_text}\n")
+            continue
+
+        parent_msg_id = msg.get("parent_header", {}).get("msg_id")
+        if parent_msg_id != msg_id:
+            continue
+
+        if msg_type == "stream":
+            text = content.get("text", "")
+            if content.get("name", "stdout") == "stdout":
+                stdout_buf.write(text)
+            else:
+                stderr_buf.write(text)
+        elif msg_type == "execute_result":
+            text = content.get("data", {}).get("text/plain", "")
+            if text:
+                stdout_buf.write(text + "\n")
+        elif msg_type == "display_data":
+            text = content.get("data", {}).get("text/plain", "")
+            if text:
+                stdout_buf.write(text + "\n")
+        elif msg_type == "error":
+            had_error = True
+            ename = content.get("ename", "Error")
+            evalue = content.get("evalue", "")
+            stderr_buf.write(f"\n{ename}: {evalue}\n")
+            for line in content.get("traceback", []):
+                stderr_buf.write(re.sub(r"\x1b\[[0-9;]*m", "", line) + "\n")
+        elif msg_type == "status":
+            if content.get("execution_state") == "idle":
+                saw_idle = True
+                break
+
+    ws.close()
+    if not saw_idle:
+        stderr_buf.write("[colab-gpu] ERROR: Timed out waiting for kernel execution to finish.\n")
+        had_error = True
+
+    return stdout_buf.getvalue(), stderr_buf.getvalue(), 1 if had_error else 0

mcp_colab_gpu-0.2.0/mcp_colab_gpu/server.py

@@ -0,0 +1,295 @@
+"""MCP server for executing Python code on Google Colab GPU/TPU runtimes.
+
+Exposes three tools via the FastMCP API:
+- colab_execute: Run inline Python code on a Colab GPU/TPU
+- colab_execute_file: Run a local .py file on a Colab GPU/TPU
+- colab_execute_notebook: Run code and collect generated artifacts
+
+Original: mcp-server-colab-exec v0.1.0 by Paritosh Dwivedi (MIT License)
+Extended: mcp-colab-gpu by Masaya Hirano --- all Colab Pro GPUs, high-memory, security fixes.
+"""
+
+import base64
+import json
+import os
+import pathlib
+import re
+import zipfile
+
+from mcp.server.fastmcp import FastMCP
+
+from .colab_runtime import (
+    allocate_runtime,
+    create_session,
+    execute_code,
+    get_credentials,
+    start_keepalive,
+    unassign_runtime,
+    validate_params,
+)
+
+mcp = FastMCP("colab-gpu")
+
+CELL_START = "===CELL_START_{n}==="
+CELL_END = "===CELL_END_{n}==="
+ARTIFACT_B64_START = "ARTIFACT_BASE64_START"
+ARTIFACT_B64_END = "ARTIFACT_BASE64_END"
+
+
+def _wrap_cells(code: str) -> tuple[str, int]:
+    """Wrap code in cell-boundary markers so output can be parsed per-cell."""
+    raw_cells = re.split(r"\n{2,}", code.strip())
+    cells = [c.strip() for c in raw_cells if c.strip()]
+    if not cells:
+        cells = [code]
+    wrapped_parts = []
+    for i, cell in enumerate(cells):
+        marker_start = CELL_START.format(n=i)
+        marker_end = CELL_END.format(n=i)
+        wrapped_parts.append(
+            f'print("{marker_start}", flush=True)\n'
+            f"{cell}\n"
+            f'print("{marker_end}", flush=True)'
+        )
+    return "\n\n".join(wrapped_parts), len(cells)
+
+
+def _parse_cell_output(stdout: str, num_cells: int) -> list[dict]:
+    cells = []
+    for i in range(num_cells):
+        start_marker = CELL_START.format(n=i)
+        end_marker = CELL_END.format(n=i)
+        pattern = re.escape(start_marker) + r"\n?(.*?)\n?" + re.escape(end_marker)
+        match = re.search(pattern, stdout, re.DOTALL)
+        cell_stdout = match.group(1).strip() if match else ""
+        cells.append({"cell_num": i, "stdout": cell_stdout, "status": "ok" if match else "no_output"})
+    return cells
+
+
+def _extract_artifact_b64(stdout: str) -> str | None:
+    pattern = re.escape(ARTIFACT_B64_START) + r"\n(.*?)\n" + re.escape(ARTIFACT_B64_END)
+    match = re.search(pattern, stdout, re.DOTALL)
+    return match.group(1).strip() if match else None
+
+
+def _run_on_colab(code: str, accelerator: str, high_memory: bool, timeout: int) -> tuple[str, str, int]:
+    validate_params(accelerator, timeout)
+    creds = get_credentials()
+    access_token = creds.token
+    assignment = allocate_runtime(access_token, accelerator, high_memory)
+    stop_event = start_keepalive(access_token, assignment["endpoint"])
+    try:
+        kernel_id = create_session(assignment["proxy_url"], assignment["proxy_token"])
+        stdout, stderr, exit_code = execute_code(
+            assignment["proxy_url"], assignment["proxy_token"], kernel_id, code,
+            timeout=timeout, access_token=access_token, endpoint=assignment["endpoint"],
+        )
+    finally:
+        stop_event.set()
+        unassign_runtime(access_token, assignment["endpoint"])
+    return stdout, stderr, exit_code
+
+
+def _validate_file_path(raw: str) -> pathlib.Path:
+    """Validate file path: must exist, must be .py."""
+    resolved = pathlib.Path(os.path.expanduser(raw)).resolve()
+    if resolved.suffix != ".py":
+        raise ValueError(f"Only .py files are allowed, got: '{resolved.suffix}'")
+    if not resolved.is_file():
+        raise FileNotFoundError(f"File not found: {resolved}")
+    return resolved
+
+
+def _safe_extract_zip(zip_path: str, output_dir: str) -> list[str]:
+    """Extract zip with zip-slip protection. Returns list of extracted filenames."""
+    resolved_output = pathlib.Path(output_dir).resolve()
+    with zipfile.ZipFile(zip_path, "r") as zf:
+        for member in zf.infolist():
+            member_path = (resolved_output / member.filename).resolve()
+            if not str(member_path).startswith(str(resolved_output)):
+                raise ValueError(f"Zip slip detected: {member.filename}")
+        artifact_files = zf.namelist()
+        zf.extractall(output_dir)
+    return artifact_files
+
+
+@mcp.tool(annotations={"readOnlyHint": False, "destructiveHint": False})
+def colab_execute(
+    code: str,
+    accelerator: str = "T4",
+    high_memory: bool = False,
+    timeout: int = 300,
+) -> str:
+    """Execute Python code on a Google Colab GPU/TPU runtime.
+
+    Allocates a GPU or TPU, runs the code, and returns structured JSON
+    with per-cell output, errors, and stderr.
+
+    Args:
+        code: Python code to execute on the Colab runtime.
+        accelerator: Hardware accelerator type. Default: "T4".
+            GPU types:
+              "T4"   - NVIDIA Tesla T4 (16 GB, free-tier)
+              "L4"   - NVIDIA L4 (24 GB, Colab Pro)
+              "A100"  - NVIDIA A100 (40 GB, Colab Pro/Pro+)
+              "H100"  - NVIDIA H100 (80 GB, Colab Pro+)
+              "G4"   - NVIDIA G4 (Colab Pro+)
+            TPU types:
+              "V5E1" - TPU v5e-1 (Colab Pro+)
+              "V6E1" - TPU v6e-1 (Colab Pro+)
+        high_memory: Enable high-memory runtime (more RAM). Default: False.
+        timeout: Max execution time in seconds. Default: 300.
+    """
+    wrapped, num_cells = _wrap_cells(code)
+    stdout, stderr, rc = _run_on_colab(wrapped, accelerator, high_memory, timeout)
+    cells = _parse_cell_output(stdout, num_cells)
+    errors = [c for c in cells if c["status"] != "ok"] if rc != 0 else []
+    return json.dumps({"cells": cells, "errors": errors, "stderr": stderr, "exit_code": rc}, indent=2)
+
+
+@mcp.tool(annotations={"readOnlyHint": False, "destructiveHint": False})
+def colab_execute_file(
+    file_path: str,
+    accelerator: str = "T4",
+    high_memory: bool = False,
+    timeout: int = 300,
+) -> str:
+    """Execute a local Python file on a Google Colab GPU/TPU runtime.
+
+    Reads the file contents and sends them for execution on a Colab runtime.
+    Only .py files are allowed for security.
+
+    Args:
+        file_path: Path to a local .py file to execute on Colab.
+        accelerator: Hardware accelerator type. Default: "T4".
+            GPU types:
+              "T4"   - NVIDIA Tesla T4 (16 GB, free-tier)
+              "L4"   - NVIDIA L4 (24 GB, Colab Pro)
+              "A100"  - NVIDIA A100 (40 GB, Colab Pro/Pro+)
+              "H100"  - NVIDIA H100 (80 GB, Colab Pro+)
+              "G4"   - NVIDIA G4 (Colab Pro+)
+            TPU types:
+              "V5E1" - TPU v5e-1 (Colab Pro+)
+              "V6E1" - TPU v6e-1 (Colab Pro+)
+        high_memory: Enable high-memory runtime (more RAM). Default: False.
+        timeout: Max execution time in seconds. Default: 300.
+    """
+    try:
+        resolved = _validate_file_path(file_path)
+    except (ValueError, FileNotFoundError) as e:
+        return json.dumps({"error": str(e)})
+
+    code = resolved.read_text()
+    wrapped, num_cells = _wrap_cells(code)
+    stdout, stderr, rc = _run_on_colab(wrapped, accelerator, high_memory, timeout)
+    cells = _parse_cell_output(stdout, num_cells)
+    errors = [c for c in cells if c["status"] != "ok"] if rc != 0 else []
+    return json.dumps({"cells": cells, "errors": errors, "stderr": stderr, "exit_code": rc}, indent=2)
+
+
+@mcp.tool(annotations={"readOnlyHint": False, "destructiveHint": False})
+def colab_execute_notebook(
+    code: str,
+    output_dir: str,
+    accelerator: str = "T4",
+    high_memory: bool = False,
+    timeout: int = 300,
+) -> str:
+    """Execute Python code on Colab GPU/TPU and collect generated artifacts.
+
+    Runs the code, then scans the runtime for output files (images, CSVs,
+    models, etc.), zips them, and downloads to a local directory.
+
+    Args:
+        code: Python code to execute on the Colab runtime.
+        output_dir: Local directory to save the artifacts zip and extracted files.
+        accelerator: Hardware accelerator type. Default: "T4".
+            GPU types:
+              "T4"   - NVIDIA Tesla T4 (16 GB, free-tier)
+              "L4"   - NVIDIA L4 (24 GB, Colab Pro)
+              "A100"  - NVIDIA A100 (40 GB, Colab Pro/Pro+)
+              "H100"  - NVIDIA H100 (80 GB, Colab Pro+)
+              "G4"   - NVIDIA G4 (Colab Pro+)
+            TPU types:
+              "V5E1" - TPU v5e-1 (Colab Pro+)
+              "V6E1" - TPU v6e-1 (Colab Pro+)
+        high_memory: Enable high-memory runtime (more RAM). Default: False.
+        timeout: Max execution time in seconds. Default: 300.
+    """
+    output_dir = os.path.expanduser(output_dir)
+    os.makedirs(output_dir, exist_ok=True)
+
+    artifact_code = '''
+
+# --- colab-exec artifact collection ---
+import os, zipfile, base64, glob
+
+_artifact_dir = "/tmp/colab_artifacts"
+os.makedirs(_artifact_dir, exist_ok=True)
+
+_scan_dirs = ["/tmp", os.getcwd(), "/content"]
+_skip_prefixes = ["/tmp/colab_artifacts", "/tmp/."]
+_collected = []
+for _sd in _scan_dirs:
+    if not os.path.isdir(_sd):
+        continue
+    for _root, _dirs, _files in os.walk(_sd):
+        _dirs[:] = [d for d in _dirs if not d.startswith('.')]
+        if any(_root.startswith(p) for p in _skip_prefixes):
+            continue
+        for _f in _files:
+            _fp = os.path.join(_root, _f)
+            if _f.endswith(('.png', '.jpg', '.jpeg', '.gif', '.svg', '.csv', '.json',
+                           '.txt', '.pt', '.pth', '.h5', '.pkl', '.npy', '.npz',
+                           '.onnx', '.mp4', '.wav', '.mp3', '.pdf')):
+                try:
+                    if os.path.getsize(_fp) < 50 * 1024 * 1024:
+                        _collected.append(_fp)
+                except OSError:
+                    pass
+
+if _collected:
+    _zip_path = "/tmp/colab_artifacts.zip"
+    with zipfile.ZipFile(_zip_path, 'w', zipfile.ZIP_DEFLATED) as _zf:
+        for _fp in _collected:
+            _zf.write(_fp, os.path.basename(_fp))
+    with open(_zip_path, 'rb') as _zfh:
+        _b64 = base64.b64encode(_zfh.read()).decode('ascii')
+    print("ARTIFACT_BASE64_START")
+    print(_b64)
+    print("ARTIFACT_BASE64_END")
+    print(f"[colab-gpu] Collected {len(_collected)} artifact(s)", flush=True)
+else:
+    print("[colab-gpu] No artifacts found to collect", flush=True)
+'''
+    full_code = code + "\n\n" + artifact_code
+    wrapped, num_cells = _wrap_cells(full_code)
+    stdout, stderr, rc = _run_on_colab(wrapped, accelerator, high_memory, timeout)
+    cells = _parse_cell_output(stdout, num_cells)
+    errors = [c for c in cells if c["status"] != "ok"] if rc != 0 else []
+
+    artifact_files = []
+    artifacts_zip_path = None
+    b64_data = _extract_artifact_b64(stdout)
+    if b64_data:
+        try:
+            zip_bytes = base64.b64decode(b64_data)
+            artifacts_zip_path = os.path.join(output_dir, "colab_artifacts.zip")
+            with open(artifacts_zip_path, "wb") as f:
+                f.write(zip_bytes)
+            artifact_files = _safe_extract_zip(artifacts_zip_path, output_dir)
+        except Exception as e:
+            errors.append({"artifact_error": str(e)})
+
+    return json.dumps({
+        "cells": cells, "errors": errors, "artifacts_zip": artifacts_zip_path,
+        "artifact_files": artifact_files, "stderr": stderr, "exit_code": rc,
+    }, indent=2)
+
+
+def main():
+    mcp.run(transport="stdio")
+
+
+if __name__ == "__main__":
+    main()

mcp_colab_gpu-0.2.0/pyproject.toml

@@ -0,0 +1,37 @@
+[project]
+name = "mcp-colab-gpu"
+version = "0.2.0"
+description = "MCP server for executing Python on Google Colab GPU/TPU runtimes — supports all Colab Pro accelerators (T4/L4/A100/H100/G4) with high-memory option and security hardening. Fork of mcp-server-colab-exec by Paritosh Dwivedi."
+requires-python = ">=3.10"
+license = "MIT"
+authors = [
+    {name = "Masaya Hirano"},
+]
+maintainers = [
+    {name = "Masaya Hirano"},
+]
+keywords = ["colab", "cuda", "gpu", "tpu", "machine-learning", "mcp", "colab-pro"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "google-auth>=2.0",
+    "google-auth-oauthlib>=1.0",
+    "mcp[cli]>=1.6.0",
+    "requests>=2.28",
+    "websocket-client>=1.5",
+]
+
+[project.urls]
+Homepage = "https://github.com/mio-github/mcp-colab-gpu"
+Repository = "https://github.com/mio-github/mcp-colab-gpu"
+"Original Project" = "https://github.com/pdwi2020/mcp-server-colab-exec"
+
+[project.scripts]
+mcp-colab-gpu = "mcp_colab_gpu.server:main"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"