This is an automated email from the ASF dual-hosted git repository.
eladkal 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 e25f413ff61 Let the Modal sandbox backend read credentials from a
modal connection (#74302)
e25f413ff61 is described below
commit e25f413ff61440c0abae1080f8ed733ab8c9fb92
Author: Kaxil Naik <[email protected]>
AuthorDate: Tue Oct 6 08:10:19 2026 +0100
Let the Modal sandbox backend read credentials from a modal connection
(#74302)
* Let the Modal sandbox backend read credentials from a modal connection
ModalSandboxBackend could only use ambient Modal credentials
(MODAL_TOKEN_ID /
MODAL_TOKEN_SECRET or ~/.modal.toml), so a deployment that keeps secrets in
a
secrets backend could not use the sandbox toolset on Modal without putting a
token in the worker environment. It now takes modal_conn_id (default
modal_default) and builds its Modal client through the Modal provider's
ModalHook, which falls back to the ambient credentials when the default
connection does not exist.
The modal extra now installs apache-airflow-providers-modal, which needs
Airflow 3, so the Modal backend is gated on Airflow 3 the way the skills and
git extras already are.
* Route Modal credentials through one accessor and trim repeated docs
Build the Modal client in one place, under a lock, so concurrent first
calls on a
shared backend cannot each open a client, and translate credential errors
there
instead of at every call site. Point the docs at the Modal connection page
for the
credential precedence rather than restating it, and drop tests that
re-checked
ModalHook rules the Modal provider already tests.
---------
Co-authored-by: Rahul Vats <[email protected]>
---
providers/common/ai/docs/index.rst | 3 +-
providers/common/ai/docs/installation.rst | 5 +-
providers/common/ai/docs/sandbox/backends.rst | 12 +-
providers/common/ai/docs/sandbox/configuration.rst | 8 +-
providers/common/ai/docs/sandbox/index.rst | 33 ++---
providers/common/ai/pyproject.toml | 8 +-
.../airflow/providers/common/ai/sandbox/modal.py | 80 ++++++++---
.../ai/tests/unit/common/ai/sandbox/fake_modal.py | 47 ++++++-
.../ai/tests/unit/common/ai/sandbox/test_modal.py | 148 +++++++++++++++++++--
uv.lock | 4 +
10 files changed, 293 insertions(+), 55 deletions(-)
diff --git a/providers/common/ai/docs/index.rst
b/providers/common/ai/docs/index.rst
index fb2d2d62caa..d6a32a58176 100644
--- a/providers/common/ai/docs/index.rst
+++ b/providers/common/ai/docs/index.rst
@@ -202,6 +202,7 @@ Dependent package
============================================================================================================
==============
`apache-airflow-providers-common-sql
<https://airflow.apache.org/docs/apache-airflow-providers-common-sql>`_
``common.sql``
`apache-airflow-providers-git
<https://airflow.apache.org/docs/apache-airflow-providers-git>`_
``git``
+`apache-airflow-providers-modal
<https://airflow.apache.org/docs/apache-airflow-providers-modal>`_
``modal``
============================================================================================================
==============
Optional dependencies
@@ -224,7 +225,7 @@ Extra Dependencies
``openai`` ``pydantic-ai-slim[openai]>=2.33.0``, ``openai>=2.47.0``
``typesafe`` ``typesafe-sdk>=0.6.0``
``mcp`` ``pydantic-ai-slim[mcp]>=2.33.0``
-``modal`` ``modal>=1.5.2``
+``modal`` ``apache-airflow-providers-modal``, ``modal>=1.5.2``
``opensandbox`` ``opensandbox>=1.1.0``
``code-mode`` ``pydantic-ai-harness[codemode]>=0.3.0``
``shields`` ``pydantic-ai-shields>=0.3.4``
diff --git a/providers/common/ai/docs/installation.rst
b/providers/common/ai/docs/installation.rst
index a7227e724d6..70b2193e0d3 100644
--- a/providers/common/ai/docs/installation.rst
+++ b/providers/common/ai/docs/installation.rst
@@ -73,6 +73,9 @@ The provider runs on Airflow 2.11, but some features need a
newer Airflow versio
- Needs
* - The ``skills`` and ``git`` extras (``apache-airflow-providers-git``
needs Airflow 3)
- Airflow 3.0
+ * - The ``modal`` extra and the Modal sandbox backend
(``apache-airflow-providers-modal``
+ needs Airflow 3)
+ - Airflow 3.0
* - :doc:`Approval gates <approval_gates>` and :doc:`HITL review
<hitl_review>`
- Airflow 3.1
* - The **Model** field in the connection form; on older Airflow versions
put the model in
@@ -104,7 +107,7 @@ On Airflow 2.11 the operators, decorators, hooks and
toolsets run as they do on
``apache-airflow-providers-common-sql`` to releases older than this provider
needs.
Installing Airflow 2.11.0 without its constraints can also pull in a
``universal-pathlib``
0.3 release, which Airflow 2's ``ObjectStoragePath`` rejects; 2.11.1 and
later cap it.
- Leave out the ``skills`` and ``git`` extras: they need Airflow 3, and without
+ Leave out the ``skills``, ``git`` and ``modal`` extras: they need Airflow 3,
and without
constraints ``pip`` upgrades Airflow to satisfy them.
* Python 3.10 to 3.12: the provider needs 3.10 or later, and Airflow 2.11
supports up to 3.12.
diff --git a/providers/common/ai/docs/sandbox/backends.rst
b/providers/common/ai/docs/sandbox/backends.rst
index a5a7b40472d..acf3d453938 100644
--- a/providers/common/ai/docs/sandbox/backends.rst
+++ b/providers/common/ai/docs/sandbox/backends.rst
@@ -34,11 +34,19 @@ provider, **this is the managed one to use in production**,
and with
:ref:`OpenSandbox <sandbox-backend-opensandbox>` one of the two that run on
Kubernetes: nothing has to be installed on the worker, model-written code never
executes on the worker host, and Modal reclaims a sandbox at its own lifetime
-whether or not the worker survives. It needs the ``modal`` extra and ambient
-credentials, as under :ref:`Quick start <sandbox-quick-start>`.
+whether or not the worker survives. It needs the ``modal`` extra and Modal
+credentials, from a ``modal`` connection or the worker environment, as under
+:ref:`Quick start <sandbox-quick-start>`.
Constructor parameters:
+- ``modal_conn_id``: ``modal`` connection the token and, optionally, the Modal
+ ``environment`` come from. Default ``"modal_default"``; ``None`` uses the
worker's
+ credentials without looking for a connection. How a missing or partial
connection
+ resolves is on the :ref:`Modal connection page <howto/connection:modal>`. A
credential
+ problem fails the task, except in the toolset's own teardown, which logs it
so a
+ finished run is not failed. The connection type comes from the Modal
provider, which
+ the ``modal`` extra installs and which needs Airflow 3.
- ``image``: Registry tag for the sandbox image, or a prepared ``modal.Image``
carrying pre-installed packages. Default ``"python:3.12-slim"``.
- ``app_name``: Modal app the sandboxes are created under. Default
diff --git a/providers/common/ai/docs/sandbox/configuration.rst
b/providers/common/ai/docs/sandbox/configuration.rst
index 2647defde51..2a5652720d2 100644
--- a/providers/common/ai/docs/sandbox/configuration.rst
+++ b/providers/common/ai/docs/sandbox/configuration.rst
@@ -137,8 +137,9 @@ Credentials
``SandboxSpec.env`` is the only way in. Airflow never populates it: no
connection,
variable or worker environment variable reaches a sandbox unless you name it
there, and the credential that *provisions* the sandbox never enters it either.
-Modal's token stays on the worker and is used by the client, so code running
-inside cannot call Modal as you or create further sandboxes.
+Modal's token, whether it comes from the ``modal`` connection or the worker
+environment, stays in the task and is used by the client, so code running
inside
+cannot call Modal as you or create further sandboxes.
Before you put a real secret in ``env``, four things are true of it.
@@ -266,7 +267,8 @@ 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. Two things cannot be
-done inside that shape: a credential cannot come from a connection, and a
second
+done inside that shape: a credential for the code inside the sandbox
+(``SandboxSpec.env``) 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
diff --git a/providers/common/ai/docs/sandbox/index.rst
b/providers/common/ai/docs/sandbox/index.rst
index 1d87a27f20e..9022e591f7b 100644
--- a/providers/common/ai/docs/sandbox/index.rst
+++ b/providers/common/ai/docs/sandbox/index.rst
@@ -107,15 +107,19 @@ is the only change between this Dag and production:
:start-after: [START howto_sandbox_agent_local]
:end-before: [END howto_sandbox_agent_local]
-Install the Modal extra and authenticate as you would for Modal's own CLI. On a
-worker, set ``MODAL_TOKEN_ID`` and ``MODAL_TOKEN_SECRET`` in the environment
-instead; nothing is read until the first sandbox is created, so a Dag file that
-constructs the backend parses without credentials present:
+Install the Modal extra, which also installs the Modal provider, and create a
+``modal`` connection with the Modal token id as its login and the token secret
as
+its password. The backend reads ``modal_default`` unless you pass
``modal_conn_id``;
+without that connection it uses the worker's own Modal credentials, as the
+:ref:`Modal connection page <howto/connection:modal>` describes. Nothing is
read until
+the first sandbox is created, so a Dag file that constructs the backend parses
without
+credentials present:
.. code-block:: bash
pip install 'apache-airflow-providers-common-ai[modal]'
- modal token new # writes ~/.modal.toml
+ airflow connections add modal_default --conn-type modal \
+ --conn-login "$MODAL_TOKEN_ID" --conn-password "$MODAL_TOKEN_SECRET"
What it is for, in practice
---------------------------
@@ -425,14 +429,13 @@ system tests, ``example_sandbox_toolset_sbx.py`` and
``example_sandbox_toolset_modal.py``, run against a real backend and are
reachable from the System Tests entry in the sidebar.
-**Credentials and where it runs.** This route does not end at an
-Airflow connection. Airflow puts none of its context, connections, variables or
-worker environment into the sandbox; only what you pass through
+**Credentials and where it runs.** Airflow puts none of its context,
connections,
+variables or worker environment into the sandbox; only what you pass through
:class:`~airflow.providers.common.ai.sandbox.SandboxSpec` goes in, and the
-credential that provisions the sandbox never enters it. Authorization to the
-backend sits outside Airflow: ``sbx login`` and ``sbx policy init`` on the
-machine for ``sbx``, or ``MODAL_TOKEN_ID`` and ``MODAL_TOKEN_SECRET`` on the
-worker for Modal. Work runs in a per-run microVM on the worker host with
+credential that provisions the sandbox never enters it. For Modal that
credential
+is a ``modal`` connection (``modal_conn_id``; see the
+:ref:`Modal connection page <howto/connection:modal>`). For ``sbx`` it sits
outside Airflow:
+``sbx login`` on the machine. Work runs in a per-run microVM on the worker
host with
``sbx``, or off the worker entirely in Modal's infrastructure. Its tool calls
act as barriers, as they do for the other
routes that build their own tools; see :ref:`toolset-call-barriers`.
@@ -486,9 +489,9 @@ is the list to read before designing a Dag around an agent
with a sandbox.
- **A file the agent built leaves only through a task**, never through the
model's context, which is text-only and capped.
:ref:`Getting a result out <sandbox-results>`.
-- **A credential comes from a connection only when a task provisions the
- sandbox**; the toolset's own spec is fixed at parse time, and anything
injected
- is readable by the model. :ref:`Credentials <sandbox-credentials>`.
+- **A credential handed to the code inside the sandbox comes from a connection
+ only when a task provisions the sandbox**; the toolset's own spec is fixed at
+ parse time, and anything injected is readable by the model.
:ref:`Credentials <sandbox-credentials>`.
- **Cannot be combined with** ``durable=True``, and with
``enable_hitl_review=True``
only when the sandbox is task-owned; ``AgentOperator`` raises otherwise.
:ref:`Lifecycle <sandbox-lifecycle>`.
diff --git a/providers/common/ai/pyproject.toml
b/providers/common/ai/pyproject.toml
index 18d50b3033e..c1663c2a81b 100644
--- a/providers/common/ai/pyproject.toml
+++ b/providers/common/ai/pyproject.toml
@@ -104,8 +104,11 @@ dependencies = [
"mcp" = ["pydantic-ai-slim[mcp]>=2.33.0"]
# Modal-hosted sandboxes for SandboxToolset, the backend to use where the
worker host
# cannot run microVMs -- Kubernetes, most notably. Needs no host-level
installation,
-# only Modal credentials.
-"modal" = ["modal>=1.5.2"]
+# only Modal credentials, which the Modal provider's `modal` connection
supplies.
+"modal" = [
+ "apache-airflow-providers-modal",
+ "modal>=1.5.2",
+]
"opensandbox" = ["opensandbox>=1.1.0"]
# Code mode: collapse tool calls into a single `run_code` tool that the model
# drives by writing Python, executed in the Monty sandbox (pydantic-monty).
@@ -163,6 +166,7 @@ dev = [
"apache-airflow-providers-common-compat",
"apache-airflow-providers-common-sql",
"apache-airflow-providers-git",
+ "apache-airflow-providers-modal",
"apache-airflow-providers-standard",
# Additional devel dependencies (do not remove this line and add extra
development dependencies)
"sqlglot>=30.0.0",
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/sandbox/modal.py
b/providers/common/ai/src/airflow/providers/common/ai/sandbox/modal.py
index c67ae080857..96de4ccd917 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/sandbox/modal.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/sandbox/modal.py
@@ -31,6 +31,9 @@ from typing import TYPE_CHECKING, Literal, TypedDict
try:
import modal
+
+ from airflow.providers.modal.exceptions import ModalConnectionError
+ from airflow.providers.modal.hooks.modal import ModalHook
except ModuleNotFoundError as e:
from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
@@ -54,6 +57,7 @@ from airflow.providers.common.ai.sandbox.base import (
encode_network_policy,
is_sandbox_handle,
)
+from airflow.providers.common.compat.sdk import AirflowNotFoundException
if TYPE_CHECKING:
from collections.abc import Mapping, Sequence
@@ -106,6 +110,15 @@ _ALLOWED_HOSTNAME = re.compile(
EgressEnforcement = Literal["strict", "sni"]
+# What reading credentials from the connection can raise before Modal is
reached at all: a
+# named connection that does not exist, or one that sets only half of the
token.
+_CREDENTIAL_ERRORS = (AirflowNotFoundException, ModalConnectionError)
+# Serializes building a backend's Modal client. ModalHook resolves its
connection and builds
+# its client without a lock, so concurrent first calls on a backend shared by
several agent
+# runs would each open a gRPC client and keep one. Module-level rather than
per instance
+# because a lock cannot be deep-copied, and operators holding a backend are.
+_CLIENT_BUILD_LOCK = threading.Lock()
+
@dataclass
class _TrackedSandbox:
@@ -196,11 +209,12 @@ class ModalSandboxBackend(AttachableSandboxBackend):
tags; keys starting with ``airflow_`` are reserved for them and overwrite
any
``tags`` of the same name.
- **Credentials are ambient.** Modal is authenticated the same way its CLI
is: run
- ``modal token new`` once to write ``~/.modal.toml``, or set
``MODAL_TOKEN_ID`` and
- ``MODAL_TOKEN_SECRET`` in the worker environment. Nothing is read until
the first
- sandbox is created, so a Dag file that constructs this backend parses
without
- credentials present.
+ **Credentials come from a** ``modal`` **connection**, resolved by
+ :class:`~airflow.providers.modal.hooks.modal.ModalHook` in the order the
+ :ref:`Modal connection page <howto/connection:modal>` sets out; without a
+ ``modal_default`` connection that is the worker's own Modal credentials.
Nothing is
+ read until the first sandbox is created or attached to, so a Dag file that
constructs
+ this backend parses without the connection or any credentials present.
**What the network policy can and cannot promise.**
``SandboxSpec(block_network=True)``
maps to Modal's own ``block_network``, which drops all outbound traffic
including DNS.
@@ -235,6 +249,9 @@ class ModalSandboxBackend(AttachableSandboxBackend):
:meth:`read_file`, which deliberately stays on the base class's shell
implementation
(see that method). Any Debian or Ubuntu based image, including
``python:*-slim``, does.
+ :param modal_conn_id: The ``modal`` connection to authenticate with.
Default
+ ``"modal_default"``. ``None`` uses the worker's Modal credentials
without looking
+ for a connection.
:param image: Registry tag for the sandbox image, or a prepared
``modal.Image``.
Default ``"python:3.12-slim"``. An image carrying the packages an
agent needs is
the alternative to opening egress so it can install them:
@@ -287,6 +304,7 @@ class ModalSandboxBackend(AttachableSandboxBackend):
def __init__(
self,
*,
+ modal_conn_id: str | None = ModalHook.default_conn_name,
image: str | modal.Image = DEFAULT_IMAGE,
app_name: str = DEFAULT_APP_NAME,
create_app_if_missing: bool = True,
@@ -330,6 +348,9 @@ class ModalSandboxBackend(AttachableSandboxBackend):
_validate_positive_finite(memory, "memory")
if egress_enforcement not in ("strict", "sni"):
raise ValueError(f"egress_enforcement must be 'strict' or 'sni',
got {egress_enforcement!r}.")
+ # Constructing the hook reads no connection and opens no client, so
this is safe at
+ # Dag-parse time; both wait for _connected_hook.
+ self._hook = ModalHook(modal_conn_id=modal_conn_id)
self._image = image
self._app_name = app_name
self._create_app_if_missing = create_app_if_missing
@@ -347,6 +368,25 @@ class ModalSandboxBackend(AttachableSandboxBackend):
# runs never has two of them touching the same entry.
self._tracked: dict[str, _TrackedSandbox] = {}
+ def _connected_hook(self) -> ModalHook:
+ """
+ Return the hook with its client built, the one place every Modal call
goes through.
+
+ Nothing the model does can fix a missing or half-filled connection, so
it fails the
+ task, naming the connection the author set. That includes ``destroy``:
the sandbox
+ would otherwise bill until its lifetime ends, and the toolset's own
teardown catches
+ and logs it, so it still never fails a finished agent run. A Modal
error from
+ building the client propagates unchanged for the caller to classify.
+ """
+ with _CLIENT_BUILD_LOCK:
+ try:
+ self._hook.get_conn()
+ except _CREDENTIAL_ERRORS as e:
+ raise SandboxTerminalError(
+ f"Could not read Modal credentials from connection
{self._hook.modal_conn_id!r}: {e}"
+ ) from e
+ return self._hook
+
# ------------------------------------------------------------------
# Lifecycle.
# ------------------------------------------------------------------
@@ -359,16 +399,16 @@ class ModalSandboxBackend(AttachableSandboxBackend):
self._check_attachable(spec)
name = _new_sandbox_name()
try:
- # Looked up per create rather than memoized on the instance: the
lookup is
- # idempotent server-side, and one round trip per sandbox is
cheaper than the
- # locking a shared cache would need, given the toolset shares one
backend
- # across concurrent runs.
- app = modal.App.lookup(self._app_name,
create_if_missing=self._create_app_if_missing)
# A registry tag is the common case; a prepared Image is how
packages get
# into the sandbox without opening egress to fetch them.
image = modal.Image.from_registry(self._image) if
isinstance(self._image, str) else self._image
- sandbox = modal.Sandbox.create(
- app=app,
+ # The hook looks the app up in the connection's environment on
every create
+ # rather than memoizing it: the lookup is idempotent server-side,
and one round
+ # trip per sandbox is cheaper than the locking a shared cache
would need, given
+ # the toolset shares one backend across concurrent runs.
+ sandbox = self._connected_hook().create_sandbox(
+ app_name=self._app_name,
+ create_app_if_missing=self._create_app_if_missing,
image=image,
name=name,
tags=self._create_tags(spec, name),
@@ -497,7 +537,9 @@ class ModalSandboxBackend(AttachableSandboxBackend):
self._check_handle(sandbox)
tracked = self._tracked.pop(sandbox, None)
try:
- target = tracked.handle if tracked and tracked.handle else
modal.Sandbox.from_id(sandbox)
+ target = (
+ tracked.handle if tracked and tracked.handle else
self._connected_hook().get_sandbox(sandbox)
+ )
target.terminate()
except modal.exception.NotFoundError:
# Already gone, which is what destroy is for. Idempotent by
contract.
@@ -984,7 +1026,7 @@ class ModalSandboxBackend(AttachableSandboxBackend):
try:
# Reachable when the handle came from another backend instance, or
from a
# caller that created the sandbox elsewhere.
- found = modal.Sandbox.from_id(sandbox)
+ found = self._connected_hook().get_sandbox(sandbox)
except modal.exception.Error as e:
raise self._as_sandbox_error(e, sandbox=sandbox) from e
if tracked is None:
@@ -1050,9 +1092,15 @@ class ModalSandboxBackend(AttachableSandboxBackend):
# error. Only reachable if it escapes from somewhere that cannot
report it.
return SandboxError(f"The command hit its deadline: {error}")
if isinstance(error, modal.exception.AuthError):
+ if self._hook.credentials is not None:
+ # The token came from the connection, so worker settings would
change nothing.
+ return SandboxTerminalError(
+ f"Modal rejected the token on connection
{self._hook.modal_conn_id!r}. Update its "
+ f"login and password: {error}"
+ )
return SandboxTerminalError(
- "Modal rejected the credentials. Run 'modal token new', or set
"
- f"MODAL_TOKEN_ID and MODAL_TOKEN_SECRET on the worker: {error}"
+ "Modal rejected the credentials. Create a 'modal' connection,
run 'modal token "
+ f"new', or set MODAL_TOKEN_ID and MODAL_TOKEN_SECRET on the
worker: {error}"
)
if isinstance(error, modal.exception.PermissionDeniedError):
# A sibling of AuthError rather than a subclass, and just as
unfixable by
diff --git a/providers/common/ai/tests/unit/common/ai/sandbox/fake_modal.py
b/providers/common/ai/tests/unit/common/ai/sandbox/fake_modal.py
index 6c4def36349..1ac92b0db21 100644
--- a/providers/common/ai/tests/unit/common/ai/sandbox/fake_modal.py
+++ b/providers/common/ai/tests/unit/common/ai/sandbox/fake_modal.py
@@ -262,6 +262,7 @@ class FakeSandboxFactory:
self.from_id_error: Exception | None = None
self.by_id: dict[str, FakeSandbox] = {}
self.next_process: FakeProcess | None = None
+ self.from_id_clients: list[FakeClient | None] = []
def create(self, **kwargs: Any) -> FakeSandbox:
if self.create_error is not None:
@@ -273,7 +274,8 @@ class FakeSandboxFactory:
self.by_id[sandbox.object_id] = sandbox
return sandbox
- def from_id(self, object_id: str) -> FakeSandbox:
+ def from_id(self, object_id: str, *, client: FakeClient | None = None) ->
FakeSandbox:
+ self.from_id_clients.append(client)
if self.from_id_error is not None:
raise self.from_id_error
try:
@@ -282,18 +284,54 @@ class FakeSandboxFactory:
raise NotFoundError(f"no such sandbox: {object_id}") from None
+class FakeClient:
+ """One authenticated client, remembering where its credentials came
from."""
+
+ def __init__(self, credentials: tuple[str, str] | None) -> None:
+ self.credentials = credentials
+
+
+class FakeClientFactory:
+ """Stands in for ``modal.Client``: ``from_env`` is the ambient path,
``from_credentials`` the connection."""
+
+ def __init__(self) -> None:
+ self.built: list[FakeClient] = []
+
+ def from_env(self) -> FakeClient:
+ client = FakeClient(None)
+ self.built.append(client)
+ return client
+
+ def from_credentials(self, token_id: str, token_secret: str) -> FakeClient:
+ client = FakeClient((token_id, token_secret))
+ self.built.append(client)
+ return client
+
+
class FakeApp:
- def __init__(self, name: str) -> None:
+ def __init__(self, name: str, *, client: FakeClient | None = None,
environment_name: str | None = None):
self.name = name
+ self.client = client
+ self.environment_name = environment_name
class FakeAppFactory:
def __init__(self) -> None:
self.lookups: list[tuple[str, bool]] = []
+ self.apps: list[FakeApp] = []
- def lookup(self, name: str, *, create_if_missing: bool = False) -> FakeApp:
+ def lookup(
+ self,
+ name: str,
+ *,
+ create_if_missing: bool = False,
+ client: FakeClient | None = None,
+ environment_name: str | None = None,
+ ) -> FakeApp:
self.lookups.append((name, create_if_missing))
- return FakeApp(name)
+ app = FakeApp(name, client=client, environment_name=environment_name)
+ self.apps.append(app)
+ return app
class FakeImageFactory:
@@ -334,6 +372,7 @@ def build_fake_modal() -> types.ModuleType:
("Sandbox", FakeSandboxFactory()),
("App", FakeAppFactory()),
("Image", FakeImageFactory()),
+ ("Client", FakeClientFactory()),
):
setattr(module, name, attr)
return module
diff --git a/providers/common/ai/tests/unit/common/ai/sandbox/test_modal.py
b/providers/common/ai/tests/unit/common/ai/sandbox/test_modal.py
index af7c9a3feb9..be57f1db80b 100644
--- a/providers/common/ai/tests/unit/common/ai/sandbox/test_modal.py
+++ b/providers/common/ai/tests/unit/common/ai/sandbox/test_modal.py
@@ -27,6 +27,7 @@ import pytest
import time_machine
import airflow.providers.common.ai.sandbox as sandbox_package
+from airflow.models import Connection
from airflow.providers.common.ai.sandbox.base import (
EXPIRES_AT_TAG,
HOLDER_TAG,
@@ -45,6 +46,11 @@ from unit.common.ai.sandbox.fake_modal import (
build_fake_modal,
)
+_BOUND_TO_MODAL = (
+ "airflow.providers.common.ai.sandbox.modal",
+ "airflow.providers.modal.hooks.modal",
+)
+
@pytest.fixture
def modal_module(monkeypatch):
@@ -55,15 +61,26 @@ def modal_module(monkeypatch):
fake is in ``sys.modules``. Re-importing also keeps the tests honest about
not sharing
state, and lets the suite run whether or not the real SDK is installed.
"""
+ # The backend needs the Modal provider, which needs Airflow 3; the Airflow
2 compatibility
+ # job removes it. Its exceptions module is checked rather than the hook,
which would
+ # import the real SDK this suite replaces.
+ pytest.importorskip("airflow.providers.modal.exceptions")
fake = build_fake_modal()
monkeypatch.setitem(sys.modules, "modal", fake)
monkeypatch.setitem(sys.modules, "modal.exception", fake.exception)
- monkeypatch.delitem(sys.modules,
"airflow.providers.common.ai.sandbox.modal", raising=False)
+ # ModalHook binds ``modal`` at import too, and the backend builds its
clients through it.
+ for name in _BOUND_TO_MODAL:
+ monkeypatch.delitem(sys.modules, name, raising=False)
+ # Re-importing the hook also rebinds it on its package, which sys.modules
restoration does
+ # not undo; left in place, the Modal provider's own tests, which patch
targets by their
+ # dotted path under that package, would patch this fake-bound copy on
Python 3.10.
+
monkeypatch.delattr(importlib.import_module("airflow.providers.modal.hooks"),
"modal", raising=False)
import airflow.providers.common.ai.sandbox.modal as backend_module
yield fake, backend_module
# The next test re-imports from scratch, so nothing here should linger.
- sys.modules.pop("airflow.providers.common.ai.sandbox.modal", None)
+ for name in _BOUND_TO_MODAL:
+ sys.modules.pop(name, None)
@pytest.fixture
@@ -157,16 +174,17 @@ class TestOptionalExtra:
"""The module has to be importable without ``modal`` installed, not merely
unusable."""
@staticmethod
- def _without_modal(monkeypatch):
- """Make ``import modal`` fail the way a missing extra does."""
+ def _block_import(monkeypatch, blocked="modal"):
+ """Make importing ``blocked`` fail the way a missing extra does."""
for name in list(sys.modules):
- if name == "modal" or name.startswith("modal."):
+ if name == blocked or name.startswith(f"{blocked}."):
monkeypatch.delitem(sys.modules, name, raising=False)
- monkeypatch.delitem(sys.modules,
"airflow.providers.common.ai.sandbox.modal", raising=False)
+ for name in _BOUND_TO_MODAL:
+ monkeypatch.delitem(sys.modules, name, raising=False)
real_import = builtins.__import__
def guarded(name, *args, **kwargs):
- if name == "modal" or name.startswith("modal."):
+ if name == blocked or name.startswith(f"{blocked}."):
raise ModuleNotFoundError(f"No module named {name!r}")
return real_import(name, *args, **kwargs)
@@ -175,19 +193,26 @@ class TestOptionalExtra:
def
test_importing_the_module_directly_raises_the_optional_feature_error(self,
monkeypatch):
# The provider verifier walks every submodule of the distribution and
imports it
# directly, so this path is not reached through the package's
__getattr__.
- self._without_modal(monkeypatch)
+ self._block_import(monkeypatch)
with pytest.raises(AirflowOptionalProviderFeatureException):
importlib.import_module("airflow.providers.common.ai.sandbox.modal")
sys.modules.pop("airflow.providers.common.ai.sandbox.modal", None)
def test_package_attribute_raises_the_optional_feature_error(self,
monkeypatch):
- self._without_modal(monkeypatch)
+ self._block_import(monkeypatch)
with pytest.raises(AirflowOptionalProviderFeatureException):
getattr(sandbox_package, "ModalSandboxBackend")
sys.modules.pop("airflow.providers.common.ai.sandbox.modal", None)
+ def test_missing_modal_provider_raises_the_optional_feature_error(self,
monkeypatch):
+ """The extra brings both the SDK and the Modal provider; an install
with only one is incomplete."""
+ monkeypatch.setitem(sys.modules, "modal", build_fake_modal())
+ self._block_import(monkeypatch, blocked="airflow.providers.modal")
+ with pytest.raises(AirflowOptionalProviderFeatureException,
match="airflow.providers.modal"):
+
importlib.import_module("airflow.providers.common.ai.sandbox.modal")
+
def test_package_still_exports_everything_that_needs_no_extra(self,
monkeypatch):
- self._without_modal(monkeypatch)
+ self._block_import(monkeypatch)
assert sandbox_package.SandboxSpec is SandboxSpec
assert sandbox_package.SbxSandboxBackend.__name__ ==
"SbxSandboxBackend"
@@ -493,7 +518,9 @@ class TestAddressAllowlist:
_created(backend, fake,
SandboxSpec(allow_egress_to_cidrs=["1.1.1.1/32"]))
- assert not caplog.records
+ # The hook reports at INFO that it fell back to ambient credentials;
only a warning
+ # would mean the network policy was weakened.
+ assert not [record for record in caplog.records if record.levelno >=
logging.WARNING]
def test_surrounding_whitespace_is_tolerated(self, backend, fake):
_, sandbox = _created(backend, fake,
SandboxSpec(allow_egress_to_cidrs=[" 10.20.0.0/16 "]))
@@ -632,6 +659,105 @@ class TestCreate:
backend.create(spec=SandboxSpec())
+class TestConnection:
+ """Credentials come from a ``modal`` connection through ModalHook, or from
the worker."""
+
+ @pytest.fixture
+ def modal_connection(self, create_connection_without_db):
+ def _create(conn_id="my_modal", login="ak-test", password="as-test",
extra=None):
+ create_connection_without_db(
+ Connection(conn_id=conn_id, conn_type="modal", login=login,
password=password, extra=extra)
+ )
+
+ return _create
+
+ def test_reads_the_token_and_environment_from_the_connection(self,
backend_class, fake, modal_connection):
+ modal_connection(extra={"environment": "staging"})
+ backend = backend_class(modal_conn_id="my_modal")
+
+ _, sandbox = _created(backend, fake, SandboxSpec())
+
+ (client,) = fake.Client.built
+ assert client.credentials == ("ak-test", "as-test")
+ assert sandbox.create_kwargs["client"] is client
+ assert fake.App.apps[0].client is client
+ assert fake.App.apps[0].environment_name == "staging"
+
+ def test_construction_reads_no_connection(self, backend_class, fake):
+ """Constructors run at Dag-parse time, where a missing connection must
not fail the parse."""
+ backend_class(modal_conn_id="does_not_exist")
+
+ assert fake.Client.built == []
+
+ def test_a_named_connection_that_does_not_exist_fails_the_task(self,
backend_class, fake):
+ backend = backend_class(modal_conn_id="does_not_exist")
+
+ with pytest.raises(SandboxTerminalError, match="connection
'does_not_exist'"):
+ backend.create(spec=SandboxSpec())
+ assert fake.Sandbox.created == []
+
+ def test_a_half_filled_connection_fails_the_task(self, backend_class,
fake, modal_connection):
+ modal_connection(password=None)
+ backend = backend_class(modal_conn_id="my_modal")
+
+ with pytest.raises(SandboxTerminalError, match="connection
'my_modal'"):
+ backend.create(spec=SandboxSpec())
+ assert fake.Sandbox.created == []
+
+ def test_a_sandbox_from_another_instance_is_reached_through_the_connection(
+ self, backend_class, fake, modal_connection
+ ):
+ """The collecting task holds only the handle, so the lookup must carry
its credentials."""
+ modal_connection()
+ handle, _ = _created(backend_class(modal_conn_id="my_modal"), fake,
SandboxSpec())
+ collector = backend_class(modal_conn_id="my_modal")
+
+ collector.destroy(handle)
+
+ assert fake.Sandbox.from_id_clients[-1].credentials == ("ak-test",
"as-test")
+
+ def test_reaching_a_sandbox_with_a_missing_connection_fails_the_task(self,
backend, backend_class, fake):
+ handle, _ = _created(backend, fake, SandboxSpec())
+ stranger = backend_class(modal_conn_id="does_not_exist")
+
+ with pytest.raises(SandboxTerminalError, match="connection
'does_not_exist'"):
+ stranger.run_command(handle, "true", timeout=5,
max_output_bytes=1024)
+
+ def
test_destroy_with_a_missing_connection_fails_instead_of_leaving_it_billing(
+ self, backend, backend_class, fake
+ ):
+ """
+ A cleanup task with a misspelled connection must not succeed while the
sandbox runs on.
+
+ The toolset's own teardown catches this and logs it, so an agent run
still never fails
+ on teardown; a task that calls ``destroy`` itself hears about it.
+ """
+ handle, sandbox = _created(backend, fake, SandboxSpec())
+ stranger = backend_class(modal_conn_id="does_not_exist")
+
+ with pytest.raises(SandboxTerminalError, match="connection
'does_not_exist'"):
+ stranger.destroy(handle)
+ assert not sandbox.terminated
+
+ def test_a_rejected_connection_token_points_at_the_connection(
+ self, backend_class, fake, modal_connection
+ ):
+ """Worker settings are ignored while the connection carries a token,
so advice to set them misleads."""
+ modal_connection()
+ fake.Sandbox.create_error = fake.exception.AuthError("token revoked")
+
+ with pytest.raises(SandboxTerminalError, match="token on connection
'my_modal'") as raised:
+ backend_class(modal_conn_id="my_modal").create(spec=SandboxSpec())
+ assert "MODAL_TOKEN_ID" not in str(raised.value)
+
+ def test_rejected_ambient_credentials_point_at_both_ways_in(self, backend,
fake):
+ fake.Sandbox.create_error = fake.exception.AuthError("token missing")
+
+ with pytest.raises(SandboxTerminalError, match="Create a 'modal'
connection") as raised:
+ backend.create(spec=SandboxSpec())
+ assert "MODAL_TOKEN_ID" in str(raised.value)
+
+
class TestRunCommand:
def test_runs_through_a_shell_and_returns_both_streams(self, backend,
fake):
handle, sandbox = _created(backend, fake)
diff --git a/uv.lock b/uv.lock
index 484e7e66ef8..f4bc231f9d3 100644
--- a/uv.lock
+++ b/uv.lock
@@ -4684,6 +4684,7 @@ mcp = [
{ name = "pydantic-ai-slim", extra = ["mcp"] },
]
modal = [
+ { name = "apache-airflow-providers-modal" },
{ name = "modal" },
]
openai = [
@@ -4722,6 +4723,7 @@ dev = [
{ name = "apache-airflow-providers-common-compat" },
{ name = "apache-airflow-providers-common-sql", extra = ["datafusion"] },
{ name = "apache-airflow-providers-git" },
+ { name = "apache-airflow-providers-modal" },
{ name = "apache-airflow-providers-standard" },
{ name = "apache-airflow-task-sdk" },
{ name = "langchain" },
@@ -4747,6 +4749,7 @@ requires-dist = [
{ name = "apache-airflow-providers-common-sql", marker = "extra == 'sql'",
editable = "providers/common/sql" },
{ name = "apache-airflow-providers-git", marker = "extra == 'git'",
editable = "providers/git" },
{ name = "apache-airflow-providers-git", marker = "extra == 'skills'",
editable = "providers/git" },
+ { name = "apache-airflow-providers-modal", marker = "extra == 'modal'",
editable = "providers/modal" },
{ name = "apache-airflow-providers-standard", editable =
"providers/standard" },
{ name = "dataclasses-json", marker = "extra == 'llamaindex'", specifier =
">=0.6.7" },
{ name = "fastavro", marker = "python_full_version >= '3.14' and extra ==
'avro'", specifier = ">=1.12.1" },
@@ -4786,6 +4789,7 @@ dev = [
{ name = "apache-airflow-providers-common-sql", editable =
"providers/common/sql" },
{ name = "apache-airflow-providers-common-sql", extras = ["datafusion"],
editable = "providers/common/sql" },
{ name = "apache-airflow-providers-git", editable = "providers/git" },
+ { name = "apache-airflow-providers-modal", editable = "providers/modal" },
{ name = "apache-airflow-providers-standard", editable =
"providers/standard" },
{ name = "apache-airflow-task-sdk", editable = "task-sdk" },
{ name = "langchain", specifier = ">=1.0.0" },
