This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 4d304d0876c Export files an agent built in a sandbox to object storage
(#73990)
4d304d0876c is described below
commit 4d304d0876cbe7d615502d5c90575a24b9fb95e4
Author: Kaxil Naik <[email protected]>
AuthorDate: Thu Oct 1 11:09:48 2026 +0100
Export files an agent built in a sandbox to object storage (#73990)
SandboxToolset(exports={sandbox_path: storage_url}) copies each file out
when the run ends and before the sandbox is destroyed, streaming through the
worker without passing the model's context or XCom. A promised file that
cannot be exported fails the task; a failed run exports nothing; a teardown
failure after a good export only logs.
Backends gain SandboxBackend.export_file, a concrete default that reads the
file in bounded slices through run_command, overridden to stream natively on
sbx and OpenSandbox.
* Leave export destinations untouched when an export fails
Each file is now copied to a sibling staging key and moved into place only
after every file has been copied, so a failed export no longer deletes or
truncates what was already at a destination. Exports are also bounded by a
total deadline that scales with max_export_bytes, so a guest trickling bytes
cannot hold the task open past the stall timeout.
---
providers/common/ai/docs/sandbox/backends.rst | 11 +-
providers/common/ai/docs/sandbox/configuration.rst | 59 +++-
providers/common/ai/docs/sandbox/index.rst | 5 +-
.../ai/example_dags/example_sandbox_toolset.py | 70 +++-
.../airflow/providers/common/ai/sandbox/base.py | 167 +++++++++-
.../providers/common/ai/sandbox/opensandbox.py | 71 +++-
.../src/airflow/providers/common/ai/sandbox/sbx.py | 90 +++++-
.../providers/common/ai/toolsets/sandbox.py | 209 +++++++++++-
.../tests/unit/common/ai/operators/test_agent.py | 15 +
.../ai/tests/unit/common/ai/sandbox/test_base.py | 163 +++++++++-
.../unit/common/ai/sandbox/test_opensandbox.py | 137 ++++++++
.../ai/tests/unit/common/ai/sandbox/test_sbx.py | 106 ++++++
.../tests/unit/common/ai/toolsets/test_sandbox.py | 356 +++++++++++++++++++++
13 files changed, 1414 insertions(+), 45 deletions(-)
diff --git a/providers/common/ai/docs/sandbox/backends.rst
b/providers/common/ai/docs/sandbox/backends.rst
index 1934ecd742a..a5a7b40472d 100644
--- a/providers/common/ai/docs/sandbox/backends.rst
+++ b/providers/common/ai/docs/sandbox/backends.rst
@@ -321,9 +321,14 @@ in. Subclass
:class:`~airflow.providers.common.ai.sandbox.SandboxBackend` in you
own package and pass an instance to ``SandboxToolset``.
**Three methods are required**: ``create``, ``run_command`` and ``destroy``.
The
-three file operations ship as defaults implemented over ``run_command``,
because
-reading, writing and listing a file are all expressible as shell commands.
-Override them only when the vendor has a native file API:
+file operations ship as defaults implemented over ``run_command``, because
+reading, writing, listing and exporting a file are all expressible as shell
+commands. The default ``export_file``, behind ``SandboxToolset(exports=...)``,
+copies a file in 4 MiB slices, one command each, and needs ``stat``, ``tail``,
+``head`` and ``base64`` in the guest. It relies on ``run_command`` returning
each
+slice's output intact, or setting ``stdout_truncated`` when it could not.
Override
+it when the vendor can stream a download, as the ``sbx`` and OpenSandbox
backends do. Override the others
+only when the vendor has a native file API:
.. code-block:: python
diff --git a/providers/common/ai/docs/sandbox/configuration.rst
b/providers/common/ai/docs/sandbox/configuration.rst
index d4e347f27cb..2647defde51 100644
--- a/providers/common/ai/docs/sandbox/configuration.rst
+++ b/providers/common/ai/docs/sandbox/configuration.rst
@@ -203,9 +203,49 @@ capped, 50 KiB per stream for ``run_command`` and
``max_read_bytes`` (5 MiB) per
``read_file``. Those caps are a budget rather than a transport limit: a 200 MB
file
reads out of a live sandbox in under ten seconds on the same path, but raising
the
cap costs roughly three times the file size in worker memory to show the model
50
-KiB of it. A file leaves through a task instead: give the agent a sandbox a
task
-owns, and read the file out through the backend after the run
-(:ref:`sandbox-attach`).
+KiB of it.
+
+A file leaves through ``exports`` instead. Name the paths the agent will write
and
+where each should land, and when the run ends the toolset copies them out
before
+it destroys the sandbox. The copy streams through the worker in bounded
pieces, so
+a file far over ``max_read_bytes`` is never held in worker memory whole, and
none
+of it passes through the model's context or XCom. The destinations are
templated
+when the toolset is passed through ``AgentOperator``, so each run writes its
own,
+and the ``run_command`` tool tells the model which files will be collected so
it
+writes them where they are expected. A local destination's missing directories
+are created; on object storage a key needs none:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_sandbox_toolset.py
+ :language: python
+ :start-after: [START howto_sandbox_agent_export]
+ :end-before: [END howto_sandbox_agent_export]
+
+A destination is anything :class:`~airflow.sdk.ObjectStoragePath` can open,
with
+credentials from ``export_conn_id``. Only a regular file is exported, up to
+``max_export_bytes`` (1 GiB by default). Two outcomes are kept apart:
+
+* **A promised file that cannot be exported fails the task.** The file is
missing,
+ is a directory, is over the limit, changed size while it was copied, took
longer
+ than the limit allows at 1 MiB/s, or the storage refused the write. Every
file is
+ first copied to a staging key beside its destination, named after it with a
+ ``.partial`` suffix, and the files are moved into place only once all of them
+ have been copied, so a failed export leaves every destination as it was. The
+ staging keys are removed, and one that cannot be is logged and named in the
+ error. The sandbox is still destroyed.
+* **A sandbox that cannot be destroyed afterwards does not.** The file is
delivered,
+ the failure is logged with the sandbox's name, and the backend's lifetime or
an
+ operator's sweep reclaims it.
+
+A run that fails exports nothing, and a run that never called a tool has no
+sandbox to export from, which fails the task the same way a missing file does.
+None of these touches the destinations, so a file an earlier try exported is
+still there: a consumer that runs whatever the outcome, such as one with
+``trigger_rule=TriggerRule.ALL_DONE``, should check that the task succeeded
before
+it trusts a file. A destination that renders to something other than a storage
+URL fails the task before the model runs.
+``exports`` cannot be combined with ``attach_to``: a sandbox another task owns
is
+read out by that task (:ref:`sandbox-attach`), which is also the shape to use
when
+the files must survive a failed run.
When the deliverable is a file and the Dag already knows the job, do not use an
agent for it at all. Drive a backend from a ``@task``: the input goes in
through
@@ -225,12 +265,13 @@ A sandbox another task owns
---------------------------
The toolset's own sandbox is provisioned from a spec fixed in the Dag file, on
the
-model's first tool call, and destroyed when the run ends. Three things cannot
be
-done inside that shape: a credential cannot come from a connection, a file the
-agent built cannot leave, and a second run against the same agent, which is
what
-:ref:`HITL review <howto:hitl_review>` does when a reviewer asks for changes,
cannot
-find the first run's files. All three have the same answer. Let a task create
the
-sandbox and hand the agent only the handle.
+model's first tool call, and destroyed when the run ends. Two things cannot be
+done inside that shape: a credential cannot come from a connection, and a
second
+run against the same agent, which is what :ref:`HITL review
<howto:hitl_review>`
+does when a reviewer asks for changes, cannot find the first run's files. A
file
+the agent built can leave through ``exports`` (:ref:`sandbox-results`), but
only
+from a run that succeeded. All of these have the same answer. Let a task create
+the sandbox and hand the agent only the handle.
.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_sandbox_toolset.py
:language: python
diff --git a/providers/common/ai/docs/sandbox/index.rst
b/providers/common/ai/docs/sandbox/index.rst
index e21f8112897..1d87a27f20e 100644
--- a/providers/common/ai/docs/sandbox/index.rst
+++ b/providers/common/ai/docs/sandbox/index.rst
@@ -229,9 +229,8 @@ actual isolation. Choose the smallest boundary that fits,
then configure it.
- ``SandboxToolset``
* - Produce a large artifact for a downstream task
- A ``@task`` driving a backend directly when the Dag knows the job. When
- the agent has to produce it, a ``@task`` provisions the sandbox, the
agent
- attaches, and a ``@task`` reads the file out; see
- :ref:`A sandbox another task owns <sandbox-attach>`.
+ the agent has to produce it, ``SandboxToolset(exports=...)`` copies the
file
+ to object storage when the run ends; see :ref:`sandbox-results`.
* - A whole task's worth of untrusted work isolated, with no agent involved
- ``KubernetesPodOperator``
* - Airflow's own credentials kept away from the agent
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_sandbox_toolset.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_sandbox_toolset.py
index 15ab35396c1..456287c23aa 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_sandbox_toolset.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_sandbox_toolset.py
@@ -18,7 +18,7 @@
"""
Example Dags for ``SandboxToolset``.
-Three shapes, each a job a data team actually runs:
+Five shapes, each a job a data team actually runs:
1. An agent investigates a revenue anomaly. It queries the warehouse through a
``SQLToolset`` (the credential stays in the task, the model only sees rows)
and
@@ -34,6 +34,9 @@ Three shapes, each a job a data team actually runs:
``@task`` reads the report the agent wrote and destroys the sandbox. The
task
that creates the sandbox decides what goes in, in ordinary Python at run
time,
and a file the agent built comes out without crossing the model's context.
+5. An agent writes a file and the toolset exports it to object storage when the
+ run ends, for a downstream task to load. No task has to own the sandbox, so
+ this is the shorter shape whenever the file is all that must come out.
"""
from __future__ import annotations
@@ -229,10 +232,10 @@ def example_sandbox_task_artifact():
Convert a CSV to parquet inside a sandbox and land the result in object
storage.
No agent is involved. The Dag knows exactly what to run, so a model would
add
- nothing, and a file produced inside an *agent's* sandbox could only come
back
- through the model's context, which is text-only and capped. Driving the
- backend from a task has neither limit: the bytes move through the worker,
and
- the caller owns the sandbox's lifetime.
+ nothing. Driving the backend from a task also gives the caller the
sandbox's
+ whole lifetime: it puts the input in, runs the conversion, and reads the
result
+ out, with the bytes moving through the worker. When an agent has to
produce the
+ file instead, see example 5.
"""
# [START howto_sandbox_task_artifact]
@@ -372,3 +375,60 @@ if modal is not None:
# [END howto_sandbox_attach]
example_sandbox_attach()
+
+
+# ---------------------------------------------------------------------------
+# 5. An agent builds a file, and the toolset exports it when the run ends.
+# ---------------------------------------------------------------------------
+
+# [START howto_sandbox_agent_export]
+# Templated per run. Point it at ``s3://`` or ``gs://`` and pass
``export_conn_id``
+# in a real deployment; ``file://`` keeps the example runnable on a laptop.
+STAGING_URI = "file:///tmp/airflow-sandbox-example/{{ run_id }}/staging.csv"
+
+
+@dag(
+ schedule=None,
+ start_date=datetime(2024, 1, 1, tzinfo=timezone.utc),
+ catchup=False,
+ tags=["example", "sandbox"],
+)
+def example_sandbox_agent_export():
+ """Have an agent normalize a vendor file, export the result, and load it
downstream."""
+ normalize = AgentOperator(
+ task_id="normalize",
+ prompt=(
+ "Here is a sample of this month's vendor export:\n\n"
+ f"{VENDOR_SAMPLE}\n"
+ "Write it to a file, then write a Python script that turns it into
staging.csv with "
+ "the columns order_id, customer_name, ordered_at (ISO date),
amount_usd (decimal) "
+ "and country_code. Run the script and fix it until every row
parses, then report "
+ "the column mapping you settled on."
+ ),
+ system_prompt=(
+ "You have a sandbox with Python 3.12 and the standard library, no
network, and "
+ "an empty working directory. Read tracebacks and fix the code
rather than guessing."
+ ),
+ llm_conn_id="pydanticai_default",
+ output_type=ColumnMapping,
+ toolsets=[
+ SandboxToolset(
+ SbxSandboxBackend(host_network_policy="deny-all"),
+ # Copied out before the sandbox is destroyed. If the agent
never wrote
+ # staging.csv, the task fails rather than leaving ``load`` to
find nothing.
+ exports={"staging.csv": STAGING_URI},
+ ),
+ ],
+ )
+
+ @task
+ def load(staging_uri: str) -> int:
+ with ObjectStoragePath(staging_uri).open() as staging:
+ return sum(1 for _ in staging) - 1
+
+ normalize >> load(STAGING_URI)
+
+
+# [END howto_sandbox_agent_export]
+
+example_sandbox_agent_export()
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/sandbox/base.py
b/providers/common/ai/src/airflow/providers/common/ai/sandbox/base.py
index 0a69c54c48c..09f91b4cb24 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/sandbox/base.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/sandbox/base.py
@@ -23,6 +23,7 @@ import binascii
import json
import logging
import math
+import re
import shlex
import time
import uuid
@@ -32,6 +33,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, TypeGuard
if TYPE_CHECKING:
from collections.abc import Mapping, Sequence
+ from typing import BinaryIO
log = logging.getLogger(__name__)
@@ -90,7 +92,7 @@ class SandboxFileTooLargeError(SandboxError):
self.path = path
self.size_bytes = size_bytes
self.max_bytes = max_bytes
- super().__init__(f"{path!r} is {size_bytes} bytes, over the
{max_bytes} byte read limit.")
+ super().__init__(f"{path!r} is {size_bytes} bytes, over the
{max_bytes} byte limit.")
# Bounded budget for the shell helpers behind the default file operations.
@@ -98,6 +100,34 @@ _FILE_OP_TIMEOUT = 120.0
# They return a status or a listing, never bulk content, so a small cap bounds
# what a hostile guest can push into worker memory.
_FILE_OP_OUTPUT_CAP = 1024 * 1024
+# Slice the default export reads per command. Each slice is held in worker
memory
+# about three times over (base64 text, then decoded bytes), so this bounds the
+# export's footprint whatever the file's size.
+_EXPORT_CHUNK_BYTES = 4 * 1024 * 1024
+# Prefix on the line an export check prints the file's size on.
+_EXPORT_SIZE_TAG = "airflow-export-size:"
+# The slowest an export may run on average before it is ended, so a guest that
trickles
+# a byte at a time cannot hold the task open: 1 GiB, the toolset's default
budget, gets
+# about 17 minutes.
+_EXPORT_MIN_RATE = 1024 * 1024
+
+
+def _export_allowance(max_bytes: int) -> float:
+ """Seconds a whole export of a file of up to ``max_bytes`` may take."""
+ return max(_FILE_OP_TIMEOUT, max_bytes / _EXPORT_MIN_RATE)
+
+
+def _export_deadline(max_bytes: int) -> float:
+ """When an export starting now must have finished, on the
``time.monotonic`` clock."""
+ return time.monotonic() + _export_allowance(max_bytes)
+
+
+def _check_export_deadline(path: str, deadline: float, max_bytes: int) -> None:
+ if time.monotonic() > deadline:
+ raise SandboxError(
+ f"Exporting {path!r} took longer than the
{_export_allowance(max_bytes):g}s allowed for a "
+ f"file of up to {max_bytes} bytes; the sandbox is sending it too
slowly."
+ )
def _validate_positive_finite(value: float, name: str) -> None:
@@ -273,7 +303,8 @@ class SandboxBackend(ABC):
:class:`~airflow.providers.common.ai.toolsets.sandbox.SandboxToolset`.
The four operation methods are named after the four tools the toolset
exposes, so the mapping from a model-facing tool to the backend call behind
- it is literal; ``create`` and ``destroy`` are lifecycle and have no tool. A
+ it is literal; ``create`` and ``destroy`` are lifecycle and have no tool,
and
+ neither has ``export_file``, which hands a finished file to the Dag
author. A
backend whose sandboxes can be found again from another process implements
:class:`AttachableSandboxBackend` instead, which adds the ownership rules a
task-provisioned sandbox needs.
@@ -336,15 +367,18 @@ class SandboxBackend(ABC):
#
# Concrete, not abstract: every one of these is expressible as a shell
# command, so a backend only has to implement ``run_command`` to get all
- # three. Override them when the vendor exposes a native file API, which
+ # four. Override them when the vendor exposes a native file API, which
# avoids base64 expansion, the command-line length ceiling, and the guest
# needing coreutils at all.
# ------------------------------------------------------------------
- # Reserved exit statuses for "the path is not readable" and "the path is a
- # directory", distinct from any status the guest's own command might
return.
+ # Reserved exit statuses for "the path is not readable", "the path is a
+ # directory", "the path is not a regular file" and "the file is over the
+ # export budget", distinct from any status the guest's own command might
return.
_MISSING_PATH_STATUS = 66
_IS_DIRECTORY_STATUS = 67
+ _NOT_REGULAR_FILE_STATUS = 68
+ _TOO_LARGE_STATUS = 69
def read_file(self, sandbox: str, path: str, *, max_bytes: int) -> bytes:
"""
@@ -440,6 +474,129 @@ class SandboxBackend(ABC):
entries.append((name, kind == "d"))
return entries
+ def export_file(self, sandbox: str, path: str, dest: BinaryIO, *,
max_bytes: int) -> int:
+ """
+ Copy a regular file out of the sandbox into ``dest`` and return the
bytes written.
+
+ ``dest`` is a writable binary stream, typically an object-storage
file, and the
+ copy goes through it without the whole file ever being held in worker
memory,
+ so a file far over :meth:`read_file`'s budget can leave the sandbox.
Only a
+ regular file is exported: a directory, a device, or a FIFO is refused,
since
+ none of them has a size to promise a caller. Raise
+ :class:`SandboxFileTooLargeError` instead of copying a file larger than
+ ``max_bytes``, and :class:`SandboxError` when the file changed size
while it
+ was being copied, which means a process in the sandbox is still
writing it.
+
+ This default reads the file in slices through :meth:`run_command`, one
+ command per slice, and needs ``stat``, ``tail``, ``head`` and
``base64`` in the
+ guest. It relies on ``run_command`` returning each slice's output
intact, or
+ setting ``stdout_truncated`` when it could not, and on nothing but the
command's
+ own output reaching stdout. Override it when the vendor can stream a
file out,
+ and bound the whole copy by :func:`_export_deadline` as this one does,
since a
+ guest that keeps sending a byte now and then never trips a stall
timeout.
+ """
+ deadline = _export_deadline(max_bytes)
+ quoted = shlex.quote(path)
+ check = self.run_command(
+ sandbox,
+ f"{self._export_checks(quoted, max_bytes)}
{self._print_export_size()}",
+ timeout=_FILE_OP_TIMEOUT,
+ max_output_bytes=_FILE_OP_OUTPUT_CAP,
+ )
+ self._raise_for_export_status(path, check, max_bytes)
+ size = self._parse_export_size(check.stdout)
+ written = 0
+ while True:
+ _check_export_deadline(path, deadline, max_bytes)
+ # ``tail -c +N`` seeks on a regular file, so each slice costs its
own
+ # length rather than a read from the start.
+ result = self.run_command(
+ sandbox,
+ f"tail -c +{written + 1} -- {quoted} | head -c
{_EXPORT_CHUNK_BYTES} | base64",
+ timeout=max(1.0, min(_FILE_OP_TIMEOUT, deadline -
time.monotonic())),
+ max_output_bytes=_EXPORT_CHUNK_BYTES * 2 + 4096,
+ )
+ if result.sandbox_terminated:
+ raise SandboxTerminalError(f"The sandbox ended while {path!r}
was being exported.")
+ if result.exit_code or result.stdout_truncated:
+ # A truncated slice decodes cleanly into the wrong bytes, so
it is an
+ # error, not a short read.
+ raise SandboxError(result.stderr.strip() or f"Could not export
{path!r}.")
+ try:
+ data = base64.b64decode(result.stdout, validate=False)
+ except (binascii.Error, ValueError) as e:
+ raise SandboxError(f"Could not decode {path!r} from the
sandbox.") from e
+ if written + len(data) > max_bytes:
+ raise SandboxFileTooLargeError(path, written + len(data),
max_bytes)
+ dest.write(data)
+ written += len(data)
+ if len(data) < _EXPORT_CHUNK_BYTES:
+ break
+ self._check_export_size(path, expected=size, written=written)
+ return written
+
+ @classmethod
+ def _export_checks(cls, quoted: str, max_bytes: int) -> str:
+ """
+ Shell that refuses what :meth:`export_file` will not copy, leaving the
size in ``$sz``.
+
+ For a backend overriding :meth:`export_file` that still drives the
guest
+ through a shell, so its refusals mean what the default's do.
+ """
+ return (
+ f"sz=$(stat -Lc %s -- {quoted} 2>/dev/null) || exit
{cls._MISSING_PATH_STATUS}; "
+ f"[ -d {quoted} ] && exit {cls._IS_DIRECTORY_STATUS}; "
+ f"[ -f {quoted} ] || exit {cls._NOT_REGULAR_FILE_STATUS}; "
+ f'[ "$sz" -gt {max_bytes} ] && {{
{cls._print_export_size(to_stderr=True)} exit {cls._TOO_LARGE_STATUS}; }};'
+ )
+
+ @staticmethod
+ def _print_export_size(*, to_stderr: bool = False) -> str:
+ """
+ Shell that prints ``$sz`` on a line of its own, tagged.
+
+ Tagged because the stream it lands on may also carry the vendor CLI's
own notices.
+ """
+ redirect = " >&2" if to_stderr else ""
+ return f'printf "\\n{_EXPORT_SIZE_TAG}%s\\n" "$sz"{redirect};'
+
+ @staticmethod
+ def _parse_export_size(text: str) -> int:
+ sizes = re.findall(rf"^{_EXPORT_SIZE_TAG}(\d+)$", text,
flags=re.MULTILINE)
+ if not sizes:
+ raise SandboxError("The sandbox did not report the size of the
file being exported.")
+ return int(sizes[-1])
+
+ @classmethod
+ def _raise_for_export_status(cls, path: str, result: SandboxExecResult,
max_bytes: int) -> None:
+ """Turn a failed :meth:`_export_checks` into the error it stands
for."""
+ if result.sandbox_terminated:
+ raise SandboxTerminalError(f"The sandbox ended while {path!r} was
being exported.")
+ if result.exit_code == cls._MISSING_PATH_STATUS:
+ raise SandboxError(f"{path!r} does not exist in the sandbox, or is
not readable.")
+ if result.exit_code == cls._IS_DIRECTORY_STATUS:
+ raise SandboxError(f"{path!r} is a directory; only a regular file
can be exported.")
+ if result.exit_code == cls._NOT_REGULAR_FILE_STATUS:
+ raise SandboxError(f"{path!r} is not a regular file; only a
regular file can be exported.")
+ if result.exit_code == cls._TOO_LARGE_STATUS:
+ try:
+ size = cls._parse_export_size(result.stderr)
+ except SandboxError:
+ size = max_bytes + 1
+ raise SandboxFileTooLargeError(path, size, max_bytes)
+ if result.exit_code:
+ raise SandboxError(result.stderr.strip() or f"Could not export
{path!r}.")
+
+ @staticmethod
+ def _check_export_size(path: str, *, expected: int, written: int) -> None:
+ if written != expected:
+ # The run is over, so something the agent started in the
background is
+ # still writing it. What was copied is not the file the run
produced.
+ raise SandboxError(
+ f"{path!r} changed while it was exported: it was {expected}
bytes when the export "
+ f"began and {written} bytes were copied. A process in the
sandbox may still be writing it."
+ )
+
@abstractmethod
def destroy(self, sandbox: str) -> None:
"""Tear down the sandbox. Must be idempotent."""
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/sandbox/opensandbox.py
b/providers/common/ai/src/airflow/providers/common/ai/sandbox/opensandbox.py
index 07aaaed61d7..f8a36d5496e 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/sandbox/opensandbox.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/sandbox/opensandbox.py
@@ -22,7 +22,7 @@ import logging
import posixpath
import threading
import time
-from contextlib import contextmanager, suppress
+from contextlib import closing, contextmanager, suppress
from datetime import timedelta
from typing import TYPE_CHECKING, Any
@@ -32,13 +32,16 @@ from airflow.providers.common.ai.sandbox.base import (
SandboxExecResult,
SandboxFileTooLargeError,
SandboxTerminalError,
+ _check_export_deadline,
+ _export_deadline,
_new_sandbox_name,
_validate_positive_finite,
)
from airflow.providers.common.compat.sdk import BaseHook
if TYPE_CHECKING:
- from collections.abc import Callable, Iterator
+ from collections.abc import Callable, Generator, Iterator
+ from typing import BinaryIO
from opensandbox import SandboxSync
from opensandbox.config import ConnectionConfigSync
@@ -459,21 +462,68 @@ class OpenSandboxBackend(SandboxBackend):
def read_file(self, sandbox: str, path: str, *, max_bytes: int) -> bytes:
_validate_positive_finite(max_bytes, "max_bytes")
sandbox_client = self._get_sandbox(sandbox)
- chunks = None
data = bytearray()
+ with closing(self._download(sandbox_client, path,
max_bytes=max_bytes)) as chunks:
+ for chunk in chunks:
+ data.extend(chunk[: max_bytes + 1 - len(data)])
+ if len(data) > max_bytes:
+ size = self._get_file_size(sandbox_client, path,
at_least=len(data))
+ raise SandboxFileTooLargeError(path, size, max_bytes)
+ return bytes(data)
+
+ def export_file(self, sandbox: str, path: str, dest: BinaryIO, *,
max_bytes: int) -> int:
+ """Override: stream the file through the SDK's ranged download, one
chunk at a time."""
+ _validate_positive_finite(max_bytes, "max_bytes")
+ sandbox_client = self._get_sandbox(sandbox)
+ try:
+ info = sandbox_client.files.get_file_info([path])
+ entry = info.get(path) or next(iter(info.values()))
+ except Exception as e:
+ if _get_status_code(e) == 404 or isinstance(e, StopIteration):
+ self._confirm_sandbox_exists(sandbox_client)
+ raise SandboxError(f"{path!r} does not exist in the sandbox,
or is not readable.") from e
+ with _translate_opensandbox_errors(
+ "inspect a sandbox file", recoverable_statuses=frozenset({400})
+ ):
+ raise
+ if entry.entry_type == "directory":
+ raise SandboxError(f"{path!r} is a directory; only a regular file
can be exported.")
+ if entry.entry_type == "other":
+ raise SandboxError(f"{path!r} is not a regular file; only a
regular file can be exported.")
+ if entry.size > max_bytes:
+ raise SandboxFileTooLargeError(path, entry.size, max_bytes)
+ deadline = _export_deadline(max_bytes)
+ written = 0
+ with closing(self._download(sandbox_client, path,
max_bytes=max_bytes)) as chunks:
+ for chunk in chunks:
+ _check_export_deadline(path, deadline, max_bytes)
+ written += len(chunk)
+ if written > max_bytes:
+ raise SandboxFileTooLargeError(path, written, max_bytes)
+ dest.write(chunk)
+ if entry.entry_type == "file":
+ # A symlink reports the link's own size, not its target's, so
there is
+ # nothing to compare against for one.
+ self._check_export_size(path, expected=entry.size, written=written)
+ return written
+
+ def _download(
+ self, sandbox_client: SandboxSync, path: str, *, max_bytes: int
+ ) -> Generator[bytes, None, None]:
+ """
+ Yield the first ``max_bytes + 1`` bytes of a sandbox file, with the
SDK's errors translated.
+
+ A generator, so a failure in whatever the caller does with a chunk is
raised in
+ the caller and never mistaken for a failed download.
+ """
+ chunks = None
try:
chunks = sandbox_client.files.read_bytes_stream(
path,
chunk_size=min(65536, max_bytes + 1),
range_header=f"bytes=0-{max_bytes}",
)
- for chunk in chunks:
- data.extend(chunk[: max_bytes + 1 - len(data)])
- if len(data) > max_bytes:
- size = self._get_file_size(sandbox_client, path,
at_least=len(data))
- raise SandboxFileTooLargeError(path, size, max_bytes)
- except SandboxFileTooLargeError:
- raise
+ yield from chunks
except Exception as e:
if _get_status_code(e) == 404:
self._confirm_sandbox_exists(sandbox_client)
@@ -485,7 +535,6 @@ class OpenSandboxBackend(SandboxBackend):
if close is not None:
with suppress(Exception):
close()
- return bytes(data)
def write_file(self, sandbox: str, path: str, content: bytes) -> None:
sandbox_client = self._get_sandbox(sandbox)
diff --git a/providers/common/ai/src/airflow/providers/common/ai/sandbox/sbx.py
b/providers/common/ai/src/airflow/providers/common/ai/sandbox/sbx.py
index 02e39430a7f..dc11059ef2d 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/sandbox/sbx.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/sandbox/sbx.py
@@ -36,12 +36,17 @@ from airflow.providers.common.ai.sandbox.base import (
SandboxBackend,
SandboxError,
SandboxExecResult,
+ SandboxFileTooLargeError,
SandboxTerminalError,
+ _check_export_deadline,
+ _export_deadline,
_new_sandbox_name,
_validate_positive_finite,
)
if TYPE_CHECKING:
+ from typing import BinaryIO
+
from airflow.providers.common.ai.sandbox.base import SandboxSpec
# Extra wall-clock beyond the per-command budget to absorb CLI and microVM
@@ -55,6 +60,10 @@ _FILE_OP_TIMEOUT = 120.0
# Helpers return a status or a directory listing, never bulk file content, so a
# small cap is enough to bound what a hostile guest can push into worker
memory.
_HELPER_OUTPUT_CAP = 1024 * 1024
+# Seconds an export may go without a byte arriving before it is ended. The
whole copy
+# is bounded separately, by ``_export_deadline``; measured, 200 MB streams out
of a
+# local microVM in under three seconds.
+_EXPORT_STALL_TIMEOUT = 120.0
log = logging.getLogger(__name__)
@@ -108,7 +117,7 @@ class SbxSandboxBackend(SandboxBackend):
can find and remove them; budget for that sweep before running this at
scale.
The template image must provide GNU coreutils ``timeout``, ``base64``,
``stat``,
- ``find``, ``mkdir`` and ``dirname``, which the command and file tools use.
Any Debian or Ubuntu based
+ ``head``, ``find``, ``mkdir`` and ``dirname``, which the command and file
tools use. Any Debian or Ubuntu based
image, including ``python:*-slim``, does.
:param image: Container image for the sandbox (``sbx --template``).
@@ -458,6 +467,85 @@ class SbxSandboxBackend(SandboxBackend):
if code:
raise SandboxError(stderr.decode(errors="replace").strip() or
f"Could not write {path!r}.")
+ def export_file(self, sandbox: str, path: str, dest: BinaryIO, *,
max_bytes: int) -> int:
+ """
+ Override: stream the file out of one ``sbx exec`` straight into
``dest``.
+
+ ``sbx exec`` carries the guest's stdout byte for byte, so the default's
+ base64 round trip per slice buys nothing here. The guest reports the
size it
+ is about to send on stderr, which is how a file still being written is
caught.
+ """
+ quoted = shlex.quote(path)
+ script = (
+ f"{self._export_checks(quoted, max_bytes)} "
+ f"{self._print_export_size(to_stderr=True)} exec head -c
{max_bytes + 1} -- {quoted}"
+ )
+ deadline = _export_deadline(max_bytes)
+ written = 0
+ stderr = bytearray()
+ stalled = threading.Event()
+ expired = threading.Event()
+ finished = threading.Event()
+ last_progress = [time.monotonic()]
+ with subprocess.Popen(
+ [self._sbx_path, "exec", sandbox, "sh", "-c", script],
+ stdin=subprocess.DEVNULL,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.PIPE,
+ start_new_session=True,
+ ) as proc:
+
+ def kill() -> None:
+ with suppress(OSError):
+ os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
+ with suppress(OSError):
+ proc.kill()
+
+ def watch() -> None:
+ # Two clocks: the stall one restarts with every chunk, so a
large file that
+ # keeps moving is not cut off; the deadline does not, so a
guest that
+ # trickles a byte now and then is ended too.
+ while not finished.wait(min(1.0, _EXPORT_STALL_TIMEOUT / 4)):
+ now = time.monotonic()
+ if now - last_progress[0] > _EXPORT_STALL_TIMEOUT or now >
deadline:
+ (stalled if now <= deadline else expired).set()
+ kill()
+ return
+
+ def drain_stderr() -> None:
+ with suppress(ValueError, OSError):
+ for chunk in iter(lambda: proc.stderr.read(4096), b""): #
type: ignore[union-attr]
+ stderr.extend(chunk[: _HELPER_OUTPUT_CAP -
len(stderr)])
+
+ drainer = threading.Thread(target=drain_stderr, daemon=True)
+ watchdog = threading.Thread(target=watch, daemon=True)
+ drainer.start()
+ watchdog.start()
+ try:
+ for chunk in iter(lambda: proc.stdout.read(65536), b""): #
type: ignore[union-attr]
+ written += len(chunk)
+ if written > max_bytes:
+ raise SandboxFileTooLargeError(path, written,
max_bytes)
+ dest.write(chunk)
+ last_progress[0] = time.monotonic()
+ proc.wait()
+ finally:
+ finished.set()
+ if proc.poll() is None:
+ kill()
+ proc.wait()
+ drainer.join(timeout=5.0)
+ if expired.is_set():
+ _check_export_deadline(path, deadline, max_bytes)
+ if stalled.is_set():
+ raise SandboxError(f"Exporting {path!r} stalled: no data arrived
for {_EXPORT_STALL_TIMEOUT:g}s.")
+ report = stderr.decode(errors="replace")
+ self._raise_for_export_status(
+ path, SandboxExecResult(exit_code=proc.returncode, stdout="",
stderr=report), max_bytes
+ )
+ self._check_export_size(path,
expected=self._parse_export_size(report), written=written)
+ return written
+
def destroy(self, sandbox: str) -> None:
# Already-gone is fine -- 'sbx rm -f' exits nonzero for a missing
sandbox,
# which we ignore so destroy stays idempotent.
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py
b/providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py
index 3ae301b8584..75762aa1a6a 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/sandbox.py
@@ -23,10 +23,13 @@ import concurrent.futures
import contextvars
import logging
import math
+import sys
import threading
import time
+import uuid
from typing import TYPE_CHECKING, Any, NamedTuple
+from fsspec.implementations.local import LocalFileSystem
from pydantic_ai.exceptions import ModelRetry
from pydantic_ai.tools import ToolDefinition
from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool
@@ -54,10 +57,10 @@ from airflow.providers.common.ai.utils.tool_definition
import (
return_schema_kwargs,
)
from airflow.providers.common.ai.utils.toolset_base import AirflowToolset
-from airflow.providers.common.compat.sdk import get_current_context
+from airflow.providers.common.compat.sdk import ObjectStoragePath,
get_current_context
if TYPE_CHECKING:
- from collections.abc import Sequence
+ from collections.abc import Mapping, Sequence
from pydantic_ai._run_context import RunContext
@@ -76,6 +79,16 @@ _provisioning =
concurrent.futures.ThreadPoolExecutor(thread_name_prefix="sandbo
RUN_COMMAND = "run_command"
+def _sentence(error: BaseException) -> str:
+ """Return an exception's message ending in exactly one full stop, to join
into a longer one."""
+ return f"{str(error).rstrip('.')}."
+
+
+# Ends the name of the staging key an export is copied to before it is moved
into place,
+# so a consumer globbing its destination's own extension never matches one.
+_PARTIAL_SUFFIX = ".partial"
+
+
class _Identity(NamedTuple):
"""Who an attaching run presents as: the owner it expects on the sandbox,
and the holder it claims to be."""
@@ -186,6 +199,15 @@ class SandboxToolset(AirflowToolset):
the handle travels from the provisioning task:
``attach_to="{{ ti.xcom_pull('provision') }}"``.
+ **Files the agent builds can leave.** ``exports`` maps a path in the
sandbox to
+ an object-storage destination, and when the run ends the toolset copies
each
+ file there before it destroys the sandbox. The copy streams through the
worker
+ in bounded pieces and never passes through the model's context or XCom, so
a
+ parquet file, a chart, or a trained model is as easy to hand downstream as
a
+ line of text. The destinations are templated the same way ``attach_to``
is. A
+ promised file that cannot be exported fails the task, and a failed run
exports
+ nothing; a sandbox that cannot be destroyed afterwards does not fail the
task.
+
A non-zero exit or a timeout is normal tool output -- the model reads it
and
corrects itself. A recoverable sandbox failure becomes a bounded retry.
Only
a terminal failure (credentials rejected, daemon unreachable) fails the
task,
@@ -213,6 +235,17 @@ class SandboxToolset(AirflowToolset):
which ``ModalSandboxBackend`` is and ``SbxSandboxBackend`` is not), and
cannot be combined with ``spec``, since the sandbox is already
provisioned.
The toolset never destroys an attached sandbox.
+ :param exports: Files to copy out of the sandbox when the run ends, as a
+ mapping from a path in the sandbox (relative paths resolve the way the
+ ``read_file`` tool resolves them) to an object-storage URL such as
+ ``"s3://bucket/{{ run_id }}/report.parquet"``, anything
+ :class:`~airflow.sdk.ObjectStoragePath` can open. Only a regular file
is
+ exported. Cannot be combined with ``attach_to``: the task that owns an
+ attached sandbox collects its files itself.
+ :param export_conn_id: Airflow connection for the export destinations, or
+ ``None`` for the storage's default credentials. Only meaningful with
+ ``exports``.
+ :param max_export_bytes: Largest file an export will copy. Default 1 GiB.
:param owner: The owner the attached sandbox must carry. Defaults to the
Dag
run the task is part of, which is what a provisioning task in the same
run
stamps with ``SandboxSpec(owner=dag_run_owner(context))``. Set it only
when
@@ -222,7 +255,7 @@ class SandboxToolset(AirflowToolset):
# Rendered, on a copy, by AgentOperator. Deliberately not
``template_fields``, which
# Airflow's templater would render in place wherever the toolset is nested.
- agent_template_fields: Sequence[str] = ("attach_to",)
+ agent_template_fields: Sequence[str] = ("attach_to", "_exports",
"_export_conn_id")
def __init__(
self,
@@ -237,12 +270,16 @@ class SandboxToolset(AirflowToolset):
tool_prefix: str = "",
attach_to: str | None = None,
owner: str | None = None,
+ exports: Mapping[str, str] | None = None,
+ export_conn_id: str | None = None,
+ max_export_bytes: int = 1024 * 1024 * 1024,
) -> None:
_validate_positive_finite(default_command_timeout,
"default_command_timeout")
_validate_positive_finite(max_command_timeout, "max_command_timeout")
_validate_positive_finite(max_output_lines, "max_output_lines")
_validate_positive_finite(max_output_bytes, "max_output_bytes")
_validate_positive_finite(max_read_bytes, "max_read_bytes")
+ _validate_positive_finite(max_export_bytes, "max_export_bytes")
if default_command_timeout > max_command_timeout:
raise ValueError(
f"default_command_timeout ({default_command_timeout}) must not
exceed "
@@ -267,6 +304,11 @@ class SandboxToolset(AirflowToolset):
)
if not attach_to:
raise ValueError("attach_to must be a sandbox handle, not an
empty string.")
+ if exports:
+ raise ValueError(
+ "exports cannot be combined with attach_to: the sandbox
belongs to the task that "
+ "created it, and that task reads out whatever the agent
left behind."
+ )
elif owner is not None:
raise ValueError("owner only applies together with attach_to.")
elif spec is not None and spec.owner is not None:
@@ -277,6 +319,14 @@ class SandboxToolset(AirflowToolset):
"sandbox it provisions itself when the run ends, so an owner
on it would mean nothing; "
"provision the sandbox in a task and pass its handle as
attach_to instead."
)
+ for path, destination in (exports or {}).items():
+ if not path or not destination:
+ raise ValueError(
+ f"exports maps a sandbox path to a destination, and
neither may be empty; got "
+ f"{path!r} -> {destination!r}."
+ )
+ if export_conn_id is not None and not exports:
+ raise ValueError("export_conn_id only applies together with
exports.")
self._backend = backend
self.attach_to = attach_to
# Fixed here, not re-derived from ``attach_to`` later: the templater
rewrites
@@ -296,6 +346,9 @@ class SandboxToolset(AirflowToolset):
self._max_output_lines = int(max_output_lines)
self._max_output_bytes = int(max_output_bytes)
self._max_read_bytes = int(max_read_bytes)
+ self._exports = dict(exports or {})
+ self._export_conn_id = export_conn_id
+ self._max_export_bytes = int(max_export_bytes)
self._tool_prefix = tool_prefix
self._sandbox: str | None = None
self._create_lock = threading.Lock()
@@ -310,6 +363,13 @@ class SandboxToolset(AirflowToolset):
# Tools run only between enter and exit, so nothing outside a run can
provision a
# sandbox that no exit will destroy, or use one it has not claimed.
self._open = False
+ # The exception already being handled when a run entered, if any. The
run was
+ # started from inside that handler, so finding it again at exit says
nothing
+ # about how the run went (see ``_run_failed``).
+ self._handling_at_enter: BaseException | None = None
+ # Set on an instance that has handed its runs to copies (see
``for_run``). It
+ # never holds a sandbox itself, so it has nothing to export or to fail
over.
+ self._forked = False
@property
def id(self) -> str:
@@ -335,6 +395,7 @@ class SandboxToolset(AirflowToolset):
# other's. The backend keys all state by unique sandbox handle, so
# sharing the backend itself is safe. ``type(self)`` so a subclass does
# not silently degrade to this class on every run.
+ self._forked = True
return type(self)(
self._backend,
# Attach mode refuses a spec, and the default one filled in above
is
@@ -348,6 +409,9 @@ class SandboxToolset(AirflowToolset):
tool_prefix=self._tool_prefix,
attach_to=self._attached_handle if self._attach_mode else None,
owner=self._owner,
+ exports=self._exports or None,
+ export_conn_id=self._export_conn_id,
+ max_export_bytes=self._max_export_bytes,
)
@property
@@ -375,14 +439,33 @@ class SandboxToolset(AirflowToolset):
# nothing leaks if the run fails before any tool executes. An attached
one is
# claimed now, so a wrong handle or a held sandbox fails the run
before the
# model has spent anything, and the tool descriptions can state the
lifetime.
+ # Read here, not in the thread: the exception being handled belongs to
this frame.
+ self._handling_at_enter = sys.exc_info()[1]
await asyncio.to_thread(self._open_run)
return self
def _open_run(self) -> None:
if self._attach_mode:
self._attach(self._attached_handle)
+ else:
+ self._check_export_destinations()
self._open = True
+ def _check_export_destinations(self) -> None:
+ """
+ Refuse a rendered destination that is not a URL, before the model
spends anything.
+
+ The destinations are templated, and a value that rendered to nothing
reads as a
+ path relative to the worker's working directory, where the file would
land and
+ the export report success. The same failure shape ``_attached_handle``
guards.
+ """
+ for path, destination in self._exports.items():
+ if not isinstance(destination, str) or "://" not in destination:
+ raise SandboxTerminalError(
+ f"exports[{path!r}] rendered to {destination!r}, which is
not a storage URL such as "
+ "s3://bucket/key or file:///path. Check the template or
the value it reads."
+ )
+
def _attach(self, handle: str) -> None:
owner, holder = self._identity()
backend = self._attachable_backend
@@ -469,19 +552,36 @@ class SandboxToolset(AirflowToolset):
return _Identity(owner=self._owner if self._owner is not None else
run, holder=holder)
async def __aexit__(self, *args: Any) -> bool | None:
- await asyncio.to_thread(self._close)
+ # Read here, not in the thread: the exception being handled belongs to
this frame.
+ run_failed = args[0] is not None or self._run_failed
+ await asyncio.to_thread(self._close, run_failed=run_failed)
return None
+ @property
+ def _run_failed(self) -> bool:
+ """
+ Whether the run this toolset served ended in an exception.
+
+ pydantic-ai exits its toolsets through an exit stack that passes no
exception to
+ them, whether the run succeeded or not, so the exception being handled
at exit is
+ the only record of a failure. One that was already being handled when
the run
+ entered is not it: that run was started from inside an ``except``
block, and it
+ may well have succeeded.
+ """
+ handling = sys.exc_info()[1]
+ return handling is not None and handling is not self._handling_at_enter
+
def __enter__(self) -> Self:
"""Own the sandbox's lifetime from synchronous code, such as a task
running a native agent."""
self._open_run()
return self
def __exit__(self, *args: Any) -> bool | None:
- self._close()
+ # A with statement passes the exception that ended the block, so it is
the answer.
+ self._close(run_failed=args[0] is not None)
return None
- def _close(self) -> None:
+ def _close(self, *, run_failed: bool) -> None:
with self._create_lock:
self._open = False
pending = self._create_future
@@ -506,8 +606,92 @@ class SandboxToolset(AirflowToolset):
if holder is not None:
self._release(self._attached_handle, holder)
return
+ export = bool(self._exports) and not run_failed
if sandbox is None:
+ if export and not self._forked:
+ # The run never called a tool, so it never provisioned a
sandbox, or the
+ # sandbox ended under its last command and nothing replaced
it. Either
+ # way the files it was to leave behind do not exist.
+ raise SandboxTerminalError(
+ "The run ended with no sandbox standing, so none of the
files it was to export "
+ f"exist: {', '.join(repr(path) for path in
self._exports)}."
+ )
return
+ try:
+ if export:
+ self._export(sandbox)
+ finally:
+ self._destroy(sandbox)
+
+ def _export(self, sandbox: str) -> None:
+ """
+ Copy every file in ``exports`` out of the sandbox, failing the task on
the first that cannot be.
+
+ Before teardown, and never best effort: a task that promised a file
and did not
+ deliver it must fail, or its downstream task finds nothing and cannot
tell why.
+
+ Each file is copied to a staging key next to its destination, and only
once every
+ file has been copied are they moved into place. A failed export
therefore leaves
+ every destination as it was, including one an earlier run filled,
rather than
+ holding a truncated copy or a part of this run's set.
+ """
+ staged: list[tuple[ObjectStoragePath, ObjectStoragePath]] = []
+ for path, destination in self._exports.items():
+ target = self._export_target(destination)
+ partial =
target.with_name(f"{target.name}.{uuid.uuid4().hex[:12]}{_PARTIAL_SUFFIX}")
+ try:
+ if isinstance(target.fs, LocalFileSystem):
+ # A per-run destination names a directory nothing has
created yet. Only
+ # locally: on object storage a key needs no parent, and an
s3fs mkdir
+ # can create a bucket.
+ target.parent.mkdir(parents=True, exist_ok=True)
+ staged.append((partial, target))
+ with partial.open("wb") as stream:
+ size = self._backend.export_file(sandbox, path, stream,
max_bytes=self._max_export_bytes)
+ except Exception as e:
+ raise self._export_failed(
+ staged,
+ f"Could not export {path!r} from sandbox {sandbox} "
+ f"on backend {self._backend.name!r} to {destination}:
{_sentence(e)}",
+ ) from e
+ log.info("Exported %s from sandbox %s (%s)", path, sandbox,
format_size(size))
+ published: list[str] = []
+ for partial, target in staged:
+ try:
+ # Within one store: a rename on a local disk, a server-side
copy on object
+ # storage, so the bytes do not pass through the worker a
second time.
+ partial.move(target)
+ except Exception as e:
+ done = f" Already in place: {', '.join(published)}." if
published else ""
+ raise self._export_failed(
+ staged, f"Could not move the export into {target}:
{_sentence(e)}{done}"
+ ) from e
+ published.append(str(target))
+ log.info("Published %s", target)
+
+ @staticmethod
+ def _export_failed(
+ staged: list[tuple[ObjectStoragePath, ObjectStoragePath]], message: str
+ ) -> SandboxTerminalError:
+ """Remove the staging keys of a failed export, and build the error
that names any left behind."""
+ left_behind = []
+ for partial, _ in staged:
+ try:
+ partial.unlink(missing_ok=True)
+ except Exception:
+ log.warning("Could not remove %s after a failed export",
partial, exc_info=True)
+ left_behind.append(str(partial))
+ note = f" Left behind, to delete by hand: {', '.join(left_behind)}."
if left_behind else ""
+ return SandboxTerminalError(f"{message}{note}")
+
+ def _export_target(self, destination: str) -> ObjectStoragePath:
+ # ``conn_id`` only when one is set: on Airflow 3.0 and 3.1 an explicit
None
+ # discards a connection named in the URL itself
(``s3://conn@bucket/key``).
+ if self._export_conn_id is None:
+ return ObjectStoragePath(destination)
+ return ObjectStoragePath(destination, conn_id=self._export_conn_id)
+
+ def _destroy(self, sandbox: str) -> None:
try:
self._backend.destroy(sandbox)
except Exception:
@@ -622,6 +806,17 @@ class SandboxToolset(AirflowToolset):
return self._attach_note or "This sandbox was set up by an earlier
task."
return self._describe_network(self._spec)
+ @property
+ def _export_note(self) -> str:
+ """Which files leave the sandbox when the run ends, so the model
writes them where they are collected."""
+ if not self._exports:
+ return ""
+ paths = ", ".join(self._exports)
+ return (
+ f" When this run ends, these files are copied out of the sandbox
and everything else is "
+ f"discarded, so write your results to them before you finish:
{paths}."
+ )
+
@staticmethod
def _describe_network(spec: SandboxSpec) -> str:
if not spec.block_network:
@@ -663,7 +858,7 @@ class SandboxToolset(AirflowToolset):
name = self._tool_name(base)
description = _DESCRIPTIONS[base]
if base == RUN_COMMAND:
- description = f"{description} {self._network_note}"
+ description = f"{description}
{self._network_note}{self._export_note}"
tool_def = ToolDefinition(
name=name,
description=description,
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
index b49bcaf1a30..7d7baedbcb2 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
@@ -414,6 +414,21 @@ class TestAgentOperatorToolsetTemplating:
assert op.capabilities[0] is capability
+ def test_sandbox_export_destinations_are_rendered_on_a_copy(self):
+ toolset = SandboxToolset(
+ _NoopBackend(),
+ exports={"out/report.csv": "s3://reports/{{ params.customer
}}/report.csv"},
+ export_conn_id="s3_{{ params.customer }}",
+ )
+ op = AgentOperator(task_id="t", prompt="p", llm_conn_id="llm",
toolsets=[toolset])
+
+ op.render_template_fields(self.CONTEXT)
+
+ (rendered,) = op.toolsets
+ assert rendered._exports == {"out/report.csv":
"s3://reports/acme/report.csv"}
+ assert rendered._export_conn_id == "s3_acme"
+ assert toolset._exports == {"out/report.csv": "s3://reports/{{
params.customer }}/report.csv"}
+
def test_only_connection_ids_are_templated(self):
"""allowed_tables is validated and canonicalised in __init__, so
rendering it later
would bypass the fail-closed empty-list check."""
diff --git a/providers/common/ai/tests/unit/common/ai/sandbox/test_base.py
b/providers/common/ai/tests/unit/common/ai/sandbox/test_base.py
index 43dac4b5b06..82e43e34c2a 100644
--- a/providers/common/ai/tests/unit/common/ai/sandbox/test_base.py
+++ b/providers/common/ai/tests/unit/common/ai/sandbox/test_base.py
@@ -17,12 +17,15 @@
from __future__ import annotations
import inspect
+import io
+import os
import subprocess
from types import SimpleNamespace
import pytest
import time_machine
+from airflow.providers.common.ai.sandbox import base
from airflow.providers.common.ai.sandbox.base import (
EXPIRES_AT_TAG,
HOLDER_TAG,
@@ -124,7 +127,7 @@ class TestBackendContract:
assert required == {"create", "run_command", "destroy"}
def test_file_operations_are_overridable_defaults(self):
- for name in ("read_file", "write_file", "list_directory"):
+ for name in ("read_file", "write_file", "list_directory",
"export_file"):
method = getattr(SandboxBackend, name)
assert not getattr(method, "__isabstractmethod__", False)
@@ -253,6 +256,164 @@ class TestDefaultFileOperations:
local.list_directory("s", str(tmp_path / "nope"))
+class TestDefaultExport:
+ """The inherited export_file, against a real shell, with slices small
enough to need several."""
+
+ @pytest.fixture(autouse=True)
+ def _small_slices(self, monkeypatch):
+ monkeypatch.setattr(base, "_EXPORT_CHUNK_BYTES", 64)
+
+ @pytest.mark.parametrize("size", [0, 1, 63, 64, 128, 1000], ids=lambda n:
f"{n}-bytes")
+ def test_copies_the_file_byte_for_byte_across_slices(self, local,
tmp_path, size):
+ blob = os.urandom(size)
+ (tmp_path / "out.bin").write_bytes(blob)
+ dest = io.BytesIO()
+
+ written = local.export_file("s", str(tmp_path / "out.bin"), dest,
max_bytes=10_000)
+
+ assert written == size
+ assert dest.getvalue() == blob
+
+ def test_a_relative_path_resolves_like_read_file(self, local, tmp_path):
+ (tmp_path / "report.csv").write_bytes(b"a,b\n1,2\n")
+ dest = io.BytesIO()
+
+ local.export_file("s", "report.csv", dest, max_bytes=100)
+
+ assert dest.getvalue() == local.read_file("s", "report.csv",
max_bytes=100)
+
+ @pytest.mark.parametrize("name", ["with space.bin", "with'quote.bin",
"semi;colon.bin", "$dollar.bin"])
+ def test_hostile_filenames_are_quoted(self, local, tmp_path, name):
+ (tmp_path / name).write_bytes(b"payload")
+ dest = io.BytesIO()
+
+ local.export_file("s", str(tmp_path / name), dest, max_bytes=100)
+
+ assert dest.getvalue() == b"payload"
+
+ def test_a_missing_file_is_an_error(self, local, tmp_path):
+ with pytest.raises(SandboxError, match="does not exist"):
+ local.export_file("s", str(tmp_path / "nope.bin"), io.BytesIO(),
max_bytes=100)
+
+ def test_a_directory_is_refused(self, local, tmp_path):
+ (tmp_path / "sub").mkdir()
+
+ with pytest.raises(SandboxError, match="is a directory"):
+ local.export_file("s", str(tmp_path / "sub"), io.BytesIO(),
max_bytes=100)
+
+ def
test_a_stream_with_no_size_is_refused_rather_than_read_without_end(self, local):
+ with pytest.raises(SandboxError, match="not a regular file"):
+ local.export_file("s", "/dev/zero", io.BytesIO(), max_bytes=100)
+
+ def
test_a_file_over_the_budget_is_refused_before_anything_is_written(self, local,
tmp_path):
+ (tmp_path / "big.bin").write_bytes(b"x" * 500)
+ dest = io.BytesIO()
+
+ with pytest.raises(SandboxFileTooLargeError) as error:
+ local.export_file("s", str(tmp_path / "big.bin"), dest,
max_bytes=100)
+
+ assert error.value.size_bytes == 500
+ assert dest.getvalue() == b""
+
+ def test_a_file_that_grows_while_it_is_exported_is_an_error(self, local,
tmp_path):
+ target = tmp_path / "growing.bin"
+ target.write_bytes(b"x" * 100)
+ run_command = local.run_command
+
+ def grow_after_the_check(sandbox, command, **kwargs):
+ result = run_command(sandbox, command, **kwargs)
+ if "stat" in command:
+ with target.open("ab") as f:
+ f.write(b"y" * 10)
+ return result
+
+ local.run_command = grow_after_the_check
+
+ with pytest.raises(SandboxError, match="changed while it was
exported"):
+ local.export_file("s", str(target), io.BytesIO(), max_bytes=1000)
+
+ def test_a_file_that_grows_past_the_budget_mid_copy_is_refused(self,
local, tmp_path):
+ target = tmp_path / "growing.bin"
+ target.write_bytes(b"x" * 100)
+ run_command = local.run_command
+
+ def grow_after_the_check(sandbox, command, **kwargs):
+ result = run_command(sandbox, command, **kwargs)
+ if "stat" in command:
+ with target.open("ab") as f:
+ f.write(b"y" * 500)
+ return result
+
+ local.run_command = grow_after_the_check
+
+ with pytest.raises(SandboxFileTooLargeError):
+ local.export_file("s", str(target), io.BytesIO(), max_bytes=200)
+
+ def test_an_export_slower_than_its_deadline_is_ended(self, local,
tmp_path, monkeypatch):
+ # Each slice has its own command timeout, so only the whole-copy
deadline stops a
+ # guest that sends every slice just inside it.
+ (tmp_path / "out.bin").write_bytes(b"x" * 100)
+ monkeypatch.setattr(base, "_export_allowance", lambda max_bytes: 0.0)
+
+ with pytest.raises(SandboxError, match="took longer than"):
+ local.export_file("s", str(tmp_path / "out.bin"), io.BytesIO(),
max_bytes=1000)
+
+ def test_the_deadline_scales_with_the_budget(self):
+ assert base._export_allowance(1024**3) == 1024
+ assert base._export_allowance(1024) == base._FILE_OP_TIMEOUT
+
+ def test_a_truncated_slice_is_an_error_not_a_short_file(self, local,
tmp_path):
+ # A slice cut short decodes cleanly into the wrong bytes, so it must
not be
+ # mistaken for the end of the file.
+ (tmp_path / "out.bin").write_bytes(b"x" * 100)
+ run_command = local.run_command
+
+ def truncate_slices(sandbox, command, **kwargs):
+ result = run_command(sandbox, command, **kwargs)
+ if "tail" in command:
+ return SandboxExecResult(
+ exit_code=0, stdout=result.stdout[:8], stderr="",
stdout_truncated=True
+ )
+ return result
+
+ local.run_command = truncate_slices
+
+ with pytest.raises(SandboxError, match="Could not export"):
+ local.export_file("s", str(tmp_path / "out.bin"), io.BytesIO(),
max_bytes=1000)
+
+ def test_a_sandbox_that_ends_mid_export_is_terminal(self, local, tmp_path):
+ (tmp_path / "out.bin").write_bytes(b"x" * 100)
+ run_command = local.run_command
+
+ def end_on_first_slice(sandbox, command, **kwargs):
+ if "tail" in command:
+ return SandboxExecResult(exit_code=-1, stdout="", stderr="",
sandbox_terminated=True)
+ return run_command(sandbox, command, **kwargs)
+
+ local.run_command = end_on_first_slice
+
+ with pytest.raises(SandboxTerminalError, match="ended"):
+ local.export_file("s", str(tmp_path / "out.bin"), io.BytesIO(),
max_bytes=1000)
+
+ def test_the_size_is_found_among_other_output(self, local, tmp_path):
+ # A vendor CLI can print notices on the same stream as the guest.
+ (tmp_path / "out.bin").write_bytes(b"x" * 10)
+ run_command = local.run_command
+
+ def noisy(sandbox, command, **kwargs):
+ result = run_command(sandbox, command, **kwargs)
+ if "stat" in command:
+ return SandboxExecResult(
+ exit_code=0, stdout=f"Starting
daemon...\n{result.stdout}", stderr=""
+ )
+ return result
+
+ local.run_command = noisy
+ dest = io.BytesIO()
+
+ assert local.export_file("s", str(tmp_path / "out.bin"), dest,
max_bytes=100) == 10
+
+
class TestDagRunOwner:
def test_names_the_dag_and_the_run(self):
context = {"ti": SimpleNamespace(dag_id="my_dag",
run_id="manual__2026-01-01T00:00:00+00:00")}
diff --git
a/providers/common/ai/tests/unit/common/ai/sandbox/test_opensandbox.py
b/providers/common/ai/tests/unit/common/ai/sandbox/test_opensandbox.py
index 0b2dcdc6910..3eaad62682a 100644
--- a/providers/common/ai/tests/unit/common/ai/sandbox/test_opensandbox.py
+++ b/providers/common/ai/tests/unit/common/ai/sandbox/test_opensandbox.py
@@ -17,6 +17,7 @@
from __future__ import annotations
import builtins
+import io
import threading
import time
from types import SimpleNamespace
@@ -29,6 +30,7 @@ pytest.importorskip("opensandbox")
from opensandbox.exceptions import SandboxApiException
from opensandbox.models.sandboxes import NetworkPolicy, NetworkRule
+from airflow.providers.common.ai.sandbox import base
from airflow.providers.common.ai.sandbox.base import (
SandboxError,
SandboxFileTooLargeError,
@@ -670,6 +672,141 @@ class TestFileOperations:
backend.read_file("box-1", "/w/a", max_bytes=0)
+def _stream(*chunks: bytes) -> mock.MagicMock:
+ stream = mock.MagicMock(spec=["__iter__", "close"])
+ stream.__iter__.return_value = iter(chunks)
+ return stream
+
+
+def _entry(entry_type: str | None, size: int) -> SimpleNamespace:
+ return SimpleNamespace(entry_type=entry_type, size=size)
+
+
+class TestExportFile:
+ def test_streams_the_ranged_download_into_the_destination(self):
+ backend, sandbox = _backend_with_sandbox()
+ sandbox.files.get_file_info.return_value = {"/w/out.bin":
_entry("file", 5)}
+ stream = _stream(b"he", b"llo")
+ sandbox.files.read_bytes_stream.return_value = stream
+ dest = io.BytesIO()
+
+ written = backend.export_file("box-1", "/w/out.bin", dest,
max_bytes=100)
+
+ assert written == 5
+ assert dest.getvalue() == b"hello"
+ sandbox.files.read_bytes_stream.assert_called_once_with(
+ "/w/out.bin", chunk_size=101, range_header="bytes=0-100"
+ )
+ stream.close.assert_called_once()
+
+ @pytest.mark.parametrize(
+ ("entry_type", "match"),
+ [("directory", "is a directory"), ("other", "not a regular file")],
+ )
+ def test_refuses_what_is_not_a_regular_file_without_downloading(self,
entry_type, match):
+ backend, sandbox = _backend_with_sandbox()
+ sandbox.files.get_file_info.return_value = {"/w/x": _entry(entry_type,
0)}
+
+ with pytest.raises(SandboxError, match=match):
+ backend.export_file("box-1", "/w/x", io.BytesIO(), max_bytes=100)
+
+ sandbox.files.read_bytes_stream.assert_not_called()
+
+ def test_a_file_over_the_budget_is_refused_without_downloading(self):
+ backend, sandbox = _backend_with_sandbox()
+ sandbox.files.get_file_info.return_value = {"/w/big": _entry("file",
1_000)}
+
+ with pytest.raises(SandboxFileTooLargeError) as error:
+ backend.export_file("box-1", "/w/big", io.BytesIO(), max_bytes=100)
+
+ assert error.value.size_bytes == 1_000
+ sandbox.files.read_bytes_stream.assert_not_called()
+
+ def test_a_missing_file_is_an_error_when_the_sandbox_exists(self):
+ backend, sandbox = _backend_with_sandbox()
+ sandbox.files.get_file_info.side_effect = _api_error(404)
+
+ with pytest.raises(SandboxError, match="does not exist"):
+ backend.export_file("box-1", "/w/missing", io.BytesIO(),
max_bytes=100)
+
+ sandbox.get_info.assert_called_once()
+
+ def test_an_empty_lookup_means_the_file_does_not_exist(self):
+ backend, sandbox = _backend_with_sandbox()
+ sandbox.files.get_file_info.return_value = {}
+
+ with pytest.raises(SandboxError, match="does not exist"):
+ backend.export_file("box-1", "/w/missing", io.BytesIO(),
max_bytes=100)
+
+ def test_a_lookup_that_fails_on_the_server_is_terminal(self):
+ backend, sandbox = _backend_with_sandbox()
+ sandbox.files.get_file_info.side_effect = _api_error(500)
+
+ with pytest.raises(SandboxTerminalError, match="inspect a sandbox
file"):
+ backend.export_file("box-1", "/w/out.bin", io.BytesIO(),
max_bytes=100)
+
+ def
test_a_relative_path_uses_the_entry_the_server_keyed_by_absolute_path(self):
+ backend, sandbox = _backend_with_sandbox()
+ sandbox.files.get_file_info.return_value = {"/workspace/out.bin":
_entry("file", 5)}
+ sandbox.files.read_bytes_stream.return_value = _stream(b"hello")
+ dest = io.BytesIO()
+
+ backend.export_file("box-1", "out.bin", dest, max_bytes=100)
+
+ assert dest.getvalue() == b"hello"
+
+ def test_a_file_that_grows_past_the_budget_mid_copy_is_refused(self):
+ backend, sandbox = _backend_with_sandbox()
+ sandbox.files.get_file_info.return_value = {"/w/out.bin":
_entry("file", 5)}
+ sandbox.files.read_bytes_stream.return_value = _stream(b"x" * 8, b"x"
* 8)
+
+ with pytest.raises(SandboxFileTooLargeError):
+ backend.export_file("box-1", "/w/out.bin", io.BytesIO(),
max_bytes=10)
+
+ def test_a_download_slower_than_its_deadline_is_ended(self, monkeypatch):
+ backend, sandbox = _backend_with_sandbox()
+ sandbox.files.get_file_info.return_value = {"/w/out.bin":
_entry("file", 5)}
+ stream = _stream(b"he", b"llo")
+ sandbox.files.read_bytes_stream.return_value = stream
+ monkeypatch.setattr(base, "_export_allowance", lambda max_bytes: 0.0)
+
+ with pytest.raises(SandboxError, match="took longer than"):
+ backend.export_file("box-1", "/w/out.bin", io.BytesIO(),
max_bytes=100)
+
+ stream.close.assert_called_once()
+
+ def test_a_file_that_changed_size_is_an_error(self):
+ backend, sandbox = _backend_with_sandbox()
+ sandbox.files.get_file_info.return_value = {"/w/out.bin":
_entry("file", 5)}
+ sandbox.files.read_bytes_stream.return_value = _stream(b"hello", b"!")
+
+ with pytest.raises(SandboxError, match="changed while it was
exported"):
+ backend.export_file("box-1", "/w/out.bin", io.BytesIO(),
max_bytes=100)
+
+ def test_a_symlink_is_exported_without_comparing_the_links_own_size(self):
+ backend, sandbox = _backend_with_sandbox()
+ sandbox.files.get_file_info.return_value = {"/w/link":
_entry("symlink", 9)}
+ sandbox.files.read_bytes_stream.return_value = _stream(b"target
contents")
+ dest = io.BytesIO()
+
+ backend.export_file("box-1", "/w/link", dest, max_bytes=100)
+
+ assert dest.getvalue() == b"target contents"
+
+ def test_a_destination_failure_is_not_reported_as_a_sandbox_error(self):
+ backend, sandbox = _backend_with_sandbox()
+ sandbox.files.get_file_info.return_value = {"/w/out.bin":
_entry("file", 5)}
+ stream = _stream(b"hello")
+ sandbox.files.read_bytes_stream.return_value = stream
+ dest = mock.MagicMock(spec=["write"])
+ dest.write.side_effect = OSError("No space left on device")
+
+ with pytest.raises(OSError, match="No space left"):
+ backend.export_file("box-1", "/w/out.bin", dest, max_bytes=100)
+
+ stream.close.assert_called_once()
+
+
class TestGetSandbox:
@mock.patch("opensandbox.SandboxSync.connect", autospec=True)
def test_uncached_sandbox_is_reconnected_and_then_cached(self, connect):
diff --git a/providers/common/ai/tests/unit/common/ai/sandbox/test_sbx.py
b/providers/common/ai/tests/unit/common/ai/sandbox/test_sbx.py
index a34bd6d098b..9cc04b40e48 100644
--- a/providers/common/ai/tests/unit/common/ai/sandbox/test_sbx.py
+++ b/providers/common/ai/tests/unit/common/ai/sandbox/test_sbx.py
@@ -17,13 +17,19 @@
from __future__ import annotations
import base64
+import io
+import os
import subprocess
+import textwrap
+import time
from unittest.mock import patch
import pytest
+from airflow.providers.common.ai.sandbox import base, sbx
from airflow.providers.common.ai.sandbox.base import (
SandboxError,
+ SandboxFileTooLargeError,
SandboxSpec,
SandboxTerminalError,
)
@@ -327,6 +333,106 @@ class TestWriteFileOverride:
backend.write_file("box", "/w/a", b"x")
+class TestExportFileOverride:
+ """
+ sbx overrides export_file to stream one ``exec`` into the destination.
+
+ Driven through a stand-in ``sbx`` binary that runs the guest command on
this host,
+ so the real subprocess, the stream and the watchdog are what is under
test. It
+ prints a daemon notice first, as the real CLI does when it starts its
daemon.
+ """
+
+ @pytest.fixture
+ def local_sbx(self, tmp_path):
+ fake = tmp_path / "sbx"
+ fake.write_text(
+ textwrap.dedent(
+ """\
+ #!/bin/sh
+ [ "$1" = exec ] || exit 2
+ shift 2
+ echo "Starting sandboxd daemon..." >&2
+ [ -n "$FAKE_SBX_STALL" ] && sleep "$FAKE_SBX_STALL"
+ "$@"
+ rc=$?
+ # A later size report stands in for a file that changed
mid-copy.
+ [ -n "$FAKE_SBX_SIZE" ] && printf
'\\nairflow-export-size:%s\\n' "$FAKE_SBX_SIZE" >&2
+ exit $rc
+ """
+ )
+ )
+ fake.chmod(0o755)
+ return SbxSandboxBackend(sbx_path=str(fake),
host_network_policy="deny-all")
+
+ def test_streams_the_file_byte_for_byte(self, local_sbx, tmp_path):
+ blob = os.urandom(3 * 1024 * 1024 + 7)
+ (tmp_path / "model.bin").write_bytes(blob)
+ dest = io.BytesIO()
+
+ written = local_sbx.export_file("box", str(tmp_path / "model.bin"),
dest, max_bytes=10 * 1024 * 1024)
+
+ assert written == len(blob)
+ assert dest.getvalue() == blob
+
+ def test_a_missing_file_is_an_error(self, local_sbx, tmp_path):
+ with pytest.raises(SandboxError, match="does not exist"):
+ local_sbx.export_file("box", str(tmp_path / "nope.bin"),
io.BytesIO(), max_bytes=100)
+
+ def test_a_file_over_the_budget_reports_its_size_despite_the_notice(self,
local_sbx, tmp_path):
+ (tmp_path / "big.bin").write_bytes(b"x" * 500)
+ dest = io.BytesIO()
+
+ with pytest.raises(SandboxFileTooLargeError) as error:
+ local_sbx.export_file("box", str(tmp_path / "big.bin"), dest,
max_bytes=100)
+
+ assert error.value.size_bytes == 500
+ assert dest.getvalue() == b""
+
+ def test_a_destination_that_fails_stops_the_copy(self, local_sbx,
tmp_path):
+ (tmp_path / "out.bin").write_bytes(os.urandom(1024 * 1024))
+
+ class FullDisk(io.RawIOBase):
+ def writable(self):
+ return True
+
+ def write(self, b):
+ raise OSError("No space left on device")
+
+ with pytest.raises(OSError, match="No space left"):
+ local_sbx.export_file("box", str(tmp_path / "out.bin"),
FullDisk(), max_bytes=10 * 1024 * 1024)
+
+ def test_a_file_that_changed_size_while_copied_is_an_error(self,
local_sbx, tmp_path, monkeypatch):
+ (tmp_path / "out.bin").write_bytes(b"x" * 10)
+ monkeypatch.setenv("FAKE_SBX_SIZE", "11")
+
+ with pytest.raises(SandboxError, match="changed while it was
exported"):
+ local_sbx.export_file("box", str(tmp_path / "out.bin"),
io.BytesIO(), max_bytes=100)
+
+ def test_a_transfer_that_keeps_trickling_is_ended_by_the_deadline(self,
local_sbx, tmp_path, monkeypatch):
+ # The stall clock is far off, so only the whole-copy deadline can end
this one.
+ (tmp_path / "out.bin").write_bytes(b"x")
+ monkeypatch.setenv("FAKE_SBX_STALL", "30")
+ monkeypatch.setattr(sbx, "_EXPORT_STALL_TIMEOUT", 3600.0)
+ monkeypatch.setattr(base, "_export_allowance", lambda max_bytes: 0.5)
+ start = time.monotonic()
+
+ with pytest.raises(SandboxError, match="took longer than"):
+ local_sbx.export_file("box", str(tmp_path / "out.bin"),
io.BytesIO(), max_bytes=100)
+
+ assert time.monotonic() - start < 10
+
+ def test_a_stalled_transfer_is_ended(self, local_sbx, tmp_path,
monkeypatch):
+ (tmp_path / "out.bin").write_bytes(b"x")
+ monkeypatch.setenv("FAKE_SBX_STALL", "30")
+ monkeypatch.setattr(sbx, "_EXPORT_STALL_TIMEOUT", 0.5)
+ start = time.monotonic()
+
+ with pytest.raises(SandboxError, match="stalled"):
+ local_sbx.export_file("box", str(tmp_path / "out.bin"),
io.BytesIO(), max_bytes=100)
+
+ assert time.monotonic() - start < 10
+
+
class TestDestroy:
def test_is_idempotent_when_the_sandbox_is_already_gone(self, backend):
with patch.object(backend, "_run_cli", autospec=True,
return_value=_completed(returncode=1)):
diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py
b/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py
index 9759b10ec96..48d9f7edc5c 100644
--- a/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py
+++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py
@@ -49,6 +49,7 @@ from airflow.providers.common.ai.sandbox.base import (
from airflow.providers.common.ai.tools import ToolCallError
from airflow.providers.common.ai.toolsets import sandbox as sandbox_module
from airflow.providers.common.ai.toolsets.sandbox import SandboxToolset
+from airflow.providers.common.compat.sdk import ObjectStoragePath
from unit.common.ai.sandbox.fake_tags import InMemoryTagStore
@@ -79,6 +80,12 @@ class _RecordingBackend(SandboxBackend):
self.create_error = create_error
self.read_payload = b""
self.read_error: Exception | None = None
+ # What the agent left in each sandbox, for export_file to hand back.
+ self.files: dict[str, bytes] = {}
+ # (sandbox, path, max_bytes, whether the sandbox was already destroyed)
+ self.exported: list[tuple[str, str, int, bool]] = []
+ # Paths whose copy writes some bytes and then fails, as a copy cut off
partway does.
+ self.fail_partway: set[str] = set()
def create(self, *, spec: SandboxSpec | None = None) -> str:
self.created.append(spec)
@@ -105,6 +112,16 @@ class _RecordingBackend(SandboxBackend):
def list_directory(self, sandbox, path):
return list(self.entries)
+ def export_file(self, sandbox, path, dest, *, max_bytes):
+ self.exported.append((sandbox, path, max_bytes, sandbox in
self.destroyed))
+ if path not in self.files:
+ raise SandboxError(f"{path!r} does not exist in the sandbox, or is
not readable.")
+ if path in self.fail_partway:
+ dest.write(self.files[path][:1])
+ raise SandboxError(f"the sandbox stopped sending {path!r}")
+ dest.write(self.files[path])
+ return len(self.files[path])
+
def destroy(self, sandbox):
self.destroyed.append(sandbox)
if self.destroy_error is not None:
@@ -157,6 +174,7 @@ class TestInit:
"max_output_lines",
"max_output_bytes",
"max_read_bytes",
+ "max_export_bytes",
],
)
def test_rejects_non_positive_bounds(self, field, bad):
@@ -779,6 +797,280 @@ class TestForRun:
assert isinstance(forked, CustomToolset)
+class TestExports:
+ @pytest.mark.parametrize(
+ ("kwargs", "match"),
+ [
+ ({"exports": {"": "file:///tmp/x"}}, "neither may be empty"),
+ ({"exports": {"out.bin": ""}}, "neither may be empty"),
+ ({"export_conn_id": "aws"}, "export_conn_id only applies together
with exports"),
+ ],
+ )
+ def test_constructor_refuses_a_shape_that_cannot_work(self, kwargs, match):
+ with pytest.raises(ValueError, match=match):
+ SandboxToolset(_RecordingBackend(), **kwargs)
+
+ def test_an_attached_sandbox_cannot_export(self):
+ # The task that owns an attached sandbox reads out what the agent left.
+ with pytest.raises(ValueError, match="exports cannot be combined with
attach_to"):
+ SandboxToolset(_AttachableRecordingBackend(), attach_to="sb-1",
exports={"a": "file:///tmp/a"})
+
+ @pytest.mark.asyncio
+ async def
test_a_finished_run_copies_each_file_out_before_the_sandbox_is_destroyed(self,
tmp_path):
+ backend = _RecordingBackend()
+ backend.files = {"out/report.parquet": b"PAR1...", "chart.png":
b"\x89PNG"}
+ ts = SandboxToolset(
+ backend,
+ exports={
+ "out/report.parquet": f"file://{tmp_path}/report.parquet",
+ "chart.png": f"file://{tmp_path}/chart.png",
+ },
+ max_export_bytes=1234,
+ )
+
+ async with ts:
+ await _call(ts, "run_command", {"command": "make report"})
+
+ assert backend.exported == [
+ ("box-1", "out/report.parquet", 1234, False),
+ ("box-1", "chart.png", 1234, False),
+ ]
+ assert backend.destroyed == ["box-1"]
+ assert (tmp_path / "report.parquet").read_bytes() == b"PAR1..."
+ assert (tmp_path / "chart.png").read_bytes() == b"\x89PNG"
+ assert sorted(f.name for f in tmp_path.iterdir()) == ["chart.png",
"report.parquet"]
+
+ @pytest.mark.asyncio
+ async def
test_a_failed_run_exports_nothing_and_still_destroys_the_sandbox(self,
tmp_path):
+ backend = _RecordingBackend()
+ backend.files = {"out.bin": b"half done"}
+ ts = SandboxToolset(backend, exports={"out.bin":
f"file://{tmp_path}/out.bin"})
+
+ async def run_that_fails():
+ async with ts:
+ await _call(ts, "run_command", {"command": "x"})
+ raise RuntimeError("agent blew up")
+
+ with pytest.raises(RuntimeError, match="agent blew up"):
+ await run_that_fails()
+
+ assert backend.exported == []
+ assert backend.destroyed == ["box-1"]
+ assert not (tmp_path / "out.bin").exists()
+
+ @pytest.mark.asyncio
+ async def test_a_file_that_cannot_be_exported_fails_the_task(self,
tmp_path):
+ backend = _RecordingBackend()
+ ts = SandboxToolset(backend, exports={"out.bin":
f"file://{tmp_path}/out.bin"})
+
+ with pytest.raises(SandboxTerminalError, match="Could not export
'out.bin' from sandbox box-1"):
+ async with ts:
+ await _call(ts, "run_command", {"command": "forgot to write
it"})
+
+ assert backend.destroyed == ["box-1"]
+ assert list(tmp_path.iterdir()) == []
+
+ @pytest.mark.asyncio
+ async def
test_a_copy_that_fails_partway_leaves_the_destination_as_it_was(self, tmp_path):
+ (tmp_path / "out.bin").write_bytes(b"ORIGINAL")
+ backend = _RecordingBackend()
+ backend.files = {"out.bin": b"new contents"}
+ backend.fail_partway = {"out.bin"}
+ ts = SandboxToolset(backend, exports={"out.bin":
f"file://{tmp_path}/out.bin"})
+
+ with pytest.raises(SandboxTerminalError, match="stopped sending"):
+ async with ts:
+ await _call(ts, "run_command", {"command": "x"})
+
+ assert (tmp_path / "out.bin").read_bytes() == b"ORIGINAL"
+ assert [f.name for f in tmp_path.iterdir()] == ["out.bin"]
+
+ @pytest.mark.asyncio
+ async def test_for_run_carries_the_export_connection_across(self):
+ base = SandboxToolset(
+ _RecordingBackend(), exports={"out.bin": "s3://b/out.bin"},
export_conn_id="reports"
+ )
+
+ forked = await base.for_run(_ctx())
+
+ with patch.object(sandbox_module, "ObjectStoragePath", autospec=True)
as path_cls:
+ forked._export_target("s3://b/out.bin")
+
+ path_cls.assert_called_once_with("s3://b/out.bin", conn_id="reports")
+
+ @pytest.mark.asyncio
+ async def test_a_failed_export_changes_none_of_the_destinations(self,
tmp_path):
+ # Nothing is moved into place until every file has been copied, so a
consumer
+ # that runs whatever the outcome never finds half of a set.
+ (tmp_path / "first.bin").write_bytes(b"ORIGINAL")
+ backend = _RecordingBackend()
+ backend.files = {"first.bin": b"one"}
+ ts = SandboxToolset(
+ backend,
+ exports={
+ "first.bin": f"file://{tmp_path}/first.bin",
+ "second.bin": f"file://{tmp_path}/second.bin",
+ },
+ )
+
+ with pytest.raises(SandboxTerminalError, match="Could not export
'second.bin'"):
+ async with ts:
+ await _call(ts, "run_command", {"command": "x"})
+
+ assert (tmp_path / "first.bin").read_bytes() == b"ORIGINAL"
+ assert [f.name for f in tmp_path.iterdir()] == ["first.bin"]
+
+ @pytest.mark.asyncio
+ async def test_a_publish_that_fails_names_what_is_already_in_place(self,
tmp_path):
+ backend = _RecordingBackend()
+ backend.files = {"first.bin": b"one", "second.bin": b"two"}
+ ts = SandboxToolset(
+ backend,
+ exports={
+ "first.bin": f"file://{tmp_path}/first.bin",
+ "second.bin": f"file://{tmp_path}/second.bin",
+ },
+ )
+ move = ObjectStoragePath.move
+
+ def fail_the_second(self, path, **kwargs):
+ if str(path).endswith("second.bin"):
+ raise OSError("permission denied")
+ return move(self, path, **kwargs)
+
+ with patch.object(ObjectStoragePath, "move", autospec=True,
side_effect=fail_the_second):
+ with pytest.raises(
+ SandboxTerminalError, match=r"second\.bin: permission
denied.*Already in place: .*first\.bin"
+ ):
+ async with ts:
+ await _call(ts, "run_command", {"command": "x"})
+
+ assert sorted(f.name for f in tmp_path.iterdir()) == ["first.bin"]
+
+ @pytest.mark.asyncio
+ async def test_a_copy_that_cannot_be_removed_is_named_in_the_error(self,
tmp_path, caplog):
+ backend = _RecordingBackend()
+ ts = SandboxToolset(backend, exports={"out.bin":
f"file://{tmp_path}/out.bin"})
+
+ with patch.object(ObjectStoragePath, "unlink", autospec=True,
side_effect=OSError("denied")):
+ with pytest.raises(
+ SandboxTerminalError, match=r"Left behind, to delete by hand:
.*out\.bin\.\w+\.partial"
+ ):
+ async with ts:
+ await _call(ts, "run_command", {"command": "x"})
+
+ assert "Could not remove" in caplog.text
+
+ @pytest.mark.asyncio
+ async def test_a_local_destination_gets_its_directories(self, tmp_path):
+ backend = _RecordingBackend()
+ backend.files = {"out.bin": b"payload"}
+ ts = SandboxToolset(backend, exports={"out.bin":
f"file://{tmp_path}/run-1/nested/out.bin"})
+
+ async with ts:
+ await _call(ts, "run_command", {"command": "x"})
+
+ assert (tmp_path / "run-1" / "nested" / "out.bin").read_bytes() ==
b"payload"
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("rendered", [None, "", "None",
"relative/out.bin"])
+ async def test_a_destination_that_is_not_a_url_fails_before_the_run(self,
rendered):
+ backend = _RecordingBackend()
+ ts = SandboxToolset(backend, exports={"out.bin": "{{ var.value.dest
}}"})
+ # What the templater writes after construction.
+ ts._exports = {"out.bin": rendered}
+
+ with pytest.raises(SandboxTerminalError, match="not a storage URL"):
+ async with ts:
+ pass
+
+ assert backend.created == []
+
+ def test_the_connection_is_not_passed_when_none_is_set(self):
+ # An explicit conn_id=None discards a connection named in the URL on
Airflow 3.0/3.1.
+ ts = SandboxToolset(_RecordingBackend(), exports={"out.bin":
"s3://conn@bucket/out.bin"})
+
+ with patch.object(sandbox_module, "ObjectStoragePath", autospec=True)
as path_cls:
+ ts._export_target("s3://conn@bucket/out.bin")
+
+ path_cls.assert_called_once_with("s3://conn@bucket/out.bin")
+
+ @pytest.mark.asyncio
+ async def
test_a_teardown_failure_after_a_good_export_does_not_fail_the_task(self,
tmp_path, caplog):
+ backend = _RecordingBackend(destroy_error=RuntimeError("daemon went
away"))
+ backend.files = {"out.bin": b"payload"}
+ ts = SandboxToolset(backend, exports={"out.bin":
f"file://{tmp_path}/out.bin"})
+
+ async with ts:
+ await _call(ts, "run_command", {"command": "x"})
+
+ assert (tmp_path / "out.bin").read_bytes() == b"payload"
+ assert "Failed to destroy sandbox box-1" in caplog.text
+
+ @pytest.mark.asyncio
+ async def
test_a_run_that_never_used_its_sandbox_cannot_deliver_its_files(self, tmp_path):
+ backend = _RecordingBackend()
+ ts = SandboxToolset(backend, exports={"out.bin":
f"file://{tmp_path}/out.bin"})
+
+ with pytest.raises(SandboxTerminalError, match="no sandbox
standing.*'out.bin'"):
+ async with ts:
+ pass
+
+ assert backend.created == []
+
+ @pytest.mark.asyncio
+ async def test_the_model_is_told_which_files_are_collected(self):
+ with_exports = SandboxToolset(_RecordingBackend(),
exports={"out/report.csv": "file:///tmp/r.csv"})
+ without = SandboxToolset(_RecordingBackend())
+
+ exported = (await
with_exports.get_tools(_ctx()))["run_command"].tool_def.description
+ plain = (await
without.get_tools(_ctx()))["run_command"].tool_def.description
+
+ assert "copied out of the sandbox" in exported
+ assert "out/report.csv" in exported
+ assert "copied out" not in plain
+
+ @pytest.mark.asyncio
+ async def test_for_run_carries_the_exports_across(self, tmp_path):
+ backend = _RecordingBackend()
+ backend.files = {"out.bin": b"payload"}
+ base = SandboxToolset(backend, exports={"out.bin":
f"file://{tmp_path}/out.bin"}, max_export_bytes=99)
+
+ forked = await base.for_run(_ctx())
+ async with forked:
+ await _call(forked, "run_command", {"command": "x"})
+
+ assert backend.exported == [("box-1", "out.bin", 99, False)]
+ assert (tmp_path / "out.bin").read_bytes() == b"payload"
+
+ def test_a_with_block_exports_for_a_native_agent(self, tmp_path):
+ backend = _RecordingBackend()
+ backend.files = {"out.bin": b"payload"}
+
+ with SandboxToolset(backend, exports={"out.bin":
f"file://{tmp_path}/out.bin"}) as sandbox:
+ run_command = {tool.name: tool for tool in
sandbox.airflow_tools()}["run_command"]
+ asyncio.run(run_command.call({"command": "ls"}))
+
+ assert (tmp_path / "out.bin").read_bytes() == b"payload"
+ assert backend.destroyed == ["box-1"]
+
+ def test_a_native_agent_that_raises_exports_nothing(self, tmp_path):
+ backend = _RecordingBackend()
+ backend.files = {"out.bin": b"payload"}
+
+ def native_run_that_fails():
+ with SandboxToolset(backend, exports={"out.bin":
f"file://{tmp_path}/out.bin"}) as sandbox:
+ run_command = {tool.name: tool for tool in
sandbox.airflow_tools()}["run_command"]
+ asyncio.run(run_command.call({"command": "ls"}))
+ raise RuntimeError("native agent failed")
+
+ with pytest.raises(RuntimeError, match="native agent failed"):
+ native_run_that_fails()
+
+ assert backend.exported == []
+ assert backend.destroyed == ["box-1"]
+
+
class TestAttachMode:
"""
``attach_to``: use a sandbox another task provisioned, and leave it
standing.
@@ -1205,6 +1497,70 @@ class TestOutsideAnAgentRun:
class TestInsideAPydanticAIRun:
+ @staticmethod
+ def _model_calling(tool: str):
+ def model(messages, info):
+ if any(isinstance(p, ToolReturnPart) for m in messages for p in
m.parts):
+ return ModelResponse(parts=[TextPart("done")])
+ return ModelResponse(parts=[ToolCallPart(tool, {"command": "ls"},
tool_call_id="c")])
+
+ return FunctionModel(model)
+
+ def test_a_successful_run_exports_its_files(self, tmp_path):
+ backend = _RecordingBackend()
+ backend.files = {"out.bin": b"payload"}
+ sandbox = SandboxToolset(backend, exports={"out.bin":
f"file://{tmp_path}/out.bin"})
+
+ Agent(self._model_calling("box_run_command"),
toolsets=[sandbox.prefixed("box")]).run_sync("go")
+
+ assert backend.exported == [("box-1", "out.bin", 1024**3, False)]
+ assert (tmp_path / "out.bin").read_bytes() == b"payload"
+
+ def test_a_run_started_inside_an_except_block_still_exports(self,
tmp_path):
+ # The exception already being handled when the run began is not the
run failing.
+ backend = _RecordingBackend()
+ backend.files = {"out.bin": b"payload"}
+ sandbox = SandboxToolset(backend, exports={"out.bin":
f"file://{tmp_path}/out.bin"})
+ agent = Agent(self._model_calling("run_command"), toolsets=[sandbox])
+
+ try:
+ raise KeyError("cache miss")
+ except KeyError:
+ result = agent.run_sync("go")
+
+ assert result.output == "done"
+ assert (tmp_path / "out.bin").read_bytes() == b"payload"
+
+ def
test_an_agent_held_open_exports_from_its_run_and_not_again_at_exit(self,
tmp_path):
+ # ``async with agent`` enters the toolset the author built; runs use
copies of it.
+ backend = _RecordingBackend()
+ backend.files = {"out.bin": b"payload"}
+ sandbox = SandboxToolset(backend, exports={"out.bin":
f"file://{tmp_path}/out.bin"})
+ agent = Agent(self._model_calling("run_command"), toolsets=[sandbox])
+
+ async def held_open():
+ async with agent:
+ await agent.run("go")
+
+ asyncio.run(held_open())
+
+ assert backend.exported == [("box-1", "out.bin", 1024**3, False)]
+ assert (tmp_path / "out.bin").read_bytes() == b"payload"
+
+ def test_a_failed_run_exports_nothing(self, tmp_path):
+ # pydantic-ai closes its toolsets through an exit stack that passes
them no
+ # exception, so the toolset has to find the failure itself.
+ backend =
_RecordingBackend(run_error=SandboxTerminalError("credentials rejected"))
+ backend.files = {"out.bin": b"payload"}
+ sandbox = SandboxToolset(backend, exports={"out.bin":
f"file://{tmp_path}/out.bin"})
+
+ with pytest.raises(SandboxTerminalError, match="credentials rejected"):
+ Agent(self._model_calling("run_command"),
toolsets=[sandbox]).run_sync("go")
+
+ assert backend.exported == []
+ assert backend.destroyed == ["box-1"]
+ assert not (tmp_path / "out.bin").exists()
+
def test_the_run_opens_the_sandbox_and_destroys_it_when_it_ends(self):
backend = _RecordingBackend()
