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 c8d8df3353d Add `@task.modal` to run Python functions on Modal without
Airflow in the image (#76502)
c8d8df3353d is described below
commit c8d8df3353d0c588310087222b18e3c60f9270a5
Author: Kaxil Naik <[email protected]>
AuthorDate: Thu Nov 5 18:14:28 2026 +0000
Add `@task.modal` to run Python functions on Modal without Airflow in the
image (#76502)
The decorator runs the function's body in a Modal sandbox built from any
modal.Image, for code whose dependencies conflict with Airflow's own pins or
that needs a GPU the workers lack. The task stays on a normal worker, so it
works on any deployment once the provider is installed.
A short standard-library bootstrap reads the function source and pickled
arguments from stdin, calls the function, and prints the pickled return
value
as the last stdout line. The image needs only python, and its Python version
need not match the worker's. Function output and tracebacks stream to the
task
log.
deferrable=True defers to ModalSandboxTrigger, which follows the sandbox's
output from the triggerer and terminates it on user kill (Airflow 3.3+) or
when
following it fails. ModalHook gains get_async_client and aget_sandbox for
the
triggerer; secret and volume names resolve through the connection at run
time.
The function's output, including what its subprocesses write, now goes to
the sandbox's stderr, and stdout carries only the pickled return value. This
drops the in-band result line and the reader-side filtering it needed, and
execute_complete no longer re-reads the function's whole output to find the
result.
Also moves stdin chunking and the exit-code constants next to the rest of
the
sandbox protocol, and resolves async credentials without writing into the
hook's cached connection. The docs page now shows output from a real run.
* Tighten the @task.modal docs against a live run
Quotes the real timeout and out-of-memory messages, notes the
KeyboardInterrupt traceback a timeout leaves in the task log, states the
Python 3.8 floor for the image (checked against a 3.10 image from a 3.12
worker), and moves the argument reference to the public modal_task
docstring.
* Fix @task.modal deferral on Airflow 3.0 and cap sandboxes at
execution_timeout
The deferred path no longer overrides execute_complete to call
_handle_output, whose signature differs on Airflow 3.0. Its only job was to
record returned Airflow assets, which a function running without Airflow
cannot return.
The sandbox's Modal timeout is now capped at the task's execution_timeout,
which is also passed to defer(), so a deferred task's sandbox cannot outlive
the task. Losing the output stream now logs a warning and keeps waiting
instead of terminating a sandbox that is still working; only losing track of
the sandbox itself terminates it. A return value that Modal kept only part
of
fails with a clear error, a single secret name is accepted, and the async
connection lookup no longer needs common-compat 1.17.
* Frame @task.modal return values and run functions in a registered module
The return value is now written as <length>:<base64>, so a value Modal kept
only the end of is always reported as incomplete instead of occasionally
decoding into a wrong value. The function now runs in a module registered in
sys.modules, which dataclasses defined inside the function need.
The timeout docs and error now describe what execution_timeout does: it is
the sandbox's timeout when timeout is unset, and the shorter of the two wins
when both are set.
On Airflow 3.0, DecoratedOperator.execute reads self.xcom_push to pass it to
_handle_output, and the task SDK's BaseOperator has no such attribute, so
every @task.modal task that returned a value failed there. Expose it as an
alias of do_xcom_push, as the common.sql decorators already do.
---
providers/modal/docs/index.rst | 13 +-
providers/modal/docs/task_decorator.rst | 179 +++++++++
providers/modal/docs/usage.rst | 13 +-
providers/modal/provider.yaml | 12 +-
providers/modal/pyproject.toml | 2 +-
.../{exceptions.py => decorators/__init__.py} | 17 -
.../airflow/providers/modal/decorators/modal.py | 308 +++++++++++++++
.../src/airflow/providers/modal/exceptions.py | 9 +
.../airflow/providers/modal/get_provider_info.py | 8 +-
.../src/airflow/providers/modal/hooks/modal.py | 80 +++-
.../modal/{exceptions.py => triggers/__init__.py} | 17 -
.../src/airflow/providers/modal/triggers/modal.py | 100 +++++
.../modal/{exceptions.py => utils/__init__.py} | 17 -
.../providers/modal/utils/python_sandbox.py | 156 ++++++++
.../exceptions.py => tests/system/__init__.py} | 18 +-
.../system/modal/__init__.py} | 17 -
.../modal/tests/system/modal/example_modal_task.py | 80 ++++
.../unit/modal/decorators/__init__.py} | 17 -
.../tests/unit/modal/decorators/test_modal.py | 428 +++++++++++++++++++++
.../modal/tests/unit/modal/hooks/test_modal.py | 152 ++++++++
providers/modal/tests/unit/modal/local_sandbox.py | 132 +++++++
.../modal/tests/unit/modal/test_exceptions.py | 4 +-
.../unit/modal/triggers/__init__.py} | 17 -
.../modal/tests/unit/modal/triggers/test_modal.py | 171 ++++++++
.../unit/modal/utils/__init__.py} | 17 -
.../tests/unit/modal/utils/test_python_sandbox.py | 163 ++++++++
26 files changed, 1990 insertions(+), 157 deletions(-)
diff --git a/providers/modal/docs/index.rst b/providers/modal/docs/index.rst
index 70f0170e1a0..7ebe68d33e5 100644
--- a/providers/modal/docs/index.rst
+++ b/providers/modal/docs/index.rst
@@ -39,6 +39,8 @@ piece of Airflow code resolves it through:
* ``ModalHook.lookup_app``, ``get_secret``, ``get_volume`` resolve the other
named resources.
* ``ModalHook.client`` is the underlying ``modal.Client`` for anything not
covered above;
``client_kwargs`` splats it into SDK calls that take only ``client=``.
+* ``@task.modal`` runs a function's body in a sandbox, in an image without
Airflow; see
+ :doc:`task_decorator`.
Every hook method applies both the connection's credentials and its
``environment``. Raw SDK
calls made with ``client_kwargs`` carry the credentials only; see :doc:`usage`
for what that
@@ -64,6 +66,14 @@ resolution, so a worker that already ran ``modal token new``
needs no Airflow se
Connection types <connections>
Using Modal from Airflow <usage>
+ Running a Python function on Modal <task_decorator>
+
+.. toctree::
+ :hidden:
+ :maxdepth: 1
+ :caption: System tests
+
+ System Tests <_api/tests/system/modal/index>
.. toctree::
:hidden:
@@ -90,7 +100,8 @@ apache-airflow-providers-modal package
`Modal <https://modal.com/>`__ provider for Apache Airflow.
Owns the ``modal`` connection type and a hook that turns it into an
authenticated
-Modal client, so operators, toolsets and executors share one credential path.
+Modal client, so operators, toolsets and executors share one credential path,
and
+``@task.modal``, which runs a Python function in a Modal sandbox.
Release: 0.1.0
diff --git a/providers/modal/docs/task_decorator.rst
b/providers/modal/docs/task_decorator.rst
new file mode 100644
index 00000000000..a9f05eb4645
--- /dev/null
+++ b/providers/modal/docs/task_decorator.rst
@@ -0,0 +1,179 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+.. _howto/decorator:modal:
+
+Running a Python function on Modal
+==================================
+
+``@task.modal`` runs the body of a Python function in a Modal
+`sandbox <https://modal.com/docs/guide/sandboxes>`__, in an image you define,
while the task
+itself stays on an Airflow worker. Use it for code whose dependencies cannot
be installed next
+to Airflow, such as ``pydantic<2`` or a specific vLLM build, or that needs a
GPU the workers do
+not have. The image does not need Airflow or the Modal client, only ``python``.
+
+If the code is already deployed as a Modal function, call it from a plain
``@task`` with
+``ModalHook.get_function`` instead; see :ref:`howto/modal-usage`.
+
+Running a function
+------------------
+
+Decorate the function and give it an image. Everything the function imports is
installed in
+the image, not on the worker:
+
+.. exampleinclude:: /../tests/system/modal/example_modal_task.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_decorator_modal]
+ :end-before: [END howto_decorator_modal]
+
+The task starts a sandbox under the Modal app ``airflow`` (``app_name``
changes it; the app is
+created on first use), sends it the function's source and arguments, and logs
the function's
+output as it runs:
+
+.. code-block:: text
+
+ Started Modal sandbox sb-01M6MBWQQ3DSHZMVF1W677BDAT in app 'airflow'.
+ Validated 2 records with pydantic 1.10.27
+
+The return value goes to XCom like any task's, here with the scores coerced to
floats by
+pydantic 1:
+
+.. code-block:: python
+
+ {
+ "pydantic_version": "1.10.27",
+ "records": [
+ {"id": 1, "email": "[email protected]", "score": 0.91},
+ {"id": 2, "email": "[email protected]", "score": 0.47},
+ ],
+ }
+
+Each sandbox carries tags with the task instance's Dag id, task id, run id,
map index and try
+number, so you can find it in the Modal dashboard from the task, and the other
way round.
+
+Writing the function
+--------------------
+
+The function's source is sent to the sandbox on its own, as with
``@task.docker`` and
+``@task.kubernetes``:
+
+* Import what the function needs inside its body. Module-level imports,
globals and helper
+ functions in the Dag file do not exist in the sandbox.
+* Take every input as an argument. There is no task context, connection,
variable or hook
+ inside the function.
+* Arguments and the return value are pickled. They must load on the other
side, so pass
+ plain Python values, not objects from packages that only the worker or only
the image has.
+ The image needs Python 3.8 or later, and its minor version does not have to
match the
+ worker's. Dates that come from another task's XCom arrive as ``pendulum``
objects, which
+ need ``pendulum`` in the image; pass ``.isoformat()`` strings instead to
avoid that.
+* Keep the return value small. It comes back through the sandbox's output, and
Modal keeps a
+ limited amount of that; the task fails with a clear error when the value was
cut short.
+ Write large results to a volume and return their path.
+* Pass credentials as Modal secrets, not as arguments: arguments are rendered
in the Airflow UI.
+
+Choosing resources
+------------------
+
+``image``, ``gpu``, ``cpu``, ``memory``, ``timeout``, ``secrets`` and
``volumes`` go to
+`modal.Sandbox.create <https://modal.com/docs/reference/modal.Sandbox>`__, and
anything left
+unset gets Modal's default;
:func:`~airflow.providers.modal.decorators.modal.modal_task`
+lists every argument. Before a long or memory-heavy run:
+
+* Set ``timeout`` for anything that runs longer than five minutes. Modal stops
a sandbox after
+ 300 seconds by default. When only the task's ``execution_timeout`` is set,
the sandbox gets
+ that instead; when both are set, the shorter wins. Either way a sandbox
stops no later than
+ ``execution_timeout`` after it starts.
+* A single ``cpu`` or ``memory`` number is a reservation, which the sandbox
can exceed. To cap
+ it, pass a ``(reservation, limit)`` pair; a function that goes over a memory
limit is
+ stopped.
+* ``secrets`` and ``volumes`` take names as well as Modal objects. Names are
looked up when the
+ task runs, through the connection, so they resolve in its ``environment``
and parsing the
+ Dag file makes no Modal call. Each secret's keys become environment
variables in the
+ sandbox.
+
+The default image is ``modal.Image.debian_slim()`` built for the worker's
Python version.
+
+A GPU job that reads a Hugging Face token from a secret, named ``huggingface``
and holding
+``HF_TOKEN``, and caches the weights on a volume. It needs an H100, so it is a
sketch rather
+than an example the docs build runs:
+
+.. code-block:: python
+
+ @task.modal(
+ image=modal.Image.debian_slim().pip_install("vllm"),
+ gpu="H100",
+ timeout=3600,
+ secrets=["huggingface"],
+ volumes={"/weights": "model-weights"},
+ )
+ def batch_infer(prompts: list[str]) -> list[str]:
+ from vllm import LLM
+
+ llm = LLM("meta-llama/Llama-3.1-8B-Instruct", download_dir="/weights")
+ return [output.outputs[0].text for output in llm.generate(prompts)]
+
+Not holding a worker slot
+-------------------------
+
+With ``deferrable=True`` the task starts the sandbox and then defers, so the
worker slot is
+free while the function runs. The triggerer follows the sandbox's output into
the task log,
+and when the sandbox exits the task resumes on a worker to read the return
value:
+
+.. exampleinclude:: /../tests/system/modal/example_modal_task.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_decorator_modal_deferrable]
+ :end-before: [END howto_decorator_modal_deferrable]
+
+On a T4 the task pushes ``"Tesla T4, 15360 MiB"`` to XCom.
+
+The triggerer talks to Modal through the same ``modal_conn_id``. If you rely
on ambient
+credentials (``MODAL_TOKEN_ID`` and ``MODAL_TOKEN_SECRET``, or
``~/.modal.toml``) rather than
+a connection, the triggerer needs them as well as the workers.
+
+If the triggerer restarts while it waits, the trigger that takes over reads
the output again
+from the oldest part Modal still holds, so the task log can repeat lines
logged before the
+restart. ``[operators] default_deferrable`` sets the default for
``deferrable``.
+
+When the function fails
+-----------------------
+
+The task fails with ``ModalSandboxError``, whose ``returncode`` and
``sandbox_id`` say what
+happened and where:
+
+* The function raised. Its traceback is in the task log, and the error reads
+ ``The function failed in Modal sandbox sb-... with exit code 1. Its
traceback is in the log
+ above.``
+* The sandbox ran past ``timeout`` (exit code 124). Modal interrupts the
function, so the task
+ log shows a ``KeyboardInterrupt`` traceback from wherever it was, followed
by ``Modal sandbox
+ sb-... reached its timeout (10s) before the function returned. Raise the
decorator's
+ timeout, or the task's execution_timeout if that is shorter.``
+* The sandbox was stopped from outside (exit code 137), for example by going
over a ``memory``
+ limit: ``Modal sandbox sb-... was terminated before the function returned,
for example
+ because it ran out of memory under a memory limit, or from the Modal
dashboard.``
+
+If the task loses the function's output stream, it logs a warning and keeps
waiting for the
+sandbox to exit. If it loses track of the sandbox itself, it tries to
terminate the sandbox
+before failing, so a retry does not run beside the earlier attempt.
+
+A retry starts a new sandbox and runs the function from the beginning. Marking
a running task
+failed or success, or clearing it, terminates its sandbox. A worker that dies
without running
+the task's ``on_kill`` leaves the sandbox running until it exits or hits its
``timeout``.
+
+A deferred task's sandbox is terminated on those actions from Airflow 3.3,
which added
+``on_kill`` for triggers. On earlier versions it runs until it exits or hits
its ``timeout``.
diff --git a/providers/modal/docs/usage.rst b/providers/modal/docs/usage.rst
index 7458eb1e48c..1afbbfd76ea 100644
--- a/providers/modal/docs/usage.rst
+++ b/providers/modal/docs/usage.rst
@@ -20,11 +20,11 @@
Using Modal from Airflow
========================
-This provider ships a hook, not operators. A task gets a
-:class:`~airflow.providers.modal.hooks.modal.ModalHook`, asks it for a handle
to a deployed
-function, class or sandbox, and then uses the Modal SDK on that handle exactly
as it would
-outside Airflow. What the hook adds is that every handle carries the
connection's credentials
-and environment.
+A task gets a :class:`~airflow.providers.modal.hooks.modal.ModalHook`, asks it
for a handle to a
+deployed function, class or sandbox, and then uses the Modal SDK on that
handle exactly as it
+would outside Airflow. What the hook adds is that every handle carries the
connection's
+credentials and environment. To run a function's body on Modal without
deploying it first, use
+``@task.modal`` instead; see :ref:`howto/decorator:modal`.
Invoking a deployed function
----------------------------
@@ -150,7 +150,8 @@ Execution semantics
What this provider does not do yet
----------------------------------
-* No operators, sensors or deferrable triggers; blocking calls in a ``@task``
are the pattern.
+* No operators or sensors for deployed functions; blocking calls in a
``@task`` are the
+ pattern. ``@task.modal`` (:ref:`howto/decorator:modal`) is the one
deferrable path.
* No executor. Tasks run on Airflow workers and call out to Modal; they do not
run inside Modal.
* Token authentication only. Modal's OAuth client credentials are not modeled
on the
connection.
diff --git a/providers/modal/provider.yaml b/providers/modal/provider.yaml
index 0687aa9999b..14e78cdf704 100644
--- a/providers/modal/provider.yaml
+++ b/providers/modal/provider.yaml
@@ -23,7 +23,8 @@ name: Modal
description: |
`Modal <https://modal.com/>`__ provider for Apache Airflow.
Owns the ``modal`` connection type and a hook that turns it into an
authenticated
- Modal client, so operators, toolsets and executors share one credential
path.
+ Modal client, so operators, toolsets and executors share one credential
path, and
+ ``@task.modal``, which runs a Python function in a Modal sandbox.
state: ready
lifecycle: incubation
@@ -45,6 +46,15 @@ hooks:
python-modules:
- airflow.providers.modal.hooks.modal
+triggers:
+ - integration-name: Modal
+ python-modules:
+ - airflow.providers.modal.triggers.modal
+
+task-decorators:
+ - class-name: airflow.providers.modal.decorators.modal.modal_task
+ name: modal
+
connection-types:
- hook-class-name: airflow.providers.modal.hooks.modal.ModalHook
hook-name: "Modal"
diff --git a/providers/modal/pyproject.toml b/providers/modal/pyproject.toml
index 0378e7633bd..8c231c8a557 100644
--- a/providers/modal/pyproject.toml
+++ b/providers/modal/pyproject.toml
@@ -59,7 +59,7 @@ requires-python = ">=3.11"
# After you modify the dependencies, and rebuild your Breeze CI image with
``breeze ci-image build``
dependencies = [
"apache-airflow>=3.0.0",
- "apache-airflow-providers-common-compat>=1.12.0",
+ "apache-airflow-providers-common-compat>=1.12.0", # use next version
"modal>=1.5.0",
]
diff --git a/providers/modal/src/airflow/providers/modal/exceptions.py
b/providers/modal/src/airflow/providers/modal/decorators/__init__.py
similarity index 59%
copy from providers/modal/src/airflow/providers/modal/exceptions.py
copy to providers/modal/src/airflow/providers/modal/decorators/__init__.py
index a1f0e53a97a..13a83393a91 100644
--- a/providers/modal/src/airflow/providers/modal/exceptions.py
+++ b/providers/modal/src/airflow/providers/modal/decorators/__init__.py
@@ -14,20 +14,3 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Exceptions raised by the Modal provider."""
-
-from __future__ import annotations
-
-
-class ModalError(Exception):
- """
- Base class for errors raised by the Modal provider.
-
- Deliberately not an ``AirflowException`` subclass, matching the other AI
providers: callers
- catch provider errors through this base, and the exceptions carry no
Airflow-specific
- task-state semantics.
- """
-
-
-class ModalConnectionError(ModalError):
- """A ``modal`` connection is misconfigured, for example only one of the
two token fields is set."""
diff --git a/providers/modal/src/airflow/providers/modal/decorators/modal.py
b/providers/modal/src/airflow/providers/modal/decorators/modal.py
new file mode 100644
index 00000000000..2a579232284
--- /dev/null
+++ b/providers/modal/src/airflow/providers/modal/decorators/modal.py
@@ -0,0 +1,308 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""``@task.modal``: run a Python function in a Modal sandbox, in an image
without Airflow."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Callable, Sequence
+from typing import TYPE_CHECKING, Any
+
+import modal
+
+from airflow.providers.common.compat.sdk import (
+ BaseOperator,
+ DecoratedOperator,
+ conf,
+ task_decorator_factory,
+)
+from airflow.providers.modal.exceptions import ModalSandboxError
+from airflow.providers.modal.hooks.modal import ModalHook
+from airflow.providers.modal.triggers.modal import ModalSandboxTrigger
+from airflow.providers.modal.utils.python_sandbox import (
+ ENTRYPOINT,
+ TERMINATED_RETURNCODE,
+ TIMEOUT_RETURNCODE,
+ build_payload,
+ decode_result,
+ send_payload,
+ wait_for_exit,
+)
+
+if TYPE_CHECKING:
+ from airflow.providers.common.compat.sdk import Context, TaskDecorator
+
+
+class _ModalSandboxOperator(BaseOperator):
+ """
+ Run a prepared Python payload in a Modal sandbox and return the function's
return value.
+
+ The base half of ``@task.modal``. ``DecoratedOperator.execute`` hands off
to the next
+ ``execute`` in the MRO, so the Modal half lives here, under it, as
``DockerOperator`` does
+ for ``@task.docker``. The subclass supplies the payload through
:meth:`_build_payload`.
+ """
+
+ template_fields: Sequence[str] = ("modal_conn_id",)
+
+ def __init__(
+ self,
+ *,
+ image: modal.Image | None = None,
+ gpu: str | None = None,
+ cpu: float | tuple[float, float] | None = None,
+ memory: int | tuple[int, int] | None = None,
+ timeout: int | None = None,
+ secrets: Sequence[modal.Secret | str] | str | None = None,
+ volumes: dict[str, modal.Volume | str] | None = None,
+ modal_conn_id: str | None = ModalHook.default_conn_name,
+ app_name: str = "airflow",
+ deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
+ **kwargs: Any,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.image = image
+ self.gpu = gpu
+ self.cpu = cpu
+ self.memory = memory
+ self.timeout = timeout
+ self.secrets = secrets
+ self.volumes = volumes
+ self.modal_conn_id = modal_conn_id
+ self.app_name = app_name
+ self.deferrable = deferrable
+ self._sandbox: modal.Sandbox | None = None
+
+ def _build_payload(self) -> str:
+ raise NotImplementedError
+
+ def _sandbox_options(self, hook: ModalHook) -> dict[str, Any]:
+ """``Sandbox.create`` keyword arguments, leaving out the unset ones so
Modal's defaults apply."""
+ options: dict[str, Any] = {
+ "gpu": self.gpu,
+ "cpu": self.cpu,
+ "memory": self.memory,
+ "timeout": self._sandbox_timeout(),
+ }
+ options = {key: value for key, value in options.items() if value is
not None}
+ # Names resolve here, through the hook, so they use the connection's
credentials and
+ # environment, and a Dag file that names them makes no Modal call when
it is parsed.
+ if self.secrets is not None:
+ secrets = [self.secrets] if isinstance(self.secrets, str) else
self.secrets
+ options["secrets"] = [
+ hook.get_secret(secret) if isinstance(secret, str) else secret
for secret in secrets
+ ]
+ if self.volumes is not None:
+ options["volumes"] = {
+ path: hook.get_volume(volume) if isinstance(volume, str) else
volume
+ for path, volume in self.volumes.items()
+ }
+ # Built when the task runs, so debian_slim() takes the Python version
of the worker,
+ # not of whichever process imported this module.
+ options["image"] = self.image or modal.Image.debian_slim()
+ return options
+
+ def _sandbox_timeout(self) -> int | None:
+ """
+ Modal's timeout for the sandbox: the shorter of ``timeout`` and
``execution_timeout``.
+
+ Either one alone is used as is; with neither, Modal's 300-second
default applies.
+
+ The cap is what stops a deferred task's sandbox when
``execution_timeout`` runs out: the
+ triggerer gives up on the trigger without calling ``on_kill``, so only
Modal can stop it.
+ """
+ if self.execution_timeout is None:
+ return self.timeout
+ cap = math.ceil(self.execution_timeout.total_seconds())
+ return cap if self.timeout is None else min(self.timeout, cap)
+
+ def _sandbox_tags(self, context: Context) -> dict[str, str]:
+ """Tags that find the sandbox from a task instance in the Modal
dashboard, and back."""
+ ti = context["ti"]
+ return {
+ "airflow_dag_id": ti.dag_id,
+ "airflow_task_id": ti.task_id,
+ "airflow_run_id": ti.run_id,
+ "airflow_map_index": str(ti.map_index),
+ "airflow_try_number": str(ti.try_number),
+ }
+
+ def execute(self, context: Context) -> Any:
+ payload = self._build_payload()
+ hook = ModalHook(modal_conn_id=self.modal_conn_id)
+ sandbox = hook.create_sandbox(
+ *ENTRYPOINT,
+ app_name=self.app_name,
+ create_app_if_missing=True,
+ tags=self._sandbox_tags(context),
+ **self._sandbox_options(hook),
+ )
+ self._sandbox = sandbox
+ self.log.info("Started Modal sandbox %s in app %r.",
sandbox.object_id, self.app_name)
+ # Deferring raises TaskDeferred, a BaseException, so it passes through
and leaves the
+ # sandbox running for the trigger. Any other error would leave a
sandbox nobody follows,
+ # and a retry would start a second one beside it.
+ try:
+ send_payload(sandbox, payload)
+ if self.deferrable:
+ self.defer(
+ trigger=ModalSandboxTrigger(
+ sandbox_id=sandbox.object_id,
modal_conn_id=self.modal_conn_id
+ ),
+ method_name="execute_complete",
+ timeout=self.execution_timeout,
+ )
+ self._log_output(sandbox)
+ wait_for_exit(sandbox)
+ except Exception:
+ self.on_kill()
+ raise
+ self._raise_for_returncode(sandbox.object_id, sandbox.returncode)
+ return decode_result(sandbox.stdout.read())
+
+ def _log_output(self, sandbox: modal.Sandbox) -> None:
+ try:
+ for line in sandbox.stderr:
+ self.log.info("%s", line.rstrip("\n"))
+ except Exception as e:
+ # Modal's reader raises once its retries for stream errors run
out. Losing the
+ # output is no reason to stop a function that is still running.
+ self.log.warning(
+ "Lost the output of Modal sandbox %s (%s); waiting for it to
exit.", sandbox.object_id, e
+ )
+
+ def execute_complete(self, context: Context, event: dict[str, Any]) -> Any:
+ """Read the return value once the trigger reports that the sandbox
exited."""
+ sandbox_id = event["sandbox_id"]
+ self._raise_for_returncode(sandbox_id, event["returncode"])
+ sandbox =
ModalHook(modal_conn_id=self.modal_conn_id).get_sandbox(sandbox_id)
+ return decode_result(sandbox.stdout.read())
+
+ def _raise_for_returncode(self, sandbox_id: str, returncode: int | None)
-> None:
+ if returncode == 0:
+ return
+ if returncode == TIMEOUT_RETURNCODE:
+ timeout = self._sandbox_timeout()
+ limit = f"{timeout}s" if timeout is not None else "300s, Modal's
default"
+ message = (
+ f"Modal sandbox {sandbox_id} reached its timeout ({limit})
before the function "
+ "returned. Raise the decorator's timeout, or the task's
execution_timeout if that "
+ "is shorter."
+ )
+ elif returncode == TERMINATED_RETURNCODE:
+ message = (
+ f"Modal sandbox {sandbox_id} was terminated before the
function returned, for "
+ "example because it ran out of memory under a memory limit, or
from the Modal "
+ "dashboard."
+ )
+ else:
+ message = (
+ f"The function failed in Modal sandbox {sandbox_id} with exit
code {returncode}. "
+ "Its traceback is in the log above."
+ )
+ raise ModalSandboxError(message, sandbox_id=sandbox_id,
returncode=returncode)
+
+ def on_kill(self) -> None:
+ if self._sandbox is None:
+ return
+ self.log.info("Terminating Modal sandbox %s.", self._sandbox.object_id)
+ try:
+ self._sandbox.terminate()
+ except Exception as e:
+ # execute() also calls this while another error propagates; a
failed terminate must
+ # not replace that error.
+ self.log.warning("Failed to terminate Modal sandbox %s: %s",
self._sandbox.object_id, e)
+
+
+class _ModalDecoratedOperator(DecoratedOperator, _ModalSandboxOperator):
+ """
+ Run the decorated function's body in a Modal sandbox.
+
+ The function's source is extracted and sent to the sandbox, as
``@task.docker`` and
+ ``@task.kubernetes`` do, so the body must be self-contained: import what
it needs inside
+ the function, and take every input as an argument. Airflow is not needed
+ in the image, and there is no task context, connection or hook inside the
function. Pass
+ credentials in as Modal secrets.
+
+ Arguments and the return value travel as pickles, so they must load on the
other side:
+ use plain Python values, not objects from packages that only one side has
installed.
+
+ Its arguments are documented on :func:`modal_task`.
+ """
+
+ custom_operator_name = "@task.modal"
+
+ template_fields: Sequence[str] = (*DecoratedOperator.template_fields,
"modal_conn_id")
+
+ # Modal objects hold SDK state that is not meant to be deep-copied, and
are never mutated here.
+ shallow_copy_attrs: Sequence[str] = ("python_callable", "image",
"secrets", "volumes")
+
+ @property
+ def xcom_push(self) -> bool:
+ """Read by ``DecoratedOperator.execute`` on Airflow 3.0, which passes
it to ``_handle_output``."""
+ return self.do_xcom_push
+
+ @xcom_push.setter
+ def xcom_push(self, value: bool) -> None:
+ self.do_xcom_push = value
+
+ def _build_payload(self) -> str:
+ return build_payload(
+ source=self.get_python_source(),
+ function_name=self.python_callable.__name__,
+ args=self.op_args,
+ kwargs=dict(self.op_kwargs),
+ )
+
+
+def modal_task(
+ python_callable: Callable | None = None,
+ multiple_outputs: bool | None = None,
+ **kwargs: Any,
+) -> TaskDecorator:
+ """
+ Wrap a function into a task that runs its body in a Modal sandbox, as
``@task.modal``.
+
+ Takes the arguments below as keywords, plus any ``BaseOperator`` argument.
+
+ :param python_callable: Function to decorate.
+ :param multiple_outputs: If set, a returned dict is unrolled to multiple
XCom values, one
+ per key. When unset, it is inferred from the function's return
annotation.
+ :param image: The ``modal.Image`` to run in. Defaults to
``modal.Image.debian_slim()`` built
+ for the worker's Python version. Any image with Python 3.8 or later on
its ``PATH`` as
+ ``python`` works; the function's dependencies go here, not on the
worker.
+ :param gpu: GPU to request, for example ``"H100"`` or ``"A10G:2"``.
+ :param cpu: CPU cores to reserve, or a ``(reservation, limit)`` pair to
also cap usage.
+ :param memory: Memory to reserve in MiB, or a ``(reservation, limit)``
pair to also cap it.
+ :param timeout: Seconds the sandbox may run before Modal stops it. Modal's
default is 300.
+ When both this and the task's ``execution_timeout`` are set, the
shorter wins; when only
+ ``execution_timeout`` is set, the sandbox gets that.
+ :param secrets: Modal secrets whose keys become environment variables
inside the sandbox,
+ as a name, or a list of names or ``modal.Secret`` objects. Names are
looked up when the
+ task runs, in the connection's environment.
+ :param volumes: Mount points mapped to Modal volumes, as names or
``modal.Volume`` objects.
+ Names are looked up when the task runs, in the connection's
environment.
+ :param modal_conn_id: :ref:`Modal connection id <howto/connection:modal>`.
+ :param app_name: Modal app the sandboxes are created under, ``airflow`` by
default; created
+ on first use.
+ :param deferrable: Wait for the sandbox on the triggerer instead of
holding a worker slot.
+ """
+ return task_decorator_factory(
+ python_callable=python_callable,
+ multiple_outputs=multiple_outputs,
+ decorated_operator_class=_ModalDecoratedOperator,
+ **kwargs,
+ )
diff --git a/providers/modal/src/airflow/providers/modal/exceptions.py
b/providers/modal/src/airflow/providers/modal/exceptions.py
index a1f0e53a97a..2d333a1291d 100644
--- a/providers/modal/src/airflow/providers/modal/exceptions.py
+++ b/providers/modal/src/airflow/providers/modal/exceptions.py
@@ -31,3 +31,12 @@ class ModalError(Exception):
class ModalConnectionError(ModalError):
"""A ``modal`` connection is misconfigured, for example only one of the
two token fields is set."""
+
+
+class ModalSandboxError(ModalError):
+ """A ``@task.modal`` function did not finish successfully in its Modal
sandbox."""
+
+ def __init__(self, message: str, *, sandbox_id: str, returncode: int |
None) -> None:
+ super().__init__(message)
+ self.sandbox_id = sandbox_id
+ self.returncode = returncode
diff --git a/providers/modal/src/airflow/providers/modal/get_provider_info.py
b/providers/modal/src/airflow/providers/modal/get_provider_info.py
index 429e0ff99d7..830da50312c 100644
--- a/providers/modal/src/airflow/providers/modal/get_provider_info.py
+++ b/providers/modal/src/airflow/providers/modal/get_provider_info.py
@@ -25,7 +25,7 @@ def get_provider_info():
return {
"package-name": "apache-airflow-providers-modal",
"name": "Modal",
- "description": "`Modal <https://modal.com/>`__ provider for Apache
Airflow.\nOwns the ``modal`` connection type and a hook that turns it into an
authenticated\nModal client, so operators, toolsets and executors share one
credential path.\n",
+ "description": "`Modal <https://modal.com/>`__ provider for Apache
Airflow.\nOwns the ``modal`` connection type and a hook that turns it into an
authenticated\nModal client, so operators, toolsets and executors share one
credential path, and\n``@task.modal``, which runs a Python function in a Modal
sandbox.\n",
"integrations": [
{
"integration-name": "Modal",
@@ -35,6 +35,12 @@ def get_provider_info():
}
],
"hooks": [{"integration-name": "Modal", "python-modules":
["airflow.providers.modal.hooks.modal"]}],
+ "triggers": [
+ {"integration-name": "Modal", "python-modules":
["airflow.providers.modal.triggers.modal"]}
+ ],
+ "task-decorators": [
+ {"class-name":
"airflow.providers.modal.decorators.modal.modal_task", "name": "modal"}
+ ],
"connection-types": [
{
"hook-class-name":
"airflow.providers.modal.hooks.modal.ModalHook",
diff --git a/providers/modal/src/airflow/providers/modal/hooks/modal.py
b/providers/modal/src/airflow/providers/modal/hooks/modal.py
index d4b05c56984..e5cafda2e1a 100644
--- a/providers/modal/src/airflow/providers/modal/hooks/modal.py
+++ b/providers/modal/src/airflow/providers/modal/hooks/modal.py
@@ -18,18 +18,27 @@
from __future__ import annotations
+import asyncio
import os
+import weakref
from functools import cached_property
from typing import TYPE_CHECKING, Any
import modal
+from airflow.providers.common.compat.connection import get_async_connection
from airflow.providers.common.compat.sdk import AirflowNotFoundException,
BaseHook
from airflow.providers.modal.exceptions import ModalConnectionError
if TYPE_CHECKING:
from airflow.sdk import Connection
+# Async clients shared across hook instances in this process, keyed on (pid,
token id, token secret).
+_async_clients: dict[tuple[int, str, str], modal.Client] = {}
+_async_client_locks: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop,
asyncio.Lock] = (
+ weakref.WeakKeyDictionary()
+)
+
class ModalHook(BaseHook):
"""
@@ -89,20 +98,26 @@ class ModalHook(BaseHook):
return None
try:
return self.get_connection(self.modal_conn_id)
- except AirflowNotFoundException:
- if self.modal_conn_id != self.default_conn_name:
- raise
- self.log.info(
- "Connection %r not found; using ambient Modal credentials
(MODAL_TOKEN_ID / "
- "MODAL_TOKEN_SECRET or ~/.modal.toml).",
- self.modal_conn_id,
- )
+ except AirflowNotFoundException as e:
+ self._fall_back_to_ambient(e)
return None
+ def _fall_back_to_ambient(self, error: AirflowNotFoundException) -> None:
+ """Accept a missing default connection, meaning "use ambient
credentials"; re-raise any other."""
+ if self.modal_conn_id != self.default_conn_name:
+ raise error
+ self.log.info(
+ "Connection %r not found; using ambient Modal credentials
(MODAL_TOKEN_ID / "
+ "MODAL_TOKEN_SECRET or ~/.modal.toml).",
+ self.modal_conn_id,
+ )
+
@cached_property
def credentials(self) -> tuple[str, str] | None:
"""Token id and secret from the connection, or ``None`` to defer to
the Modal SDK."""
- conn = self._connection
+ return self._credentials_from(self._connection)
+
+ def _credentials_from(self, conn: Connection | None) -> tuple[str, str] |
None:
if conn is None:
return None
token_id, token_secret = conn.login or "", conn.password or ""
@@ -152,6 +167,53 @@ class ModalHook(BaseHook):
"""Return the Modal client for this connection (same cached object as
:attr:`client`)."""
return self.client
+ async def get_async_client(self) -> modal.Client:
+ """
+ Return a Modal client for async callers, such as a trigger, without
blocking the event loop.
+
+ The connection is looked up asynchronously, and the client is built
with the Modal SDK's
+ ``.aio`` constructors. Use it with the ``.aio`` form of SDK calls, for
example
+ ``await modal.Sandbox.from_id.aio(sandbox_id, client=client)``.
+
+ Unlike :attr:`client`, which is cached per hook instance, async
clients are shared per
+ process and per token pair. The triggerer is long-lived and runs many
triggers, the Modal
+ SDK has no public way to close a client, and it shares its own ambient
``from_env``
+ client in the same way. The number of clients is bounded by the number
of distinct Modal
+ tokens the connections hold.
+ """
+ connection: Connection | None = None
+ if self.modal_conn_id is not None:
+ try:
+ connection = await get_async_connection(self.modal_conn_id)
+ except AirflowNotFoundException as e:
+ self._fall_back_to_ambient(e)
+ credentials = self._credentials_from(connection)
+ if credentials is None:
+ return await modal.Client.from_env.aio()
+ key = (os.getpid(), *credentials)
+ # One lock per event loop, held across construction: without it,
triggers that start
+ # together all miss the cache and each build a client that Modal keeps
open until exit.
+ loop = asyncio.get_running_loop()
+ lock = _async_client_locks.setdefault(loop, asyncio.Lock())
+ async with lock:
+ client = _async_clients.get(key)
+ if client is None or client.is_closed():
+ client = await modal.Client.from_credentials.aio(*credentials)
+ _async_clients[key] = client
+ return client
+
+ async def aget_sandbox(self, sandbox_id: str) -> modal.Sandbox:
+ """
+ Reattach to an existing sandbox by id without blocking the event loop.
+
+ The async counterpart of :meth:`get_sandbox`; call the ``.aio`` form
of methods on the
+ returned handle.
+
+ :param sandbox_id: The ``sb-...`` id returned by ``Sandbox.object_id``.
+ """
+ client = await self.get_async_client()
+ return await modal.Sandbox.from_id.aio(sandbox_id, client=client)
+
@property
def client_kwargs(self) -> dict[str, Any]:
"""
diff --git a/providers/modal/src/airflow/providers/modal/exceptions.py
b/providers/modal/src/airflow/providers/modal/triggers/__init__.py
similarity index 59%
copy from providers/modal/src/airflow/providers/modal/exceptions.py
copy to providers/modal/src/airflow/providers/modal/triggers/__init__.py
index a1f0e53a97a..13a83393a91 100644
--- a/providers/modal/src/airflow/providers/modal/exceptions.py
+++ b/providers/modal/src/airflow/providers/modal/triggers/__init__.py
@@ -14,20 +14,3 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Exceptions raised by the Modal provider."""
-
-from __future__ import annotations
-
-
-class ModalError(Exception):
- """
- Base class for errors raised by the Modal provider.
-
- Deliberately not an ``AirflowException`` subclass, matching the other AI
providers: callers
- catch provider errors through this base, and the exceptions carry no
Airflow-specific
- task-state semantics.
- """
-
-
-class ModalConnectionError(ModalError):
- """A ``modal`` connection is misconfigured, for example only one of the
two token fields is set."""
diff --git a/providers/modal/src/airflow/providers/modal/triggers/modal.py
b/providers/modal/src/airflow/providers/modal/triggers/modal.py
new file mode 100644
index 00000000000..001f35bf1be
--- /dev/null
+++ b/providers/modal/src/airflow/providers/modal/triggers/modal.py
@@ -0,0 +1,100 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Trigger that waits for a ``@task.modal`` sandbox to exit."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.modal.hooks.modal import ModalHook
+from airflow.providers.modal.utils.python_sandbox import await_exit
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+if TYPE_CHECKING:
+ import modal
+
+
+class ModalSandboxTrigger(BaseTrigger):
+ """
+ Wait for a ``@task.modal`` sandbox to exit, writing its output to the task
log.
+
+ The trigger logs the function's output (the sandbox's stderr) line by
line, waits for the
+ sandbox to exit, and then fires one event with the exit code. The operator
reads the return
+ value in ``execute_complete``. Losing the output stream only stops the
logging.
+
+ If the triggerer restarts, the trigger that takes over reads the output
again from the
+ oldest part Modal still holds, so the task log can repeat lines logged
before the restart.
+
+ :param sandbox_id: The ``sb-...`` id of the sandbox to wait for.
+ :param modal_conn_id: :ref:`Modal connection id <howto/connection:modal>`
the sandbox was
+ created through.
+ """
+
+ def __init__(self, sandbox_id: str, modal_conn_id: str | None =
ModalHook.default_conn_name) -> None:
+ super().__init__()
+ self.sandbox_id = sandbox_id
+ self.modal_conn_id = modal_conn_id
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ return (
+ f"{type(self).__module__}.{type(self).__qualname__}",
+ {"sandbox_id": self.sandbox_id, "modal_conn_id":
self.modal_conn_id},
+ )
+
+ async def run(self) -> AsyncIterator[TriggerEvent]:
+ try:
+ sandbox = await
ModalHook(modal_conn_id=self.modal_conn_id).aget_sandbox(self.sandbox_id)
+ await self._log_output(sandbox)
+ await await_exit(sandbox)
+ except Exception:
+ # The task fails with this error and a retry starts a new sandbox,
so stop this one.
+ # Cancellation is not an Exception and passes through. On
triggerer shutdown or
+ # reassignment another triggerer takes the sandbox over; on the
trigger's timeout,
+ # the sandbox's Modal timeout, never longer than
execution_timeout, stops it.
+ await self._terminate()
+ raise
+ yield TriggerEvent({"sandbox_id": self.sandbox_id, "returncode":
sandbox.returncode})
+
+ async def _log_output(self, sandbox: modal.Sandbox) -> None:
+ try:
+ async for line in sandbox.stderr:
+ self.log.info("%s", line.rstrip("\n"))
+ except Exception as e:
+ # Modal's reader raises once its retries for stream errors run
out. Losing the
+ # output is no reason to stop a function that is still running.
+ self.log.warning(
+ "Lost the output of Modal sandbox %s (%s); waiting for it to
exit.", self.sandbox_id, e
+ )
+
+ async def on_kill(self) -> None:
+ """
+ Terminate the sandbox when a user kills the deferred task.
+
+ Runs on Airflow 3.3+, which calls it when a deferred task is marked
failed, marked
+ success or cleared; older versions never call it. Without it, a
sandbox whose task was
+ killed keeps running, and billing, until its own ``timeout``.
+ """
+ await self._terminate()
+
+ async def _terminate(self) -> None:
+ try:
+ sandbox = await
ModalHook(modal_conn_id=self.modal_conn_id).aget_sandbox(self.sandbox_id)
+ await sandbox.terminate.aio()
+ self.log.info("Terminated Modal sandbox %s", self.sandbox_id)
+ except Exception as e:
+ self.log.warning("Failed to terminate Modal sandbox %s: %s",
self.sandbox_id, e)
diff --git a/providers/modal/src/airflow/providers/modal/exceptions.py
b/providers/modal/src/airflow/providers/modal/utils/__init__.py
similarity index 59%
copy from providers/modal/src/airflow/providers/modal/exceptions.py
copy to providers/modal/src/airflow/providers/modal/utils/__init__.py
index a1f0e53a97a..13a83393a91 100644
--- a/providers/modal/src/airflow/providers/modal/exceptions.py
+++ b/providers/modal/src/airflow/providers/modal/utils/__init__.py
@@ -14,20 +14,3 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Exceptions raised by the Modal provider."""
-
-from __future__ import annotations
-
-
-class ModalError(Exception):
- """
- Base class for errors raised by the Modal provider.
-
- Deliberately not an ``AirflowException`` subclass, matching the other AI
providers: callers
- catch provider errors through this base, and the exceptions carry no
Airflow-specific
- task-state semantics.
- """
-
-
-class ModalConnectionError(ModalError):
- """A ``modal`` connection is misconfigured, for example only one of the
two token fields is set."""
diff --git
a/providers/modal/src/airflow/providers/modal/utils/python_sandbox.py
b/providers/modal/src/airflow/providers/modal/utils/python_sandbox.py
new file mode 100644
index 00000000000..fbc479f42f6
--- /dev/null
+++ b/providers/modal/src/airflow/providers/modal/utils/python_sandbox.py
@@ -0,0 +1,156 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+How ``@task.modal`` runs a function in a Modal sandbox and gets its return
value back.
+
+The sandbox's command is a short bootstrap program (:data:`BOOTSTRAP`) run
with ``python -c``.
+It reads one JSON payload from stdin holding the function's source and its
pickled arguments,
+defines the function and calls it. Everything the function prints, including
what its
+subprocesses write and any traceback, goes to the sandbox's stderr, in order;
the task log
+follows that stream. Stdout carries only the return value, pickled,
base64-encoded and
+prefixed with its length (``<length>:<base64>``), written once the function
has returned, so a
+function that raises leaves stdout empty and exits nonzero.
+
+The bootstrap needs only the Python standard library, so the image needs
neither Airflow nor
+the Modal client. The shared virtualenv script template is deliberately not
used: it
+reconnects to an Airflow supervisor when it finds Airflow installed, and a
sandbox has none.
+
+Both streams outlive the sandbox and can be read through a new handle from
+``Sandbox.from_id``, which is what lets the deferrable path follow the output
on the triggerer
+and collect the return value back on a worker. Both the decorator, on a
worker, and the
+trigger, on the triggerer, use this module.
+"""
+
+from __future__ import annotations
+
+import base64
+import contextlib
+import json
+import pickle
+from collections.abc import Iterable
+from typing import Any
+
+import modal
+
+BOOTSTRAP = r"""
+import __future__, base64, json, os, pickle, sys, tempfile, types
+
+# Keep the real stdout for the result, and point fd 1 at stderr so everything
else the
+# function prints, from Python or from a subprocess, lands on stderr.
+result = os.fdopen(os.dup(1), "wb")
+os.dup2(2, 1)
+payload = json.loads(sys.stdin.read())
+# Compiling from a file, rather than a string, lets tracebacks show the
function's source lines.
+path = os.path.join(tempfile.mkdtemp(prefix="airflow-modal-task-"),
payload["function_name"] + ".py")
+with open(path, "w") as f:
+ f.write(payload["source"])
+# A real, registered module: dataclasses and other tools look a class's module
up in sys.modules.
+module = types.ModuleType("__airflow_modal_task__")
+sys.modules[module.__name__] = module
+# Postponed annotations, as in a Dag file with `from __future__ import
annotations`: the
+# function's annotations may name types that only the Dag file imported.
+code = compile(
+ payload["source"], path, "exec",
flags=__future__.annotations.compiler_flag, dont_inherit=True
+)
+exec(code, module.__dict__)
+arguments = pickle.loads(base64.b64decode(payload["arguments"]))
+value = getattr(module, payload["function_name"])(*arguments["args"],
**arguments["kwargs"])
+# Length-prefixed, so a value Modal kept only the end of is detected instead
of misread.
+encoded = base64.b64encode(pickle.dumps(value))
+result.write(b"%d:" % len(encoded) + encoded)
+result.close()
+"""
+
+#: The sandbox command: run :data:`BOOTSTRAP` with the image's ``python``,
unbuffered so
+#: output reaches the task log as it is printed.
+ENTRYPOINT: tuple[str, ...] = ("python", "-u", "-c", BOOTSTRAP)
+
+#: Exit codes the Modal SDK reports for a sandbox that did not exit on its own.
+TIMEOUT_RETURNCODE = 124
+TERMINATED_RETURNCODE = 137
+
+#: Bytes written to sandbox stdin between drains: half of the Modal SDK's 2
MiB stdin buffer,
+#: which raises ``BufferError`` when a single write would overflow it.
+STDIN_CHUNK_BYTES = 1024 * 1024
+
+
+def build_payload(*, source: str, function_name: str, args: Iterable[Any],
kwargs: dict[str, Any]) -> str:
+ """
+ Return the stdin payload for :data:`BOOTSTRAP`.
+
+ :param source: Source of the function, without its decorators.
+ :param function_name: Name the source defines the function under.
+ :param args: Positional arguments to call it with.
+ :param kwargs: Keyword arguments to call it with.
+ """
+ arguments = pickle.dumps({"args": list(args), "kwargs": kwargs})
+ return json.dumps(
+ {
+ "source": source,
+ "function_name": function_name,
+ "arguments": base64.b64encode(arguments).decode(),
+ }
+ )
+
+
+def send_payload(sandbox: modal.Sandbox, payload: str) -> None:
+ """Write the payload to the sandbox's stdin in pieces that fit the SDK's
buffer, then close it."""
+ data = memoryview(payload.encode())
+ for start in range(0, len(data), STDIN_CHUNK_BYTES):
+ sandbox.stdin.write(data[start : start + STDIN_CHUNK_BYTES])
+ sandbox.stdin.drain()
+ sandbox.stdin.write_eof()
+ sandbox.stdin.drain()
+
+
+def wait_for_exit(sandbox: modal.Sandbox) -> None:
+ """
+ Wait for the sandbox to exit, leaving its exit code on
``sandbox.returncode``.
+
+ ``Sandbox.wait`` raises ``SandboxTimeoutError`` for a sandbox that hit its
timeout even with
+ ``raise_on_termination=False``. The exit code (124) is recorded before it
raises, so the
+ caller reports a timeout the same way as any other exit code.
+ """
+ with contextlib.suppress(modal.exception.SandboxTimeoutError):
+ sandbox.wait(raise_on_termination=False)
+
+
+async def await_exit(sandbox: modal.Sandbox) -> None:
+ """Async counterpart of :func:`wait_for_exit`, for the triggerer's event
loop."""
+ with contextlib.suppress(modal.exception.SandboxTimeoutError):
+ await sandbox.wait.aio(raise_on_termination=False)
+
+
+def decode_result(stdout: str) -> Any:
+ """
+ Return the function's return value from the sandbox's stdout.
+
+ :raises ValueError: When stdout is empty, which happens when the function
did not return,
+ or does not hold a whole value, which happens when Modal kept only the
end of a large one.
+ """
+ if not stdout:
+ raise ValueError("The sandbox exited cleanly but wrote no return
value.")
+ # Modal drops the start of output it cannot keep, which takes the length
header with it
+ # or leaves a header that no longer matches.
+ header, separator, encoded = stdout.partition(":")
+ if not separator or not header.isdigit() or int(header) != len(encoded):
+ raise ValueError(
+ f"The return value ({len(stdout)} characters read back) is
incomplete, most likely "
+ "because it was too large for Modal to keep. Write large results
to a Modal volume "
+ "and return their path instead."
+ )
+ return pickle.loads(base64.b64decode(encoded))
diff --git a/providers/modal/src/airflow/providers/modal/exceptions.py
b/providers/modal/tests/system/__init__.py
similarity index 59%
copy from providers/modal/src/airflow/providers/modal/exceptions.py
copy to providers/modal/tests/system/__init__.py
index a1f0e53a97a..5966d6b1d52 100644
--- a/providers/modal/src/airflow/providers/modal/exceptions.py
+++ b/providers/modal/tests/system/__init__.py
@@ -14,20 +14,4 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Exceptions raised by the Modal provider."""
-
-from __future__ import annotations
-
-
-class ModalError(Exception):
- """
- Base class for errors raised by the Modal provider.
-
- Deliberately not an ``AirflowException`` subclass, matching the other AI
providers: callers
- catch provider errors through this base, and the exceptions carry no
Airflow-specific
- task-state semantics.
- """
-
-
-class ModalConnectionError(ModalError):
- """A ``modal`` connection is misconfigured, for example only one of the
two token fields is set."""
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)
diff --git a/providers/modal/src/airflow/providers/modal/exceptions.py
b/providers/modal/tests/system/modal/__init__.py
similarity index 59%
copy from providers/modal/src/airflow/providers/modal/exceptions.py
copy to providers/modal/tests/system/modal/__init__.py
index a1f0e53a97a..13a83393a91 100644
--- a/providers/modal/src/airflow/providers/modal/exceptions.py
+++ b/providers/modal/tests/system/modal/__init__.py
@@ -14,20 +14,3 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Exceptions raised by the Modal provider."""
-
-from __future__ import annotations
-
-
-class ModalError(Exception):
- """
- Base class for errors raised by the Modal provider.
-
- Deliberately not an ``AirflowException`` subclass, matching the other AI
providers: callers
- catch provider errors through this base, and the exceptions carry no
Airflow-specific
- task-state semantics.
- """
-
-
-class ModalConnectionError(ModalError):
- """A ``modal`` connection is misconfigured, for example only one of the
two token fields is set."""
diff --git a/providers/modal/tests/system/modal/example_modal_task.py
b/providers/modal/tests/system/modal/example_modal_task.py
new file mode 100644
index 00000000000..2d165505260
--- /dev/null
+++ b/providers/modal/tests/system/modal/example_modal_task.py
@@ -0,0 +1,80 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Run Python functions on Modal with ``@task.modal``, in images that do not
contain Airflow."""
+
+from __future__ import annotations
+
+import modal
+
+from airflow.sdk import dag, task
+
+
+@dag(tags=["example", "modal"])
+def example_modal_task():
+ @task
+ def extract() -> list[dict]:
+ return [
+ {"id": 1, "email": "[email protected]", "score": "0.91"},
+ {"id": 2, "email": "[email protected]", "score": "0.47"},
+ ]
+
+ # [START howto_decorator_modal]
+ @task.modal(
+ image=modal.Image.debian_slim().pip_install("pydantic<2"),
+ memory=1024,
+ timeout=600,
+ )
+ def validate(records: list[dict]) -> dict:
+ # pydantic 1 cannot be installed next to Airflow; inside the sandbox
it is the only one.
+ import pydantic
+
+ class Record(pydantic.BaseModel):
+ id: int
+ email: str
+ score: float
+
+ validated = [Record(**record).dict() for record in records]
+ print(f"Validated {len(validated)} records with pydantic
{pydantic.VERSION}")
+ return {"pydantic_version": pydantic.VERSION, "records": validated}
+
+ # [END howto_decorator_modal]
+
+ # [START howto_decorator_modal_deferrable]
+ @task.modal(gpu="T4", timeout=900, deferrable=True)
+ def gpu_report() -> str:
+ import subprocess
+
+ return subprocess.run(
+ ["nvidia-smi", "--query-gpu=name,memory.total",
"--format=csv,noheader"],
+ capture_output=True,
+ text=True,
+ check=True,
+ ).stdout.strip()
+
+ # [END howto_decorator_modal_deferrable]
+
+ validate(extract())
+ gpu_report()
+
+
+example_dag = example_modal_task()
+
+
+from tests_common.test_utils.system_tests import get_test_run # noqa: E402
+
+# Needed to run the example Dag with pytest (see:
contributing-docs/testing/system_tests.rst)
+test_run = get_test_run(example_dag)
diff --git a/providers/modal/src/airflow/providers/modal/exceptions.py
b/providers/modal/tests/unit/modal/decorators/__init__.py
similarity index 59%
copy from providers/modal/src/airflow/providers/modal/exceptions.py
copy to providers/modal/tests/unit/modal/decorators/__init__.py
index a1f0e53a97a..13a83393a91 100644
--- a/providers/modal/src/airflow/providers/modal/exceptions.py
+++ b/providers/modal/tests/unit/modal/decorators/__init__.py
@@ -14,20 +14,3 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Exceptions raised by the Modal provider."""
-
-from __future__ import annotations
-
-
-class ModalError(Exception):
- """
- Base class for errors raised by the Modal provider.
-
- Deliberately not an ``AirflowException`` subclass, matching the other AI
providers: callers
- catch provider errors through this base, and the exceptions carry no
Airflow-specific
- task-state semantics.
- """
-
-
-class ModalConnectionError(ModalError):
- """A ``modal`` connection is misconfigured, for example only one of the
two token fields is set."""
diff --git a/providers/modal/tests/unit/modal/decorators/test_modal.py
b/providers/modal/tests/unit/modal/decorators/test_modal.py
new file mode 100644
index 00000000000..ece0d9e2413
--- /dev/null
+++ b/providers/modal/tests/unit/modal/decorators/test_modal.py
@@ -0,0 +1,428 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import timedelta
+from unittest import mock
+
+import modal
+import pytest
+
+from airflow.providers.modal.decorators.modal import _ModalDecoratedOperator
+from airflow.providers.modal.exceptions import ModalSandboxError
+from airflow.providers.modal.hooks.modal import ModalHook
+from airflow.providers.modal.triggers.modal import ModalSandboxTrigger
+from airflow.providers.modal.utils.python_sandbox import ENTRYPOINT,
send_payload
+from airflow.sdk import DAG, task
+from airflow.sdk.api.datamodels._generated import TaskInstanceState
+from airflow.sdk.execution_time.comms import DeferTask
+
+from unit.modal.local_sandbox import LocalSandbox
+
+
[email protected]
+def sandboxes(tmp_path):
+ """
+ Patch ``ModalHook.create_sandbox`` to run each sandbox locally; yields the
sandboxes created.
+
+ ``get_sandbox`` returns the sandbox with the requested id, as reattaching
does on Modal.
+ """
+ created: list[LocalSandbox] = []
+
+ def create_sandbox(self, *entrypoint, app_name,
create_app_if_missing=False, **sandbox_kwargs):
+ sandbox = LocalSandbox(*entrypoint, tmp_path=tmp_path,
timeout=sandbox_kwargs.get("timeout"))
+ sandbox.create_kwargs = {
+ "app_name": app_name,
+ "create_app_if_missing": create_app_if_missing,
+ **sandbox_kwargs,
+ }
+ created.append(sandbox)
+ return sandbox
+
+ def get_sandbox(self, sandbox_id):
+ return next(sandbox for sandbox in created if sandbox.object_id ==
sandbox_id)
+
+ with (
+ mock.patch.object(ModalHook, "create_sandbox", autospec=True,
side_effect=create_sandbox),
+ mock.patch.object(ModalHook, "get_sandbox", autospec=True,
side_effect=get_sandbox),
+ ):
+ yield created
+
+
+def _operator(xcom_arg) -> _ModalDecoratedOperator:
+ operator = xcom_arg.operator
+ assert isinstance(operator, _ModalDecoratedOperator)
+ return operator
+
+
+class TestModalTask:
+ def test_return_value_reaches_xcom(self, run_task, sandboxes):
+ with DAG(dag_id="modal_dag"):
+
+ @task.modal(image=modal.Image.debian_slim(python_version="3.12"))
+ def summarize(values: list[int], scale: int) -> dict:
+ import statistics
+
+ return {"mean": statistics.mean(values) * scale}
+
+ operator = _operator(summarize([1, 2, 3], scale=10))
+
+ run_task(operator)
+
+ assert run_task.state == TaskInstanceState.SUCCESS, run_task.error
+ assert run_task.xcom.get(key="return_value") == {"mean": 20}
+
+ def test_multiple_outputs_unrolls_the_returned_dict(self, run_task,
sandboxes):
+ with DAG(dag_id="modal_dag"):
+
+ @task.modal(multiple_outputs=True)
+ def split() -> dict[str, int]:
+ return {"rows": 3, "bytes": 1024}
+
+ operator = _operator(split())
+
+ run_task(operator)
+
+ assert run_task.state == TaskInstanceState.SUCCESS, run_task.error
+ assert run_task.xcom.get(key="rows") == 3
+ assert run_task.xcom.get(key="bytes") == 1024
+
+ def test_function_error_fails_the_task_with_its_exit_code(self, run_task,
sandboxes):
+ with DAG(dag_id="modal_dag"):
+
+ @task.modal
+ def fail():
+ raise ValueError("boom")
+
+ operator = _operator(fail())
+
+ run_task(operator)
+
+ assert run_task.state == TaskInstanceState.FAILED
+ assert isinstance(run_task.error, ModalSandboxError)
+ assert run_task.error.returncode == 1
+ assert run_task.error.sandbox_id == "sb-local"
+
+ def test_function_output_goes_to_the_task_log(self, run_task, sandboxes):
+ with DAG(dag_id="modal_dag"):
+
+ @task.modal
+ def chatty():
+ print("loading the model")
+ return "done"
+
+ operator = _operator(chatty())
+
+ with mock.patch.object(type(operator), "log",
new_callable=mock.PropertyMock) as log:
+ run_task(operator)
+
+ logged = [call.args for call in log.return_value.info.call_args_list]
+ output_lines = [args for args in logged if args[0] == "%s"]
+ assert output_lines == [("%s", "loading the model")], "only the
function's output is logged"
+
+ def test_sandbox_gets_the_task_options_and_identifying_tags(self,
run_task, sandboxes):
+ image =
modal.Image.debian_slim(python_version="3.12").pip_install("pydantic<2")
+ secret = modal.Secret.from_dict({"API_KEY": "x"})
+ volume = mock.sentinel.volume
+ with DAG(dag_id="modal_dag"):
+
+ @task.modal(
+ image=image,
+ gpu="H100",
+ cpu=(2.0, 4.0),
+ memory=8192,
+ timeout=3600,
+ secrets=[secret],
+ volumes={"/data": volume},
+ app_name="ml-jobs",
+ )
+ def infer():
+ return None
+
+ operator = _operator(infer())
+
+ run_task(operator, run_id="manual__2026-11-04T10:00:00+00:00",
try_number=2)
+
+ (sandbox,) = sandboxes
+ assert sandbox.entrypoint == ENTRYPOINT
+ assert sandbox.create_kwargs == {
+ "app_name": "ml-jobs",
+ "create_app_if_missing": True,
+ "image": image,
+ "gpu": "H100",
+ "cpu": (2.0, 4.0),
+ "memory": 8192,
+ "timeout": 3600,
+ "secrets": [secret],
+ "volumes": {"/data": volume},
+ "tags": {
+ "airflow_dag_id": "modal_dag",
+ "airflow_task_id": "infer",
+ "airflow_run_id": "manual__2026-11-04T10:00:00+00:00",
+ "airflow_map_index": "-1",
+ "airflow_try_number": "2",
+ },
+ }
+
+ def
test_secret_and_volume_names_resolve_through_the_connection_at_run_time(self,
run_task, sandboxes):
+ named_secret = modal.Secret.from_dict({"HF_TOKEN": "x"})
+ named_volume = modal.Volume.from_name("weights")
+ inline_secret = modal.Secret.from_dict({"API_KEY": "y"})
+ with DAG(dag_id="modal_dag"):
+
+ @task.modal(
+ modal_conn_id="modal_gpu",
+ secrets=["huggingface", inline_secret],
+ volumes={"/weights": "weights"},
+ )
+ def infer():
+ return None
+
+ operator = _operator(infer())
+
+ with (
+ mock.patch.object(
+ ModalHook, "get_secret", autospec=True,
return_value=named_secret
+ ) as get_secret,
+ mock.patch.object(
+ ModalHook, "get_volume", autospec=True,
return_value=named_volume
+ ) as get_volume,
+ ):
+ run_task(operator)
+
+ (sandbox,) = sandboxes
+ assert sandbox.create_kwargs["secrets"] == [named_secret,
inline_secret]
+ assert sandbox.create_kwargs["volumes"] == {"/weights": named_volume}
+ (hook, name), _ = get_secret.call_args
+ assert (hook.modal_conn_id, name) == ("modal_gpu", "huggingface")
+ (hook, name), _ = get_volume.call_args
+ assert (hook.modal_conn_id, name) == ("modal_gpu", "weights")
+
+ def
test_unset_options_are_left_to_modal_and_image_defaults_to_debian_slim(self,
run_task, sandboxes):
+ with DAG(dag_id="modal_dag"):
+
+ @task.modal
+ def noop():
+ return None
+
+ operator = _operator(noop())
+
+ with mock.patch.object(modal.Image, "debian_slim", autospec=True) as
debian_slim:
+ run_task(operator)
+
+ (sandbox,) = sandboxes
+ assert set(sandbox.create_kwargs) == {"app_name",
"create_app_if_missing", "image", "tags"}
+ assert sandbox.create_kwargs["app_name"] == "airflow"
+ assert sandbox.create_kwargs["image"] is debian_slim.return_value
+
+ def test_a_single_secret_name_is_one_secret(self, run_task, sandboxes):
+ secret = modal.Secret.from_dict({"HF_TOKEN": "x"})
+ with DAG(dag_id="modal_dag"):
+
+ @task.modal(secrets="huggingface")
+ def infer():
+ return None
+
+ operator = _operator(infer())
+
+ with mock.patch.object(ModalHook, "get_secret", autospec=True,
return_value=secret) as get_secret:
+ run_task(operator)
+
+ (sandbox,) = sandboxes
+ assert sandbox.create_kwargs["secrets"] == [secret]
+ assert get_secret.call_args.args[1:] == ("huggingface",)
+
+ @pytest.mark.parametrize(
+ ("timeout", "execution_timeout", "sandbox_timeout"),
+ [
+ pytest.param(3600, timedelta(seconds=30), 30,
id="execution-timeout-caps-timeout"),
+ pytest.param(None, timedelta(seconds=30.5), 31,
id="execution-timeout-alone"),
+ pytest.param(10, timedelta(seconds=30), 10,
id="shorter-timeout-wins"),
+ pytest.param(10, None, 10, id="timeout-alone"),
+ ],
+ )
+ def test_execution_timeout_caps_the_sandbox_timeout(
+ self, run_task, sandboxes, timeout, execution_timeout, sandbox_timeout
+ ):
+ with DAG(dag_id="modal_dag"):
+
+ @task.modal(timeout=timeout, execution_timeout=execution_timeout)
+ def noop():
+ return None
+
+ operator = _operator(noop())
+
+ run_task(operator)
+
+ (sandbox,) = sandboxes
+ assert sandbox.create_kwargs["timeout"] == sandbox_timeout
+
+
+class TestFailures:
+ def test_timeout_fails_the_task_with_exit_code_124(self, run_task,
sandboxes):
+ with DAG(dag_id="modal_dag"):
+
+ @task.modal(timeout=1)
+ def nap():
+ import time
+
+ time.sleep(30)
+
+ operator = _operator(nap())
+
+ run_task(operator)
+
+ assert run_task.state == TaskInstanceState.FAILED
+ assert isinstance(run_task.error, ModalSandboxError)
+ assert run_task.error.returncode == 124
+ assert "reached its timeout (1s)" in str(run_task.error)
+
+ def test_losing_the_output_stream_keeps_waiting_for_the_result(self,
run_task, sandboxes, monkeypatch):
+ def broken_stream(self):
+ raise modal.exception.ConnectionError("stream reset")
+
+ monkeypatch.setattr(LocalSandbox, "stderr", property(broken_stream))
+ with DAG(dag_id="modal_dag"):
+
+ @task.modal
+ def train():
+ return {"loss": 0.1}
+
+ operator = _operator(train())
+
+ run_task(operator)
+
+ assert run_task.state == TaskInstanceState.SUCCESS, run_task.error
+ assert run_task.xcom.get(key="return_value") == {"loss": 0.1}
+ (sandbox,) = sandboxes
+ assert not sandbox.terminated
+
+ def test_losing_track_of_the_sandbox_terminates_it(self, run_task,
sandboxes, monkeypatch):
+ def lost(self, raise_on_termination=True):
+ raise modal.exception.ConnectionError("connection lost")
+
+ monkeypatch.setattr(LocalSandbox, "wait", lost)
+ with DAG(dag_id="modal_dag"):
+
+ @task.modal
+ def train():
+ return None
+
+ operator = _operator(train())
+
+ run_task(operator)
+
+ assert run_task.state == TaskInstanceState.FAILED
+ assert isinstance(run_task.error, modal.exception.ConnectionError)
+ (sandbox,) = sandboxes
+ assert sandbox.terminated
+
+
+class TestDeferrable:
+ def test_defers_to_a_trigger_for_the_started_sandbox(self, run_task,
sandboxes):
+ with DAG(dag_id="modal_dag"):
+
+ @task.modal(deferrable=True, modal_conn_id="modal_gpu",
execution_timeout=timedelta(hours=2))
+ def train():
+ return {"loss": 0.1}
+
+ operator = _operator(train())
+
+ run_task(operator)
+
+ assert run_task.state == TaskInstanceState.DEFERRED
+ assert isinstance(run_task.msg, DeferTask)
+ classpath, kwargs = ModalSandboxTrigger(sandbox_id="sb-local",
modal_conn_id="modal_gpu").serialize()
+ assert run_task.msg.classpath == classpath
+ assert run_task.msg.trigger_kwargs == kwargs
+ assert run_task.msg.next_method == "execute_complete"
+ assert run_task.msg.trigger_timeout == timedelta(hours=2)
+ (sandbox,) = sandboxes
+ assert not sandbox.terminated, "deferring must leave the sandbox
running for the trigger"
+
+ def test_execute_complete_reads_the_return_value_from_the_sandbox(self,
sandboxes):
+ with DAG(dag_id="modal_dag"):
+
+ @task.modal(deferrable=True)
+ def train():
+ return {"loss": 0.1}
+
+ operator = _operator(train())
+ sandbox = ModalHook().create_sandbox(*ENTRYPOINT, app_name="airflow")
+ send_payload(sandbox, operator._build_payload())
+
+ result = operator.execute_complete(context={}, event={"sandbox_id":
"sb-local", "returncode": 0})
+
+ assert result == {"loss": 0.1}
+
+ @pytest.mark.parametrize(
+ ("returncode", "message"),
+ [
+ pytest.param(124, r"reached its timeout \(300s, Modal's
default\)", id="timeout"),
+ pytest.param(137, "was terminated", id="terminated"),
+ pytest.param(1, "failed in Modal sandbox sb-1 with exit code 1",
id="function-error"),
+ ],
+ )
+ def test_execute_complete_raises_for_a_failed_sandbox(self, returncode,
message):
+ with DAG(dag_id="modal_dag"):
+
+ @task.modal(deferrable=True)
+ def train():
+ return None
+
+ operator = _operator(train())
+
+ with pytest.raises(ModalSandboxError, match=message) as exc_info:
+ operator.execute_complete(context={}, event={"sandbox_id": "sb-1",
"returncode": returncode})
+
+ assert exc_info.value.returncode == returncode
+
+
+class TestOnKill:
+ def test_terminates_the_running_sandbox(self, tmp_path):
+ with DAG(dag_id="modal_dag"):
+
+ @task.modal
+ def train():
+ return None
+
+ operator = _operator(train())
+ sandbox = LocalSandbox(*ENTRYPOINT, tmp_path=tmp_path)
+ operator._sandbox = sandbox
+
+ operator.on_kill()
+
+ assert sandbox.terminated
+
+ def test_a_failed_terminate_is_logged_not_raised(self, tmp_path,
monkeypatch):
+ def gone(self):
+ raise modal.exception.NotFoundError("sandbox gone")
+
+ monkeypatch.setattr(LocalSandbox, "terminate", gone)
+ with DAG(dag_id="modal_dag"):
+
+ @task.modal
+ def train():
+ return None
+
+ operator = _operator(train())
+ operator._sandbox = LocalSandbox(*ENTRYPOINT, tmp_path=tmp_path)
+
+ with mock.patch.object(type(operator), "log",
new_callable=mock.PropertyMock) as log:
+ operator.on_kill()
+
+ assert "Failed to terminate" in
log.return_value.warning.call_args.args[0]
diff --git a/providers/modal/tests/unit/modal/hooks/test_modal.py
b/providers/modal/tests/unit/modal/hooks/test_modal.py
index 7a444e10c08..de4cf4e9702 100644
--- a/providers/modal/tests/unit/modal/hooks/test_modal.py
+++ b/providers/modal/tests/unit/modal/hooks/test_modal.py
@@ -16,6 +16,7 @@
# under the License.
from __future__ import annotations
+import asyncio
import os
import re
from unittest import mock
@@ -370,3 +371,154 @@ class TestResourceLookups:
sandbox_cls.from_id.assert_called_once_with(
"sb-123", client=modal_client_cls.from_credentials.return_value
)
+
+
+class _FakeClient:
+ """A Modal client as the hook's async path sees it: only ``is_closed`` is
read."""
+
+ def __init__(self, *credentials: str) -> None:
+ self.credentials = credentials
+ self.closed = False
+
+ def is_closed(self) -> bool:
+ return self.closed
+
+
+class _AioConstructor:
+ """A ``modal.Client`` constructor as async code calls it, through
``.aio``; records the calls."""
+
+ def __init__(self) -> None:
+ self.calls: list[tuple[str, ...]] = []
+
+ async def aio(self, *credentials: str) -> _FakeClient:
+ self.calls.append(credentials)
+ return _FakeClient(*credentials)
+
+
+class _FakeClientClass:
+ def __init__(self) -> None:
+ self.from_credentials = _AioConstructor()
+ self.from_env = _AioConstructor()
+
+
[email protected]
+def async_client_cls():
+ """Patch ``modal.Client`` with async constructors and start from an empty
shared-client cache."""
+ client_cls = _FakeClientClass()
+ with (
+ mock.patch("airflow.providers.modal.hooks.modal.modal.Client",
new=client_cls),
+ mock.patch.dict("airflow.providers.modal.hooks.modal._async_clients",
clear=True),
+ ):
+ yield client_cls
+
+
+def _async_lookup(*, returns=None, raises=None):
+ """Stand-in for ``get_async_connection``, as the hook module imports it."""
+
+ # The signature at the provider's floor, common-compat 1.12.0, which has
no ``hook=``.
+ async def get_async_connection(conn_id):
+ if raises is not None:
+ raise raises
+ return returns
+
+ return
mock.patch("airflow.providers.modal.hooks.modal.get_async_connection",
new=get_async_connection)
+
+
[email protected]
+class TestGetAsyncClient:
+ async def test_connection_tokens_build_a_client_shared_across_hooks(self,
async_client_cls):
+ connection = Connection(
+ conn_id="modal_full", conn_type="modal", login=TOKEN_ID,
password=TOKEN_SECRET
+ )
+
+ with _async_lookup(returns=connection):
+ first = await
ModalHook(modal_conn_id="modal_full").get_async_client()
+ second = await
ModalHook(modal_conn_id="modal_full").get_async_client()
+
+ assert first is second
+ assert async_client_cls.from_credentials.calls == [(TOKEN_ID,
TOKEN_SECRET)]
+ assert async_client_cls.from_env.calls == []
+
+ async def test_closed_shared_client_is_replaced(self, async_client_cls):
+ connection = Connection(
+ conn_id="modal_full", conn_type="modal", login=TOKEN_ID,
password=TOKEN_SECRET
+ )
+
+ with _async_lookup(returns=connection):
+ first = await
ModalHook(modal_conn_id="modal_full").get_async_client()
+ first.closed = True
+ second = await
ModalHook(modal_conn_id="modal_full").get_async_client()
+
+ assert second is not first
+ assert len(async_client_cls.from_credentials.calls) == 2
+
+ async def test_connection_without_tokens_uses_ambient_client(self,
async_client_cls):
+ with _async_lookup(returns=Connection(conn_id="modal_ambient",
conn_type="modal")):
+ await ModalHook(modal_conn_id="modal_ambient").get_async_client()
+
+ assert async_client_cls.from_env.calls == [()]
+ assert async_client_cls.from_credentials.calls == []
+
+ async def test_missing_default_connection_falls_back_to_ambient(self,
async_client_cls):
+ with _async_lookup(raises=AirflowNotFoundException("nope")):
+ await ModalHook().get_async_client()
+
+ assert async_client_cls.from_env.calls == [()]
+
+ async def test_missing_named_connection_raises(self, async_client_cls):
+ with _async_lookup(raises=AirflowNotFoundException("nope")),
pytest.raises(AirflowNotFoundException):
+ await ModalHook(modal_conn_id="modal_typo").get_async_client()
+
+ assert async_client_cls.from_env.calls == []
+
+ async def test_half_filled_connection_raises(self, async_client_cls):
+ connection = Connection(conn_id="modal_half", conn_type="modal",
login=TOKEN_ID)
+
+ with _async_lookup(returns=connection),
pytest.raises(ModalConnectionError):
+ await ModalHook(modal_conn_id="modal_half").get_async_client()
+
+ async def test_concurrent_callers_share_one_new_client(self,
async_client_cls):
+ connection = Connection(
+ conn_id="modal_full", conn_type="modal", login=TOKEN_ID,
password=TOKEN_SECRET
+ )
+ constructor = async_client_cls.from_credentials
+ construct = constructor.aio
+
+ async def slow_construct(*credentials):
+ await asyncio.sleep(0.01)
+ return await construct(*credentials)
+
+ constructor.aio = slow_construct
+
+ with _async_lookup(returns=connection):
+ clients = await asyncio.gather(
+ *(ModalHook(modal_conn_id="modal_full").get_async_client() for
_ in range(5))
+ )
+
+ assert len({id(client) for client in clients}) == 1
+ assert constructor.calls == [(TOKEN_ID, TOKEN_SECRET)]
+
+ async def test_aget_sandbox_reattaches_through_the_async_client(self,
async_client_cls):
+ connection = Connection(
+ conn_id="modal_full", conn_type="modal", login=TOKEN_ID,
password=TOKEN_SECRET
+ )
+ reattached: list[tuple[str, object]] = []
+
+ class _FromId:
+ async def aio(self, sandbox_id, *, client):
+ reattached.append((sandbox_id, client))
+ return mock.sentinel.sandbox
+
+ class _SandboxClass:
+ from_id = _FromId()
+
+ with (
+ _async_lookup(returns=connection),
+ mock.patch("airflow.providers.modal.hooks.modal.modal.Sandbox",
new=_SandboxClass),
+ ):
+ hook = ModalHook(modal_conn_id="modal_full")
+ sandbox = await hook.aget_sandbox("sb-1")
+ client = await hook.get_async_client()
+
+ assert sandbox is mock.sentinel.sandbox
+ assert reattached == [("sb-1", client)]
diff --git a/providers/modal/tests/unit/modal/local_sandbox.py
b/providers/modal/tests/unit/modal/local_sandbox.py
new file mode 100644
index 00000000000..0f4dfeb589e
--- /dev/null
+++ b/providers/modal/tests/unit/modal/local_sandbox.py
@@ -0,0 +1,132 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""A stand-in for ``modal.Sandbox`` that runs the sandbox command in a local
subprocess."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import modal
+
+from airflow.providers.modal.utils.python_sandbox import TIMEOUT_RETURNCODE
+
+# Modal's stdin buffer for a sandbox entrypoint, ``MAX_BUFFER_SIZE`` in
``modal/io_streams.py``
+# (not in the package's type stubs, so not imported).
+STDIN_BUFFER_BYTES = 2 * 1024 * 1024
+
+
+class _Stdin:
+ def __init__(self, sandbox: LocalSandbox) -> None:
+ self._sandbox = sandbox
+ self._pending = b""
+ self._received = b""
+
+ def write(self, data: str | bytes | memoryview) -> None:
+ data = data.encode() if isinstance(data, str) else bytes(data)
+ if len(self._pending) + len(data) > STDIN_BUFFER_BYTES:
+ raise BufferError("Buffer size exceed limit. Call drain to flush
the buffer.")
+ self._pending += data
+
+ def write_eof(self) -> None:
+ self._sandbox._stdin_closed = True
+
+ def drain(self) -> None:
+ self._received += self._pending
+ self._pending = b""
+ if self._sandbox._stdin_closed:
+ self._sandbox._run(self._received)
+
+
+class _Stream:
+ """An output stream as Modal's reader serves it: line by line with
newlines, or whole."""
+
+ def __init__(self, text: str) -> None:
+ self._text = text
+
+ def __iter__(self):
+ return iter(self._text.splitlines(keepends=True))
+
+ def read(self) -> str:
+ return self._text
+
+
+class LocalSandbox:
+ """
+ Run a sandbox's command locally, the way Modal runs it in a container.
+
+ Behaves like the parts of ``modal.Sandbox`` that ``@task.modal`` uses:
stdin holds at most
+ 2 MiB between drains and reaches the command once it is closed and
drained, ``stdout`` and
+ ``stderr`` are separate streams read line by line or whole, ``returncode``
stays ``None``
+ until ``wait`` has seen the command exit, and a command that outlives
``timeout`` gets exit
+ code 124 and makes ``wait`` raise ``SandboxTimeoutError`` whatever
``raise_on_termination``
+ says. ``python`` in
+ the command is replaced by the interpreter running the tests, and
``TMPDIR`` points into the
+ test's temporary directory.
+ """
+
+ object_id = "sb-local"
+
+ def __init__(self, *entrypoint: str, tmp_path: Path, timeout: int | None =
None) -> None:
+ self.entrypoint = entrypoint
+ self.tmp_path = tmp_path
+ self.timeout = timeout
+ self.stdin = _Stdin(self)
+ self.returncode: int | None = None
+ self.terminated = False
+ self._stdin_closed = False
+ self._stdout = self._stderr = ""
+ self._exit_code: int | None = None
+
+ def _run(self, stdin: bytes) -> None:
+ program, *args = self.entrypoint
+ command = [sys.executable if program == "python" else program, *args]
+ try:
+ completed = subprocess.run(
+ command,
+ input=stdin,
+ capture_output=True,
+ env={**os.environ, "TMPDIR": str(self.tmp_path)},
+ timeout=self.timeout,
+ check=False,
+ )
+ except subprocess.TimeoutExpired as e:
+ self._stdout = (e.stdout or b"").decode()
+ self._stderr = (e.stderr or b"").decode()
+ self._exit_code = TIMEOUT_RETURNCODE
+ return
+ self._stdout = completed.stdout.decode()
+ self._stderr = completed.stderr.decode()
+ self._exit_code = completed.returncode
+
+ @property
+ def stdout(self) -> _Stream:
+ return _Stream(self._stdout)
+
+ @property
+ def stderr(self) -> _Stream:
+ return _Stream(self._stderr)
+
+ def wait(self, raise_on_termination: bool = True) -> None:
+ self.returncode = self._exit_code
+ if self.returncode == TIMEOUT_RETURNCODE:
+ raise modal.exception.SandboxTimeoutError()
+
+ def terminate(self) -> None:
+ self.terminated = True
diff --git a/providers/modal/tests/unit/modal/test_exceptions.py
b/providers/modal/tests/unit/modal/test_exceptions.py
index 5309e12a0d2..b4ce9c169a7 100644
--- a/providers/modal/tests/unit/modal/test_exceptions.py
+++ b/providers/modal/tests/unit/modal/test_exceptions.py
@@ -18,10 +18,10 @@ from __future__ import annotations
import pytest
-from airflow.providers.modal.exceptions import ModalConnectionError, ModalError
+from airflow.providers.modal.exceptions import ModalConnectionError,
ModalError, ModalSandboxError
[email protected]("exc", [ModalError, ModalConnectionError])
[email protected]("exc", [ModalError, ModalConnectionError,
ModalSandboxError])
def test_provider_errors_share_base_and_are_not_airflow_exceptions(exc):
# Every provider error is catchable via the ModalError base and is
deliberately
# NOT an AirflowException subclass, matching the other AI providers.
diff --git a/providers/modal/src/airflow/providers/modal/exceptions.py
b/providers/modal/tests/unit/modal/triggers/__init__.py
similarity index 59%
copy from providers/modal/src/airflow/providers/modal/exceptions.py
copy to providers/modal/tests/unit/modal/triggers/__init__.py
index a1f0e53a97a..13a83393a91 100644
--- a/providers/modal/src/airflow/providers/modal/exceptions.py
+++ b/providers/modal/tests/unit/modal/triggers/__init__.py
@@ -14,20 +14,3 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Exceptions raised by the Modal provider."""
-
-from __future__ import annotations
-
-
-class ModalError(Exception):
- """
- Base class for errors raised by the Modal provider.
-
- Deliberately not an ``AirflowException`` subclass, matching the other AI
providers: callers
- catch provider errors through this base, and the exceptions carry no
Airflow-specific
- task-state semantics.
- """
-
-
-class ModalConnectionError(ModalError):
- """A ``modal`` connection is misconfigured, for example only one of the
two token fields is set."""
diff --git a/providers/modal/tests/unit/modal/triggers/test_modal.py
b/providers/modal/tests/unit/modal/triggers/test_modal.py
new file mode 100644
index 00000000000..e7ad9a2429f
--- /dev/null
+++ b/providers/modal/tests/unit/modal/triggers/test_modal.py
@@ -0,0 +1,171 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import asyncio
+from unittest import mock
+
+import modal
+import pytest
+
+from airflow.providers.modal.hooks.modal import ModalHook
+from airflow.providers.modal.triggers.modal import ModalSandboxTrigger
+from airflow.triggers.base import TriggerEvent
+
+
+async def _lines(*lines: str):
+ for line in lines:
+ yield line
+
+
+class _AioMethod:
+ """A Modal SDK method as async code calls it, through ``.aio``; records
the calls."""
+
+ def __init__(self, raises: BaseException | None = None) -> None:
+ self.calls: list[tuple[tuple, dict]] = []
+ self.raises = raises
+
+ async def aio(self, *args, **kwargs) -> None:
+ self.calls.append((args, kwargs))
+ if self.raises is not None:
+ raise self.raises
+
+
+class _AsyncSandbox:
+ """The parts of a ``modal.Sandbox`` handle the trigger uses from the event
loop."""
+
+ def __init__(self, *, stderr_lines: list[str], returncode: int | None) ->
None:
+ self.stderr = _lines(*stderr_lines)
+ self.wait = _AioMethod()
+ self.terminate = _AioMethod()
+ self.returncode = returncode
+
+
+def test_serialize_round_trips():
+ trigger = ModalSandboxTrigger(sandbox_id="sb-1", modal_conn_id="modal_gpu")
+
+ classpath, kwargs = trigger.serialize()
+
+ assert classpath ==
"airflow.providers.modal.triggers.modal.ModalSandboxTrigger"
+ assert kwargs == {"sandbox_id": "sb-1", "modal_conn_id": "modal_gpu"}
+ assert ModalSandboxTrigger(**kwargs).serialize() == (classpath, kwargs)
+
+
[email protected]
[email protected]("returncode", [0, 1, 124])
+async def test_run_logs_output_and_fires_once_with_the_exit_code(returncode):
+ sandbox = _AsyncSandbox(stderr_lines=["epoch 1\n", "epoch 2\n"],
returncode=returncode)
+ trigger = ModalSandboxTrigger(sandbox_id="sb-1", modal_conn_id="modal_gpu")
+
+ with (
+ mock.patch.object(ModalHook, "aget_sandbox", autospec=True,
return_value=sandbox) as aget_sandbox,
+ mock.patch.object(type(trigger), "log",
new_callable=mock.PropertyMock) as log,
+ ):
+ events = [event async for event in trigger.run()]
+
+ assert events == [TriggerEvent({"sandbox_id": "sb-1", "returncode":
returncode})]
+ assert aget_sandbox.call_args.args[1:] == ("sb-1",)
+ assert aget_sandbox.call_args.args[0].modal_conn_id == "modal_gpu"
+ assert [call.args for call in log.return_value.info.call_args_list] == [
+ ("%s", "epoch 1"),
+ ("%s", "epoch 2"),
+ ]
+ assert sandbox.wait.calls == [((), {"raise_on_termination": False})]
+
+
[email protected]
+async def test_on_kill_terminates_the_sandbox():
+ sandbox = _AsyncSandbox(stderr_lines=[], returncode=None)
+ trigger = ModalSandboxTrigger(sandbox_id="sb-1")
+
+ with mock.patch.object(ModalHook, "aget_sandbox", autospec=True,
return_value=sandbox):
+ await trigger.on_kill()
+
+ assert sandbox.terminate.calls == [((), {})]
+
+
[email protected]
+async def test_on_kill_swallows_errors_so_the_triggerer_keeps_going():
+ trigger = ModalSandboxTrigger(sandbox_id="sb-gone")
+
+ with mock.patch.object(
+ ModalHook, "aget_sandbox", autospec=True,
side_effect=modal.exception.NotFoundError("gone")
+ ):
+ await trigger.on_kill()
+
+
[email protected]
+async def test_timeout_fires_an_event_with_exit_code_124():
+ # Modal's wait raises SandboxTimeoutError for a timed-out sandbox even with
+ # raise_on_termination=False, after recording exit code 124 on the handle.
+ sandbox = _AsyncSandbox(stderr_lines=["epoch 1\n"], returncode=124)
+ sandbox.wait = _AioMethod(raises=modal.exception.SandboxTimeoutError())
+ trigger = ModalSandboxTrigger(sandbox_id="sb-1")
+
+ with mock.patch.object(ModalHook, "aget_sandbox", autospec=True,
return_value=sandbox):
+ events = [event async for event in trigger.run()]
+
+ assert events == [TriggerEvent({"sandbox_id": "sb-1", "returncode": 124})]
+ assert sandbox.terminate.calls == []
+
+
+async def _stream_ending_with(error: BaseException):
+ yield "epoch 1\n"
+ raise error
+
+
[email protected]
+async def test_losing_the_output_stream_keeps_waiting_for_the_sandbox():
+ sandbox = _AsyncSandbox(stderr_lines=[], returncode=0)
+ sandbox.stderr =
_stream_ending_with(modal.exception.ConnectionError("stream reset"))
+ trigger = ModalSandboxTrigger(sandbox_id="sb-1")
+
+ with mock.patch.object(ModalHook, "aget_sandbox", autospec=True,
return_value=sandbox):
+ events = [event async for event in trigger.run()]
+
+ assert events == [TriggerEvent({"sandbox_id": "sb-1", "returncode": 0})]
+ assert sandbox.terminate.calls == []
+
+
[email protected]
+async def test_losing_track_of_the_sandbox_terminates_it_and_fails():
+ sandbox = _AsyncSandbox(stderr_lines=[], returncode=None)
+ sandbox.wait =
_AioMethod(raises=modal.exception.ConnectionError("connection lost"))
+ trigger = ModalSandboxTrigger(sandbox_id="sb-1")
+
+ with (
+ mock.patch.object(ModalHook, "aget_sandbox", autospec=True,
return_value=sandbox),
+ pytest.raises(modal.exception.ConnectionError),
+ ):
+ [event async for event in trigger.run()]
+
+ assert sandbox.terminate.calls == [((), {})]
+
+
[email protected]
+async def test_cancellation_leaves_the_sandbox_to_the_next_triggerer():
+ sandbox = _AsyncSandbox(stderr_lines=[], returncode=None)
+ sandbox.stderr = _stream_ending_with(asyncio.CancelledError())
+ trigger = ModalSandboxTrigger(sandbox_id="sb-1")
+
+ with (
+ mock.patch.object(ModalHook, "aget_sandbox", autospec=True,
return_value=sandbox),
+ pytest.raises(asyncio.CancelledError),
+ ):
+ [event async for event in trigger.run()]
+
+ assert sandbox.terminate.calls == []
diff --git a/providers/modal/src/airflow/providers/modal/exceptions.py
b/providers/modal/tests/unit/modal/utils/__init__.py
similarity index 59%
copy from providers/modal/src/airflow/providers/modal/exceptions.py
copy to providers/modal/tests/unit/modal/utils/__init__.py
index a1f0e53a97a..13a83393a91 100644
--- a/providers/modal/src/airflow/providers/modal/exceptions.py
+++ b/providers/modal/tests/unit/modal/utils/__init__.py
@@ -14,20 +14,3 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Exceptions raised by the Modal provider."""
-
-from __future__ import annotations
-
-
-class ModalError(Exception):
- """
- Base class for errors raised by the Modal provider.
-
- Deliberately not an ``AirflowException`` subclass, matching the other AI
providers: callers
- catch provider errors through this base, and the exceptions carry no
Airflow-specific
- task-state semantics.
- """
-
-
-class ModalConnectionError(ModalError):
- """A ``modal`` connection is misconfigured, for example only one of the
two token fields is set."""
diff --git a/providers/modal/tests/unit/modal/utils/test_python_sandbox.py
b/providers/modal/tests/unit/modal/utils/test_python_sandbox.py
new file mode 100644
index 00000000000..ed790a2405b
--- /dev/null
+++ b/providers/modal/tests/unit/modal/utils/test_python_sandbox.py
@@ -0,0 +1,163 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import textwrap
+
+import pytest
+
+from airflow.providers.modal.utils.python_sandbox import (
+ ENTRYPOINT,
+ TIMEOUT_RETURNCODE,
+ build_payload,
+ decode_result,
+ send_payload,
+ wait_for_exit,
+)
+
+from unit.modal.local_sandbox import STDIN_BUFFER_BYTES, LocalSandbox
+
+
+def _run(tmp_path, source: str, function_name: str, *args, timeout: int | None
= None, **kwargs):
+ sandbox = LocalSandbox(*ENTRYPOINT, tmp_path=tmp_path, timeout=timeout)
+ payload = build_payload(
+ source=textwrap.dedent(source), function_name=function_name,
args=args, kwargs=kwargs
+ )
+ # LocalSandbox stands in for modal.Sandbox with the same surface.
+ send_payload(sandbox, payload) # type: ignore[arg-type]
+ wait_for_exit(sandbox) # type: ignore[arg-type]
+ return sandbox
+
+
+class TestBootstrap:
+ def test_round_trips_arguments_and_return_value(self, tmp_path):
+ source = """
+ def summarize(values, *, scale):
+ import statistics
+
+ return {"mean": statistics.mean(values) * scale, "n": len(values)}
+ """
+
+ sandbox = _run(tmp_path, source, "summarize", [1, 2, 3], scale=10)
+
+ assert sandbox.returncode == 0
+ assert decode_result(sandbox.stdout.read()) == {"mean": 20, "n": 3}
+
+ def test_none_return_value(self, tmp_path):
+ sandbox = _run(tmp_path, "def noop():\n pass\n", "noop")
+
+ assert sandbox.returncode == 0
+ assert decode_result(sandbox.stdout.read()) is None
+
+ def
test_everything_the_function_prints_goes_to_stderr_and_stdout_holds_only_the_result(self,
tmp_path):
+ source = """
+ def chatty():
+ import subprocess
+ import sys
+
+ print("from print")
+ print("from stderr", file=sys.stderr)
+ subprocess.run([sys.executable, "-c", "print('from a
subprocess')"], check=True)
+ sys.stdout.write("no newline at the end")
+ return 42
+ """
+
+ sandbox = _run(tmp_path, source, "chatty")
+
+ assert list(sandbox.stderr) == [
+ "from print\n",
+ "from stderr\n",
+ "from a subprocess\n",
+ "no newline at the end",
+ ]
+ assert decode_result(sandbox.stdout.read()) == 42
+
+ def
test_exception_exits_nonzero_with_traceback_on_stderr_and_no_result(self,
tmp_path):
+ source = """
+ def fail():
+ print("before the error")
+ raise ValueError("boom")
+ """
+
+ sandbox = _run(tmp_path, source, "fail")
+
+ assert sandbox.returncode == 1
+ lines = [line.rstrip("\n") for line in sandbox.stderr]
+ assert lines[0] == "before the error"
+ assert lines[-1] == "ValueError: boom"
+ assert ' raise ValueError("boom")' in lines, "traceback should
quote the function's source"
+ assert sandbox.stdout.read() == ""
+
+ def test_classes_defined_in_the_function_work(self, tmp_path):
+ source = """
+ def records(rows):
+ from dataclasses import asdict, dataclass
+
+ @dataclass
+ class Row:
+ id: int
+ name: str
+
+ return [asdict(Row(**row)) for row in rows]
+ """
+
+ sandbox = _run(tmp_path, source, "records", [{"id": 1, "name": "ada"}])
+
+ assert sandbox.returncode == 0, sandbox.stderr.read()
+ assert decode_result(sandbox.stdout.read()) == [{"id": 1, "name":
"ada"}]
+
+ def test_arguments_larger_than_the_stdin_buffer_arrive_whole(self,
tmp_path):
+ big = "x" * (STDIN_BUFFER_BYTES + 1)
+
+ sandbox = _run(tmp_path, "def size(text):\n return len(text)\n",
"size", big)
+
+ assert sandbox.returncode == 0
+ assert decode_result(sandbox.stdout.read()) == len(big)
+
+ def
test_annotations_naming_types_the_sandbox_lacks_are_not_evaluated(self,
tmp_path):
+ source = """
+ def identity(frame: pd.DataFrame, rows: Sequence[Row]) -> pd.DataFrame:
+ return frame
+ """
+
+ sandbox = _run(tmp_path, source, "identity", {"a": 1}, [])
+
+ assert sandbox.returncode == 0, sandbox.stderr.read()
+ assert decode_result(sandbox.stdout.read()) == {"a": 1}
+
+
+def test_wait_for_exit_turns_a_timeout_into_its_exit_code(tmp_path):
+ # Modal's wait raises SandboxTimeoutError here even with
raise_on_termination=False.
+ sandbox = _run(tmp_path, "def nap():\n import time\n
time.sleep(30)\n", "nap", timeout=1)
+
+ assert sandbox.returncode == TIMEOUT_RETURNCODE
+
+
+class TestDecodeResult:
+ def test_empty_stdout_raises(self):
+ with pytest.raises(ValueError, match="wrote no return value"):
+ decode_result("")
+
+ @pytest.mark.parametrize("dropped", [1, 2, 3, 4, 5, 6, 7, 1000, 1001,
1002, 1003])
+ def test_a_value_modal_kept_only_the_end_of_raises_a_clear_error(self,
tmp_path, dropped):
+ # Modal drops the start of output it cannot keep. Cuts at every base64
alignment and
+ # inside the length header must all be caught, not decoded into a
wrong value.
+ sandbox = _run(tmp_path, "def big():\n return list(range(5000))\n",
"big")
+ whole = sandbox.stdout.read()
+
+ with pytest.raises(ValueError, match="too large for Modal to keep"):
+ decode_result(whole[dropped:])
