uranusjr commented on code in PR #70805:
URL: https://github.com/apache/airflow/pull/70805#discussion_r3827704844
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -385,23 +432,146 @@ 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
+ the version current when the task starts; the task's own bundle uses
the
+ run's version. Either way the resolved version is pinned for the whole
task.
"""
task_startup_timeout: float = 10.0
+ dag_bundle_name: str | None = None
+
+ _artifact_source: _ArtifactSource = attrs.field(init=False)
+ # 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)
+ _active_bundle_info: BundleInfo | None = attrs.field(init=False,
default=None)
+ _active_scan_roots: tuple[pathlib.Path, ...] | None =
attrs.field(init=False, default=None)
+
+ @property
+ def _explicit_artifact_roots(self) -> tuple[str, Sequence[pathlib.Path]]:
+ """
+ The subclass's explicit-root kwarg name and its configured value.
+
+ The name is only used in error messages. An empty value — the default,
for a
+ subclass that does not override this — selects task-bundle mode rather
than
+ failing at execute time.
+ """
+ return "root", ()
+
+ def __attrs_post_init__(self) -> None:
+ self._classify_artifact_source()
+
+ def _classify_artifact_source(self) -> None:
+ """
+ Classify and validate how this coordinator locates artifacts
(construction time).
+
+ 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.
+ """
+ root_kwarg, configured = self._explicit_artifact_roots
+ 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
+ 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
+ log.debug(
+ "Coordinator artifact source selected",
+ mode=source.name,
+ dag_bundle_name=self.dag_bundle_name,
+ configured_roots=[str(root) for root in self._configured_roots],
+ )
+
+ def _init_root_source(
+ self, logger: FilteringBoundLogger
+ ) -> tuple[list[pathlib.Path], BaseDagBundle | None]:
+ """
+ Resolve the directories to scan for artifacts, dispatched on the
classified mode.
+
+ Returns ``(roots, bundle)``: an explicit root yields no bundle
(``None``);
+ a Dag-bundle mode returns the materialized path and the resolved
bundle so
+ :meth:`execute_task` can hold a version lock over it. *logger* is the
task
+ logger, so materialization failures surface in the task log.
+ """
+ if self._artifact_source is _ArtifactSource.EXPLICIT_ROOT:
+ return self._configured_roots, None
+
+ if self._artifact_source is _ArtifactSource.NAMED_BUNDLE:
+ # NAMED_BUNDLE implies dag_bundle_name is set.
+ target = BundleInfo(name=cast("str", self.dag_bundle_name))
+ else:
+ if self._active_bundle_info is None:
+ raise RuntimeError("_init_root_source requires an active task;
call it during execute_task.")
Review Comment:
This guard only exists because bundle info is stored on the instance instead
of being passed in. `_init_root_source` is not supposed to be overridden or
manually called by a subclass, so we can simply pass the bundle info into this
function as an argument instead.
This is separate from the `_get_scan_roots` compatibility question. No
external contract constrains this method's signature, so keeping the subclass
hook stable doesn't require the bundle to be stashed too.
--
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]
