jason810496 commented on code in PR #70805:
URL: https://github.com/apache/airflow/pull/70805#discussion_r3711102499
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -385,23 +399,134 @@ class SubprocessCoordinator(BaseCoordinator):
:param task_startup_timeout: Maximum time the coordinator waits for the
subprocess to connect to both servers, in seconds. The default is 10
seconds.
+ :param dag_bundle_name: Locate artifacts through a configured Dag bundle
rather
+ than an explicit root. Mutually exclusive with the subclass's explicit
root;
+ if neither is set, the task's own bundle is used. A named bundle
resolves to
+ its latest version; the task's own bundle is pinned to the run's
version.
"""
task_startup_timeout: float = 10.0
+ dag_bundle_name: str | None = None
+
+ # Classified once at construction by :meth:`_classify_artifact_source` and
+ # dispatched on by :meth:`_init_root_source` at execute time.
+ _artifact_source: _ArtifactSource | None = attrs.field(init=False,
default=None)
+ # The subclass's explicit root, recorded at construction so the base can
+ # resolve roots without knowing the subclass field name.
+ _configured_roots: list[pathlib.Path] = attrs.field(init=False,
factory=list)
+ # The task's own bundle, bound for the duration of a single
:meth:`execute_task`
+ # call by :meth:`_set_current_bundle` so :meth:`_init_root_source` can
resolve
+ # co-located artifacts.
+ _active_bundle_info: BundleInfo | None = attrs.field(init=False,
default=None)
+
+ def _classify_artifact_source(self, configured: Sequence[pathlib.Path], *,
root_kwarg: str) -> None:
+ """
+ Classify and validate how this coordinator locates artifacts
(construction time).
+
+ Subclasses call this from ``__attrs_post_init__`` with their own root
+ field. It rejects setting both an explicit root and
``dag_bundle_name``,
+ fails fast when ``dag_bundle_name`` names a bundle that is not
configured,
+ and records the resulting :class:`_ArtifactSource` and explicit root.
+ """
+ if configured and self.dag_bundle_name is not None:
+ raise ValueError(
+ f"Set at most one of {root_kwarg!r} or 'dag_bundle_name':
{root_kwarg!r} for an "
+ f"explicit path, 'dag_bundle_name' for a configured Dag
bundle, or leave both "
+ f"unset to scan the task's own bundle."
+ )
+ if configured:
+ source = _ArtifactSource.EXPLICIT_ROOT
+ self._configured_roots = list(configured)
+ elif self.dag_bundle_name is not None:
+ source = _ArtifactSource.NAMED_BUNDLE
+ from airflow.dag_processing.bundles.manager import
DagBundlesManager # noqa: SDK002
+
+ if not
DagBundlesManager.is_bundle_configured(self.dag_bundle_name):
+ raise ValueError(
+ f"Coordinator 'dag_bundle_name' references unconfigured
Dag bundle "
+ f"{self.dag_bundle_name!r}."
+ )
+ else:
+ source = _ArtifactSource.TASK_BUNDLE
+
+ self._artifact_source = source
+ details: dict[str, str | list[str]] = {"mode": source.name}
+ if self.dag_bundle_name is not None:
+ details["dag_bundle_name"] = self.dag_bundle_name
+ if self._configured_roots:
+ details["configured_roots"] = [str(root) for root in
self._configured_roots]
+ log.debug("Coordinator artifact source selected", **details)
+
+ def _init_root_source(self) -> list[pathlib.Path]:
+ """
+ Resolve the directories to scan for artifacts for the current task.
+
+ Dispatches on the :class:`_ArtifactSource` classified at construction.
+ An explicit root is returned as-is (no Dag bundle is resolved);
otherwise
+ the root is a Dag bundle's materialized path — the named
+ ``dag_bundle_name`` bundle at its latest version, or the task's own
+ bundle pinned to the run's version. Called by :meth:`execute_task`,
which
+ forwards the result to :meth:`_build_execute_task_command`.
+ """
+ if self._artifact_source is _ArtifactSource.EXPLICIT_ROOT:
+ return self._configured_roots
+
+ if self._artifact_source is _ArtifactSource.NAMED_BUNDLE:
+ from airflow.sdk.api.datamodels._generated import BundleInfo
+
+ # NAMED_BUNDLE implies dag_bundle_name is set.
+ target = BundleInfo(name=cast("str", self.dag_bundle_name))
+ elif self._artifact_source is _ArtifactSource.TASK_BUNDLE:
+ if self._active_bundle_info is None:
+ raise RuntimeError("_init_root_source requires an active task;
call it during execute_task.")
+ target = self._active_bundle_info
+ else:
+ raise RuntimeError(
+ "Coordinator artifact source was not classified; call
_classify_artifact_source first."
+ )
+
+ # Lazy import: task_runner is a heavy module and importing it at module
+ # load would risk an import cycle through the supervisor.
+ from airflow.sdk.execution_time.task_runner import initialize_ti_bundle
- def _build_execute_task_command(self, *, what: TaskInstance) ->
tuple[list[str], str | None]:
+ bundle = initialize_ti_bundle(target, log)
Review Comment:
Follow-up is up as #71075.
It moves the RUNNING transition to the top of
`SubprocessCoordinator.execute_task`, ahead of `_build_execute_task_command`,
so artifact resolution and the bundle materialization this PR adds are charged
to the task's runtime instead of to `[scheduler] task_queued_timeout`. The pid
reported there has to be the supervisor's rather than the runtime's: the server
keeps whatever pid `task_instances.start` carried and 409s a heartbeat that
disagrees with it, and it rejects a second `start` carrying a new pid, so there
is no other way for the transition to precede the spawn.
Two things fall out of that and are in the same PR. A heartbeat now runs
across the launch window, because nothing else heartbeats until `wait()` starts
monitoring and a slow launch would otherwise be reaped by
`task_instance_heartbeat_timeout`. And a runtime that never connects gets a
terminal state reported for it, rather than the run sitting QUEUED for the
stuck-in-queued handler to requeue; the pre-handshake stdout/stderr that
`_accept_connections` used to discard goes to the task log, which for "could
not find or load main class" is the only explanation there is.
The logger half of this comment is already addressed in this PR:
`execute_task` resolves `logger or log` and threads it through
`_init_root_source` into `initialize_ti_bundle`.
---
Drafted-by: Claude Code (Opus 5) (no human review before posting)
--
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]
