From 4049a825f694697a05dcfbd4a0a5795680db0fd7 Mon Sep 17 00:00:00 2001 From: Eric Charles Date: Sat, 8 Aug 2026 17:33:29 +0200 Subject: [PATCH 1/3] add kaggle and colab --- README.md | 75 ++++ code_sandboxes/__init__.py | 17 +- code_sandboxes/__version__.py | 2 +- code_sandboxes/client.py | 168 +++++++- code_sandboxes/colab.py | 208 ++++++++++ code_sandboxes/colab_sandbox.py | 21 +- code_sandboxes/interfaces.py | 30 +- code_sandboxes/jupyter_sandbox.py | 2 +- code_sandboxes/kaggle.py | 162 ++++++++ code_sandboxes/kaggle_execute.py | 591 +++++++++++++++++++++++++++ code_sandboxes/kaggle_sandbox.py | 27 +- docs/docs/api-reference/index.mdx | 60 +++ docs/docs/installation/index.mdx | 5 +- docs/docs/sandboxes/google-colab.mdx | 92 +++++ docs/docs/sandboxes/index.mdx | 18 +- docs/docs/sandboxes/kaggle.mdx | 170 ++++++++ pyproject.toml | 5 +- tests/test_client.py | 99 ++++- tests/test_colab.py | 126 ++++++ tests/test_jupyter.py | 4 + tests/test_kaggle.py | 127 ++++++ tests/test_kaggle_execute.py | 268 ++++++++++++ tests/test_modal_colab_sandbox.py | 35 +- 23 files changed, 2219 insertions(+), 93 deletions(-) create mode 100644 code_sandboxes/colab.py create mode 100644 code_sandboxes/kaggle.py create mode 100644 code_sandboxes/kaggle_execute.py create mode 100644 docs/docs/sandboxes/google-colab.mdx create mode 100644 docs/docs/sandboxes/kaggle.mdx create mode 100644 tests/test_colab.py create mode 100644 tests/test_kaggle.py create mode 100644 tests/test_kaggle_execute.py diff --git a/README.md b/README.md index 6ebc427..bd5bf13 100644 --- a/README.md +++ b/README.md @@ -90,6 +90,81 @@ export KAGGLE_API_KEY="" sandbox repl --variant kaggle ``` +### Kaggle + +Kaggle supports both batch execution and interactive connections through the +`kaggle` sandbox. Install its optional dependency first: + +```bash +pip install "code-sandboxes[kaggle]" +``` + +For batch execution, configure Kaggle credentials and create the sandbox +without a runtime URL: + +```python +from code_sandboxes import Sandbox + +with Sandbox.create(variant="kaggle") as sandbox: + result = sandbox.run_code("print('hello from kaggle')") + print(result.stdout) +``` + +The lower-level batch API is also available directly: + +```python +from code_sandboxes import KaggleKernelExecutor + +executor = KaggleKernelExecutor() +result = executor.execute( + "print('hello from kaggle')", + title="code-sandboxes-demo", + accelerator="NvidiaTeslaT4", + wait=True, +) +print(result.status, result.stdout) +print(result.to_kernel_reply()) +``` + +For interactive execution, copy the WebSocket channels URL from an active +Kaggle notebook session and pass it to the sandbox or client: + +```python +from code_sandboxes import KaggleKernelClient + +with KaggleKernelClient.from_channels_url(channels_url, token=None) as kernel: + print(kernel.execute("x = 1 + 1; print(x)")) +``` + +See the [complete Kaggle guide](docs/docs/sandboxes/kaggle.mdx) for authentication, +accelerators, channels URL retrieval, and execution options. + +### Google Colab + +Google Colab exposes an already-running kernel through an authenticating proxy. +Copy its WebSocket channels URL from the browser's Network tools, then pass it +directly to the sandbox: + +```python +from code_sandboxes import Sandbox + +with Sandbox.create(variant="colab", channels_url=channels_url) as sandbox: + print(sandbox.run_code("x = 1 + 1; print(x)").stdout) +``` + +The lower-level client and parser are owned by Code Sandboxes as well: + +```python +from code_sandboxes import ColabKernelClient, parse_colab_channels_url + +server_url, kernel_id, proxy_token = parse_colab_channels_url(channels_url) +with ColabKernelClient.from_channels_url(channels_url) as kernel: + print(kernel.execute("print('hello from colab')")) +``` + +See the [complete Google Colab guide](docs/docs/sandboxes/google-colab.mdx) for +proxy authentication, explicit connection values, and channels URL retrieval. + For full setup and parameters for all variants, see: - [https://code-sandboxes.datalayer.tech/sandboxes](https://code-sandboxes.datalayer.tech/sandboxes) diff --git a/code_sandboxes/__init__.py b/code_sandboxes/__init__.py index 2113a3f..550fcf5 100644 --- a/code_sandboxes/__init__.py +++ b/code_sandboxes/__init__.py @@ -58,7 +58,8 @@ """ from .base import Sandbox -from .client import CodeExecutionOutcome, CodeSandboxClient +from .client import CodeExecutionOutcome, CodeSandboxClient, execution_result_to_reply +from .colab import ColabKernelClient, parse_colab_channels_url from .colab_sandbox import ColabSandbox from .commands import CommandResult, ProcessHandle, SandboxCommands from .datalayer_sandbox import DatalayerSandbox @@ -86,8 +87,10 @@ SandboxFileHandle, SandboxFilesystem, ) -from .interfaces import IJupyterKernelClient, ISandboxClient +from .interfaces import ISandboxClient from .jupyter_sandbox import JupyterSandbox +from .kaggle import KAGGLE_API_TOKEN_ENV, KaggleKernelClient, parse_kaggle_channels_url +from .kaggle_execute import KaggleExecutionResult, KaggleKernelExecutor from .kaggle_sandbox import KaggleSandbox from .modal_sandbox import ModalSandbox from .models import ( @@ -112,10 +115,12 @@ from .monty_sandbox import MontySandbox __all__ = [ + "KAGGLE_API_TOKEN_ENV", # Models "CodeError", "CodeExecutionOutcome", "CodeSandboxClient", + "ColabKernelClient", "ColabSandbox", "CommandResult", "Context", @@ -130,9 +135,11 @@ "FileWatchEvent", "FileWatchEventType", "GPUType", - "IJupyterKernelClient", "ISandboxClient", "JupyterSandbox", + "KaggleExecutionResult", + "KaggleKernelClient", + "KaggleKernelExecutor", "KaggleSandbox", "Logs", "MIMEType", @@ -166,8 +173,10 @@ "SandboxStatus", "SandboxTimeoutError", "SandboxVariant", - "SandboxVariant", "SnapshotInfo", "TunnelInfo", "VariableNotFoundError", + "execution_result_to_reply", + "parse_colab_channels_url", + "parse_kaggle_channels_url", ] diff --git a/code_sandboxes/__version__.py b/code_sandboxes/__version__.py index 11db5c9..ebdbba6 100644 --- a/code_sandboxes/__version__.py +++ b/code_sandboxes/__version__.py @@ -3,4 +3,4 @@ """Code Sandboxes.""" -__version__ = "0.17.0" +__version__ = "1.0.0" diff --git a/code_sandboxes/client.py b/code_sandboxes/client.py index b64fe4f..514d355 100644 --- a/code_sandboxes/client.py +++ b/code_sandboxes/client.py @@ -45,17 +45,73 @@ from collections.abc import AsyncIterator, Iterator from dataclasses import dataclass, field -from typing import Union +from typing import Any, Callable, Union from .base import Sandbox from .commands import CommandResult -from .models import CodeError, ExecutionResult, OutputMessage, Result, SandboxConfig, SandboxVariant - -__all__ = ["CodeExecutionOutcome", "CodeSandboxClient"] +from .models import ( + CodeError, + ExecutionResult, + OutputMessage, + Result, + SandboxConfig, + SandboxInfo, + SandboxVariant, +) + +__all__ = ["CodeExecutionOutcome", "CodeSandboxClient", "execution_result_to_reply"] StreamingItem = Union[OutputMessage, Result, CodeError] +def execution_result_to_reply(execution: ExecutionResult) -> dict[str, Any]: + """Convert a variant-neutral execution result to a Jupyter-shaped reply.""" + outputs: list[dict[str, Any]] = [] + + if execution.logs.stdout: + outputs.append( + { + "output_type": "stream", + "name": "stdout", + "text": "\n".join(message.line for message in execution.logs.stdout) + "\n", + } + ) + if execution.logs.stderr: + outputs.append( + { + "output_type": "stream", + "name": "stderr", + "text": "\n".join(message.line for message in execution.logs.stderr) + "\n", + } + ) + + for result in execution.results: + outputs.append( + { + "output_type": "execute_result" if result.is_main_result else "display_data", + "data": result.data, + "metadata": result.extra, + } + ) + + if execution.code_error is not None: + traceback = execution.code_error.traceback or "" + outputs.append( + { + "output_type": "error", + "ename": execution.code_error.name, + "evalue": execution.code_error.value, + "traceback": traceback.splitlines(), + } + ) + + return { + "execution_count": execution.execution_count, + "outputs": outputs, + "status": "ok" if execution.success else "error", + } + + @dataclass class CodeExecutionOutcome: """Normalized result of a code execution, independent of sandbox variant. @@ -183,6 +239,39 @@ def sandbox(self) -> Sandbox: """The wrapped sandbox instance.""" return self._sandbox + @property + def config(self) -> SandboxConfig: + """Variant-neutral configuration for the wrapped sandbox.""" + return self._sandbox.config + + @property + def info(self) -> SandboxInfo | None: + """Runtime information for the wrapped sandbox, when started.""" + return self._sandbox.info + + @property + def id(self) -> str | None: + """Stable execution-backend identifier when the variant exposes one.""" + info = getattr(self._sandbox, "info", None) + metadata = getattr(info, "metadata", None) or {} + kernel_id = metadata.get("kernel_id") + if kernel_id: + return str(kernel_id) + backend = getattr(self._sandbox, "kernel_client", None) + backend_id = getattr(backend, "id", None) + return str(backend_id) if backend_id else self._sandbox.sandbox_id + + @property + def kernel_info(self) -> dict[str, Any]: + """Language metadata without exposing a variant's underlying client.""" + backend = getattr(self._sandbox, "kernel_client", None) + info = getattr(backend, "kernel_info", None) + if isinstance(info, dict): + return info + environments = self._sandbox.list_environments() + language = environments[0].language if environments else "python" + return {"language_info": {"name": language}} + @property def variant(self) -> SandboxVariant | None: """The variant of the wrapped sandbox, if known.""" @@ -219,6 +308,16 @@ def close(self) -> None: if self._owns_sandbox and callable(stop_fn) and self.is_started: stop_fn() + def stop(self, shutdown_kernel: bool = True) -> None: + """Release the client, optionally preserving a borrowed remote backend.""" + if shutdown_kernel: + self.close() + return + backend = getattr(self._sandbox, "kernel_client", None) + backend_stop = getattr(backend, "stop", None) + if callable(backend_stop): + backend_stop(shutdown_kernel=False) + async def close_async(self) -> None: """Async variant of :meth:`close`.""" if not (self._owns_sandbox and self.is_started): @@ -244,6 +343,67 @@ def execute_code( execution = self._sandbox.run_code(code, language=language, timeout=timeout, envs=envs) return CodeExecutionOutcome.from_execution_result(execution) + def execute( + self, + code: str, + silent: bool = False, + timeout: float | None = None, + **kwargs: Any, + ) -> dict[str, Any]: + """Execute code and return a backend-neutral Jupyter-shaped reply.""" + del silent, kwargs + self.start() + execution = self._sandbox.run_code(code, timeout=timeout) + return execution_result_to_reply(execution) + + def execute_interactive( + self, + code: str, + silent: bool = False, + timeout: float | None = None, + output_hook: Callable[[dict[str, Any]], None] | None = None, + **kwargs: Any, + ) -> dict[str, Any]: + """Execute code and optionally emit each normalized output to a callback.""" + reply = self.execute(code, silent=silent, timeout=timeout, **kwargs) + if output_hook is not None: + for output in reply["outputs"]: + output_hook( + { + "msg_type": output.get("output_type", "display_data"), + "content": output, + } + ) + return reply + + def get_variable(self, name: str) -> Any: + """Read a variable through the wrapped sandbox.""" + self.start() + return self._sandbox.get_variable(name) + + def set_variable(self, name: str, value: Any) -> None: + """Set a variable through the wrapped sandbox.""" + self.start() + self._sandbox.set_variable(name, value) + + def set_variables(self, variables: dict[str, Any]) -> None: + """Set multiple variables through the wrapped sandbox.""" + self.start() + self._sandbox.set_variables(variables) + + def interrupt(self) -> bool: + """Interrupt the active execution when supported by the variant.""" + return self._sandbox.interrupt() + + def is_alive(self) -> bool: + """Whether the sandbox is started and available for execution.""" + return self.is_started + + def restart(self) -> None: + """Restart the wrapped sandbox through its public lifecycle.""" + self._sandbox.stop() + self._sandbox.start() + async def execute_code_async( self, code: str, diff --git a/code_sandboxes/colab.py b/code_sandboxes/colab.py new file mode 100644 index 0000000..35f86e5 --- /dev/null +++ b/code_sandboxes/colab.py @@ -0,0 +1,208 @@ +# Copyright (c) 2025-2026 Datalayer, Inc. +# +# BSD 3-Clause License + +"""Google Colab kernel client. + +This module provides :class:`ColabKernelClient`, a thin specialization of +:class:`~jupyter_kernel_client.client.JupyterKernelClient` that connects to an +**already-running** Google Colab kernel. + +A Colab runtime is reached through a per-session proxy. Compared to a vanilla +Jupyter Server connection, it requires: + +* a ``colab-runtime-proxy-token`` query parameter on the websocket URL, and +* the ``X-Colab-Client-Agent`` / ``X-Colab-Runtime-Proxy-Token`` HTTP headers on + both the REST and websocket requests. + +.. note:: + This client **reuses an existing kernel**; it does not create a Colab runtime + from scratch. Consumer Colab has no public API to provision a runtime from a + standalone process (authentication lives in the browser session). Start a + runtime from the Colab UI, then connect to it here using the values taken + from the websocket *channels* URL. + +The ``server_url``, ``kernel_id`` and ``proxy_token`` are the parts of the +websocket *channels* URL that Colab's own frontend uses to reach your assigned +runtime, for example:: + + wss:///api/kernels//channels?session_id=<...>&colab-runtime-proxy-token=&colab-client-agent=web + +Use :func:`parse_colab_channels_url` (or +:meth:`ColabKernelClient.from_channels_url`) to turn that URL into the +``server_url``, ``kernel_id`` and ``proxy_token`` expected by the client. + +Example: + >>> from code_sandboxes import ColabKernelClient + >>> kernel = ColabKernelClient( + ... server_url="https://", + ... kernel_id="", + ... proxy_token="", + ... ) + >>> kernel.start() + >>> reply = kernel.execute("print('hey')") + >>> print(reply) + >>> kernel.stop(shutdown_kernel=False) # disconnect only +""" + +from __future__ import annotations + +import logging +import re +import typing as t +from urllib.parse import parse_qs, urlsplit + +from jupyter_kernel_client.client import JupyterKernelClient +from jupyter_kernel_client.wsclient import JupyterSubprotocol + +#: HTTP header identifying the client agent to the Colab proxy. +COLAB_CLIENT_AGENT_HEADER = "X-Colab-Client-Agent" +#: HTTP header carrying the Colab runtime proxy token. +COLAB_RUNTIME_PROXY_TOKEN_HEADER = "X-Colab-Runtime-Proxy-Token" # noqa: S105 +#: Websocket query parameter carrying the Colab runtime proxy token. +COLAB_RUNTIME_PROXY_TOKEN_PARAM = "colab-runtime-proxy-token" # noqa: S105 +#: Default value advertised through :data:`COLAB_CLIENT_AGENT_HEADER`. +DEFAULT_COLAB_CLIENT_AGENT = "code-sandboxes" + +#: Regular expression extracting the kernel id from a Colab channels URL. +_COLAB_KERNEL_RE = re.compile(r"/api/kernels/([^/]+)/channels", re.IGNORECASE) + + +def parse_colab_channels_url(channels_url: str) -> tuple[str, str, str]: + """Extract ``server_url``, ``kernel_id`` and ``proxy_token`` from a URL. + + Parses the websocket *channels* URL of a running Colab kernel session, as + seen in the browser network tab. + + Args: + channels_url: The websocket *channels* URL of a running Colab kernel + session, e.g. + ``wss:///api/kernels//channels?...&colab-runtime-proxy-token=``. + + Returns: + A ``(server_url, kernel_id, proxy_token)`` tuple where ``server_url`` is + the HTTP(S) base *before* ``/api/kernels``, ``kernel_id`` is the kernel + identifier, and ``proxy_token`` is the Colab runtime proxy token. + + Raises: + ValueError: If the URL does not look like a Colab channels URL or is + missing the proxy token. + """ + split = urlsplit(channels_url) + if not split.netloc: + raise ValueError( + f"Could not parse a Colab channels URL from: {channels_url!r}. " + "Expected a websocket URL of the form " + "'wss:///api/kernels//channels?...&colab-runtime-proxy-token='." + ) + + scheme = "https" if split.scheme.lower() in ("wss", "https") else "http" + + marker = "/api/kernels/" + idx = split.path.find(marker) + if idx == -1: + raise ValueError( + f"Could not find '/api/kernels/' in: {channels_url!r}. " + "Expected the URL to contain 'api/kernels//channels'." + ) + prefix = split.path[:idx] + server_url = f"{scheme}://{split.netloc}{prefix}" + + kernel_match = _COLAB_KERNEL_RE.search(split.path) + if kernel_match is None: + raise ValueError( + f"Could not parse a Colab kernel id from: {channels_url!r}. " + "Expected the URL to contain 'api/kernels//channels'." + ) + kernel_id = kernel_match.group(1) + + query = parse_qs(split.query) + proxy_tokens = query.get(COLAB_RUNTIME_PROXY_TOKEN_PARAM) + if not proxy_tokens or not proxy_tokens[0]: + raise ValueError( + f"Could not find the '{COLAB_RUNTIME_PROXY_TOKEN_PARAM}' query " + f"parameter in: {channels_url!r}." + ) + proxy_token = proxy_tokens[0] + + return server_url, kernel_id, proxy_token + + +class ColabKernelClient(JupyterKernelClient): + """Kernel client connected to an existing Google Colab runtime. + + This client connects to a kernel that is **already running** on a Colab + runtime. It does not create a runtime from scratch. + + Args: + server_url: The Colab runtime proxy URL (the HTTP(S) base *before* + ``/api/kernels``, derived from the session's channels URL). + proxy_token: The Colab runtime proxy token (the + ``colab-runtime-proxy-token`` value from the channels URL). + kernel_id: The identifier of the existing Colab kernel to connect to. + client_agent: Value sent through the ``X-Colab-Client-Agent`` header. + subprotocol: Websocket subprotocol to use; Colab uses the default one. + log: Optional logger. + **kwargs: Forwarded to :class:`~jupyter_kernel_client.client.JupyterKernelClient`. + ``client_kwargs`` and ``headers`` may be provided and are merged with + the Colab-specific values. + """ + + def __init__( + self, + server_url: str, + proxy_token: str, + *, + kernel_id: str, + client_agent: str = DEFAULT_COLAB_CLIENT_AGENT, + subprotocol: JupyterSubprotocol | None = JupyterSubprotocol.DEFAULT, + log: logging.Logger | None = None, + **kwargs: t.Any, + ) -> None: + client_kwargs: dict[str, t.Any] = dict(kwargs.pop("client_kwargs", None) or {}) + client_kwargs.setdefault("subprotocol", subprotocol) + + extra_params: dict[str, t.Any] = dict(client_kwargs.get("extra_params", None) or {}) + extra_params[COLAB_RUNTIME_PROXY_TOKEN_PARAM] = proxy_token + client_kwargs["extra_params"] = extra_params + + headers: dict[str, t.Any] = dict(kwargs.pop("headers", None) or {}) + headers.setdefault(COLAB_CLIENT_AGENT_HEADER, client_agent) + headers[COLAB_RUNTIME_PROXY_TOKEN_HEADER] = proxy_token + + # Colab authenticates through the proxy token, not the Jupyter token. + # Drop any provided Jupyter token to avoid sending an Authorization + # header and a `token=` query parameter that Colab does not use. + kwargs.pop("token", None) + + super().__init__( + kernel_id=kernel_id, + log=log, + server_url=server_url, + token=None, + client_kwargs=client_kwargs, + headers=headers, + **kwargs, + ) + + @classmethod + def from_channels_url( + cls, + channels_url: str, + **kwargs: t.Any, + ) -> ColabKernelClient: + """Create a client from a Colab kernel session *channels* URL. + + Args: + channels_url: The websocket *channels* URL of a running Colab kernel + session (see :func:`parse_colab_channels_url`). + **kwargs: Forwarded to :class:`ColabKernelClient`. Values provided + here override those parsed from the URL. + + Returns: + A configured :class:`ColabKernelClient` instance. + """ + server_url, kernel_id, proxy_token = parse_colab_channels_url(channels_url) + kwargs.setdefault("kernel_id", kernel_id) + kwargs.setdefault("proxy_token", proxy_token) + return cls(server_url=server_url, **kwargs) diff --git a/code_sandboxes/colab_sandbox.py b/code_sandboxes/colab_sandbox.py index 64adaa2..f9bb934 100644 --- a/code_sandboxes/colab_sandbox.py +++ b/code_sandboxes/colab_sandbox.py @@ -5,7 +5,7 @@ """Google Colab sandbox implementation. This sandbox connects to an existing Google Colab runtime and executes code in -its kernel using ``jupyter-kernel-client``'s :class:`ColabKernelClient`. +its kernel using :class:`code_sandboxes.colab.ColabKernelClient`. Unlike the Jupyter/Docker sandboxes, this sandbox does **not** provision a runtime: a Colab runtime must already be running in a browser session. Reuse it @@ -20,6 +20,7 @@ import uuid from .base import Sandbox +from .colab import ColabKernelClient, parse_colab_channels_url from .exceptions import SandboxConfigurationError, SandboxNotStartedError from .interfaces import ISandboxClient from .models import ( @@ -95,14 +96,6 @@ def start(self) -> None: if self._channels_url and ( not self._server_url or not self._kernel_id or not self._proxy_token ): - try: - from jupyter_kernel_client import parse_colab_channels_url - except ImportError as exc: - raise SandboxConfigurationError( - "jupyter-kernel-client>=0.14.0 is required for Colab channels_url parsing. " - "Install it with: pip install jupyter-kernel-client" - ) from exc - parsed_server_url, parsed_kernel_id, parsed_proxy_token = parse_colab_channels_url( self._channels_url ) @@ -116,14 +109,6 @@ def start(self) -> None: "Provide them directly, or pass 'channels_url' from an active Colab session." ) - try: - from jupyter_kernel_client import ColabKernelClient - except ImportError as exc: - raise SandboxConfigurationError( - "jupyter-kernel-client>=0.12.0 is required for ColabSandbox. " - "Install it with: pip install jupyter-kernel-client" - ) from exc - self._client = ColabKernelClient( server_url=self._server_url, kernel_id=self._kernel_id, @@ -139,7 +124,7 @@ def start(self) -> None: status=SandboxStatus.RUNNING, created_at=time.time(), name=self.config.name, - metadata={"server_url": self._server_url, "kernel_id": self._kernel_id}, + metadata={"server_url": self._server_url, "kernel_id": self._client.id}, config=self.config, ) self._started = True diff --git a/code_sandboxes/interfaces.py b/code_sandboxes/interfaces.py index bc3a1db..84c4f03 100644 --- a/code_sandboxes/interfaces.py +++ b/code_sandboxes/interfaces.py @@ -6,15 +6,29 @@ from __future__ import annotations -from typing import Protocol, runtime_checkable - -from jupyter_kernel_client.interfaces import IJupyterKernelClient +from typing import Any, Protocol, runtime_checkable @runtime_checkable -class ISandboxClient(IJupyterKernelClient, Protocol): - """Kernel client protocol exposed by sandbox variants. +class ISandboxClient(Protocol): + """Internal execution backend used by kernel-backed sandbox variants.""" + + @property + def id(self) -> str | None: ... + + @property + def kernel_info(self) -> dict[str, Any] | None: ... + + def start(self, **kwargs: Any) -> None: ... + + def stop(self, shutdown_kernel: bool = True) -> None: ... + + def execute(self, code: str, **kwargs: Any) -> dict[str, Any]: ... + + def execute_interactive(self, code: str, **kwargs: Any) -> dict[str, Any]: ... + + def get_variable(self, name: str) -> Any: ... + + def set_variable(self, name: str, value: Any) -> None: ... - This currently matches ``IJupyterKernelClient`` exactly and acts as an extension - point for sandbox-specific client capabilities. - """ + def interrupt(self) -> bool: ... diff --git a/code_sandboxes/jupyter_sandbox.py b/code_sandboxes/jupyter_sandbox.py index 74658ce..fa565b9 100644 --- a/code_sandboxes/jupyter_sandbox.py +++ b/code_sandboxes/jupyter_sandbox.py @@ -411,7 +411,7 @@ def start(self) -> None: status=SandboxStatus.RUNNING, created_at=time.time(), name=self.config.name, - metadata={"server_url": self._server_url}, + metadata={"server_url": self._server_url, "kernel_id": self._client.id}, config=self.config, ) self._started = True diff --git a/code_sandboxes/kaggle.py b/code_sandboxes/kaggle.py new file mode 100644 index 0000000..4a8a04e --- /dev/null +++ b/code_sandboxes/kaggle.py @@ -0,0 +1,162 @@ +# Copyright (c) 2025-2026 Datalayer, Inc. +# +# BSD 3-Clause License + +"""Kaggle kernel client. + +This module provides :class:`KaggleKernelClient`, a thin specialization of +:class:`~jupyter_kernel_client.client.JupyterKernelClient` that knows how to connect to +a Kaggle interactive notebook runtime. + +There are two ways to authenticate: + +* **API token (default).** Provide a Kaggle API token, either explicitly through + the ``token`` argument or via the :data:`KAGGLE_API_TOKEN` environment + variable. The token is used to authenticate REST/websocket requests, so + omitting ``kernel_id`` lets :meth:`start` *create* a new kernel on the runtime + (``POST /api/kernels``). +* **Signed proxy URL.** When you connect to an already-running notebook session, + the signed JWT embedded in the proxied ``server_url`` path carries the + authentication and no token is required (pass ``token=None`` explicitly). + +The ``server_url`` and ``kernel_id`` can be derived from the websocket +*channels* URL of a running Kaggle notebook session (visible in the browser +network tab), for example:: + + wss://kkb-production.jupyter-proxy.kaggle.net/k/12345678/eyJhbGci.../proxy/api/kernels/11e073f0-e82d-4029-be8d-3918f7ed1a9e/channels?session_id=... + +Use :func:`parse_kaggle_channels_url` (or +:meth:`KaggleKernelClient.from_channels_url`) to turn that URL into the +``server_url`` and ``kernel_id`` expected by the client. + +Example: + >>> import os + >>> from code_sandboxes import KaggleKernelClient + >>> os.environ["KAGGLE_API_TOKEN"] = "..." + >>> with KaggleKernelClient(server_url="https://.../proxy") as kernel: + ... print("kernel_id:", kernel.id) # a new kernel was created + ... reply = kernel.execute("print('hey')") + ... print(reply) +""" + +from __future__ import annotations + +import logging +import os +import re +import typing as t + +from jupyter_kernel_client.client import JupyterKernelClient +from jupyter_kernel_client.wsclient import JupyterSubprotocol + +#: Environment variable holding the Kaggle API token used for authentication. +KAGGLE_API_TOKEN_ENV = "KAGGLE_API_TOKEN" # noqa: S105 +_TOKEN_UNSET = object() +#: Regular expression matching the proxied server base of a Kaggle channels URL. +_KAGGLE_SERVER_RE = re.compile(r"^(wss?)://(.*?)/proxy", re.IGNORECASE) +#: Regular expression extracting the kernel id from a Kaggle channels URL. +_KAGGLE_KERNEL_RE = re.compile(r"kernels/([0-9a-f-]+)/channels", re.IGNORECASE) + + +def parse_kaggle_channels_url(channels_url: str) -> tuple[str, str]: + """Extract the ``server_url`` and ``kernel_id`` from a Kaggle channels URL. + + Args: + channels_url: The websocket *channels* URL of a running Kaggle notebook + session, e.g. ``wss://.../proxy/api/kernels//channels?...``. + + Returns: + A ``(server_url, kernel_id)`` tuple where ``server_url`` is the HTTP(S) + base ending in ``/proxy`` and ``kernel_id`` is the kernel identifier. + + Raises: + ValueError: If the URL does not look like a Kaggle channels URL. + """ + server_match = _KAGGLE_SERVER_RE.match(channels_url) + if server_match is None: + raise ValueError( + f"Could not parse a Kaggle proxy server URL from: {channels_url!r}. " + "Expected a websocket URL of the form 'wss://.../proxy/api/kernels//channels'." + ) + scheme = "https" if server_match.group(1).lower() == "wss" else "http" + server_url = f"{scheme}://{server_match.group(2)}/proxy" + + kernel_match = _KAGGLE_KERNEL_RE.search(channels_url) + if kernel_match is None: + raise ValueError( + f"Could not parse a Kaggle kernel id from: {channels_url!r}. " + "Expected the URL to contain 'kernels//channels'." + ) + kernel_id = kernel_match.group(1) + + return server_url, kernel_id + + +class KaggleKernelClient(JupyterKernelClient): + """Kernel client connected to a Kaggle interactive notebook runtime. + + Args: + server_url: The Kaggle runtime proxy URL (ending in ``/proxy``). This is + the HTTP(S) base derived from the notebook session's channels URL. + kernel_id: The identifier of the Kaggle kernel to connect to. If omitted, + :meth:`start` creates a new kernel on the runtime (which requires a + valid API ``token``). + token: The Kaggle API token used to authenticate. When omitted, it + falls back to the :data:`KAGGLE_API_TOKEN` environment variable. + Pass ``token=None`` to rely solely on the signed proxy + ``server_url`` (even if the environment variable is set). + subprotocol: Websocket subprotocol to use; Kaggle uses the default one. + log: Optional logger. + **kwargs: Forwarded to :class:`~jupyter_kernel_client.client.JupyterKernelClient`. + ``client_kwargs`` may be provided and is merged with the + Kaggle-specific values. + """ + + def __init__( + self, + server_url: str, + *, + kernel_id: str | None = None, + token: str | None | object = _TOKEN_UNSET, + subprotocol: JupyterSubprotocol | None = JupyterSubprotocol.DEFAULT, + log: logging.Logger | None = None, + **kwargs: t.Any, + ) -> None: + client_kwargs: dict[str, t.Any] = dict(kwargs.pop("client_kwargs", None) or {}) + client_kwargs.setdefault("subprotocol", subprotocol) + + # Resolve the Kaggle API token from the environment only when omitted. + # Explicit ``token=None`` disables env fallback and relies on signed + # proxy authentication embedded in ``server_url``. + if token is _TOKEN_UNSET: + token = os.environ.get(KAGGLE_API_TOKEN_ENV) + + super().__init__( + kernel_id=kernel_id, + log=log, + server_url=server_url, + token=token, + client_kwargs=client_kwargs, + **kwargs, + ) + + @classmethod + def from_channels_url( + cls, + channels_url: str, + **kwargs: t.Any, + ) -> KaggleKernelClient: + """Create a client from a Kaggle notebook session *channels* URL. + + Args: + channels_url: The websocket *channels* URL of a running Kaggle + notebook session (see :func:`parse_kaggle_channels_url`). + **kwargs: Forwarded to :class:`KaggleKernelClient`. A ``kernel_id`` + provided here overrides the one parsed from the URL. + + Returns: + A configured :class:`KaggleKernelClient` instance. + """ + server_url, kernel_id = parse_kaggle_channels_url(channels_url) + kwargs.setdefault("kernel_id", kernel_id) + return cls(server_url=server_url, **kwargs) diff --git a/code_sandboxes/kaggle_execute.py b/code_sandboxes/kaggle_execute.py new file mode 100644 index 0000000..bb2ae7c --- /dev/null +++ b/code_sandboxes/kaggle_execute.py @@ -0,0 +1,591 @@ +# Copyright (c) 2025-2026 Datalayer, Inc. +# +# BSD 3-Clause License + +"""Run code on Kaggle through the official *kernels* (notebooks) API. + +Unlike :class:`~code_sandboxes.kaggle.KaggleKernelClient`, which connects +to an **already-running** interactive Kaggle notebook session over websocket, +this module drives Kaggle's **batch** kernels API: it creates (or updates) a +Kaggle notebook, runs it end to end on Kaggle's infrastructure, waits for +completion, and downloads the output. This is the supported "from zero" way to +create a Kaggle kernel and run code from a standalone process — no browser +session is required, only Kaggle API credentials. + +Authentication is delegated to the official ``kaggle`` package, which resolves +credentials from (in order) the ``KAGGLE_USERNAME`` / ``KAGGLE_KEY`` environment +variables, a ``KAGGLE_API_TOKEN`` (newer CLIs), or a ``~/.kaggle/kaggle.json`` +file. Install the optional dependency with +``pip install 'code-sandboxes[kaggle]'``. + +Example: + >>> from code_sandboxes import KaggleKernelExecutor + >>> executor = KaggleKernelExecutor() # credentials from the environment + >>> result = executor.execute( + ... "import pandas as pd\\nprint('Running on Kaggle')", + ... title="My Python Notebook", + ... ) + >>> print(result.status) # e.g. "COMPLETE" + >>> print(result.log) # captured execution log +""" + +from __future__ import annotations + +import json +import logging +import re +import tempfile +import time +import typing as t +import uuid +from dataclasses import dataclass, field +from pathlib import Path + +if t.TYPE_CHECKING: # pragma: no cover - typing only + from kaggle.api.kaggle_api_extended import KaggleApi + +#: Kernel statuses that indicate execution has finished (successfully or not). +TERMINAL_STATUSES: frozenset[str] = frozenset({"COMPLETE", "ERROR", "CANCEL_ACKNOWLEDGED"}) +#: Default number of seconds to wait for a Kaggle run to finish. +DEFAULT_EXECUTION_TIMEOUT = 3600.0 +#: Default number of seconds between status polls. +DEFAULT_POLL_INTERVAL = 10.0 + +_SLUG_CLEAN_RE = re.compile(r"[^a-z0-9]+") + +_KAGGLE_ACCELERATOR_ALIASES: dict[str, str] = { + "tesla p100": "NvidiaTeslaP100", + "nvidiateslap100": "NvidiaTeslaP100", + "p100": "NvidiaTeslaP100", + "tesla t4": "NvidiaTeslaT4", + "nvidiateslat4": "NvidiaTeslaT4", + "t4": "NvidiaTeslaT4", + "tesla t4 high memory": "NvidiaTeslaT4Highmem", + "nvidiateslat4highmem": "NvidiaTeslaT4Highmem", + "t4 high memory": "NvidiaTeslaT4Highmem", + "t4highmem": "NvidiaTeslaT4Highmem", + "l4": "NvidiaL4", + "nvidial4": "NvidiaL4", + "l4 x1": "NvidiaL4X1", + "nvidial4x1": "NvidiaL4X1", + "l4x1": "NvidiaL4X1", + "a100": "NvidiaTeslaA100", + "nvidiateslaa100": "NvidiaTeslaA100", + "h100": "NvidiaH100", + "nvidiah100": "NvidiaH100", + "rtx pro 6000": "NvidiaRtxPro6000", + "nvidiartxpro6000": "NvidiaRtxPro6000", + "rtxpro6000": "NvidiaRtxPro6000", +} + +_KAGGLE_ACCELERATOR_VALUES = sorted(set(_KAGGLE_ACCELERATOR_ALIASES.values())) + + +def _slugify(value: str) -> str: + """Turn an arbitrary title into a Kaggle-compatible slug.""" + slug = _SLUG_CLEAN_RE.sub("-", value.strip().lower()).strip("-") + return slug or f"code-sandboxes-run-{uuid.uuid4().hex[:8]}" + + +def _normalize_status(status: t.Any) -> str: + """Normalize a Kaggle status value (enum, int or string) to a plain name.""" + name = getattr(status, "name", None) + if name is None: + name = str(status) + return name.split(".")[-1].strip().upper() + + +def _build_notebook(code: str) -> dict[str, t.Any]: + """Build a minimal nbformat v4 notebook with a single code cell.""" + return { + "cells": [ + { + "cell_type": "code", + "id": uuid.uuid4().hex[:8], + "metadata": {"language": "python"}, + "execution_count": None, + "outputs": [], + "source": code, + } + ], + "metadata": { + "kernelspec": { + "name": "python3", + "display_name": "Python 3", + "language": "python", + }, + "language_info": {"name": "python"}, + }, + "nbformat": 4, + "nbformat_minor": 5, + } + + +def _normalize_accelerator(accelerator: str | None) -> str | None: + """Normalize human-friendly accelerator names to Kaggle API values.""" + if accelerator is None: + return None + normalized = accelerator.strip() + if not normalized: + return None + alias_key = normalized.lower() + resolved = _KAGGLE_ACCELERATOR_ALIASES.get(alias_key) + if resolved is None: + raise ValueError( + "Unsupported Kaggle accelerator {!r}. Supported values: {}".format( + accelerator, + ", ".join(_KAGGLE_ACCELERATOR_VALUES), + ) + ) + return resolved + + +@dataclass +class KaggleExecutionResult: + """Outcome of running code on Kaggle through the kernels API.""" + + #: Full kernel reference, ``"/"``. + slug: str + #: Normalized final status, e.g. ``"COMPLETE"`` or ``"ERROR"``. + status: str + #: Kaggle notebook URL, when known. + url: str | None = None + #: Kernel version number produced by the push, when known. + version_number: int | None = None + #: Failure message reported by Kaggle, when the run did not succeed. + failure_message: str | None = None + #: Directory the outputs were downloaded to, when ``download_output=True``. + output_dir: str | None = None + #: Paths of downloaded output files. + output_files: list[str] = field(default_factory=list) + #: Contents of the execution ``.log`` file, when present. + log: str | None = None + #: The executed notebook (nbformat dict), when present in the output. + notebook: dict[str, t.Any] | None = None + #: Jupyter-style execute reply derived from the notebook/log output. + kernel_reply: dict[str, t.Any] | None = None + + @property + def succeeded(self) -> bool: + """Whether the run finished with a ``COMPLETE`` status.""" + return self.status == "COMPLETE" + + @property + def outputs(self) -> list[dict[str, t.Any]]: + """Jupyter-like output list for the execution.""" + reply = self.kernel_reply or self.to_kernel_reply() + return list(reply.get("outputs", [])) + + @property + def stdout(self) -> str: + """Merged stdout stream extracted from kernel-like outputs.""" + chunks: list[str] = [] + for output in self.outputs: + if output.get("output_type") == "stream" and output.get("name") == "stdout": + chunks.append(str(output.get("text", ""))) + return "".join(chunks) + + @property + def stderr(self) -> str: + """Merged stderr stream extracted from kernel-like outputs.""" + chunks: list[str] = [] + for output in self.outputs: + if output.get("output_type") == "stream" and output.get("name") == "stderr": + chunks.append(str(output.get("text", ""))) + return "".join(chunks) + + def __repr__(self) -> str: + """Compact representation that avoids printing full raw logs.""" + reply = self.kernel_reply or self.to_kernel_reply() + return ( + "KaggleExecutionResult(" + f"slug={self.slug!r}, status={self.status!r}, kernel_status={reply.get('status')!r}, " + f"url={self.url!r}, version_number={self.version_number!r}, " + f"execution_count={reply.get('execution_count', 0)!r}, " + f"stdout={self.stdout.strip()!r}, stderr={self.stderr.strip()!r}, " + f"failure_message={self.failure_message!r}, output_dir={self.output_dir!r}, " + f"output_files={self.output_files!r}" + ")" + ) + + def to_kernel_reply(self) -> dict[str, t.Any]: + """Return a Jupyter-like execute reply. + + This mirrors the shape returned by ``JupyterKernelClient.execute``: + ``{"execution_count": int, "outputs": list, "status": "ok"|"error"}``. + """ + outputs: list[dict[str, t.Any]] = [] + execution_count = 0 + + if self.notebook: + outputs, execution_count = _extract_notebook_outputs(self.notebook) + + if not outputs and self.log: + outputs = _outputs_from_kaggle_log(self.log) + + has_error_output = any(output.get("output_type") == "error" for output in outputs) + status = "ok" if (self.succeeded and not has_error_output) else "error" + + return { + "execution_count": execution_count, + "outputs": outputs, + "status": status, + } + + +def _extract_notebook_outputs(notebook: dict[str, t.Any]) -> tuple[list[dict[str, t.Any]], int]: + """Extract outputs from the last executed code cell in an nbformat dict.""" + execution_count = 0 + outputs: list[dict[str, t.Any]] = [] + + for cell in notebook.get("cells", []): + if cell.get("cell_type") != "code": + continue + cell_outputs = cell.get("outputs") or [] + if not cell_outputs: + continue + outputs = [ + output + for output in cell_outputs + if isinstance(output, dict) and "output_type" in output + ] + count = cell.get("execution_count") + execution_count = count if isinstance(count, int) else execution_count + + return outputs, execution_count + + +def _outputs_from_kaggle_log(log: str) -> list[dict[str, t.Any]]: + """Convert Kaggle log stream events into Jupyter stream outputs.""" + if not log.strip(): + return [] + + events: list[dict[str, t.Any]] = [] + try: + parsed = json.loads(log) + if isinstance(parsed, list): + events = [event for event in parsed if isinstance(event, dict)] + except ValueError: + # Keep plain-text fallback for unexpected log formats. + return [{"output_type": "stream", "name": "stdout", "text": log}] + + outputs: list[dict[str, t.Any]] = [] + for event in events: + stream_name = str(event.get("stream_name", "stdout")).strip().lower() + if stream_name not in {"stdout", "stderr"}: + stream_name = "stdout" + text = event.get("data") + if text is None: + continue + outputs.append( + { + "output_type": "stream", + "name": stream_name, + "text": str(text), + } + ) + + return outputs + + +class KaggleKernelExecutor: + """Create and run Kaggle notebooks (kernels) through the official API. + + Args: + username: Kaggle username used to build the kernel reference + (``"/"``). When omitted, it is read from the + authenticated Kaggle configuration. + api: An already-authenticated ``KaggleApi`` instance. When omitted, one + is created and authenticated lazily on first use. + quiet: Whether to suppress the Kaggle client's own progress output. + log: Optional logger. + """ + + def __init__( + self, + username: str | None = None, + *, + api: KaggleApi | None = None, + quiet: bool = True, + log: logging.Logger | None = None, + ) -> None: + self._username = username + self._api = api + self._quiet = quiet + self.log = log or logging.getLogger(__name__) + + # -- Kaggle client --------------------------------------------------- + + @property + def api(self) -> KaggleApi: + """The authenticated Kaggle API client (created lazily).""" + if self._api is None: + try: + from kaggle.api.kaggle_api_extended import KaggleApi + except ImportError as exc: # pragma: no cover - optional dependency + raise ImportError( + "The 'kaggle' package is required for KaggleKernelExecutor. " + "Install it with: pip install 'code-sandboxes[kaggle]'" + ) from exc + api = KaggleApi() + api.authenticate() + self._api = api + return self._api + + def _resolve_username(self) -> str: + if self._username: + return self._username + username = self.api.get_config_value("username") + if not username: + raise ValueError( + "Could not determine the Kaggle username. Pass 'username=...' or " + "configure Kaggle credentials (KAGGLE_USERNAME / kaggle.json)." + ) + self._username = username + return username + + # -- Public API ------------------------------------------------------ + + def execute( + self, + code: str, + *, + slug: str | None = None, + title: str | None = None, + language: str = "python", + kernel_type: str = "notebook", + enable_gpu: bool = False, + accelerator: str | None = None, + enable_internet: bool = True, + is_private: bool = True, + dataset_sources: t.Sequence[str] | None = None, + competition_sources: t.Sequence[str] | None = None, + kernel_sources: t.Sequence[str] | None = None, + model_sources: t.Sequence[str] | None = None, + wait: bool = True, + timeout: float = DEFAULT_EXECUTION_TIMEOUT, + poll_interval: float = DEFAULT_POLL_INTERVAL, + download_output: bool = True, + output_dir: str | None = None, + ) -> KaggleExecutionResult: + """Create (or update) a Kaggle notebook, run it, and collect the output. + + Args: + code: The Python (or R) source to run in a single notebook cell. + slug: Kaggle notebook slug. Generated from ``title`` (or a random + value) when omitted. Reusing a slug creates a new version. + title: Human-readable notebook title. Defaults to the slug. + language: Notebook language (``"python"`` or ``"r"``). + kernel_type: ``"notebook"`` or ``"script"``. + enable_gpu: Whether to request a GPU. + accelerator: Explicit Kaggle accelerator value, for example + ``"NvidiaTeslaT4"`` or ``"NvidiaTeslaP100"``. Friendly aliases + such as ``"T4"`` and ``"P100"`` are also accepted. When set, GPU + is enabled automatically. + enable_internet: Whether the kernel may access the internet. + is_private: Whether the notebook is private. + dataset_sources: Kaggle dataset references to attach. + competition_sources: Kaggle competition references to attach. + kernel_sources: Kaggle kernel references to attach. + model_sources: Kaggle model references to attach. + wait: Whether to block until the run reaches a terminal status. + timeout: Maximum seconds to wait for completion. + poll_interval: Seconds between status polls. + download_output: Whether to download outputs after completion. + output_dir: Directory to download outputs into. A temporary + directory is used when omitted. + + Returns: + A :class:`KaggleExecutionResult` describing the run. + """ + username = self._resolve_username() + slug = slug or _slugify(title or f"jkc-run-{uuid.uuid4().hex[:8]}") + ref = f"{username}/{slug}" + title = title or slug + normalized_accelerator = _normalize_accelerator(accelerator) + resolved_enable_gpu = bool(enable_gpu or normalized_accelerator) + + with tempfile.TemporaryDirectory(prefix="jkc-kaggle-") as tmp: + folder = Path(tmp) + code_file = self._write_sources(folder, code, kernel_type, language) + self._write_metadata( + folder, + ref=ref, + title=title, + code_file=code_file, + language=language, + kernel_type=kernel_type, + enable_gpu=resolved_enable_gpu, + accelerator=normalized_accelerator, + enable_internet=enable_internet, + is_private=is_private, + dataset_sources=dataset_sources, + competition_sources=competition_sources, + kernel_sources=kernel_sources, + model_sources=model_sources, + ) + + self.log.info("Pushing Kaggle kernel %s", ref) + push_response = self._kernels_push(str(folder), normalized_accelerator) + + url = getattr(push_response, "url", None) + version_number = getattr(push_response, "version_number", None) + + result = KaggleExecutionResult( + slug=ref, + status="QUEUED", + url=url, + version_number=version_number, + ) + + if not wait: + result.status = self.status(ref) + return result + + result.status, result.failure_message = self._wait_for_completion( + ref, timeout=timeout, poll_interval=poll_interval + ) + + if download_output: + self._download_output(ref, result, output_dir) + + result.kernel_reply = result.to_kernel_reply() + + return result + + def status(self, slug: str) -> str: + """Return the normalized status of a Kaggle kernel.""" + response = self.api.kernels_status(slug) + return _normalize_status(getattr(response, "status", response)) + + def output( + self, + slug: str, + dest: str, + *, + force: bool = True, + quiet: bool | None = None, + ) -> list[str]: + """Download a kernel's output files into ``dest`` and return their paths.""" + Path(dest).mkdir(parents=True, exist_ok=True) + quiet = self._quiet if quiet is None else quiet + files, _token = self.api.kernels_output(slug, path=dest, force=force, quiet=quiet) + return list(files) + + # -- Internals ------------------------------------------------------- + + @staticmethod + def _write_sources(folder: Path, code: str, kernel_type: str, language: str) -> str: + if kernel_type == "script": + extension = "r" if language == "r" else "py" + code_file = f"script.{extension}" + (folder / code_file).write_text(code, encoding="utf-8") + else: + code_file = "notebook.ipynb" + (folder / code_file).write_text( + json.dumps(_build_notebook(code), indent=1), encoding="utf-8" + ) + return code_file + + @staticmethod + def _write_metadata( + folder: Path, + *, + ref: str, + title: str, + code_file: str, + language: str, + kernel_type: str, + enable_gpu: bool, + accelerator: str | None, + enable_internet: bool, + is_private: bool, + dataset_sources: t.Sequence[str] | None, + competition_sources: t.Sequence[str] | None, + kernel_sources: t.Sequence[str] | None, + model_sources: t.Sequence[str] | None, + ) -> None: + metadata = { + "id": ref, + "title": title, + "code_file": code_file, + "language": language, + "kernel_type": kernel_type, + "is_private": is_private, + "enable_gpu": enable_gpu, + "enable_internet": enable_internet, + "dataset_sources": list(dataset_sources or []), + "competition_sources": list(competition_sources or []), + "kernel_sources": list(kernel_sources or []), + "model_sources": list(model_sources or []), + } + if accelerator: + metadata["accelerator"] = accelerator + (folder / "kernel-metadata.json").write_text( + json.dumps(metadata, indent=2), encoding="utf-8" + ) + + def _kernels_push(self, folder: str, accelerator: str | None): + if accelerator is None: + return self.api.kernels_push(folder) + try: + return self.api.kernels_push(folder, accelerator=accelerator) + except TypeError: + # Older kaggle clients may not expose the accelerator kwarg yet. + return self.api.kernels_push(folder) + + def _wait_for_completion( + self, + slug: str, + *, + timeout: float, + poll_interval: float, + ) -> tuple[str, str | None]: + deadline = time.monotonic() + timeout + status = "QUEUED" + failure_message: str | None = None + while True: + response = self.api.kernels_status(slug) + status = _normalize_status(getattr(response, "status", response)) + failure_message = getattr(response, "failure_message", None) or None + self.log.debug("Kaggle kernel %s status: %s", slug, status) + if status in TERMINAL_STATUSES: + break + if time.monotonic() >= deadline: + self.log.warning("Timed out waiting for Kaggle kernel %s", slug) + break + time.sleep(poll_interval) + return status, failure_message + + def _download_output( + self, + slug: str, + result: KaggleExecutionResult, + output_dir: str | None, + ) -> None: + try: + if output_dir is None: + output_dir = tempfile.mkdtemp(prefix="jkc-kaggle-out-") + files = self.output(slug, output_dir) + result.output_dir = output_dir + result.output_files = files + self._populate_log_and_notebook(result, files) + except Exception as exc: # pragma: no cover - best-effort download + self.log.warning("Could not download Kaggle output for %s: %s", slug, exc) + + @staticmethod + def _populate_log_and_notebook( + result: KaggleExecutionResult, files: t.Sequence[str] + ) -> None: + for path_str in files: + path = Path(path_str) + if path.suffix == ".log" and result.log is None: + try: + result.log = path.read_text(encoding="utf-8", errors="replace") + except OSError: + pass + elif path.suffix == ".ipynb" and result.notebook is None: + try: + result.notebook = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + pass diff --git a/code_sandboxes/kaggle_sandbox.py b/code_sandboxes/kaggle_sandbox.py index a470734..4c398d0 100644 --- a/code_sandboxes/kaggle_sandbox.py +++ b/code_sandboxes/kaggle_sandbox.py @@ -5,7 +5,7 @@ """Kaggle sandbox implementation. This sandbox connects to a Kaggle interactive notebook runtime and executes code -in its kernel using ``jupyter-kernel-client``'s :class:`KaggleKernelClient`. +in its kernel using :class:`code_sandboxes.kaggle.KaggleKernelClient`. When runtime connection details are not provided, it transparently falls back to Kaggle's batch execution API via ``KaggleKernelExecutor``. This mode is useful @@ -36,6 +36,8 @@ from .base import Sandbox from .exceptions import SandboxConfigurationError, SandboxNotStartedError from .interfaces import ISandboxClient +from .kaggle import KaggleKernelClient, parse_kaggle_channels_url +from .kaggle_execute import KaggleKernelExecutor from .models import ( CodeError, Context, @@ -114,14 +116,6 @@ def start(self) -> None: # Transparent batch mode: when no interactive runtime connection details # are available, fall back to Kaggle's official job API. if not self._server_url and not self._channels_url: - try: - from jupyter_kernel_client import KaggleKernelExecutor - except ImportError as exc: - raise SandboxConfigurationError( - "jupyter-kernel-client>=0.14.0 is required for Kaggle batch execution. " - "Install it with: pip install jupyter-kernel-client" - ) from exc - self._executor = KaggleKernelExecutor( username=self._extra_kwargs.get("username"), quiet=True, @@ -140,17 +134,6 @@ def start(self) -> None: self._started = True return - try: - from jupyter_kernel_client import ( - KaggleKernelClient, - parse_kaggle_channels_url, - ) - except ImportError as exc: - raise SandboxConfigurationError( - "jupyter-kernel-client>=0.12.0 is required for KaggleSandbox. " - "Install it with: pip install jupyter-kernel-client" - ) from exc - # Derive server_url / kernel_id from a channels URL when provided. if self._channels_url and (not self._server_url or not self._kernel_id): parsed_server_url, parsed_kernel_id = parse_kaggle_channels_url(self._channels_url) @@ -178,7 +161,7 @@ def start(self) -> None: status=SandboxStatus.RUNNING, created_at=time.time(), name=self.config.name, - metadata={"server_url": self._server_url, "kernel_id": self._kernel_id}, + metadata={"server_url": self._server_url, "kernel_id": self._client.id}, config=self.config, ) self._started = True @@ -634,7 +617,7 @@ def _run_code_batch( # noqa: C901 if on_error: on_error(code_error) elif result.log: - # Backward-compatible fallback for older jupyter-kernel-client versions. + # Preserve plain log output when no normalized kernel reply is available. for line in result.log.splitlines(): msg = OutputMessage(line=line, timestamp=now, error=False) stdout_messages.append(msg) diff --git a/docs/docs/api-reference/index.mdx b/docs/docs/api-reference/index.mdx index cf53d33..9091ee9 100644 --- a/docs/docs/api-reference/index.mdx +++ b/docs/docs/api-reference/index.mdx @@ -237,6 +237,39 @@ async def execute_code_streaming_async( ) -> AsyncIterator[OutputMessage | Result | CodeError] ``` +#### Compatibility-free execution facade + +Consumers that need a Jupyter-shaped reply can stay on the sandbox client +without accessing a variant's underlying kernel implementation: + +```python +reply = client.execute(code, timeout=60) +reply = client.execute_interactive(code, output_hook=handle_output) +``` + +The client also exposes variant-neutral lifecycle and variable operations: + +```python +client.start() +client.interrupt() +client.restart() +client.set_variable("name", value) +client.set_variables({"one": 1, "two": 2}) +value = client.get_variable("name") +client.stop() +``` + +### Properties + +| Property | Description | +|----------|-------------| +| `id` | Active backend identifier, including a remote kernel ID when applicable | +| `variant` | Configured sandbox variant | +| `config` | Variant-neutral sandbox configuration | +| `info` | Runtime sandbox information after startup | +| `kernel_info` | Language metadata without exposing the underlying kernel client | +| `is_started` | Whether the sandbox has been started | + ### Outcome model `CodeExecutionOutcome` includes normalized fields such as: @@ -253,6 +286,33 @@ async def execute_code_streaming_async( --- +## Kaggle and Colab Clients + +Code Sandboxes owns the provider-specific clients used by its managed notebook +variants. They are exported from `code_sandboxes`: + +```python +from code_sandboxes import ( + ColabKernelClient, + KaggleExecutionResult, + KaggleKernelClient, + KaggleKernelExecutor, + parse_colab_channels_url, + parse_kaggle_channels_url, +) +``` + +- `KaggleKernelClient` connects to an interactive Kaggle notebook kernel. +- `KaggleKernelExecutor` submits and monitors Kaggle batch notebook jobs. +- `KaggleExecutionResult` normalizes batch output into a Jupyter-like reply. +- `ColabKernelClient` connects to an already-running Google Colab kernel. +- The parser helpers extract connection details from browser channels URLs. + +See [Kaggle](/sandboxes/kaggle) and [Google Colab](/sandboxes/google-colab) for +authentication and complete examples. + +--- + ## SandboxFilesystem File operations interface. diff --git a/docs/docs/installation/index.mdx b/docs/docs/installation/index.mdx index 25410d0..67daf99 100644 --- a/docs/docs/installation/index.mdx +++ b/docs/docs/installation/index.mdx @@ -27,9 +27,6 @@ pip install code-sandboxes[docker] # With Kaggle support pip install code-sandboxes[kaggle] -# With Google Colab support -pip install code-sandboxes[colab] - # With Monty (secure in-process interpreter) support pip install code-sandboxes[monty] @@ -46,7 +43,7 @@ pip install code-sandboxes[all] - For Docker variant: Docker installed and running - For Datalayer variant: valid `DATALAYER_API_KEY` - For Kaggle variant: `code-sandboxes[kaggle]` and Kaggle credentials (for batch mode) or runtime connection values (for interactive mode) -- For Google Colab variant: `code-sandboxes[colab]` and a Colab runtime assignment +- For Google Colab variant: a Colab runtime assignment (server URL, kernel id, proxy token) - For Monty variant: `code-sandboxes[monty]` (no credentials required) - For Modal variant: `code-sandboxes[modal]` and Modal credentials diff --git a/docs/docs/sandboxes/google-colab.mdx b/docs/docs/sandboxes/google-colab.mdx new file mode 100644 index 0000000..b250f03 --- /dev/null +++ b/docs/docs/sandboxes/google-colab.mdx @@ -0,0 +1,92 @@ +--- +sidebar_position: 5 +title: Google Colab +--- + +# Google Colab + +Google Colab exposes a Jupyter-compatible kernel behind an authenticating proxy. +Use `ColabKernelClient` to connect to a kernel that is already running in Colab. + +This client reuses an existing Colab runtime; it does not create a Colab runtime +from scratch. + +You need three values to connect: + +- `server_url` +- `kernel_id` +- `proxy_token` + +All three are available in Colab's channels WebSocket URL. + +## Option A: Connect With Explicit Values + +```python +from code_sandboxes import ColabKernelClient + +kernel = ColabKernelClient( + server_url="https://", + kernel_id="", + proxy_token="", +) +kernel.start() +reply = kernel.execute("x = 1") +print(reply) +# Disconnect only; do not shut down shared Colab runtime kernels. +kernel.stop(shutdown_kernel=False) +``` + +## Option B: Connect From Channels URL + +```python +from code_sandboxes import ColabKernelClient + +channels_url = ( + "wss:///api/kernels//channels" + "?session_id=<...>&colab-runtime-proxy-token=&colab-client-agent=web" +) + +with ColabKernelClient.from_channels_url(channels_url) as kernel: + reply = kernel.execute("x = 1 + 1; print(x)") + print(reply) +``` + +You can also parse values directly: + +```python +from code_sandboxes import parse_colab_channels_url + +server_url, kernel_id, proxy_token = parse_colab_channels_url(channels_url) +``` + +`ColabKernelClient` forwards the proxy token as both the +`X-Colab-Runtime-Proxy-Token` HTTP header and the +`colab-runtime-proxy-token` WebSocket query parameter. + +## How To Obtain The Colab Channels URL + +The `server_url`, `kernel_id`, and `proxy_token` are in the WebSocket channels URL +used by Colab itself: + +```text +wss:///api/kernels//channels?session_id=<...>&colab-runtime-proxy-token=&colab-client-agent=web +``` + +Values are tied to your browser session and are short-lived. + +1. Open your notebook on https://colab.research.google.com and connect a runtime. +2. Open DevTools (`F12`) and switch to Network with the WS filter. +3. Run a cell to generate traffic. +4. Open the `.../api/kernels//channels?...` request and copy the full URL. + +If you extract values manually: + +- `server_url`: scheme+host before `/api/kernels` (use `https://`) +- `kernel_id`: UUID path segment after `/api/kernels/` +- `proxy_token`: `colab-runtime-proxy-token` query parameter + +Consumer Colab does not provide a public API key based flow to create runtimes +from standalone scripts. + +For a true programmatic runtime provisioning flow, use Colab Enterprise on +Google Cloud. diff --git a/docs/docs/sandboxes/index.mdx b/docs/docs/sandboxes/index.mdx index 77b39a4..c36789d 100644 --- a/docs/docs/sandboxes/index.mdx +++ b/docs/docs/sandboxes/index.mdx @@ -71,8 +71,8 @@ Each section below explains how to configure each variant. ### jupyter -Runs code against a local or remote Jupyter Server and connects using -`jupyter-kernel-client`. This variant provides process isolation via the Jupyter +Runs code against a local or remote Jupyter Server and connects through its +Jupyter kernel client. This variant provides process isolation via the Jupyter kernel and persistent state across requests. - **Requirements:** `jupyter_server` and `jupyter-kernel-client` (included by default). @@ -166,7 +166,7 @@ sandbox.run_code("print(now())") ### Docker Runs a Jupyter Server inside a Docker container for local isolated execution and -connects using `jupyter-kernel-client`. +connects through its Jupyter kernel client. - **Requirements:** `code-sandboxes[docker]` plus a running Docker Engine (verify with `docker version`). @@ -188,6 +188,9 @@ with Sandbox.create(variant="docker", image="code-sandboxes-jupyter:latest") as ### Kaggle +For the complete authentication, channels URL, batch execution, and accelerator +reference, see [Kaggle](./kaggle). + Runs code on Kaggle with two transparent modes: 1. **Interactive kernel mode** using `KaggleKernelClient` (runtime proxy URL). @@ -224,10 +227,13 @@ with Sandbox.create(variant="kaggle") as sandbox: ### Colab Runs code against a Google Colab runtime. Colab exposes a Jupyter-compatible -kernel behind an authenticating proxy, so this variant connects using -`jupyter-kernel-client`'s `ColabKernelClient` under the hood. +kernel behind an authenticating proxy, using the `ColabKernelClient` implemented +by Code Sandboxes. + +For the complete proxy authentication and channels URL guide, see +[Google Colab](./google-colab). -- **Requirements:** `code-sandboxes[colab]` (installs `jupyter-kernel-client`). +- **Requirements:** the base `code-sandboxes` installation. - **Parameters:** `server_url`, `kernel_id`, `proxy_token` (pass as keyword arguments or through the sandbox configuration). You can also pass `channels_url` and let the client parse the values. diff --git a/docs/docs/sandboxes/kaggle.mdx b/docs/docs/sandboxes/kaggle.mdx new file mode 100644 index 0000000..944b8c1 --- /dev/null +++ b/docs/docs/sandboxes/kaggle.mdx @@ -0,0 +1,170 @@ +--- +sidebar_position: 4 +title: Kaggle +--- + +# Kaggle + +Kaggle notebooks expose a Jupyter-compatible kernel behind an authenticating +proxy. `KaggleKernelClient` supports two auth modes: + +- API token mode (`token` arg or `KAGGLE_API_TOKEN` env var) +- Signed proxy URL mode (`token=None` for existing browser session URLs) + +When using signed proxy URLs, authentication is carried by the JWT in the +proxied `server_url` path. + +## Batch Execution + +For notebook jobs started from code, use `KaggleKernelExecutor`. + +Unlike consumer Colab, Kaggle provides an official public API to create and run +notebooks from code, without requiring an active browser session. + +Install optional dependencies: + +```bash +pip install "code-sandboxes[kaggle]" +``` + +Run code as a Kaggle batch job: + +```python +from code_sandboxes import KaggleKernelExecutor + +executor = KaggleKernelExecutor() + +result = executor.execute( + "print('hello from kaggle')", + title="code-sandboxes-demo", + accelerator="NvidiaTeslaT4", + enable_internet=True, + wait=True, + timeout=3600, + download_output=True, +) + +print(result.status) +print(result.succeeded) +print(result.url) +print(result.stdout) +print(result.stderr) +print(result.kernel_reply) +print(result.to_kernel_reply()) +print(result.output_files) +``` + +`to_kernel_reply()` returns a Jupyter-like shape compatible with +`JupyterKernelClient.execute(...)` responses: + +```python +{"execution_count": int, "outputs": [...], "status": "ok" | "error"} +``` + +`result.kernel_reply` exposes the same normalized payload directly. + +Batch submissions now include generated notebook cell IDs in the emitted notebook +JSON metadata, which avoids schema warnings in newer notebook tooling. + +Friendly accelerator aliases are supported, for example `T4`, `P100`, `A100`, and `H100`. + +For authentication in batch mode, use `~/.kaggle/kaggle.json` or +`KAGGLE_API_TOKEN` environment credentials. + +Useful `execute(...)` options: + +- `slug` and `title` for kernel identity +- `accelerator`, `enable_gpu`, `enable_internet`, `is_private` +- `dataset_sources`, `competition_sources`, `kernel_sources`, `model_sources` +- `wait=False` to submit now and poll later via `status(...)` and `output(...)` + +Each execution runs as a batch job and reaches terminal states like +`complete`, `error`, or `cancel_acknowledged`. + +You can also submit as a script with `kernel_type="script"` when needed. + +Common free-tier accelerators are typically `P100` and `T4`; higher tiers like +`A100`, `H100`, or `L4` may be limited to specific environments. + +Supported accelerator values include: + +- `NvidiaTeslaP100` +- `NvidiaTeslaT4` +- `NvidiaTeslaT4Highmem` +- `NvidiaL4` +- `NvidiaL4X1` +- `NvidiaTeslaA100` +- `NvidiaH100` +- `NvidiaRtxPro6000` + +## Interactive Kernel + +Use `KaggleKernelClient` for interactive execution. + +- Provide a Kaggle API token to create a new kernel. +- Or connect to an existing running session from a copied channels URL. + +Create a kernel with API token credentials: + +```python +import os +from code_sandboxes import KaggleKernelClient + +os.environ["KAGGLE_API_TOKEN"] = "..." + +with KaggleKernelClient( + server_url="https://kkb-production.jupyter-proxy.kaggle.net/k/12345678/eyJhbGci.../proxy", +) as kernel: + print("kernel_id:", kernel.id) + reply = kernel.execute("x = 1 + 1; print(x)") + print(reply) +``` + +```python +from code_sandboxes import KaggleKernelClient + +channels_url = ( + "wss://kkb-production.jupyter-proxy.kaggle.net/k/12345678/eyJhbGci.../proxy" + "/api/kernels/11e073f0-e82d-4029-be8d-3918f7ed1a9e/channels?session_id=..." +) + +with KaggleKernelClient.from_channels_url(channels_url, token=None) as kernel: + reply = kernel.execute("x = 1 + 1; print(x)") + print(reply) +``` + +You can also pass explicit values: + +```python +from code_sandboxes import KaggleKernelClient + +kernel = KaggleKernelClient( + server_url="https://kkb-production.jupyter-proxy.kaggle.net/k/12345678/eyJhbGci.../proxy", + kernel_id="11e073f0-e82d-4029-be8d-3918f7ed1a9e", +) +kernel.start() +reply = kernel.execute("x = 1") +print(reply) +kernel.stop(shutdown_kernel=False) +``` + +Parser helper: + +```python +from code_sandboxes import parse_kaggle_channels_url + +server_url, kernel_id = parse_kaggle_channels_url(channels_url) +``` + +### How To Obtain The Kaggle Channels URL + +The official Kaggle API (`kaggle` CLI / `kagglehub`) is primarily for batch +kernel operations (push, pull, status, output). Interactive kernel channels URL +values come from an active browser notebook session. + +1. Open your notebook on https://www.kaggle.com and run any cell. +2. Open DevTools (`F12`) and switch to Network with the WS filter. +3. Select the `.../proxy/api/kernels//channels?...` request. +4. Copy the full URL and pass it to `from_channels_url(...)`. + +Values are tied to your active browser session and rotate when sessions reconnect. diff --git a/pyproject.toml b/pyproject.toml index 9eff009..05966d8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,14 +35,13 @@ sandbox = "code_sandboxes.cli:main" [project.optional-dependencies] datalayer = ["agent_runtimes>=1.0.16"] docker = ["docker>=6.0"] -colab = ["jupyter-kernel-client"] -kaggle = ["jupyter-kernel-client"] +kaggle = ["kaggle>=1.6"] monty = ["pydantic-monty"] modal = ["modal>=0.64"] all = [ "agent_runtimes", "docker>=6.0", - "jupyter-kernel-client", + "kaggle>=1.6", "pydantic-monty", "modal>=0.64", ] diff --git a/tests/test_client.py b/tests/test_client.py index f98924b..65a07a8 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -11,14 +11,24 @@ import pytest -from code_sandboxes.client import CodeSandboxClient -from code_sandboxes.models import CodeError, OutputMessage, Result +from code_sandboxes.client import CodeSandboxClient, execution_result_to_reply +from code_sandboxes.models import ( + CodeError, + ExecutionResult, + Logs, + OutputMessage, + Result, + SandboxEnvironment, +) class _FakeSandbox: def __init__(self): self._started = False + self._variables = {} self.config = SimpleNamespace(variant="kaggle") + self.info = SimpleNamespace(metadata={"kernel_id": "execution-id"}) + self.sandbox_id = "sandbox-id" @property def is_started(self): @@ -30,6 +40,34 @@ def start(self): async def start_async(self): self.start() + def stop(self): + self._started = False + + def run_code(self, code: str, language: str = "python", timeout=None, envs=None): + _ = (code, language, timeout, envs) + return ExecutionResult( + execution_ok=True, + execution_count=2, + logs=Logs(stdout=[OutputMessage(line="hello")]), + results=[Result(data={"text/plain": "42"}, is_main_result=True)], + ) + + def get_variable(self, name): + return self._variables[name] + + def set_variable(self, name, value): + self._variables[name] = value + + def set_variables(self, variables): + self._variables.update(variables) + + def interrupt(self): + return True + + @classmethod + def list_environments(cls): + return [SandboxEnvironment(name="fake", title="Fake", language="python")] + def run_code_streaming(self, code: str, language: str = "python", timeout=None, envs=None): _ = (code, language, timeout, envs) yield OutputMessage(line="hello", timestamp=0.0, error=False) @@ -70,3 +108,60 @@ async def test_execute_code_streaming_async_proxies_sandbox_events(): assert isinstance(events[0], OutputMessage) assert isinstance(events[1], Result) assert isinstance(events[2], CodeError) + + +def test_execute_returns_variant_neutral_reply_and_metadata(): + client = CodeSandboxClient(_FakeSandbox()) + + reply = client.execute("print('hello')") + + assert client.id == "execution-id" + assert client.info is not None + assert client.config.variant == "kaggle" + assert client.kernel_info == {"language_info": {"name": "python"}} + assert reply == { + "execution_count": 2, + "outputs": [ + {"output_type": "stream", "name": "stdout", "text": "hello\n"}, + { + "output_type": "execute_result", + "data": {"text/plain": "42"}, + "metadata": {}, + }, + ], + "status": "ok", + } + + +def test_variables_interrupt_and_restart_delegate_to_sandbox(): + sandbox = _FakeSandbox() + client = CodeSandboxClient(sandbox) + + client.set_variable("one", 1) + client.set_variables({"two": 2}) + assert client.get_variable("one") == 1 + assert client.get_variable("two") == 2 + assert client.interrupt() is True + + client.restart() + assert client.is_alive() is True + + +def test_execution_error_converts_to_error_output(): + execution = ExecutionResult( + execution_ok=True, + execution_count=3, + code_error=CodeError(name="ValueError", value="boom", traceback="line 1\nline 2"), + ) + + reply = execution_result_to_reply(execution) + + assert reply["status"] == "error" + assert reply["outputs"] == [ + { + "output_type": "error", + "ename": "ValueError", + "evalue": "boom", + "traceback": ["line 1", "line 2"], + } + ] diff --git a/tests/test_colab.py b/tests/test_colab.py new file mode 100644 index 0000000..c21d51b --- /dev/null +++ b/tests/test_colab.py @@ -0,0 +1,126 @@ +# Copyright (c) 2025-2026 Datalayer, Inc. +# +# BSD 3-Clause License + +from __future__ import annotations + +import logging + +import pytest + +from code_sandboxes.colab import ( + COLAB_CLIENT_AGENT_HEADER, + COLAB_RUNTIME_PROXY_TOKEN_HEADER, + COLAB_RUNTIME_PROXY_TOKEN_PARAM, + ColabKernelClient, + parse_colab_channels_url, +) + +CHANNELS_URL = ( + "wss://abc123.prod.colab.dev/api/kernels/" + "11e073f0-e82d-4029-be8d-3918f7ed1a9e/channels" + "?session_id=96f4a03c-e4e0-4f15-8e9f-0cd33d3edecf" + "&colab-runtime-proxy-token=proxy-abc" + "&colab-client-agent=web" +) +SERVER_URL = "https://abc123.prod.colab.dev" +KERNEL_ID = "11e073f0-e82d-4029-be8d-3918f7ed1a9e" +PROXY_TOKEN = "proxy-abc" # noqa: S105 + + + +def test_colab_kernel_client_injects_headers_and_extra_params(monkeypatch): + captured: dict = {} + + def fake_kernel_client_init(self, *args, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr( + "code_sandboxes.colab.JupyterKernelClient.__init__", fake_kernel_client_init + ) + + ColabKernelClient( + server_url="https://colab-host.example", + kernel_id="kernel-123", + proxy_token="proxy-abc", # noqa: S106 + client_agent="custom-agent", + headers={"Existing": "value"}, + client_kwargs={"extra_params": {"existing": "p"}}, + ) + + assert captured["server_url"] == "https://colab-host.example" + assert captured["kernel_id"] == "kernel-123" + assert captured["token"] is None + + headers = captured["headers"] + assert headers["Existing"] == "value" + assert headers[COLAB_CLIENT_AGENT_HEADER] == "custom-agent" + assert headers[COLAB_RUNTIME_PROXY_TOKEN_HEADER] == "proxy-abc" + + client_kwargs = captured["client_kwargs"] + assert client_kwargs["extra_params"]["existing"] == "p" + assert client_kwargs["extra_params"][COLAB_RUNTIME_PROXY_TOKEN_PARAM] == "proxy-abc" + + +def test_colab_kernel_client_drops_any_provided_jupyter_token(monkeypatch): + captured: dict = {} + + def fake_kernel_client_init(self, *args, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr( + "code_sandboxes.colab.JupyterKernelClient.__init__", fake_kernel_client_init + ) + + ColabKernelClient( + server_url="https://colab-host.example", + kernel_id="kernel-123", + proxy_token="proxy-abc", # noqa: S106 + token="should-be-ignored", # noqa: S106 + log=logging.getLogger("test"), + ) + + assert captured["token"] is None + + +def test_parse_colab_channels_url_extracts_parts(): + server_url, kernel_id, proxy_token = parse_colab_channels_url(CHANNELS_URL) + assert server_url == SERVER_URL + assert kernel_id == KERNEL_ID + assert proxy_token == PROXY_TOKEN + + +def test_parse_colab_channels_url_maps_ws_to_http(): + server_url, _, _ = parse_colab_channels_url(CHANNELS_URL.replace("wss://", "ws://")) + assert server_url.startswith("http://") + + +def test_parse_colab_channels_url_requires_proxy_token(): + without_token = CHANNELS_URL.replace("&colab-runtime-proxy-token=proxy-abc", "") + with pytest.raises(ValueError): + parse_colab_channels_url(without_token) + + +def test_parse_colab_channels_url_rejects_invalid_url(): + with pytest.raises(ValueError): + parse_colab_channels_url("https://colab.research.google.com/not-a-channels-url") + + +def test_colab_kernel_client_from_channels_url(monkeypatch): + captured: dict = {} + + def fake_kernel_client_init(self, *args, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr( + "code_sandboxes.colab.JupyterKernelClient.__init__", fake_kernel_client_init + ) + + ColabKernelClient.from_channels_url(CHANNELS_URL) + + assert captured["server_url"] == SERVER_URL + assert captured["kernel_id"] == KERNEL_ID + headers = captured["headers"] + assert headers[COLAB_RUNTIME_PROXY_TOKEN_HEADER] == PROXY_TOKEN + client_kwargs = captured["client_kwargs"] + assert client_kwargs["extra_params"][COLAB_RUNTIME_PROXY_TOKEN_PARAM] == PROXY_TOKEN diff --git a/tests/test_jupyter.py b/tests/test_jupyter.py index 25ed0df..d8b9b6d 100644 --- a/tests/test_jupyter.py +++ b/tests/test_jupyter.py @@ -23,6 +23,7 @@ def test_explicit_kernel_id_wins_over_reuse(monkeypatch): class _KernelClientStub: def __init__(self, server_url, token, kernel_id, client_kwargs=None): + self.id = kernel_id or "started-kernel" captured["server_url"] = server_url captured["token"] = token captured["kernel_id"] = kernel_id @@ -67,6 +68,7 @@ def test_reuse_kernel_false_forces_new_kernel(monkeypatch): class _KernelClientStub: def __init__(self, server_url, token, kernel_id, client_kwargs=None): + self.id = kernel_id or "started-kernel" captured["kernel_id"] = kernel_id def start(self, path=None): @@ -108,6 +110,7 @@ def test_kernel_client_forwards_client_kwargs(monkeypatch, tmp_path: Path): class _KernelClientStub: def __init__(self, server_url, token, kernel_id, client_kwargs=None): + self.id = kernel_id or "started-kernel" captured["server_url"] = server_url captured["token"] = token captured["kernel_id"] = kernel_id @@ -177,6 +180,7 @@ def _kernel_client_stub(captured: dict): class _KernelClientStub: def __init__(self, server_url, token, kernel_id, client_kwargs=None, **kwargs): + self.id = kernel_id or "started-kernel" captured["server_url"] = server_url captured["token"] = token captured["kernel_id"] = kernel_id diff --git a/tests/test_kaggle.py b/tests/test_kaggle.py new file mode 100644 index 0000000..7a36cf8 --- /dev/null +++ b/tests/test_kaggle.py @@ -0,0 +1,127 @@ +# Copyright (c) 2025-2026 Datalayer, Inc. +# +# BSD 3-Clause License + +from __future__ import annotations + +import logging + +import pytest + +from code_sandboxes.kaggle import ( + KaggleKernelClient, + parse_kaggle_channels_url, +) + +CHANNELS_URL = ( + "wss://kkb-production.jupyter-proxy.kaggle.net/k/12345678/eyJhbGciMGIgwPdMJ" + "/proxy/api/kernels/11e073f0-e82d-4029-be8d-3918f7ed1a9e/channels" + "?session_id=96f4a03c-e4e0-4f15-8e9f-0cd33d3edecf" +) +SERVER_URL = "https://kkb-production.jupyter-proxy.kaggle.net/k/12345678/eyJhbGciMGIgwPdMJ/proxy" +KERNEL_ID = "11e073f0-e82d-4029-be8d-3918f7ed1a9e" + + +def test_parse_kaggle_channels_url_extracts_server_and_kernel(): + server_url, kernel_id = parse_kaggle_channels_url(CHANNELS_URL) + assert server_url == SERVER_URL + assert kernel_id == KERNEL_ID + + +def test_parse_kaggle_channels_url_maps_ws_to_http(): + server_url, _ = parse_kaggle_channels_url(CHANNELS_URL.replace("wss://", "ws://")) + assert server_url.startswith("http://") + + +def test_parse_kaggle_channels_url_rejects_invalid_url(): + with pytest.raises(ValueError): + parse_kaggle_channels_url("https://kaggle.com/not-a-channels-url") + + +def test_kaggle_kernel_client_uses_explicit_token(monkeypatch): + captured: dict = {} + + def fake_kernel_client_init(self, *args, **kwargs): + captured.update(kwargs) + + monkeypatch.delenv("KAGGLE_API_TOKEN", raising=False) + monkeypatch.setattr( + "code_sandboxes.kaggle.JupyterKernelClient.__init__", fake_kernel_client_init + ) + + KaggleKernelClient( + server_url=SERVER_URL, + kernel_id=KERNEL_ID, + token="explicit-token", # noqa: S106 + log=logging.getLogger("test"), + ) + + assert captured["server_url"] == SERVER_URL + assert captured["kernel_id"] == KERNEL_ID + assert captured["token"] == "explicit-token" # noqa: S105 + + +def test_kaggle_kernel_client_reads_token_from_env(monkeypatch): + captured: dict = {} + + def fake_kernel_client_init(self, *args, **kwargs): + captured.update(kwargs) + + monkeypatch.setenv("KAGGLE_API_TOKEN", "env-token") + monkeypatch.setattr( + "code_sandboxes.kaggle.JupyterKernelClient.__init__", fake_kernel_client_init + ) + + KaggleKernelClient(server_url=SERVER_URL) + + assert captured["token"] == "env-token" # noqa: S105 + + +def test_kaggle_kernel_client_token_none_without_env(monkeypatch): + captured: dict = {} + + def fake_kernel_client_init(self, *args, **kwargs): + captured.update(kwargs) + + monkeypatch.delenv("KAGGLE_API_TOKEN", raising=False) + monkeypatch.setattr( + "code_sandboxes.kaggle.JupyterKernelClient.__init__", fake_kernel_client_init + ) + + KaggleKernelClient(server_url=SERVER_URL) + + assert captured["token"] is None + + +def test_kaggle_kernel_client_from_channels_url(monkeypatch): + captured: dict = {} + + def fake_kernel_client_init(self, *args, **kwargs): + captured.update(kwargs) + + monkeypatch.delenv("KAGGLE_API_TOKEN", raising=False) + monkeypatch.setattr( + "code_sandboxes.kaggle.JupyterKernelClient.__init__", fake_kernel_client_init + ) + + KaggleKernelClient.from_channels_url(CHANNELS_URL) + + assert captured["server_url"] == SERVER_URL + assert captured["kernel_id"] == KERNEL_ID + + +def test_kaggle_kernel_client_allows_missing_kernel_id_for_new_kernel(monkeypatch): + captured: dict = {} + + def fake_kernel_client_init(self, *args, **kwargs): + captured.update(kwargs) + + monkeypatch.setenv("KAGGLE_API_TOKEN", "env-token") + monkeypatch.setattr( + "code_sandboxes.kaggle.JupyterKernelClient.__init__", fake_kernel_client_init + ) + + KaggleKernelClient(server_url=SERVER_URL) + + assert captured["kernel_id"] is None + assert captured["token"] == "env-token" # noqa: S105 diff --git a/tests/test_kaggle_execute.py b/tests/test_kaggle_execute.py new file mode 100644 index 0000000..9d6e700 --- /dev/null +++ b/tests/test_kaggle_execute.py @@ -0,0 +1,268 @@ +# Copyright (c) 2025-2026 Datalayer, Inc. +# +# BSD 3-Clause License + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from code_sandboxes.kaggle_execute import ( + KaggleExecutionResult, + KaggleKernelExecutor, + _normalize_accelerator, + _normalize_status, + _slugify, +) + + +class _FakeStatus: + def __init__(self, status, failure_message=None): + self.status = status + self.failure_message = failure_message + + +class _FakePushResponse: + url = "https://www.kaggle.com/code/me/my-notebook" + version_number = 3 + + +class _FakeApi: + """Minimal stand-in for kaggle.api.kaggle_api_extended.KaggleApi.""" + + def __init__(self, statuses, *, username="me"): + self._statuses = list(statuses) + self._username = username + self.pushed_metadata = None + self.pushed_code = None + self.pushed_accelerator = None + + def get_config_value(self, name): + return self._username if name == "username" else None + + def kernels_push(self, folder, accelerator=None): + folder_path = Path(folder) + self.pushed_metadata = json.loads( + (folder_path / "kernel-metadata.json").read_text(encoding="utf-8") + ) + code_file = self.pushed_metadata["code_file"] + self.pushed_code = (folder_path / code_file).read_text(encoding="utf-8") + self.pushed_accelerator = accelerator + return _FakePushResponse() + + def kernels_status(self, slug): + # Advance through the queued statuses, holding on the last one. + status = self._statuses.pop(0) if len(self._statuses) > 1 else self._statuses[0] + return _FakeStatus(status) + + def kernels_output(self, kernel, path, force=False, quiet=True): + log_path = Path(path) / "run.log" + log_path.write_text("execution log contents", encoding="utf-8") + return ([str(log_path)], None) + + +def test_slugify_normalizes_title(): + assert _slugify("My Python Notebook!") == "my-python-notebook" + assert _slugify(" ").startswith("code-sandboxes-run-") + + +def test_normalize_status_handles_enum_like(): + class _Enum: + name = "COMPLETE" + + assert _normalize_status(_Enum()) == "COMPLETE" + assert _normalize_status("KernelWorkerStatus.ERROR") == "ERROR" + assert _normalize_status("running") == "RUNNING" + + +def test_normalize_accelerator_supports_aliases(): + assert _normalize_accelerator("T4") == "NvidiaTeslaT4" + assert _normalize_accelerator("tesla p100") == "NvidiaTeslaP100" + assert _normalize_accelerator("NvidiaH100") == "NvidiaH100" + + +def test_normalize_accelerator_rejects_unknown_value(): + with pytest.raises(ValueError): + _normalize_accelerator("SomeFutureGpu") + + +def test_execute_success_downloads_log(): + api = _FakeApi(["RUNNING", "COMPLETE"]) + executor = KaggleKernelExecutor(api=api) + + result = executor.execute( + "print('hi')", + title="My Python Notebook", + poll_interval=0, + ) + + assert result.slug == "me/my-python-notebook" + assert result.status == "COMPLETE" + assert result.succeeded is True + assert result.url == "https://www.kaggle.com/code/me/my-notebook" + assert result.version_number == 3 + assert result.log == "execution log contents" + + # Metadata was written with the expected id and a notebook code file. + assert api.pushed_metadata["id"] == "me/my-python-notebook" + assert api.pushed_metadata["code_file"] == "notebook.ipynb" + assert "print('hi')" in api.pushed_code + + notebook = json.loads(api.pushed_code) + assert notebook["cells"][0]["id"] + assert notebook["cells"][0]["metadata"]["language"] == "python" + + +def test_execute_script_kernel_writes_python_file(): + api = _FakeApi(["COMPLETE"]) + executor = KaggleKernelExecutor(api=api) + + executor.execute( + "print('hi')", + slug="my-script", + kernel_type="script", + poll_interval=0, + download_output=False, + ) + + assert api.pushed_metadata["code_file"] == "script.py" + assert api.pushed_metadata["kernel_type"] == "script" + assert api.pushed_code == "print('hi')" + + +def test_execute_forwards_accelerator_and_enables_gpu_metadata(): + api = _FakeApi(["COMPLETE"]) + executor = KaggleKernelExecutor(api=api) + + executor.execute( + "print('hi')", + slug="my-notebook", + accelerator="T4", + poll_interval=0, + download_output=False, + ) + + assert api.pushed_accelerator == "NvidiaTeslaT4" + assert api.pushed_metadata["accelerator"] == "NvidiaTeslaT4" + assert api.pushed_metadata["enable_gpu"] is True + + +def test_execute_without_wait_returns_current_status(): + api = _FakeApi(["QUEUED"]) + executor = KaggleKernelExecutor(api=api) + + result = executor.execute( + "print('hi')", + slug="my-notebook", + wait=False, + ) + + assert result.slug == "me/my-notebook" + assert result.status == "QUEUED" + + +def test_execute_times_out(): + api = _FakeApi(["RUNNING"]) + executor = KaggleKernelExecutor(api=api) + + result = executor.execute( + "print('hi')", + slug="stuck", + timeout=0, + poll_interval=0, + download_output=False, + ) + + assert result.status == "RUNNING" + + +def test_missing_username_raises(): + api = _FakeApi(["COMPLETE"], username=None) + executor = KaggleKernelExecutor(api=api) + + with pytest.raises(ValueError): + executor.execute("print('hi')", poll_interval=0) + + +def test_explicit_username_is_used(): + api = _FakeApi(["COMPLETE"], username="ignored") + executor = KaggleKernelExecutor(username="explicit", api=api) + + result = executor.execute( + "print('hi')", + slug="nb", + poll_interval=0, + download_output=False, + ) + + assert result.slug == "explicit/nb" + + +def test_to_kernel_reply_uses_notebook_outputs_when_available(): + result = { + "cells": [ + { + "cell_type": "code", + "execution_count": 3, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": "hello from kaggle\n", + } + ], + } + ] + } + + execution = KaggleExecutionResult( + slug="me/demo", + status="COMPLETE", + notebook=result, + log='[{"stream_name":"stderr","data":"infra warning\\n"}]', + ) + + reply = execution.to_kernel_reply() + assert reply["status"] == "ok" + assert reply["execution_count"] == 3 + assert reply["outputs"][0]["text"] == "hello from kaggle\n" + + +def test_to_kernel_reply_falls_back_to_log_streams(): + execution = KaggleExecutionResult( + slug="me/demo", + status="COMPLETE", + log=( + '[' + '{"stream_name":"stdout","data":"hello\\n"},' + '{"stream_name":"stderr","data":"warn\\n"}' + ']' + ), + ) + + reply = execution.to_kernel_reply() + assert reply["status"] == "ok" + assert reply["execution_count"] == 0 + assert [output["name"] for output in reply["outputs"]] == ["stdout", "stderr"] + + +def test_stdout_stderr_and_repr_are_compact(): + execution = KaggleExecutionResult( + slug="me/demo", + status="COMPLETE", + log=( + '[' + '{"stream_name":"stdout","data":"hello\\n"},' + '{"stream_name":"stderr","data":"warning\\n"}' + ']' + ), + ) + + assert execution.stdout == "hello\n" + assert execution.stderr == "warning\n" + rendered = repr(execution) + assert "kernel_status='ok'" in rendered + assert "stdout='hello'" in rendered + assert "stderr='warning'" in rendered diff --git a/tests/test_modal_colab_sandbox.py b/tests/test_modal_colab_sandbox.py index 1dd02b3..06e9ee6 100644 --- a/tests/test_modal_colab_sandbox.py +++ b/tests/test_modal_colab_sandbox.py @@ -147,10 +147,9 @@ def execute(self, code, wait=True, timeout=0.0, download_output=True, accelerato succeeded=True, ) - monkeypatch.setitem( - sys.modules, - "jupyter_kernel_client", - SimpleNamespace(KaggleKernelExecutor=_FakeKaggleExecutor), + monkeypatch.setattr( + "code_sandboxes.kaggle_sandbox.KaggleKernelExecutor", + _FakeKaggleExecutor, ) sandbox = KaggleSandbox(config=SandboxConfig(timeout=10.0)) @@ -190,10 +189,9 @@ def execute(self, code, wait=True, timeout=0.0, download_output=True, accelerato succeeded=False, ) - monkeypatch.setitem( - sys.modules, - "jupyter_kernel_client", - SimpleNamespace(KaggleKernelExecutor=_FakeKaggleExecutor), + monkeypatch.setattr( + "code_sandboxes.kaggle_sandbox.KaggleKernelExecutor", + _FakeKaggleExecutor, ) sandbox = KaggleSandbox(config=SandboxConfig(timeout=10.0)) @@ -234,10 +232,9 @@ def execute(self, code, wait=True, timeout=0.0, download_output=True, accelerato succeeded=True, ) - monkeypatch.setitem( - sys.modules, - "jupyter_kernel_client", - SimpleNamespace(KaggleKernelExecutor=_FakeKaggleExecutor), + monkeypatch.setattr( + "code_sandboxes.kaggle_sandbox.KaggleKernelExecutor", + _FakeKaggleExecutor, ) sandbox = KaggleSandbox(config=SandboxConfig(timeout=10.0, gpu="T4")) @@ -286,10 +283,9 @@ def __init__(self, username=None, quiet=True): def execute(self, code, wait=True, timeout=0.0, download_output=True, accelerator=None): return _FakeKaggleResult() - monkeypatch.setitem( - sys.modules, - "jupyter_kernel_client", - SimpleNamespace(KaggleKernelExecutor=_FakeKaggleExecutor), + monkeypatch.setattr( + "code_sandboxes.kaggle_sandbox.KaggleKernelExecutor", + _FakeKaggleExecutor, ) sandbox = KaggleSandbox(config=SandboxConfig(timeout=10.0)) @@ -358,10 +354,9 @@ def output(self, slug, dest, force=True, quiet=None): path.write_text("[]", encoding="utf-8") return [str(path)] - monkeypatch.setitem( - sys.modules, - "jupyter_kernel_client", - SimpleNamespace(KaggleKernelExecutor=_FakeKaggleExecutor), + monkeypatch.setattr( + "code_sandboxes.kaggle_sandbox.KaggleKernelExecutor", + _FakeKaggleExecutor, ) sandbox = KaggleSandbox(config=SandboxConfig(timeout=10.0), poll_interval=0.0) From 5f7eee2dd950d61ff0ff510065910b5ac28d80bc Mon Sep 17 00:00:00 2001 From: Eric Charles Date: Sat, 8 Aug 2026 17:47:34 +0200 Subject: [PATCH 2/3] docs --- code_sandboxes/client.py | 5 + code_sandboxes/kaggle_execute.py | 4 +- docs/docs/sandboxes/datalayer.mdx | 63 +++++ docs/docs/sandboxes/docker.mdx | 39 ++++ docs/docs/sandboxes/eval.mdx | 29 +++ docs/docs/sandboxes/google-colab.mdx | 57 ++++- docs/docs/sandboxes/index.mdx | 330 +-------------------------- docs/docs/sandboxes/jupyter.mdx | 56 +++++ docs/docs/sandboxes/kaggle.mdx | 40 +++- docs/docs/sandboxes/modal.mdx | 69 ++++++ docs/docs/sandboxes/monty.mdx | 48 ++++ pyproject.toml | 2 +- tests/test_client.py | 10 + tests/test_colab.py | 1 - tests/test_kaggle_execute.py | 10 +- 15 files changed, 428 insertions(+), 335 deletions(-) create mode 100644 docs/docs/sandboxes/datalayer.mdx create mode 100644 docs/docs/sandboxes/docker.mdx create mode 100644 docs/docs/sandboxes/eval.mdx create mode 100644 docs/docs/sandboxes/jupyter.mdx create mode 100644 docs/docs/sandboxes/modal.mdx create mode 100644 docs/docs/sandboxes/monty.mdx diff --git a/code_sandboxes/client.py b/code_sandboxes/client.py index 514d355..a104bd2 100644 --- a/code_sandboxes/client.py +++ b/code_sandboxes/client.py @@ -391,6 +391,11 @@ def set_variables(self, variables: dict[str, Any]) -> None: self.start() self._sandbox.set_variables(variables) + def register_tool_caller(self, caller: Callable[..., Any]) -> None: + """Register the callable used by generated tools inside the sandbox.""" + self.start() + self._sandbox.register_tool_caller(caller) + def interrupt(self) -> bool: """Interrupt the active execution when supported by the variant.""" return self._sandbox.interrupt() diff --git a/code_sandboxes/kaggle_execute.py b/code_sandboxes/kaggle_execute.py index bb2ae7c..9f17759 100644 --- a/code_sandboxes/kaggle_execute.py +++ b/code_sandboxes/kaggle_execute.py @@ -574,9 +574,7 @@ def _download_output( self.log.warning("Could not download Kaggle output for %s: %s", slug, exc) @staticmethod - def _populate_log_and_notebook( - result: KaggleExecutionResult, files: t.Sequence[str] - ) -> None: + def _populate_log_and_notebook(result: KaggleExecutionResult, files: t.Sequence[str]) -> None: for path_str in files: path = Path(path_str) if path.suffix == ".log" and result.log is None: diff --git a/docs/docs/sandboxes/datalayer.mdx b/docs/docs/sandboxes/datalayer.mdx new file mode 100644 index 0000000..089c0d9 --- /dev/null +++ b/docs/docs/sandboxes/datalayer.mdx @@ -0,0 +1,63 @@ +--- +sidebar_position: 9 +title: Datalayer +--- + +# Datalayer + +Cloud-based execution with full isolation, GPU support, snapshots, and persistence. + +- **Requirements:** `code-sandboxes[datalayer]` (installs `agent_runtimes`). +- **Parameters:** `token` (defaults to `DATALAYER_API_KEY`), `run_url`, + `snapshot_name`, plus creation options like `environment`, `gpu`, `cpu`, `memory`. + +## How To Obtain Datalayer Credentials + +1. Create an account at [datalayer.ai](https://datalayer.ai). +2. Generate an API token from your account settings (**IAM → Tokens / API Keys**). +3. Export `DATALAYER_API_KEY` (or pass it as the `token` parameter). Set + `DATALAYER_RUN_URL` only for a self-hosted / custom deployment. + +## Usage + +```python +import os +from code_sandboxes import Sandbox + +os.environ["DATALAYER_API_KEY"] = "your-datalayer-token" + +with Sandbox.create( + variant="datalayer", + gpu="A100", + environment="python-gpu-env", +) as sandbox: + sandbox.run_code("import torch; print(torch.cuda.is_available())") +``` + +The concrete implementation is available from a top-level module: + +```python +from code_sandboxes.datalayer_sandbox import DatalayerSandbox +``` + +## Snapshots + +Save and restore sandbox state (datalayer only): + +```python +with Sandbox.create(variant="datalayer") as sandbox: + # Set up environment + sandbox.run_code("import pandas as pd") + sandbox.run_code("df = pd.DataFrame({'a': [1,2,3]})") + + # Create snapshot + snapshot = sandbox.create_snapshot("my-setup") + print(f"Snapshot: {snapshot.id}") + +# Later: restore from snapshot +with Sandbox.create( + variant="datalayer", + snapshot_name="my-setup" +) as sandbox: + result = sandbox.run_code("print(df)") # State restored +``` diff --git a/docs/docs/sandboxes/docker.mdx b/docs/docs/sandboxes/docker.mdx new file mode 100644 index 0000000..e5a53c2 --- /dev/null +++ b/docs/docs/sandboxes/docker.mdx @@ -0,0 +1,39 @@ +--- +sidebar_position: 5 +title: Docker +--- + +# Docker + +Runs a Jupyter Server inside a Docker container for local isolated execution and +connects through its Jupyter kernel client. + +- **Requirements:** `code-sandboxes[docker]` plus a running Docker Engine + (verify with `docker version`). +- **Credentials:** none — the kernel `token` is generated automatically. +- **Parameters:** `image` (default `code-sandboxes-jupyter:latest`), + `container_name`, `host`/`container_port`, `auto_remove`, `workdir`. + +## Building The Image + +Build the default image used by `DockerSandbox`: + +```bash +docker build -t code-sandboxes-jupyter:latest -f docker/Dockerfile . +``` + +## Usage + +```python +from code_sandboxes import Sandbox + +with Sandbox.create(variant="docker", image="code-sandboxes-jupyter:latest") as sandbox: + result = sandbox.run_code("import sys; print(sys.version)") + print(result.stdout) +``` + +The concrete implementation is available from a top-level module: + +```python +from code_sandboxes.docker_sandbox import DockerSandbox +``` diff --git a/docs/docs/sandboxes/eval.mdx b/docs/docs/sandboxes/eval.mdx new file mode 100644 index 0000000..5d2189e --- /dev/null +++ b/docs/docs/sandboxes/eval.mdx @@ -0,0 +1,29 @@ +--- +sidebar_position: 3 +title: Eval +--- + +# Eval + +Uses Python's `exec()` for code execution. No isolation, but fast and simple for +development. + +- **Credentials / parameters:** none — nothing to configure. +- ⚠️ `eval` shares memory with the host process and provides no sandboxing. Never + run untrusted code with it. + +## Usage + +```python +from code_sandboxes import Sandbox + +with Sandbox.create(variant="eval") as sandbox: + result = sandbox.run_code("x = 1 + 1") + result = sandbox.run_code("print(x)") # prints 2 +``` + +The concrete implementation is available from a top-level module: + +```python +from code_sandboxes.eval_sandbox import EvalSandbox +``` diff --git a/docs/docs/sandboxes/google-colab.mdx b/docs/docs/sandboxes/google-colab.mdx index b250f03..3caf0fc 100644 --- a/docs/docs/sandboxes/google-colab.mdx +++ b/docs/docs/sandboxes/google-colab.mdx @@ -1,11 +1,55 @@ --- -sidebar_position: 5 +sidebar_position: 7 title: Google Colab --- # Google Colab -Google Colab exposes a Jupyter-compatible kernel behind an authenticating proxy. +Runs code against a Google Colab runtime. Colab exposes a Jupyter-compatible +kernel behind an authenticating proxy, using the `ColabKernelClient` implemented +by Code Sandboxes. + +- **Requirements:** the base `code-sandboxes` installation. +- **Parameters:** `server_url`, `kernel_id`, `proxy_token` (pass as keyword + arguments or through the sandbox configuration). You can also pass + `channels_url` and let the client parse the values. + +Consumer Colab does not expose an official third-party API to provision runtimes +from scratch. Start/connect a runtime in the Colab UI first, then reuse it here. +The values are tied to your Colab session and are short-lived — refresh them +after the runtime is reassigned or reconnected. + +## Sandbox Usage + +```python +from code_sandboxes import Sandbox + +with Sandbox.create( + variant="colab", + server_url="https://8080-m-s-kkb-...-d.us-east1-0.prod.colab.dev", + kernel_id="c9bba548-3995-4f26-8e1a-7b8fbb10c578", + proxy_token="eyJhbGci....", +) as sandbox: + sandbox.run_code("x = 40") + result = sandbox.run_code("x + 2") + print(result.text) # 42 +``` + +Or pass a channels URL directly: + +```python +with Sandbox.create( + variant="colab", + channels_url=( + "wss:///api/kernels//channels" + "?session_id=<...>&colab-runtime-proxy-token=&colab-client-agent=web" + ), +) as sandbox: + print(sandbox.run_code("print(1 + 1)").text) +``` + +## Kernel Client + Use `ColabKernelClient` to connect to a kernel that is already running in Colab. This client reuses an existing Colab runtime; it does not create a Colab runtime @@ -81,9 +125,14 @@ Values are tied to your browser session and are short-lived. If you extract values manually: -- `server_url`: scheme+host before `/api/kernels` (use `https://`) +- `server_url`: scheme+host before `/api/kernels` (use `https://`). Colab assigns + a per-session host such as + `https://8080-m-s-kkb-...-d.us-east1-0.prod.colab.dev`; there is usually **no** + `/tun/m/...` path segment. - `kernel_id`: UUID path segment after `/api/kernels/` -- `proxy_token`: `colab-runtime-proxy-token` query parameter +- `proxy_token`: `colab-runtime-proxy-token` query parameter (same value as the + `X-Colab-Runtime-Proxy-Token` request header). Ignore the `session_id` and + `colab-client-agent` parameters. Consumer Colab does not provide a public API key based flow to create runtimes from standalone scripts. diff --git a/docs/docs/sandboxes/index.mdx b/docs/docs/sandboxes/index.mdx index c36789d..b8e9352 100644 --- a/docs/docs/sandboxes/index.mdx +++ b/docs/docs/sandboxes/index.mdx @@ -56,328 +56,18 @@ from code_sandboxes.eval_sandbox import EvalSandbox from code_sandboxes.datalayer_sandbox import DatalayerSandbox ``` -Each section below explains how to configure each variant. +Each page below explains how to configure each variant. | Variant | Summary | |---------|---------| -| `jupyter` | Jupyter kernel-backed execution with persistent state | -| `eval` | In-process `exec()` for fast development-only runs | -| `monty` | Secure in-process Python subset via Monty | -| `docker` | Jupyter execution in a Docker container | -| `kaggle` | Kaggle runtime (interactive or batch) | -| `colab` | Google Colab runtime via runtime proxy | -| `modal` | Modal container execution | -| `datalayer` | Datalayer managed runtime with optional GPU | - -### jupyter - -Runs code against a local or remote Jupyter Server and connects through its -Jupyter kernel client. This variant provides process isolation via the Jupyter -kernel and persistent state across requests. - -- **Requirements:** `jupyter_server` and `jupyter-kernel-client` (included by default). -- **Parameters:** `server_url`, `token`, `host`/`port` (when the sandbox starts - its own server), `python_executable`. - -**How to obtain the `token`:** - -- If you start the server yourself, you choose the token: - ```bash - jupyter server --port 8888 --IdentityProvider.token MY_TOKEN - ``` -- For an already-running server, list servers and read the `token=...` value: - ```bash - jupyter server list - # http://localhost:8888/?token=abcd1234... :: /home/you/notebooks - ``` - You can also pass the full `http://host:port/?token=...` URL as `server_url`; - the token is parsed automatically. -- If you omit `server_url`, `JupyterSandbox` **starts and manages its own local - Jupyter Server** and generates the token for you — no configuration needed. - -```python -# Connect to an existing server: -with Sandbox.create( - variant="jupyter", - server_url="http://localhost:8888", - token="MY_TOKEN", -) as sandbox: - sandbox.run_code("x = 40") - result = sandbox.run_code("x + 2") - print(result.text) # 42 - -# Or let the sandbox manage a local server automatically: -with Sandbox.create(variant="jupyter") as sandbox: - print(sandbox.run_code("1 + 1").text) # 2 -``` - -### eval - -Uses Python's `exec()` for code execution. No isolation, but fast and simple for development. - -- **Credentials / parameters:** none — nothing to configure. -- ⚠️ `eval` shares memory with the host process and provides no sandboxing. Never - run untrusted code with it. - -```python -with Sandbox.create(variant="eval") as sandbox: - result = sandbox.run_code("x = 1 + 1") - result = sandbox.run_code("print(x)") # prints 2 -``` - -### monty - -Runs code in [Monty](https://github.com/pydantic/monty), a minimal, secure Python -interpreter written in Rust (`pydantic-monty`). Monty runs a restricted subset of -Python in-process with microsecond startup and no access to the host filesystem, -environment, or network unless explicitly granted. It is ideal for short, -LLM-generated snippets where a full container or kernel would be overkill. - -Session state (variables, imports, definitions) persists across `run_code` calls. -Note that Monty supports only a subset of Python — third-party libraries and rich -display outputs are not available. - -- **Requirements:** `code-sandboxes[monty]` (installs `pydantic-monty`). -- **Credentials:** none required (fully local, in-process). -- **Optional parameters** (on `MontySandbox`): `type_check` (type-check before - running), `type_check_stubs`, `external_functions` (`{name: callable}` host - functions the code may call), `limits` (memory / stack / time limits). - -```python -with Sandbox.create(variant="monty") as sandbox: - sandbox.run_code("x = 21") - result = sandbox.run_code("x * 2") - print(result.text) # 42 -``` - -You can expose host callables to the sandboxed code and enable type checking: - -```python -from code_sandboxes.monty_sandbox import MontySandbox - -sandbox = MontySandbox( - type_check=True, - external_functions={"now": lambda: "2026-01-01"}, -) -sandbox.start() -sandbox.run_code("print(now())") -``` - -### Docker - -Runs a Jupyter Server inside a Docker container for local isolated execution and -connects through its Jupyter kernel client. - -- **Requirements:** `code-sandboxes[docker]` plus a running Docker Engine - (verify with `docker version`). -- **Credentials:** none — the kernel `token` is generated automatically. -- **Parameters:** `image` (default `code-sandboxes-jupyter:latest`), - `container_name`, `host`/`container_port`, `auto_remove`, `workdir`. - -Build the default image used by `DockerSandbox`: - -```bash -docker build -t code-sandboxes-jupyter:latest -f docker/Dockerfile . -``` - -```python -with Sandbox.create(variant="docker", image="code-sandboxes-jupyter:latest") as sandbox: - result = sandbox.run_code("import sys; print(sys.version)") - print(result.stdout) -``` - -### Kaggle - -For the complete authentication, channels URL, batch execution, and accelerator -reference, see [Kaggle](./kaggle). - -Runs code on Kaggle with two transparent modes: - -1. **Interactive kernel mode** using `KaggleKernelClient` (runtime proxy URL). -2. **Batch job mode** using `KaggleKernelExecutor` when no runtime connection - details are supplied. - -- **Requirements:** `code-sandboxes[kaggle]`. -- **Interactive parameters:** `server_url`, optional `kernel_id`, optional - `channels_url`, optional `token`. -- **Batch credentials:** `~/.kaggle/kaggle.json` or `KAGGLE_API_KEY`. - -```python -from code_sandboxes import Sandbox - -# Interactive mode (existing session / runtime proxy) -with Sandbox.create( - variant="kaggle", - channels_url="wss://kkb-production.jupyter-proxy.kaggle.net/k/.../proxy/api/kernels/.../channels?session_id=...", -) as sandbox: - print(sandbox.run_code("print(1 + 1)").text) - -# Transparent batch mode (no runtime URL required) -with Sandbox.create(variant="kaggle") as sandbox: - result = sandbox.run_code("print('hello from kaggle batch')") - print(result.text or result.stdout) - -# Stream status updates and outputs while the Kaggle job runs. -with Sandbox.create(variant="kaggle") as sandbox: - for event in sandbox.run_code_streaming("print('hello from kaggle stream')"): - if hasattr(event, "line"): - print(event.line) -``` - -### Colab - -Runs code against a Google Colab runtime. Colab exposes a Jupyter-compatible -kernel behind an authenticating proxy, using the `ColabKernelClient` implemented -by Code Sandboxes. - -For the complete proxy authentication and channels URL guide, see -[Google Colab](./google-colab). - -- **Requirements:** the base `code-sandboxes` installation. -- **Parameters:** `server_url`, `kernel_id`, `proxy_token` (pass as keyword - arguments or through the sandbox configuration). You can also pass - `channels_url` and let the client parse the values. - -**How to obtain these values** — they are the pieces of the WebSocket URL Colab's -own frontend uses to reach your runtime: - -``` -wss:///api/kernels//channels?session_id=<...>&colab-runtime-proxy-token=&colab-client-agent=web -``` - -Read them from your browser's developer tools while a Colab runtime is connected: - -1. Open your notebook on [colab.research.google.com](https://colab.research.google.com) - and **connect to a runtime** (*Runtime → Connect*, or run any cell). -2. Open DevTools (`F12`) → **Network** tab → **WS** filter, then run a cell to - trigger kernel traffic. -3. Click the `.../api/kernels//channels?...` request and read off: - - **`server_url`** — scheme + host *before* `/api/kernels` (change `wss://` to - `https://`). Colab assigns a per-session host such as - `https://8080-m-s-kkb-...-d.us-east1-0.prod.colab.dev`; there is usually **no** - `/tun/m/...` path segment. - - **`kernel_id`** — the UUID right after `/api/kernels/`. - - **`proxy_token`** — the `colab-runtime-proxy-token` query parameter (same - value as the `X-Colab-Runtime-Proxy-Token` request header). Ignore the - `session_id` and `colab-client-agent` parameters. - -Consumer Colab does not expose an official third-party API to provision runtimes -from scratch. Start/connect a runtime in the Colab UI first, then reuse it here. -The values are tied to your Colab session and are short-lived — refresh them -after the runtime is reassigned or reconnected. - -```python -with Sandbox.create( - variant="colab", - server_url="https://8080-m-s-kkb-...-d.us-east1-0.prod.colab.dev", - kernel_id="c9bba548-3995-4f26-8e1a-7b8fbb10c578", - proxy_token="eyJhbGci....", -) as sandbox: - sandbox.run_code("x = 40") - result = sandbox.run_code("x + 2") - print(result.text) # 42 -``` - -Or pass a channels URL directly: - -```python -with Sandbox.create( - variant="colab", - channels_url=( - "wss:///api/kernels//channels" - "?session_id=<...>&colab-runtime-proxy-token=&colab-client-agent=web" - ), -) as sandbox: - print(sandbox.run_code("print(1 + 1)").text) -``` - -### Modal - -Runs code in a [Modal](https://modal.com/docs/guide) cloud sandbox, providing -fully isolated, on-demand containers with configurable images and secrets. - -Each `run_code` call executes in a fresh `python -c` process, so state does **not** -persist across calls (use a single multi-statement snippet if you need shared -state). Configure the image with additional pip packages as needed. - -- **Requirements:** `code-sandboxes[modal]` (installs `modal`). -- **Parameters:** `app_name`, `image` (a prebuilt `modal.Image`), `pip_packages`, - `python_executable`. - -**How to obtain Modal credentials:** - -1. Create a free account at [modal.com](https://modal.com). -2. Authenticate the CLI (opens a browser, writes `~/.modal.toml`): - ```bash - modal token new - ``` - This is enough for local use: the Modal SDK reads credentials from - `~/.modal.toml` automatically. -3. Alternatively, create a token in the Modal dashboard (**Settings → API Tokens**) - and export `MODAL_TOKEN_ID` and `MODAL_TOKEN_SECRET`. - -**Do you need both `MODAL_TOKEN_ID` and `MODAL_TOKEN_SECRET`?** - -- For environment-based auth (CI/CD, containers, hosted runners): **yes**, - you need both values because Modal authenticates with a token pair - (public id + secret). -- For local development with `modal token new`: **not necessarily**. The SDK can - authenticate directly from `~/.modal.toml`. - -If you need environment variables, you can export them from your local config: - -```bash -python - <<'PY' -import pathlib -import tomllib - -cfg = tomllib.loads(pathlib.Path("~/.modal.toml").expanduser().read_text()) -profile = cfg.get("default", cfg) -token_id = profile.get("token_id") -token_secret = profile.get("token_secret") -if token_id and token_secret: - print(f"export MODAL_TOKEN_ID={token_id}") - print(f"export MODAL_TOKEN_SECRET={token_secret}") -else: - raise SystemExit("Could not find token_id/token_secret in ~/.modal.toml") -PY -``` - -```python -with Sandbox.create( - variant="modal", - pip_packages=["numpy"], -) as sandbox: - result = sandbox.run_code("import numpy as np; print(np.arange(3).sum())") - print(result.stdout) # "3" -``` - -### datalayer - -Cloud-based execution with full isolation, GPU support, snapshots, and persistence. - -- **Requirements:** `code-sandboxes[datalayer]` (installs `agent_runtimes`). -- **Parameters:** `token` (defaults to `DATALAYER_API_KEY`), `run_url`, - `snapshot_name`, plus creation options like `environment`, `gpu`, `cpu`, `memory`. - -**How to obtain Datalayer credentials:** - -1. Create an account at [datalayer.ai](https://datalayer.ai). -2. Generate an API token from your account settings (**IAM → Tokens / API Keys**). -3. Export `DATALAYER_API_KEY` (or pass it as the `token` parameter). Set - `DATALAYER_RUN_URL` only for a self-hosted / custom deployment. - -```python -import os -os.environ["DATALAYER_API_KEY"] = "your-datalayer-token" - -with Sandbox.create( - variant="datalayer", - gpu="A100", - environment="python-gpu-env", -) as sandbox: - sandbox.run_code("import torch; print(torch.cuda.is_available())") -``` +| [`jupyter`](./jupyter) | Jupyter kernel-backed execution with persistent state | +| [`eval`](./eval) | In-process `exec()` for fast development-only runs | +| [`monty`](./monty) | Secure in-process Python subset via Monty | +| [`docker`](./docker) | Jupyter execution in a Docker container | +| [`kaggle`](./kaggle) | Kaggle runtime (interactive or batch) | +| [`colab`](./google-colab) | Google Colab runtime via runtime proxy | +| [`modal`](./modal) | Modal container execution | +| [`datalayer`](./datalayer) | Datalayer managed runtime with optional GPU | ## Environments @@ -426,6 +116,7 @@ x = 10 x * 2 """) print(result.text) +``` ## Async Execution @@ -451,7 +142,6 @@ value * 2 print(result.stdout) # "value: 21" print(result.text) # "42" ``` -``` ## State Persistence diff --git a/docs/docs/sandboxes/jupyter.mdx b/docs/docs/sandboxes/jupyter.mdx new file mode 100644 index 0000000..20f5882 --- /dev/null +++ b/docs/docs/sandboxes/jupyter.mdx @@ -0,0 +1,56 @@ +--- +sidebar_position: 2 +title: Jupyter +--- + +# Jupyter + +Runs code against a local or remote Jupyter Server and connects through its +Jupyter kernel client. This variant provides process isolation via the Jupyter +kernel and persistent state across requests. + +- **Requirements:** `jupyter_server` and `jupyter-kernel-client` (included by default). +- **Parameters:** `server_url`, `token`, `host`/`port` (when the sandbox starts + its own server), `python_executable`. + +## How To Obtain The Token + +- If you start the server yourself, you choose the token: + ```bash + jupyter server --port 8888 --IdentityProvider.token MY_TOKEN + ``` +- For an already-running server, list servers and read the `token=...` value: + ```bash + jupyter server list + # http://localhost:8888/?token=abcd1234... :: /home/you/notebooks + ``` + You can also pass the full `http://host:port/?token=...` URL as `server_url`; + the token is parsed automatically. +- If you omit `server_url`, `JupyterSandbox` **starts and manages its own local + Jupyter Server** and generates the token for you — no configuration needed. + +## Usage + +```python +from code_sandboxes import Sandbox + +# Connect to an existing server: +with Sandbox.create( + variant="jupyter", + server_url="http://localhost:8888", + token="MY_TOKEN", +) as sandbox: + sandbox.run_code("x = 40") + result = sandbox.run_code("x + 2") + print(result.text) # 42 + +# Or let the sandbox manage a local server automatically: +with Sandbox.create(variant="jupyter") as sandbox: + print(sandbox.run_code("1 + 1").text) # 2 +``` + +The concrete implementation is available from a top-level module: + +```python +from code_sandboxes.jupyter_sandbox import JupyterSandbox +``` diff --git a/docs/docs/sandboxes/kaggle.mdx b/docs/docs/sandboxes/kaggle.mdx index 944b8c1..8175eed 100644 --- a/docs/docs/sandboxes/kaggle.mdx +++ b/docs/docs/sandboxes/kaggle.mdx @@ -1,10 +1,48 @@ --- -sidebar_position: 4 +sidebar_position: 6 title: Kaggle --- # Kaggle +Runs code on Kaggle with two transparent modes: + +1. **Interactive kernel mode** using `KaggleKernelClient` (runtime proxy URL). +2. **Batch job mode** using `KaggleKernelExecutor` when no runtime connection + details are supplied. + +- **Requirements:** `code-sandboxes[kaggle]`. +- **Interactive parameters:** `server_url`, optional `kernel_id`, optional + `channels_url`, optional `token`. +- **Batch credentials:** `~/.kaggle/kaggle.json`, `KAGGLE_API_TOKEN`, or + `KAGGLE_USERNAME` / `KAGGLE_KEY`. + +## Sandbox Usage + +```python +from code_sandboxes import Sandbox + +# Interactive mode (existing session / runtime proxy) +with Sandbox.create( + variant="kaggle", + channels_url="wss://kkb-production.jupyter-proxy.kaggle.net/k/.../proxy/api/kernels/.../channels?session_id=...", +) as sandbox: + print(sandbox.run_code("print(1 + 1)").text) + +# Transparent batch mode (no runtime URL required) +with Sandbox.create(variant="kaggle") as sandbox: + result = sandbox.run_code("print('hello from kaggle batch')") + print(result.text or result.stdout) + +# Stream status updates and outputs while the Kaggle job runs. +with Sandbox.create(variant="kaggle") as sandbox: + for event in sandbox.run_code_streaming("print('hello from kaggle stream')"): + if hasattr(event, "line"): + print(event.line) +``` + +## Authentication + Kaggle notebooks expose a Jupyter-compatible kernel behind an authenticating proxy. `KaggleKernelClient` supports two auth modes: diff --git a/docs/docs/sandboxes/modal.mdx b/docs/docs/sandboxes/modal.mdx new file mode 100644 index 0000000..0f860fd --- /dev/null +++ b/docs/docs/sandboxes/modal.mdx @@ -0,0 +1,69 @@ +--- +sidebar_position: 8 +title: Modal +--- + +# Modal + +Runs code in a [Modal](https://modal.com/docs/guide) cloud sandbox, providing +fully isolated, on-demand containers with configurable images and secrets. + +Each `run_code` call executes in a fresh `python -c` process, so state does **not** +persist across calls (use a single multi-statement snippet if you need shared +state). Configure the image with additional pip packages as needed. + +- **Requirements:** `code-sandboxes[modal]` (installs `modal`). +- **Parameters:** `app_name`, `image` (a prebuilt `modal.Image`), `pip_packages`, + `python_executable`. + +## How To Obtain Modal Credentials + +1. Create a free account at [modal.com](https://modal.com). +2. Authenticate the CLI (opens a browser, writes `~/.modal.toml`): + ```bash + modal token new + ``` + This is enough for local use: the Modal SDK reads credentials from + `~/.modal.toml` automatically. +3. Alternatively, create a token in the Modal dashboard (**Settings → API Tokens**) + and export `MODAL_TOKEN_ID` and `MODAL_TOKEN_SECRET`. + +**Do you need both `MODAL_TOKEN_ID` and `MODAL_TOKEN_SECRET`?** + +- For environment-based auth (CI/CD, containers, hosted runners): **yes**, + you need both values because Modal authenticates with a token pair + (public id + secret). +- For local development with `modal token new`: **not necessarily**. The SDK can + authenticate directly from `~/.modal.toml`. + +If you need environment variables, you can export them from your local config: + +```bash +python - <<'PY' +import pathlib +import tomllib + +cfg = tomllib.loads(pathlib.Path("~/.modal.toml").expanduser().read_text()) +profile = cfg.get("default", cfg) +token_id = profile.get("token_id") +token_secret = profile.get("token_secret") +if token_id and token_secret: + print(f"export MODAL_TOKEN_ID={token_id}") + print(f"export MODAL_TOKEN_SECRET={token_secret}") +else: + raise SystemExit("Could not find token_id/token_secret in ~/.modal.toml") +PY +``` + +## Usage + +```python +from code_sandboxes import Sandbox + +with Sandbox.create( + variant="modal", + pip_packages=["numpy"], +) as sandbox: + result = sandbox.run_code("import numpy as np; print(np.arange(3).sum())") + print(result.stdout) # "3" +``` diff --git a/docs/docs/sandboxes/monty.mdx b/docs/docs/sandboxes/monty.mdx new file mode 100644 index 0000000..1496327 --- /dev/null +++ b/docs/docs/sandboxes/monty.mdx @@ -0,0 +1,48 @@ +--- +sidebar_position: 4 +title: Monty +--- + +# Monty + +Runs code in [Monty](https://github.com/pydantic/monty), a minimal, secure Python +interpreter written in Rust (`pydantic-monty`). Monty runs a restricted subset of +Python in-process with microsecond startup and no access to the host filesystem, +environment, or network unless explicitly granted. It is ideal for short, +LLM-generated snippets where a full container or kernel would be overkill. + +Session state (variables, imports, definitions) persists across `run_code` calls. +Note that Monty supports only a subset of Python — third-party libraries and rich +display outputs are not available. + +- **Requirements:** `code-sandboxes[monty]` (installs `pydantic-monty`). +- **Credentials:** none required (fully local, in-process). +- **Optional parameters** (on `MontySandbox`): `type_check` (type-check before + running), `type_check_stubs`, `external_functions` (`{name: callable}` host + functions the code may call), `limits` (memory / stack / time limits). + +## Usage + +```python +from code_sandboxes import Sandbox + +with Sandbox.create(variant="monty") as sandbox: + sandbox.run_code("x = 21") + result = sandbox.run_code("x * 2") + print(result.text) # 42 +``` + +## Host Callables And Type Checking + +You can expose host callables to the sandboxed code and enable type checking: + +```python +from code_sandboxes.monty_sandbox import MontySandbox + +sandbox = MontySandbox( + type_check=True, + external_functions={"now": lambda: "2026-01-01"}, +) +sandbox.start() +sandbox.run_code("print(now())") +``` diff --git a/pyproject.toml b/pyproject.toml index 05966d8..6922b7e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,8 +29,8 @@ dependencies = [ ] [project.scripts] +code-sandbox = "code_sandboxes.cli:main" code-sandboxes = "code_sandboxes.cli:main" -sandbox = "code_sandboxes.cli:main" [project.optional-dependencies] datalayer = ["agent_runtimes>=1.0.16"] diff --git a/tests/test_client.py b/tests/test_client.py index 65a07a8..8940a11 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -26,6 +26,7 @@ class _FakeSandbox: def __init__(self): self._started = False self._variables = {} + self._tool_caller = None self.config = SimpleNamespace(variant="kaggle") self.info = SimpleNamespace(metadata={"kernel_id": "execution-id"}) self.sandbox_id = "sandbox-id" @@ -64,6 +65,9 @@ def set_variables(self, variables): def interrupt(self): return True + def register_tool_caller(self, caller): + self._tool_caller = caller + @classmethod def list_environments(cls): return [SandboxEnvironment(name="fake", title="Fake", language="python")] @@ -139,8 +143,14 @@ def test_variables_interrupt_and_restart_delegate_to_sandbox(): client.set_variable("one", 1) client.set_variables({"two": 2}) + + def tool_caller(): + return None + + client.register_tool_caller(tool_caller) assert client.get_variable("one") == 1 assert client.get_variable("two") == 2 + assert sandbox._tool_caller is tool_caller assert client.interrupt() is True client.restart() diff --git a/tests/test_colab.py b/tests/test_colab.py index c21d51b..5604a99 100644 --- a/tests/test_colab.py +++ b/tests/test_colab.py @@ -28,7 +28,6 @@ PROXY_TOKEN = "proxy-abc" # noqa: S105 - def test_colab_kernel_client_injects_headers_and_extra_params(monkeypatch): captured: dict = {} diff --git a/tests/test_kaggle_execute.py b/tests/test_kaggle_execute.py index 9d6e700..f3023e5 100644 --- a/tests/test_kaggle_execute.py +++ b/tests/test_kaggle_execute.py @@ -210,7 +210,7 @@ def test_to_kernel_reply_uses_notebook_outputs_when_available(): { "output_type": "stream", "name": "stdout", - "text": "hello from kaggle\n", + "text": "hello from kaggle\n", } ], } @@ -235,10 +235,10 @@ def test_to_kernel_reply_falls_back_to_log_streams(): slug="me/demo", status="COMPLETE", log=( - '[' + "[" '{"stream_name":"stdout","data":"hello\\n"},' '{"stream_name":"stderr","data":"warn\\n"}' - ']' + "]" ), ) @@ -253,10 +253,10 @@ def test_stdout_stderr_and_repr_are_compact(): slug="me/demo", status="COMPLETE", log=( - '[' + "[" '{"stream_name":"stdout","data":"hello\\n"},' '{"stream_name":"stderr","data":"warning\\n"}' - ']' + "]" ), ) From 6a887cb519080d0584fd576ece5955a9e6e8b839 Mon Sep 17 00:00:00 2001 From: Eric Charles Date: Sat, 8 Aug 2026 17:52:35 +0200 Subject: [PATCH 3/3] lint --- code_sandboxes/base.py | 13 +++++++++++++ code_sandboxes/client.py | 6 ++++++ tests/test_client.py | 38 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 57 insertions(+) diff --git a/code_sandboxes/base.py b/code_sandboxes/base.py index 0a99a9d..f589fc5 100644 --- a/code_sandboxes/base.py +++ b/code_sandboxes/base.py @@ -25,6 +25,7 @@ SandboxConfig, SandboxEnvironment, SandboxInfo, + SandboxStatus, SandboxVariant, ) @@ -413,6 +414,18 @@ def stop(self) -> None: """ pass + def mark_stopped(self) -> None: + """Mark the sandbox as stopped after its backend was disconnected externally. + + Callers that bypass :meth:`stop` — for example to disconnect from a + borrowed remote kernel without shutting it down — use this to keep + :attr:`is_started` consistent, so a later :meth:`start` reconnects + instead of silently reusing a closed backend. + """ + self._started = False + if self._info: + self._info.status = SandboxStatus.STOPPED + async def start_async(self) -> None: """Async version of start(). Default implementation calls sync version.""" self.start() diff --git a/code_sandboxes/client.py b/code_sandboxes/client.py index a104bd2..ec40f21 100644 --- a/code_sandboxes/client.py +++ b/code_sandboxes/client.py @@ -317,6 +317,12 @@ def stop(self, shutdown_kernel: bool = True) -> None: backend_stop = getattr(backend, "stop", None) if callable(backend_stop): backend_stop(shutdown_kernel=False) + # The backend connection is now closed, so the sandbox must no longer + # report itself as started; otherwise start() would no-op and later + # executions would run against a closed backend. + mark_stopped = getattr(self._sandbox, "mark_stopped", None) + if callable(mark_stopped): + mark_stopped() async def close_async(self) -> None: """Async variant of :meth:`close`.""" diff --git a/tests/test_client.py b/tests/test_client.py index 8940a11..9b8164f 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -175,3 +175,41 @@ def test_execution_error_converts_to_error_output(): "traceback": ["line 1", "line 2"], } ] + + +class _FakeKernelClient: + """Minimal kernel client that records how it was stopped.""" + + def __init__(self): + self.stopped_with = None + + def stop(self, shutdown_kernel=True): + self.stopped_with = shutdown_kernel + + +class _FakeKernelBackedSandbox(_FakeSandbox): + """Sandbox exposing a borrowed kernel client, like colab/kaggle variants.""" + + def __init__(self): + super().__init__() + self.kernel_client = _FakeKernelClient() + + def mark_stopped(self): + self._started = False + + +def test_stop_without_shutdown_disconnects_backend_and_clears_started(): + sandbox = _FakeKernelBackedSandbox() + client = CodeSandboxClient(sandbox) + client.start() + assert client.is_started is True + + client.stop(shutdown_kernel=False) + + # The borrowed kernel is disconnected, not shut down... + assert sandbox.kernel_client.stopped_with is False + # ...and the sandbox no longer claims to be started, so a later start() + # reconnects instead of reusing the closed backend. + assert client.is_started is False + client.start() + assert client.is_started is True