kaxil commented on code in PR #68847: URL: https://github.com/apache/airflow/pull/68847#discussion_r3606254816
########## providers/common/ai/src/airflow/providers/common/ai/sandbox/sbx.py: ########## @@ -0,0 +1,246 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Docker Sandboxes (``sbx``) microVM backend for the SandboxToolset.""" + +from __future__ import annotations + +import math +import shutil +import subprocess +import tempfile +import threading +from contextlib import suppress + +from airflow.providers.common.ai.sandbox.base import ( + _MAX_OUTPUT_CHARS, + SandboxBackend, + SandboxResult, + _new_sandbox_name, + _validate_positive_finite, +) +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException + +# Extra wall-clock beyond the per-command budget to absorb CLI + microVM +# round-trip overhead before we treat the sbx call itself as hung. +_EXEC_GRACE = 30.0 +# Grace after the timeout fires before the command is force-killed (SIGKILL) if +# it ignored SIGTERM. Kept well under _EXEC_GRACE so the outer call still returns. +_KILL_AFTER = 10 + + +class SbxSandboxBackend(SandboxBackend): + """ + Sandbox backend that runs agent code in a Docker Sandboxes (``sbx``) microVM. + + Drives the ``sbx`` CLI: ``create`` provisions a per-session microVM, ``exec`` + runs commands in it (``docker exec`` semantics), and ``rm`` tears it down. + Each sandbox is a genuine microVM with its own kernel, so agent code is far + better isolated than a shared-kernel container — the same isolation grade as + :class:`~airflow.providers.common.ai.sandbox.IsloSandboxBackend`, but local. + + Requires the ``sbx`` binary on ``PATH`` (``brew install docker/tap/sbx`` / + ``winget install Docker.sbx``) and a one-time ``sbx policy init`` on the host; + it is a Deployment Manager prerequisite, not something Airflow configures. The + template image must provide the GNU coreutils ``timeout`` binary, which + enforces the per-command timeout (any Debian/Ubuntu-based image, including + ``python:*-slim``, does). + + Airflow injects none of its context, connections, or worker environment into + the sandbox; a custom template image may still bundle its own tools. Outbound + network egress is governed by the host ``sbx policy`` (a Deployment Manager + setting), not by this backend — run ``sbx policy init deny-all`` for a + no-egress default. + + :param image: Container image for the sandbox (``sbx --template``). + Default ``"python:3.12-slim"``. + :param memory: Memory limit in binary units (e.g. ``"2g"``). ``sbx`` enforces + a 1 GiB minimum. Default ``"2g"``. + :param cpus: Number of CPUs to allocate. ``None`` (default) uses the ``sbx`` + default (all host CPUs). + :param sbx_path: Path to the ``sbx`` binary. Default ``"sbx"``. + :param create_timeout: Seconds to allow for provisioning (first-run microVM + boot plus image pull can be slow). Default ``600``. + """ + + name = "sbx" + + def __init__( + self, + *, + image: str = "python:3.12-slim", + memory: str = "2g", + cpus: int | None = None, + sbx_path: str = "sbx", + create_timeout: float = 600.0, + ) -> None: + if not image: + raise ValueError("image must not be empty.") + if not memory: + raise ValueError("memory must not be empty.") + if cpus is not None: + _validate_positive_finite(cpus, "cpus") + if not sbx_path: + raise ValueError("sbx_path must not be empty.") + _validate_positive_finite(create_timeout, "create_timeout") + self._image = image + self._memory = memory + self._cpus = cpus + self._sbx_path = sbx_path + self._create_timeout = create_timeout + # Each sandbox mounts a throwaway host workspace; remember it so destroy + # can remove it. ``sbx create`` requires a workspace path but the agent + # never needs host files, so an empty temp dir keeps the host untouched. + self._workspaces: dict[str, str] = {} + + def _run_cli(self, args: list[str], *, timeout: float) -> subprocess.CompletedProcess[str]: + # Only for small, bounded output (create/rm). Command execution uses + # _exec_capped, which bounds memory against unbounded agent output. + return subprocess.run( + [self._sbx_path, *args], + capture_output=True, + text=True, + timeout=timeout, + check=False, + ) + + def _exec_capped(self, args: list[str], *, timeout: float) -> tuple[int, str, str, bool]: + """ + Run the CLI with the output retained bounded to ``_MAX_OUTPUT_CHARS`` per stream. + + Agent code can print unbounded output; ``subprocess.run`` would buffer + all of it and OOM the worker, so drain incrementally and keep only the + cap. Returns ``(returncode, stdout, stderr, truncated)``. + """ + proc = subprocess.Popen( + [self._sbx_path, *args], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + + def drain(stream, buf: bytearray, flag: list[bool]) -> None: + for chunk in iter(lambda: stream.read(65536), b""): + room = _MAX_OUTPUT_CHARS - len(buf) + if room > 0: + buf.extend(chunk[:room]) + if len(chunk) > room: + flag[0] = True + + out, err = bytearray(), bytearray() + out_trunc, err_trunc = [False], [False] + threads = [ + threading.Thread(target=drain, args=(proc.stdout, out, out_trunc)), + threading.Thread(target=drain, args=(proc.stderr, err, err_trunc)), + ] + for t in threads: + t.start() + try: + proc.wait(timeout=timeout) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait() + for t in threads: + t.join() + raise + for t in threads: + t.join() + return ( + proc.returncode, + out.decode(errors="replace"), + err.decode(errors="replace"), + out_trunc[0] or err_trunc[0], + ) + + def create(self) -> str: + if shutil.which(self._sbx_path) is None: + raise AirflowOptionalProviderFeatureException( + f"The {self._sbx_path!r} binary was not found on PATH. Install Docker Sandboxes " + "(https://docs.docker.com/ai/sandboxes/) and run 'sbx policy init' once." + ) + name = _new_sandbox_name() + workspace = tempfile.mkdtemp(prefix="airflow-sandbox-ws-") + args = ["create", "--name", name, "--memory", self._memory, "--template", self._image] + if self._cpus is not None: + args += ["--cpus", str(int(self._cpus))] + args += ["shell", workspace] + try: + result = self._run_cli(args, timeout=self._create_timeout) + if result.returncode != 0: + raise RuntimeError(f"'sbx create' failed ({result.returncode}): {result.stderr.strip()}") + except BaseException: + # A timeout, a nonzero exit, or a partial provision all leave a + # possibly-orphaned microVM and a workspace tempdir. Best-effort + # clean both so nothing leaks on the host or in sbx. + with suppress(Exception): + self._run_cli(["rm", "-f", name], timeout=120.0) + shutil.rmtree(workspace, ignore_errors=True) + raise + self._workspaces[name] = workspace + return name + + def run(self, sandbox: str, command: list[str], *, timeout: float) -> SandboxResult: + _validate_positive_finite(timeout, "timeout") + # Round up: GNU timeout treats 0 as "no timeout", so a sub-second value + # must not truncate to it. + seconds = max(1, math.ceil(timeout)) + # Default (SIGTERM) timeout with a SIGKILL fallback: GNU ``timeout`` exits + # exactly 124 when the budget is hit, and passes the command's own exit + # code through otherwise. So exit 124 is an unambiguous timeout, while an + # OOM-kill or a crash surfaces as its real nonzero code (137, …) rather + # than being mislabelled a timeout — no wall-clock guessing needed. + exec_args = [ + "exec", + sandbox, + "timeout", + f"--kill-after={_KILL_AFTER}", + str(seconds), + *command, + ] + try: + returncode, stdout, stderr, truncated = self._exec_capped( + exec_args, timeout=timeout + _EXEC_GRACE + ) + except subprocess.TimeoutExpired: + # The CLI never returned: the command may still be running in the + # shared microVM. Destroy it so it cannot continue, and tell the + # toolset to provision a fresh one. + self.destroy(sandbox) + return SandboxResult( + exit_code=-1, + stdout="", + stderr="", + timed_out=True, + sandbox_terminated=True, + ) + return SandboxResult( + exit_code=returncode, + stdout=stdout, + stderr=stderr, + timed_out=returncode == 124, Review Comment: GNU `timeout` exits 124 only when the command dies to the initial SIGTERM. If the code ignores SIGTERM (a signal handler, or being stuck in an uninterruptible syscall), `--kill-after` escalates to SIGKILL and `timeout` exits 137, so a real timeout gets reported as `timed_out=False, exit_code=137`, indistinguishable from an OOM kill. Confirmed locally: `gtimeout --kill-after=1 1 bash -c 'trap "" TERM; sleep 30'` exits 137. Default Python code takes SIGTERM so this is an edge case, but the 124-only rule misses exactly the stubborn-timeout case it exists to catch. Treating 137 (or the `--kill-after` escalation more generally) as a timeout would close the gap. ########## providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py: ########## @@ -0,0 +1,212 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Toolset that executes agent-written Python in an isolated sandbox, off the worker.""" + +from __future__ import annotations + +import asyncio +import json +from typing import TYPE_CHECKING, Any + +from pydantic_ai.tools import ToolDefinition +from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool +from pydantic_core import SchemaValidator, core_schema +from typing_extensions import Self + +from airflow.providers.common.ai.sandbox.base import _validate_positive_finite +from airflow.providers.common.ai.utils.tool_definition import return_schema_kwargs + +if TYPE_CHECKING: + from pydantic_ai._run_context import RunContext + + from airflow.providers.common.ai.sandbox.base import SandboxBackend + +_PASSTHROUGH_VALIDATOR = SchemaValidator(core_schema.any_schema()) + +# Deliberately not "run_code": that name is reserved by the Monty code_mode +# meta-tool, and the pinned pydantic-ai-harness rejects a second tool claiming +# it. A distinct name lets this toolset and code_mode coexist. +_TOOL_NAME = "run_python_in_sandbox" + +_RUN_CODE_SCHEMA: dict[str, Any] = { + "type": "object", + "properties": { + "code": {"type": "string", "description": "Python code to execute."}, + }, + "required": ["code"], +} + +_RUN_CODE_DESCRIPTION = ( + "Execute Python code in an isolated sandbox and return a JSON object with " + "exit_code, stdout, stderr, timed_out, and optionally truncated or " + "sandbox_terminated. Airflow does not inject its context, connections, or " + "worker environment. Each call runs a fresh interpreter, but files written " + "by earlier calls persist unless sandbox_terminated is true. Print anything " + "you need to see." +) + + +class SandboxToolset(AbstractToolset[Any]): + """ + Toolset that executes agent-written Python in an isolated sandbox. + + Provides one tool, ``run_python_in_sandbox``. Unlike ``code_mode`` — whose + ``run_code`` executes in-process on the worker via the Monty interpreter — + this ships the code to a disposable sandbox provisioned by the given + :class:`~airflow.providers.common.ai.sandbox.SandboxBackend`. Airflow does + not inject its context, connections, credentials, or worker environment; + custom images and hosted backends can still provide their own credentials. + + The sandbox is created lazily on the first ``run_code`` call, shared by + every call within the agent run, and destroyed when the run ends. A run + that never calls the tool never provisions one. + + A nonzero exit code or a timeout is normal tool output — the model reads + ``stderr`` and fixes its own code. A backend can terminate and transparently + replace the sandbox when that is required to stop timed-out code. Only + infrastructure exceptions from the backend (daemon unreachable, API + failure) propagate and fail the task. + + :param backend: Sandbox backend that provisions and runs the sandbox, e.g. + :class:`~airflow.providers.common.ai.sandbox.SbxSandboxBackend` or + :class:`~airflow.providers.common.ai.sandbox.IsloSandboxBackend`. + :param timeout: Timeout in seconds for a single ``run_code`` call. + Default ``300``. + :param python_command: Python executable inside the sandbox. + Default ``"python3"``. + """ + + def __init__( + self, + backend: SandboxBackend, + *, + timeout: float = 300.0, + python_command: str = "python3", + ) -> None: + _validate_positive_finite(timeout, "timeout") + if not python_command: + raise ValueError("python_command must not be empty.") + self._backend = backend + self._timeout = timeout + self._python_command = python_command + self._sandbox: str | None = None + self._create_task: asyncio.Task[str] | None = None + + @property + def id(self) -> str: + return f"sandbox-{self._backend.name}" + + async def for_run(self, ctx: RunContext[Any]) -> AbstractToolset[Any]: + # Per-run isolation: pydantic-ai shares one toolset instance across runs, + # but each run holds its own sandbox in ``_sandbox``/``_create_task``. + # Hand every run a fresh instance so concurrent runs never share a + # sandbox or destroy each other's. The backend keys all state by unique + # sandbox name, so sharing it across runs is safe. + return SandboxToolset(self._backend, timeout=self._timeout, python_command=self._python_command) + + async def __aenter__(self) -> Self: + # The sandbox is provisioned lazily in call_tool, not here: a durable + # replay that only serves cached tool results must not provision one, + # and nothing leaks if the run fails before any tool executes. + return self + + async def __aexit__(self, *args: Any) -> bool | None: + if self._sandbox is not None: + try: + await asyncio.to_thread(self._backend.destroy, self._sandbox) + finally: + # Re-entering the same instance (HITL regenerate_with_feedback + # does) must never reuse a sandbox whose cleanup was attempted. + self._sandbox = None + return None + + async def _ensure_sandbox(self) -> str: + if self._sandbox is not None: + return self._sandbox + if self._create_task is None: + self._create_task = asyncio.create_task(asyncio.to_thread(self._backend.create)) + create_task = self._create_task + try: + sandbox = await asyncio.shield(create_task) + except asyncio.CancelledError: + # A thread cannot be cancelled. Wait until it publishes the handle + # so __aexit__ can destroy a sandbox created during cancellation. + self._sandbox = await create_task + raise + else: + self._sandbox = sandbox + return sandbox + finally: + if self._create_task is create_task: + self._create_task = None + + async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]: + # sequential=True: every call shares one sandbox and may depend on + # state left by earlier calls — they must not run concurrently. + # return_schema is "string": the tool returns a JSON-encoded string + # (json.dumps), so code mode renders `-> str` instead of `-> Any`. + tool_def = ToolDefinition( + name=_TOOL_NAME, + description=_RUN_CODE_DESCRIPTION, + parameters_json_schema=_RUN_CODE_SCHEMA, + sequential=True, + **return_schema_kwargs({"type": "string"}), + ) + return { + _TOOL_NAME: ToolsetTool( + toolset=self, + tool_def=tool_def, + max_retries=2, + args_validator=_PASSTHROUGH_VALIDATOR, Review Comment: `_PASSTHROUGH_VALIDATOR` is `any_schema`, so pydantic-ai validates the tool args against nothing. A call that omits `code` (or sends a non-string) passes validation, then `tool_args["code"]` on line 194 raises `KeyError`. A `KeyError` is not a `ValidationError`/`ModelRetry`, so it propagates and fails the task instead of prompting the model to fix its call. That also defeats the `max_retries=2` set just above, which only retries on validation/ModelRetry errors. Validating `code` as a required string (a small args model / schema instead of the passthrough) would turn a malformed call into a retryable prompt. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
