Skip to content

Commit e25f413

Browse files
kaxilvatsrahul1001
andauthored
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 <43964496+vatsrahul1001@users.noreply.github.com>
1 parent dcc9ce6 commit e25f413

10 files changed

Lines changed: 293 additions & 55 deletions

File tree

‎providers/common/ai/docs/index.rst‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -202,6 +202,7 @@ Dependent package
202202
============================================================================================================ ==============
203203
`apache-airflow-providers-common-sql <https://airflow.apache.org/docs/apache-airflow-providers-common-sql>`_ ``common.sql``
204204
`apache-airflow-providers-git <https://airflow.apache.org/docs/apache-airflow-providers-git>`_ ``git``
205+
`apache-airflow-providers-modal <https://airflow.apache.org/docs/apache-airflow-providers-modal>`_ ``modal``
205206
============================================================================================================ ==============
206207

207208
Optional dependencies
@@ -224,7 +225,7 @@ Extra Dependencies
224225
``openai`` ``pydantic-ai-slim[openai]>=2.33.0``, ``openai>=2.47.0``
225226
``typesafe`` ``typesafe-sdk>=0.6.0``
226227
``mcp`` ``pydantic-ai-slim[mcp]>=2.33.0``
227-
``modal`` ``modal>=1.5.2``
228+
``modal`` ``apache-airflow-providers-modal``, ``modal>=1.5.2``
228229
``opensandbox`` ``opensandbox>=1.1.0``
229230
``code-mode`` ``pydantic-ai-harness[codemode]>=0.3.0``
230231
``shields`` ``pydantic-ai-shields>=0.3.4``

‎providers/common/ai/docs/installation.rst‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,9 @@ The provider runs on Airflow 2.11, but some features need a newer Airflow versio
7373
- Needs
7474
* - The ``skills`` and ``git`` extras (``apache-airflow-providers-git`` needs Airflow 3)
7575
- Airflow 3.0
76+
* - The ``modal`` extra and the Modal sandbox backend (``apache-airflow-providers-modal``
77+
needs Airflow 3)
78+
- Airflow 3.0
7679
* - :doc:`Approval gates <approval_gates>` and :doc:`HITL review <hitl_review>`
7780
- Airflow 3.1
7881
* - 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
104107
``apache-airflow-providers-common-sql`` to releases older than this provider needs.
105108
Installing Airflow 2.11.0 without its constraints can also pull in a ``universal-pathlib``
106109
0.3 release, which Airflow 2's ``ObjectStoragePath`` rejects; 2.11.1 and later cap it.
107-
Leave out the ``skills`` and ``git`` extras: they need Airflow 3, and without
110+
Leave out the ``skills``, ``git`` and ``modal`` extras: they need Airflow 3, and without
108111
constraints ``pip`` upgrades Airflow to satisfy them.
109112
* Python 3.10 to 3.12: the provider needs 3.10 or later, and Airflow 2.11 supports up to 3.12.
110113

‎providers/common/ai/docs/sandbox/backends.rst‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,11 +34,19 @@ provider, **this is the managed one to use in production**, and with
3434
:ref:`OpenSandbox <sandbox-backend-opensandbox>` one of the two that run on
3535
Kubernetes: nothing has to be installed on the worker, model-written code never
3636
executes on the worker host, and Modal reclaims a sandbox at its own lifetime
37-
whether or not the worker survives. It needs the ``modal`` extra and ambient
38-
credentials, as under :ref:`Quick start <sandbox-quick-start>`.
37+
whether or not the worker survives. It needs the ``modal`` extra and Modal
38+
credentials, from a ``modal`` connection or the worker environment, as under
39+
:ref:`Quick start <sandbox-quick-start>`.
3940

4041
Constructor parameters:
4142

43+
- ``modal_conn_id``: ``modal`` connection the token and, optionally, the Modal
44+
``environment`` come from. Default ``"modal_default"``; ``None`` uses the worker's
45+
credentials without looking for a connection. How a missing or partial connection
46+
resolves is on the :ref:`Modal connection page <howto/connection:modal>`. A credential
47+
problem fails the task, except in the toolset's own teardown, which logs it so a
48+
finished run is not failed. The connection type comes from the Modal provider, which
49+
the ``modal`` extra installs and which needs Airflow 3.
4250
- ``image``: Registry tag for the sandbox image, or a prepared ``modal.Image``
4351
carrying pre-installed packages. Default ``"python:3.12-slim"``.
4452
- ``app_name``: Modal app the sandboxes are created under. Default

‎providers/common/ai/docs/sandbox/configuration.rst‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -137,8 +137,9 @@ Credentials
137137
``SandboxSpec.env`` is the only way in. Airflow never populates it: no connection,
138138
variable or worker environment variable reaches a sandbox unless you name it
139139
there, and the credential that *provisions* the sandbox never enters it either.
140-
Modal's token stays on the worker and is used by the client, so code running
141-
inside cannot call Modal as you or create further sandboxes.
140+
Modal's token, whether it comes from the ``modal`` connection or the worker
141+
environment, stays in the task and is used by the client, so code running inside
142+
cannot call Modal as you or create further sandboxes.
142143

143144
Before you put a real secret in ``env``, four things are true of it.
144145

@@ -266,7 +267,8 @@ A sandbox another task owns
266267

267268
The toolset's own sandbox is provisioned from a spec fixed in the Dag file, on the
268269
model's first tool call, and destroyed when the run ends. Two things cannot be
269-
done inside that shape: a credential cannot come from a connection, and a second
270+
done inside that shape: a credential for the code inside the sandbox
271+
(``SandboxSpec.env``) cannot come from a connection, and a second
270272
run against the same agent, which is what :ref:`HITL review <howto:hitl_review>`
271273
does when a reviewer asks for changes, cannot find the first run's files. A file
272274
the agent built can leave through ``exports`` (:ref:`sandbox-results`), but only

‎providers/common/ai/docs/sandbox/index.rst‎

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -107,15 +107,19 @@ is the only change between this Dag and production:
107107
:start-after: [START howto_sandbox_agent_local]
108108
:end-before: [END howto_sandbox_agent_local]
109109

110-
Install the Modal extra and authenticate as you would for Modal's own CLI. On a
111-
worker, set ``MODAL_TOKEN_ID`` and ``MODAL_TOKEN_SECRET`` in the environment
112-
instead; nothing is read until the first sandbox is created, so a Dag file that
113-
constructs the backend parses without credentials present:
110+
Install the Modal extra, which also installs the Modal provider, and create a
111+
``modal`` connection with the Modal token id as its login and the token secret as
112+
its password. The backend reads ``modal_default`` unless you pass ``modal_conn_id``;
113+
without that connection it uses the worker's own Modal credentials, as the
114+
:ref:`Modal connection page <howto/connection:modal>` describes. Nothing is read until
115+
the first sandbox is created, so a Dag file that constructs the backend parses without
116+
credentials present:
114117

115118
.. code-block:: bash
116119
117120
pip install 'apache-airflow-providers-common-ai[modal]'
118-
modal token new # writes ~/.modal.toml
121+
airflow connections add modal_default --conn-type modal \
122+
--conn-login "$MODAL_TOKEN_ID" --conn-password "$MODAL_TOKEN_SECRET"
119123
120124
What it is for, in practice
121125
---------------------------
@@ -425,14 +429,13 @@ system tests, ``example_sandbox_toolset_sbx.py`` and
425429
``example_sandbox_toolset_modal.py``, run against a real backend and are
426430
reachable from the System Tests entry in the sidebar.
427431

428-
**Credentials and where it runs.** This route does not end at an
429-
Airflow connection. Airflow puts none of its context, connections, variables or
430-
worker environment into the sandbox; only what you pass through
432+
**Credentials and where it runs.** Airflow puts none of its context, connections,
433+
variables or worker environment into the sandbox; only what you pass through
431434
:class:`~airflow.providers.common.ai.sandbox.SandboxSpec` goes in, and the
432-
credential that provisions the sandbox never enters it. Authorization to the
433-
backend sits outside Airflow: ``sbx login`` and ``sbx policy init`` on the
434-
machine for ``sbx``, or ``MODAL_TOKEN_ID`` and ``MODAL_TOKEN_SECRET`` on the
435-
worker for Modal. Work runs in a per-run microVM on the worker host with
435+
credential that provisions the sandbox never enters it. For Modal that credential
436+
is a ``modal`` connection (``modal_conn_id``; see the
437+
:ref:`Modal connection page <howto/connection:modal>`). For ``sbx`` it sits outside Airflow:
438+
``sbx login`` on the machine. Work runs in a per-run microVM on the worker host with
436439
``sbx``, or off the worker entirely in Modal's infrastructure. Its tool calls
437440
act as barriers, as they do for the other
438441
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.
486489
- **A file the agent built leaves only through a task**, never through the
487490
model's context, which is text-only and capped.
488491
:ref:`Getting a result out <sandbox-results>`.
489-
- **A credential comes from a connection only when a task provisions the
490-
sandbox**; the toolset's own spec is fixed at parse time, and anything injected
491-
is readable by the model. :ref:`Credentials <sandbox-credentials>`.
492+
- **A credential handed to the code inside the sandbox comes from a connection
493+
only when a task provisions the sandbox**; the toolset's own spec is fixed at
494+
parse time, and anything injected is readable by the model. :ref:`Credentials <sandbox-credentials>`.
492495
- **Cannot be combined with** ``durable=True``, and with ``enable_hitl_review=True``
493496
only when the sandbox is task-owned; ``AgentOperator`` raises otherwise.
494497
:ref:`Lifecycle <sandbox-lifecycle>`.

‎providers/common/ai/pyproject.toml‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -104,8 +104,11 @@ dependencies = [
104104
"mcp" = ["pydantic-ai-slim[mcp]>=2.33.0"]
105105
# Modal-hosted sandboxes for SandboxToolset, the backend to use where the worker host
106106
# cannot run microVMs -- Kubernetes, most notably. Needs no host-level installation,
107-
# only Modal credentials.
108-
"modal" = ["modal>=1.5.2"]
107+
# only Modal credentials, which the Modal provider's `modal` connection supplies.
108+
"modal" = [
109+
"apache-airflow-providers-modal",
110+
"modal>=1.5.2",
111+
]
109112
"opensandbox" = ["opensandbox>=1.1.0"]
110113
# Code mode: collapse tool calls into a single `run_code` tool that the model
111114
# drives by writing Python, executed in the Monty sandbox (pydantic-monty).
@@ -163,6 +166,7 @@ dev = [
163166
"apache-airflow-providers-common-compat",
164167
"apache-airflow-providers-common-sql",
165168
"apache-airflow-providers-git",
169+
"apache-airflow-providers-modal",
166170
"apache-airflow-providers-standard",
167171
# Additional devel dependencies (do not remove this line and add extra development dependencies)
168172
"sqlglot>=30.0.0",

‎providers/common/ai/src/airflow/providers/common/ai/sandbox/modal.py‎

Lines changed: 64 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,9 @@
3131

3232
try:
3333
import modal
34+
35+
from airflow.providers.modal.exceptions import ModalConnectionError
36+
from airflow.providers.modal.hooks.modal import ModalHook
3437
except ModuleNotFoundError as e:
3538
from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
3639

@@ -54,6 +57,7 @@
5457
encode_network_policy,
5558
is_sandbox_handle,
5659
)
60+
from airflow.providers.common.compat.sdk import AirflowNotFoundException
5761

5862
if TYPE_CHECKING:
5963
from collections.abc import Mapping, Sequence
@@ -106,6 +110,15 @@
106110

107111
EgressEnforcement = Literal["strict", "sni"]
108112

113+
# What reading credentials from the connection can raise before Modal is reached at all: a
114+
# named connection that does not exist, or one that sets only half of the token.
115+
_CREDENTIAL_ERRORS = (AirflowNotFoundException, ModalConnectionError)
116+
# Serializes building a backend's Modal client. ModalHook resolves its connection and builds
117+
# its client without a lock, so concurrent first calls on a backend shared by several agent
118+
# runs would each open a gRPC client and keep one. Module-level rather than per instance
119+
# because a lock cannot be deep-copied, and operators holding a backend are.
120+
_CLIENT_BUILD_LOCK = threading.Lock()
121+
109122

110123
@dataclass
111124
class _TrackedSandbox:
@@ -196,11 +209,12 @@ class ModalSandboxBackend(AttachableSandboxBackend):
196209
tags; keys starting with ``airflow_`` are reserved for them and overwrite any
197210
``tags`` of the same name.
198211
199-
**Credentials are ambient.** Modal is authenticated the same way its CLI is: run
200-
``modal token new`` once to write ``~/.modal.toml``, or set ``MODAL_TOKEN_ID`` and
201-
``MODAL_TOKEN_SECRET`` in the worker environment. Nothing is read until the first
202-
sandbox is created, so a Dag file that constructs this backend parses without
203-
credentials present.
212+
**Credentials come from a** ``modal`` **connection**, resolved by
213+
:class:`~airflow.providers.modal.hooks.modal.ModalHook` in the order the
214+
:ref:`Modal connection page <howto/connection:modal>` sets out; without a
215+
``modal_default`` connection that is the worker's own Modal credentials. Nothing is
216+
read until the first sandbox is created or attached to, so a Dag file that constructs
217+
this backend parses without the connection or any credentials present.
204218
205219
**What the network policy can and cannot promise.** ``SandboxSpec(block_network=True)``
206220
maps to Modal's own ``block_network``, which drops all outbound traffic including DNS.
@@ -235,6 +249,9 @@ class ModalSandboxBackend(AttachableSandboxBackend):
235249
:meth:`read_file`, which deliberately stays on the base class's shell implementation
236250
(see that method). Any Debian or Ubuntu based image, including ``python:*-slim``, does.
237251
252+
:param modal_conn_id: The ``modal`` connection to authenticate with. Default
253+
``"modal_default"``. ``None`` uses the worker's Modal credentials without looking
254+
for a connection.
238255
:param image: Registry tag for the sandbox image, or a prepared ``modal.Image``.
239256
Default ``"python:3.12-slim"``. An image carrying the packages an agent needs is
240257
the alternative to opening egress so it can install them:
@@ -287,6 +304,7 @@ class ModalSandboxBackend(AttachableSandboxBackend):
287304
def __init__(
288305
self,
289306
*,
307+
modal_conn_id: str | None = ModalHook.default_conn_name,
290308
image: str | modal.Image = DEFAULT_IMAGE,
291309
app_name: str = DEFAULT_APP_NAME,
292310
create_app_if_missing: bool = True,
@@ -330,6 +348,9 @@ def __init__(
330348
_validate_positive_finite(memory, "memory")
331349
if egress_enforcement not in ("strict", "sni"):
332350
raise ValueError(f"egress_enforcement must be 'strict' or 'sni', got {egress_enforcement!r}.")
351+
# Constructing the hook reads no connection and opens no client, so this is safe at
352+
# Dag-parse time; both wait for _connected_hook.
353+
self._hook = ModalHook(modal_conn_id=modal_conn_id)
333354
self._image = image
334355
self._app_name = app_name
335356
self._create_app_if_missing = create_app_if_missing
@@ -347,6 +368,25 @@ def __init__(
347368
# runs never has two of them touching the same entry.
348369
self._tracked: dict[str, _TrackedSandbox] = {}
349370

371+
def _connected_hook(self) -> ModalHook:
372+
"""
373+
Return the hook with its client built, the one place every Modal call goes through.
374+
375+
Nothing the model does can fix a missing or half-filled connection, so it fails the
376+
task, naming the connection the author set. That includes ``destroy``: the sandbox
377+
would otherwise bill until its lifetime ends, and the toolset's own teardown catches
378+
and logs it, so it still never fails a finished agent run. A Modal error from
379+
building the client propagates unchanged for the caller to classify.
380+
"""
381+
with _CLIENT_BUILD_LOCK:
382+
try:
383+
self._hook.get_conn()
384+
except _CREDENTIAL_ERRORS as e:
385+
raise SandboxTerminalError(
386+
f"Could not read Modal credentials from connection {self._hook.modal_conn_id!r}: {e}"
387+
) from e
388+
return self._hook
389+
350390
# ------------------------------------------------------------------
351391
# Lifecycle.
352392
# ------------------------------------------------------------------
@@ -359,16 +399,16 @@ def create(self, *, spec: SandboxSpec | None = None) -> str:
359399
self._check_attachable(spec)
360400
name = _new_sandbox_name()
361401
try:
362-
# Looked up per create rather than memoized on the instance: the lookup is
363-
# idempotent server-side, and one round trip per sandbox is cheaper than the
364-
# locking a shared cache would need, given the toolset shares one backend
365-
# across concurrent runs.
366-
app = modal.App.lookup(self._app_name, create_if_missing=self._create_app_if_missing)
367402
# A registry tag is the common case; a prepared Image is how packages get
368403
# into the sandbox without opening egress to fetch them.
369404
image = modal.Image.from_registry(self._image) if isinstance(self._image, str) else self._image
370-
sandbox = modal.Sandbox.create(
371-
app=app,
405+
# The hook looks the app up in the connection's environment on every create
406+
# rather than memoizing it: the lookup is idempotent server-side, and one round
407+
# trip per sandbox is cheaper than the locking a shared cache would need, given
408+
# the toolset shares one backend across concurrent runs.
409+
sandbox = self._connected_hook().create_sandbox(
410+
app_name=self._app_name,
411+
create_app_if_missing=self._create_app_if_missing,
372412
image=image,
373413
name=name,
374414
tags=self._create_tags(spec, name),
@@ -497,7 +537,9 @@ def destroy(self, sandbox: str) -> None:
497537
self._check_handle(sandbox)
498538
tracked = self._tracked.pop(sandbox, None)
499539
try:
500-
target = tracked.handle if tracked and tracked.handle else modal.Sandbox.from_id(sandbox)
540+
target = (
541+
tracked.handle if tracked and tracked.handle else self._connected_hook().get_sandbox(sandbox)
542+
)
501543
target.terminate()
502544
except modal.exception.NotFoundError:
503545
# Already gone, which is what destroy is for. Idempotent by contract.
@@ -984,7 +1026,7 @@ def _handle(self, sandbox: str) -> modal.Sandbox:
9841026
try:
9851027
# Reachable when the handle came from another backend instance, or from a
9861028
# caller that created the sandbox elsewhere.
987-
found = modal.Sandbox.from_id(sandbox)
1029+
found = self._connected_hook().get_sandbox(sandbox)
9881030
except modal.exception.Error as e:
9891031
raise self._as_sandbox_error(e, sandbox=sandbox) from e
9901032
if tracked is None:
@@ -1050,9 +1092,15 @@ def _as_sandbox_error(self, error: modal.exception.Error, *, sandbox: str | None
10501092
# error. Only reachable if it escapes from somewhere that cannot report it.
10511093
return SandboxError(f"The command hit its deadline: {error}")
10521094
if isinstance(error, modal.exception.AuthError):
1095+
if self._hook.credentials is not None:
1096+
# The token came from the connection, so worker settings would change nothing.
1097+
return SandboxTerminalError(
1098+
f"Modal rejected the token on connection {self._hook.modal_conn_id!r}. Update its "
1099+
f"login and password: {error}"
1100+
)
10531101
return SandboxTerminalError(
1054-
"Modal rejected the credentials. Run 'modal token new', or set "
1055-
f"MODAL_TOKEN_ID and MODAL_TOKEN_SECRET on the worker: {error}"
1102+
"Modal rejected the credentials. Create a 'modal' connection, run 'modal token "
1103+
f"new', or set MODAL_TOKEN_ID and MODAL_TOKEN_SECRET on the worker: {error}"
10561104
)
10571105
if isinstance(error, modal.exception.PermissionDeniedError):
10581106
# A sibling of AuthError rather than a subclass, and just as unfixable by

0 commit comments

Comments
 (0)