kaxil opened a new pull request, #76502: URL: https://github.com/apache/airflow/pull/76502
Adds `@task.modal` to the Modal provider. It 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 a normal worker, the way `@task.kubernetes` runs one in a pod. It is for code whose dependencies cannot be installed next to Airflow (`pydantic<2`, a pinned vLLM or CUDA build) or that needs a GPU the workers lack. The image needs only `python`: no Airflow, no Modal client. The usual workaround, `@task.virtualenv`, still installs the dependencies on the worker and cannot give you a GPU. ```python @task.modal(image=modal.Image.debian_slim().pip_install("pydantic<2"), memory=1024, timeout=600) def validate(records: list[dict]) -> dict: import pydantic # 1.x, only inside the sandbox ... ``` Options pass through to `modal.Sandbox.create` (`image`, `gpu`, `cpu`, `memory`, `timeout`, `secrets`, `volumes`), credentials come from `ModalHook` (`modal_conn_id`), and `deferrable=True` waits on the triggerer instead of a worker slot. The new docs page also points users with already-deployed Modal functions at the existing `ModalHook.get_function(...).remote(...)` instead. ## Design rationale **Sandbox plus a source-shipping bootstrap, not Modal's serialized functions.** I started from an ephemeral `App` with `@app.function(serialized=True)`, the most direct way to run a local function on Modal. Cloudpickled functions carry bytecode, though, so the image's Python minor version would have to match the worker's; here only the arguments and return value are pickled, and a Python 3.10 image ran fine from a 3.12 worker. Deferral is the other problem. An ephemeral app stops when `App.run()` exits, so it would need detached apps that something has to clean up, whereas a sandbox has its own server-side lifetime, can be reattached by id from the triggerer, and stopped from `on_kill`. Function logs would also only reach the task through Modal's process-global rich output manager. The authoring contract is the one `@task.docker` and `@task.kubernetes` already have: the function's source is extracted, so it imports inside its body and takes every input as an argument. **Why not reuse the virtualenv script template, as `@task.docker` does.** When it finds Airflow installed, that template reconnects to an Airflow supervisor over a file descriptor. A sandbox has none, so an image that happens to contain the task SDK (easy to do if you reuse a runtime image) would crash. **Output and return value use separate streams.** The function's stdout and stderr, including what its subprocesses write, go to the sandbox's stderr, which the operator (or the trigger) follows into the task log line by line. Stdout carries only the return value, as `<length>:<base64 pickle>`. Modal keeps only the most recent part of a large output, so the length header makes a truncated value fail with a clear error instead of decoding into something wrong. **A sandbox does not outlive its task.** - Sync tasks terminate it in `on_kill`. - Deferred tasks terminate it from the trigger's `on_kill`, which runs on Airflow 3.3 and later. - The sandbox's Modal `timeout` is capped at the task's `execution_timeout` and also passed to `defer()`, because the triggerer drops a timed-out trigger without calling `on_kill`. - If the task loses track of the sandbox, it terminates it, so a retry does not run beside it. Losing only the output stream logs a warning and keeps waiting. **`ModalHook` gains `get_async_client` and `aget_sandbox` for the triggerer.** Async clients are shared per process and token pair: Modal 1.6 has no public way to close a client, and its own `from_env` client is shared the same way. ## Evidence Run on a live Airflow stack (scheduler, worker and triggerer as separate processes) against real Modal, with the example Dag from the docs plus a Dag of failure cases: | Task | Mode | Outcome | |---|---|---| | `validate`: `pydantic<2` image | sync | success. Logged `Validated 2 records with pydantic 1.10.27` and pushed the validated records to XCom. | | `gpu_report`: `gpu="T4"` | deferrable | success. Trigger fired, the task resumed, and XCom is `"Tesla T4, 15360 MiB"`. | | function raises | sync | failed: traceback in the task log, then `ModalSandboxError ... exit code 1` | | function raises | deferrable | failed: traceback logged through the triggerer, then the same error | | `timeout=10`, sleeps 60 s | sync | failed: `KeyboardInterrupt` traceback, then `reached its timeout (10s)` | | `memory=(256, 256)`, allocates 1 GiB | sync | failed: exit code 137, "terminated ... ran out of memory" | ## Gotchas - Arguments and return values are pickled, so pass plain Python values. Dates that come from another task's XCom are `pendulum` objects and need `pendulum` in the image, or pass ISO strings. - With ambient Modal credentials rather than a connection, the triggerer needs them too for `deferrable=True`. - On Airflow 3.0 to 3.2, marking a deferred task failed does not stop its sandbox; it runs until it exits or reaches its timeout. That is because trigger `on_kill` arrived in 3.3. - `common-compat` is marked `# use next version`. On Airflow before 3.3.2, the trigger's async connection lookup needs the compat release that resolves async connections off the event loop. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
