kaxil commented on code in PR #70805:
URL: https://github.com/apache/airflow/pull/70805#discussion_r3690030816
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -414,18 +539,20 @@ def execute_task(
subprocess_logs_to_stdout: bool,
**kwargs,
) -> BaseCoordinator.ExecutionResult:
- command, subprocess_schema_version =
self._build_execute_task_command(what=what)
- process = _PopenActivitySubprocess.start(
- what=what,
- dag_rel_path=dag_rel_path,
- bundle_info=bundle_info,
- client=client,
- logger=logger,
- subprocess_logs_to_stdout=subprocess_logs_to_stdout,
- sentry_integration=sentry_integration,
- command=command,
- subprocess_schema_version=subprocess_schema_version,
- startup_timeout=self.task_startup_timeout,
- )
- exit_code = process.wait()
- return self.ExecutionResult(exit_code, process.final_state)
+ with self._set_current_bundle(bundle_info):
Review Comment:
The Python path holds `BundleVersionLock` around the whole run
(`task_runner.py:2375`), but that lock lives in `main()`, not in `parse()`, so
extracting `initialize_ti_bundle` gives this path the materialization without
the protection.
Two consequences. `BundleUsageTrackingManager._remove_stale_bundle`
(`bundles/base.py:150-177`) takes `flock(LOCK_EX|LOCK_NB)` and `rmtree`s the
version directory, backing off only on `BlockingIOError`, and Celery workers
run that cleanup loop in their own subprocess (`celery_command.py:131-138`) on
the same host as the JVM. In a mixed Python + language deployment a version
that already has a tracking file from an earlier Python task can be deleted
while a JVM is still lazily loading classes off that classpath. Separately, on
a coordinator-only worker no tracking file is ever written, so
`_find_all_tracking_files` returns nothing and every version cloned here is
invisible to cleanup and never reclaimed.
Could `_init_root_source` return the bundle rather than just the path, and
this block hold `BundleVersionLock(bundle_name=bundle.name,
bundle_version=bundle.version)` across `start()` and `wait()`? Making
`_set_current_bundle` an `ExitStack` that also enters the lock would mirror
`task_runner.main()`.
##########
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))
Review Comment:
`BundleInfo` with no version resolves to the shared mutable checkout rather
than a per-version tree. `GitDagBundle.__init__` sets `repo_path = base_dir /
"tracking_repo"` when version is falsy (`providers/git/.../git.py:84-87`), and
`_initialize()` falls through to `refresh()`, which fetches and then does
`head.reset(target, index=True, working_tree=True)`.
With `worker_concurrency > 1`, task B's startup hard-resets the tree that
task A's JVM is reading its JARs from. `bundle.lock()` covers the refresh
itself but not the subsequent use, and `BundleVersionLock.acquire()` returns
early when version is falsy (`base.py:492`), so this mode cannot be protected
by the existing mechanism at all. The same exposure applies to the task-bundle
mode whenever the run's `bundle_version` is None.
Would resolving `get_current_version()` once and re-resolving with that
concrete version work? That gives a `versions/<sha>` tree a lock can actually
protect. It also pulls the two network fetches out of this path, which today
run before `_PopenActivitySubprocess.start` and so are not covered by
`task_startup_timeout`.
##########
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(
Review Comment:
Classification is opt-in per subclass, so a third-party subclass that
defines `__attrs_post_init__` for its own reasons and does not call
`_classify_artifact_source` constructs fine and then dies here on every task.
Could the base define `__attrs_post_init__` itself, calling an overridable
`_explicit_artifact_roots()` that defaults to empty, plus a `_root_kwarg`
ClassVar? A subclass that overrides nothing would then land on the task-bundle
mode, which is the documented default, instead of a runtime error.
##########
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
Review Comment:
`dag_bundle_name` does not appear in any `.rst`, and the three config tables
still list the root parameter as *(required)*: `java.rst:665`, `go.rst:427`,
`typescript.rst:272`. Those become wrong on merge. The three-mode semantics
live only in this docstring, so there is nowhere a user can read when to pick
co-located over named-bundle over explicit root.
`contributing-docs/30_new_language_sdk.rst:107` also documents the old hook
signature verbatim (`def _build_execute_task_command(self, *, what:
TaskInstanceDTO) -> tuple[list[str], str]`), which this PR changes. task-sdk
1.3.0 shipped that signature, though the docs pages do carry `|experimental|`
(java.rst:23), so this is more about whether you want to call out the break
than a blocker.
##########
task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py:
##########
@@ -370,13 +370,42 @@ def test_executables_root_accepts_list(self, tmp_path):
coordinator = ExecutableCoordinator(executables_root=[str(tmp_path),
other])
assert coordinator.executables_root == [tmp_path, other]
- def test_executables_root_required(self):
- with pytest.raises(TypeError, match="executables_root"):
- ExecutableCoordinator()
-
- def test_executables_root_must_be_non_empty(self):
- with pytest.raises(ValueError, match="executables_root"):
- ExecutableCoordinator(executables_root=None)
+ def test_executables_root_optional_defaults_to_empty(self):
+ # Neither an explicit root nor dag_bundle_name: co-located mode, valid.
+ coordinator = ExecutableCoordinator()
+ assert coordinator.executables_root == []
+ assert coordinator.dag_bundle_name is None
+
+ def test_none_executables_root_normalized_to_empty(self):
+ coordinator = ExecutableCoordinator(executables_root=None)
+ assert coordinator.executables_root == []
+
+ def test_root_and_dag_bundle_name_are_mutually_exclusive(self, tmp_path):
+ with pytest.raises(ValueError, match="at most one of
'executables_root' or 'dag_bundle_name'"):
+ ExecutableCoordinator(executables_root=[tmp_path],
dag_bundle_name="artifacts")
+
+ @patch("airflow.dag_processing.bundles.manager.DagBundlesManager")
+ def test_unconfigured_dag_bundle_name_raises(self, mock_manager):
+ mock_manager.is_bundle_configured.return_value = False
+ with pytest.raises(ValueError, match="unconfigured Dag bundle
'ghost'"):
+ ExecutableCoordinator(dag_bundle_name="ghost")
+
+ @patch("airflow.dag_processing.bundles.manager.DagBundlesManager")
+ def test_configured_dag_bundle_name_accepted(self, mock_manager):
+ mock_manager.is_bundle_configured.return_value = True
+ coordinator = ExecutableCoordinator(dag_bundle_name="artifacts")
+ assert coordinator.dag_bundle_name == "artifacts"
+ mock_manager.is_bundle_configured.assert_called_once_with("artifacts")
+
+ def test_build_command_scans_given_roots(self, tmp_path):
Review Comment:
This builds the coordinator with `executables_root=[tmp_path]` and then
hands it `roots=[tmp_path]`, so it passes whether `_build_execute_task_command`
consumes `roots` or ignores it and reads `self.executables_root`. It passes
against main-shaped code too, so the comment above it asserts something the
test does not check. Constructing in co-located mode
(`ExecutableCoordinator()`) and passing the populated directory as `roots`
would make it fail without the change. Same shape in the java and node copies.
Separately, nothing exercises the task-bundle mode through `execute_task`.
`_StubSubprocessCoordinator` is pinned to the explicit-root mode and ignores
`roots`, and the `TestInitRootSource` tests set `_artifact_source` directly, so
the wiring that binds `bundle_info` and forwards resolved roots is not covered.
##########
task-sdk/src/airflow/sdk/coordinators/java/coordinator.py:
##########
@@ -200,16 +201,21 @@ class JavaCoordinator(SubprocessCoordinator):
jvm_args: list[str] = attrs.field(factory=list)
jars_root: list[pathlib.Path] = attrs.field(
converter=convert_roots,
- validator=attrs.validators.min_len(1),
+ factory=list,
Review Comment:
Dropping `min_len(1)` means a deployment whose `jars_root` templates to
empty or None stops failing at construction and silently switches to scanning
the whole Dag bundle.
For Java that is worse than a slow scan. With `main_class` unset,
`_JarInfo.find` takes the first JAR carrying a `Main-Class` in walk order and
`_calculate_classpath` puts everything it found on the classpath.
`java.rst:678` already warns the result is non-deterministic when multiple
executable JARs are present, and co-located mode makes that the normal case.
Requiring `main_class` when there is no explicit root would cover the Java side.
More broadly, is the implicit default worth it? An empty root now means both
"use the task's bundle" and "my config rendered empty". Something like
`use_task_bundle=True` would keep the new mode while letting a broken config
keep failing fast.
##########
airflow-core/src/airflow/dag_processing/bundles/manager.py:
##########
@@ -427,6 +427,11 @@ def get_bundle(
name=name, version=version, version_data=version_data,
**cfg_bundle.kwargs
)
+ @classmethod
+ def is_bundle_configured(cls, name: str) -> bool:
+ """Return whether *name* is a configured Dag bundle, without
constructing the bundle."""
+ return name in cls()._bundle_config
Review Comment:
`cls()` runs a full `parse_config()` on every call: `__init__` always calls
it (`manager.py:228-231`) and the `if self._bundle_config: return` guard cannot
fire on a fresh instance. That `import_string`s every configured bundle class,
and with `[core] load_examples` on it also walks `ProvidersManager()`.
The part that bites is the error surface. `parse_config` raises
`AirflowConfigException` for a duplicate or reserved name and lets
`ImportError` escape, out of the coordinator's `__attrs_post_init__`.
`CoordinatorManager.for_queue` catches `ImportError` and re-raises
`InvalidCoordinatorError("Cannot import coordinator ...")`
(`execution_time/coordinator.py:289`), pointing the operator at the wrong
config key. Worth catching both here and re-raising with the bundle name so
"bundle not configured" stays distinguishable from "the bundle config list does
not load".
##########
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:
This passes the module logger, so bundle resolution failures land in the
worker log and never in the task log. `logger` is in scope in `execute_task`
and already threaded into `_PopenActivitySubprocess.start`, so `logger or log`
here would fix it.
The ordering also flipped. The TI goes RUNNING in `_on_child_started` via
`client.task_instances.start` (`supervisor.py:1388`), which now happens after
bundle materialization, so a first-time clone of a large repo runs while the TI
still shows QUEUED and counts against `[scheduler] task_queued_timeout`. The
Python path is the other way round: RUNNING first, then `parse()` clones in the
child where output is captured to the task log.
--
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]
