jason810496 commented on code in PR #70805:
URL: https://github.com/apache/airflow/pull/70805#discussion_r3746478845
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -385,23 +400,148 @@ 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
+
+ # Name of the subclass's explicit-root kwarg, used only in error messages.
+ # Subclasses that expose an explicit root set this and override
+ # :meth:`_explicit_artifact_roots`; a subclass that does neither lands on
the
+ # task-bundle default instead of failing at execute time.
+ _root_kwarg: ClassVar[str] = "root"
+
+ # 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)
+
+ @property
+ def _explicit_artifact_roots(self) -> Sequence[pathlib.Path]:
+ """Subclass's explicit artifact roots; empty (the default) selects
task-bundle mode."""
+ return []
+
+ def __attrs_post_init__(self) -> None:
+ self._classify_artifact_source(self._explicit_artifact_roots,
root_kwarg=self._root_kwarg)
Review Comment:
Nice catch. The `_classify_artifact_source` is now argument less and the
`_explicit_artifact_roots` is the only property that subclass needs to override
(no `_root_kwargs` class attribute anymore).
##########
task-sdk/src/airflow/sdk/execution_time/tracing.py:
##########
@@ -0,0 +1,60 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Span helpers shared by the modules that make up a task run."""
+
+from __future__ import annotations
+
+import functools
+import inspect
+
+from opentelemetry import trace
+
+from airflow.sdk._shared.observability.traces import get_task_span_detail_level
+
+tracer = trace.get_tracer(__name__)
+
+
Review Comment:
Please let me know whether extracting `detail_span` into `tracing` dedicated
module changes too much or not.
Another direction is to keep the `detail_span` as-is and call both
`initialize_ti_bundle` and `verify_bundle_access` sequentially within
`task_runner`.
##########
airflow-core/src/airflow/dag_processing/bundles/manager.py:
##########
@@ -107,6 +108,20 @@ def _bundle_item_exc(msg):
)
[email protected]
Review Comment:
Agreed, I made this a common layer that both `is_bundle_configured` and
`parse_config` will get the same config snapshot.
Additionally, I prefer to keep `is_bundle_configured` as a class method so
that we won't import those actual class paths during the `is_bundle_configured`
call. Only the dag bundle manager constructor (`parse_config` call ) will
import the Dag bundle class.
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -414,18 +554,45 @@ 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)
+ task_logger = logger or log
+ with contextlib.ExitStack() as stack:
+ stack.enter_context(self._set_current_bundle(bundle_info))
+ roots, resolved_bundle = self._init_root_source(task_logger)
+ if resolved_bundle is not None:
+ # Hold the version lock across start()/wait() so bundle cleanup
+ # cannot rmtree a version this task is still reading from,
+ # mirroring task_runner.main() for the Python task path.
+ from airflow.dag_processing.bundles.base import ( # noqa:
SDK002
+ BundleVersionLock,
+ unpack_bundle_version,
+ )
+
+ # NAMED_BUNDLE resolves "latest" with no pinned version, so
+ # resolved_bundle.version is None and the lock would be a
silent
+ # no-op. Pin the concrete current version so the lock protects
the read.
+ lock_version = resolved_bundle.version
+ if lock_version is None:
+ lock_version, _ = unpack_bundle_version(
+ resolved_bundle.get_current_version(), resolved_bundle
+ )
+ stack.enter_context(
+ BundleVersionLock(
+ bundle_name=resolved_bundle.name,
+ bundle_version=lock_version,
+ )
+ )
Review Comment:
Good point, I added the `unpack_bundle_version` to get the actual version
when then the versioning is available.
The reason I didn't add in the first place is to prevent the multiple
`initialize_ti_bundle` call, but the correctness is more than the efficiency.
Also, I found that the second `initialize_ti_bundle` _shouldn't_ be heavy the
GitDagBundle will only pull the resources if necessary.
##########
task-sdk/src/airflow/sdk/coordinators/_subprocess.py:
##########
@@ -385,23 +400,148 @@ 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
+
+ # Name of the subclass's explicit-root kwarg, used only in error messages.
+ # Subclasses that expose an explicit root set this and override
+ # :meth:`_explicit_artifact_roots`; a subclass that does neither lands on
the
+ # task-bundle default instead of failing at execute time.
+ _root_kwarg: ClassVar[str] = "root"
+
+ # 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)
Review Comment:
I made it as `_artifact_source: _ArtifactSource = attrs.field(init=False)`,
thanks for the catch.
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -1012,21 +1013,34 @@ def _register_deserialization_allowed_classes(dag, log:
Logger) -> None:
)
-@detail_span("parse")
-def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance:
- # TODO: Task-SDK:
- # Using BundleDagBag here is about 98% wrong, but it'll do for now
- from airflow.dag_processing.dagbag import BundleDagBag
+def initialize_ti_bundle(bundle_info: BundleInfo, log: Logger) ->
BaseDagBundle:
Review Comment:
Both `initialize_ti_bundle` and the `verify_bundle_access` are moved into
dedicated modules. Additionally, the `detail_span` is moved into a dedicated
`tracing` module to preventing to coupling while preserving the existing
behavior (having span on the `verify_bundle_access`).
--
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]
