3131
3232try :
3333 import modal
34+
35+ from airflow .providers .modal .exceptions import ModalConnectionError
36+ from airflow .providers .modal .hooks .modal import ModalHook
3437except ModuleNotFoundError as e :
3538 from airflow .providers .common .compat .sdk import AirflowOptionalProviderFeatureException
3639
5457 encode_network_policy ,
5558 is_sandbox_handle ,
5659)
60+ from airflow .providers .common .compat .sdk import AirflowNotFoundException
5761
5862if TYPE_CHECKING :
5963 from collections .abc import Mapping , Sequence
106110
107111EgressEnforcement = Literal ["strict" , "sni" ]
108112
113+ # What reading credentials from the connection can raise before Modal is reached at all: a
114+ # named connection that does not exist, or one that sets only half of the token.
115+ _CREDENTIAL_ERRORS = (AirflowNotFoundException , ModalConnectionError )
116+ # Serializes building a backend's Modal client. ModalHook resolves its connection and builds
117+ # its client without a lock, so concurrent first calls on a backend shared by several agent
118+ # runs would each open a gRPC client and keep one. Module-level rather than per instance
119+ # because a lock cannot be deep-copied, and operators holding a backend are.
120+ _CLIENT_BUILD_LOCK = threading .Lock ()
121+
109122
110123@dataclass
111124class _TrackedSandbox :
@@ -196,11 +209,12 @@ class ModalSandboxBackend(AttachableSandboxBackend):
196209 tags; keys starting with ``airflow_`` are reserved for them and overwrite any
197210 ``tags`` of the same name.
198211
199- **Credentials are ambient.** Modal is authenticated the same way its CLI is: run
200- ``modal token new`` once to write ``~/.modal.toml``, or set ``MODAL_TOKEN_ID`` and
201- ``MODAL_TOKEN_SECRET`` in the worker environment. Nothing is read until the first
202- sandbox is created, so a Dag file that constructs this backend parses without
203- credentials present.
212+ **Credentials come from a** ``modal`` **connection**, resolved by
213+ :class:`~airflow.providers.modal.hooks.modal.ModalHook` in the order the
214+ :ref:`Modal connection page <howto/connection:modal>` sets out; without a
215+ ``modal_default`` connection that is the worker's own Modal credentials. Nothing is
216+ read until the first sandbox is created or attached to, so a Dag file that constructs
217+ this backend parses without the connection or any credentials present.
204218
205219 **What the network policy can and cannot promise.** ``SandboxSpec(block_network=True)``
206220 maps to Modal's own ``block_network``, which drops all outbound traffic including DNS.
@@ -235,6 +249,9 @@ class ModalSandboxBackend(AttachableSandboxBackend):
235249 :meth:`read_file`, which deliberately stays on the base class's shell implementation
236250 (see that method). Any Debian or Ubuntu based image, including ``python:*-slim``, does.
237251
252+ :param modal_conn_id: The ``modal`` connection to authenticate with. Default
253+ ``"modal_default"``. ``None`` uses the worker's Modal credentials without looking
254+ for a connection.
238255 :param image: Registry tag for the sandbox image, or a prepared ``modal.Image``.
239256 Default ``"python:3.12-slim"``. An image carrying the packages an agent needs is
240257 the alternative to opening egress so it can install them:
@@ -287,6 +304,7 @@ class ModalSandboxBackend(AttachableSandboxBackend):
287304 def __init__ (
288305 self ,
289306 * ,
307+ modal_conn_id : str | None = ModalHook .default_conn_name ,
290308 image : str | modal .Image = DEFAULT_IMAGE ,
291309 app_name : str = DEFAULT_APP_NAME ,
292310 create_app_if_missing : bool = True ,
@@ -330,6 +348,9 @@ def __init__(
330348 _validate_positive_finite (memory , "memory" )
331349 if egress_enforcement not in ("strict" , "sni" ):
332350 raise ValueError (f"egress_enforcement must be 'strict' or 'sni', got { egress_enforcement !r} ." )
351+ # Constructing the hook reads no connection and opens no client, so this is safe at
352+ # Dag-parse time; both wait for _connected_hook.
353+ self ._hook = ModalHook (modal_conn_id = modal_conn_id )
333354 self ._image = image
334355 self ._app_name = app_name
335356 self ._create_app_if_missing = create_app_if_missing
@@ -347,6 +368,25 @@ def __init__(
347368 # runs never has two of them touching the same entry.
348369 self ._tracked : dict [str , _TrackedSandbox ] = {}
349370
371+ def _connected_hook (self ) -> ModalHook :
372+ """
373+ Return the hook with its client built, the one place every Modal call goes through.
374+
375+ Nothing the model does can fix a missing or half-filled connection, so it fails the
376+ task, naming the connection the author set. That includes ``destroy``: the sandbox
377+ would otherwise bill until its lifetime ends, and the toolset's own teardown catches
378+ and logs it, so it still never fails a finished agent run. A Modal error from
379+ building the client propagates unchanged for the caller to classify.
380+ """
381+ with _CLIENT_BUILD_LOCK :
382+ try :
383+ self ._hook .get_conn ()
384+ except _CREDENTIAL_ERRORS as e :
385+ raise SandboxTerminalError (
386+ f"Could not read Modal credentials from connection { self ._hook .modal_conn_id !r} : { e } "
387+ ) from e
388+ return self ._hook
389+
350390 # ------------------------------------------------------------------
351391 # Lifecycle.
352392 # ------------------------------------------------------------------
@@ -359,16 +399,16 @@ def create(self, *, spec: SandboxSpec | None = None) -> str:
359399 self ._check_attachable (spec )
360400 name = _new_sandbox_name ()
361401 try :
362- # Looked up per create rather than memoized on the instance: the lookup is
363- # idempotent server-side, and one round trip per sandbox is cheaper than the
364- # locking a shared cache would need, given the toolset shares one backend
365- # across concurrent runs.
366- app = modal .App .lookup (self ._app_name , create_if_missing = self ._create_app_if_missing )
367402 # A registry tag is the common case; a prepared Image is how packages get
368403 # into the sandbox without opening egress to fetch them.
369404 image = modal .Image .from_registry (self ._image ) if isinstance (self ._image , str ) else self ._image
370- sandbox = modal .Sandbox .create (
371- app = app ,
405+ # The hook looks the app up in the connection's environment on every create
406+ # rather than memoizing it: the lookup is idempotent server-side, and one round
407+ # trip per sandbox is cheaper than the locking a shared cache would need, given
408+ # the toolset shares one backend across concurrent runs.
409+ sandbox = self ._connected_hook ().create_sandbox (
410+ app_name = self ._app_name ,
411+ create_app_if_missing = self ._create_app_if_missing ,
372412 image = image ,
373413 name = name ,
374414 tags = self ._create_tags (spec , name ),
@@ -497,7 +537,9 @@ def destroy(self, sandbox: str) -> None:
497537 self ._check_handle (sandbox )
498538 tracked = self ._tracked .pop (sandbox , None )
499539 try :
500- target = tracked .handle if tracked and tracked .handle else modal .Sandbox .from_id (sandbox )
540+ target = (
541+ tracked .handle if tracked and tracked .handle else self ._connected_hook ().get_sandbox (sandbox )
542+ )
501543 target .terminate ()
502544 except modal .exception .NotFoundError :
503545 # Already gone, which is what destroy is for. Idempotent by contract.
@@ -984,7 +1026,7 @@ def _handle(self, sandbox: str) -> modal.Sandbox:
9841026 try :
9851027 # Reachable when the handle came from another backend instance, or from a
9861028 # caller that created the sandbox elsewhere.
987- found = modal . Sandbox . from_id (sandbox )
1029+ found = self . _connected_hook (). get_sandbox (sandbox )
9881030 except modal .exception .Error as e :
9891031 raise self ._as_sandbox_error (e , sandbox = sandbox ) from e
9901032 if tracked is None :
@@ -1050,9 +1092,15 @@ def _as_sandbox_error(self, error: modal.exception.Error, *, sandbox: str | None
10501092 # error. Only reachable if it escapes from somewhere that cannot report it.
10511093 return SandboxError (f"The command hit its deadline: { error } " )
10521094 if isinstance (error , modal .exception .AuthError ):
1095+ if self ._hook .credentials is not None :
1096+ # The token came from the connection, so worker settings would change nothing.
1097+ return SandboxTerminalError (
1098+ f"Modal rejected the token on connection { self ._hook .modal_conn_id !r} . Update its "
1099+ f"login and password: { error } "
1100+ )
10531101 return SandboxTerminalError (
1054- "Modal rejected the credentials. Run 'modal token new', or set "
1055- f"MODAL_TOKEN_ID and MODAL_TOKEN_SECRET on the worker: { error } "
1102+ "Modal rejected the credentials. Create a 'modal' connection, run 'modal token "
1103+ f"new', or set MODAL_TOKEN_ID and MODAL_TOKEN_SECRET on the worker: { error } "
10561104 )
10571105 if isinstance (error , modal .exception .PermissionDeniedError ):
10581106 # A sibling of AuthError rather than a subclass, and just as unfixable by
0 commit comments