You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "aoelvp94 (via GitHub)" <gi...@apache.org> on 2026/09/11 23:11:33 UTC
[PR] Add archive staging support to S3DagBundle [airflow]
aoelvp94 opened a new pull request, #73015:
URL: https://github.com/apache/airflow/pull/73015
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
closes: #73010
`S3DagBundle` stages the bundle by downloading the prefix one object at a time (`S3Hook.sync_to_local_dir`), so a bundle of N small files costs ~N GET (+N HEAD) sequential requests on every component that stages it. Under KubernetesExecutor every task instance is a fresh pod with an empty bundle directory, which puts the full staging cost on the critical path of every task start (~20+ s for a 400-file bundle against real S3), and even a steady-state "nothing changed" refresh still pays a full per-object walk.
This PR adds an opt-in `archive_key` kwarg to `S3DagBundle`: when set, `refresh()` stages the bundle from a single `.tar.gz` object instead — one `HEAD` (skipped staging entirely when the ETag matches the last staged archive), one `GET`, a safe unpack (`tarfile` with `filter="data"`), and an atomic rename swap so a partially staged tree is never observable. On any archive problem (missing object, corrupt tar, unpack error) staging falls back to the existing per-object sync of `prefix`, so the worst case is exactly the current behavior. Publishing the archive is the deployer's job (documented: `tar -C dags . | aws s3 cp` in the same CI that syncs the prefix).
Measured with the reproduction environment from the issue (fully local: Airflow 3 + MinIO, identical 400-file / ~3 MiB file set — https://github.com/aoelvp94/airflow-s3-dag-bundle-staging-demo):
| staging strategy | wall-clock (median) | requests served |
|---|---:|---:|
| per-object sync | 15.8 s | 802 |
| single archive | 1.5 s | 5 |
Design notes:
* One class, opt-in: no `archive_key` ⇒ nothing changes. `prefix` keeps its meaning as the fallback source, and `_initialize()` skips the prefix-existence check only when the archive exists (archive-only buckets work; missing archives warn and validate/stage from the prefix as today).
* The local ETag marker lives inside the bundle directory, so a fallback per-object sync (`delete_stale=True`) removes it — a later re-published archive can never be wrongly skipped against a stale marker.
* Deletions need no reconciliation logic in the archive path: the unpacked tree replaces the previous one wholesale.
* The staging strategy is provider-agnostic; if there is interest, the same shape could be extracted for the GCS bundle (and the proposed Azure Blob bundle in #67016) in a follow-up.
Covered by 6 new unit tests (moto): archive staging, ETag skip + re-staging on change, archive-only bucket, fallback on corrupt and on missing archive, and marker cleanup after fallback. Docs updated in `providers/amazon/docs/bundles/index.rst`. `prek` static checks pass locally.
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes (please specify the tool below)
Generated-by: Claude Code following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions) — all code, tests, and docs were reviewed by the author, and static checks plus the provider unit test suite were run locally.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---
* Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment: `{pr_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments). You can add this file in a follow-up commit after the PR is created so you know the PR number.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "aoelvp94 (via GitHub) via commits" <co...@airflow.apache.org>.
aoelvp94 commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4434644033
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -79,7 +123,12 @@ def _initialize(self):
if not self.s3_hook.check_for_bucket(bucket_name=self.bucket_name):
raise AirflowException(f"S3 bucket '{self.bucket_name}' does not exist.")
- if self.prefix:
+ if self.archive_key:
+ if not self.s3_hook.check_for_key(key=self.archive_key, bucket_name=self.bucket_name):
+ raise AirflowNotFoundException(
+ f"S3 archive 's3://{self.bucket_name}/{self.archive_key}' does not exist."
+ )
Review Comment:
Removed that check. A missing archive fails on the HEAD inside refresh, as the client's 404.
https://github.com/apache/airflow/pull/73015/commits/d046c82a87a055628513a3f34bac846b300de78e
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4410467577
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +193,95 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _archive_staging_prefix(self) -> str:
+ # The storage root is shared by every bundle. The lock is per bundle name, so
+ # leftover cleanup must not match another bundle's in-progress staging directory.
+ return f".{self.name}-archive-staging-"
Review Comment:
Using `-` can cause collision. `/` would be better since it is forbidden in file names.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4410443554
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +193,95 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _archive_staging_prefix(self) -> str:
+ # The storage root is shared by every bundle. The lock is per bundle name, so
+ # leftover cleanup must not match another bundle's in-progress staging directory.
Review Comment:
```suggestion
```
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +193,95 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _archive_staging_prefix(self) -> str:
+ # The storage root is shared by every bundle. The lock is per bundle name, so
+ # leftover cleanup must not match another bundle's in-progress staging directory.
+ return f".{self.name}-archive-staging-"
+
+ def _remove_stale_archive_staging_dirs(self) -> None:
+ storage_root = self.s3_dags_dir.parent
+ if not storage_root.is_dir():
+ return
+ prefix = self._archive_staging_prefix()
+ for leftover in storage_root.iterdir():
+ if leftover.is_dir() and leftover.name.startswith(prefix):
+ shutil.rmtree(leftover, ignore_errors=True)
+
+ def _download_archive(self, client, archive_path: Path, etag: str) -> str:
+ """Download the archive object named by ``etag``."""
+ # download_file HEADs again and GETs whatever is current, so a replace in between
+ # can store an ETag for a different object than the one unpacked.
Review Comment:
```suggestion
```
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +193,95 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _archive_staging_prefix(self) -> str:
+ # The storage root is shared by every bundle. The lock is per bundle name, so
+ # leftover cleanup must not match another bundle's in-progress staging directory.
+ return f".{self.name}-archive-staging-"
+
+ def _remove_stale_archive_staging_dirs(self) -> None:
+ storage_root = self.s3_dags_dir.parent
+ if not storage_root.is_dir():
+ return
+ prefix = self._archive_staging_prefix()
+ for leftover in storage_root.iterdir():
+ if leftover.is_dir() and leftover.name.startswith(prefix):
+ shutil.rmtree(leftover, ignore_errors=True)
+
+ def _download_archive(self, client, archive_path: Path, etag: str) -> str:
+ """Download the archive object named by ``etag``."""
+ # download_file HEADs again and GETs whatever is current, so a replace in between
+ # can store an ETag for a different object than the one unpacked.
+ request = {"Bucket": self.bucket_name, "Key": self.archive_key}
+ if etag:
+ request["IfMatch"] = etag
+ response = client.get_object(**request)
+ body = response["Body"]
+ try:
+ with archive_path.open("wb") as handle:
+ for chunk in body.iter_chunks():
+ handle.write(chunk)
+ finally:
+ body.close()
+ return response.get("ETag", "")
+
+ def _refresh_from_archive(self) -> None:
+ """Stage the Dag bundle by downloading and unpacking the single archive object."""
+ # A killed process never reaches the finally that removes its staging directory.
Review Comment:
```suggestion
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4410382851
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -17,14 +17,40 @@
from __future__ import annotations
import os
+import shutil
+import tarfile
+import tempfile
from pathlib import Path
import structlog
from airflow.dag_processing.bundles.base import BaseDagBundle
from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException
+
+
+def _extract_archive(archive_path: Path, unpack_dir: Path) -> None:
+ """Unpack a bundle archive without writing outside ``unpack_dir``."""
+ with tarfile.open(archive_path, "r:*") as tar:
+ # filter="data" arrived in Python 3.11.4. 3.11.0-3.11.3 still have to reject
+ # the same escapes by hand.
Review Comment:
```suggestion
filter="data" requires Python 3.11.4 or later.
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "aoelvp94 (via GitHub)" <gi...@apache.org>.
aoelvp94 commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4407459280
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +170,60 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _refresh_from_archive(self) -> None:
+ """Stage the Dag bundle by downloading and unpacking the single archive object."""
+ client = self.s3_hook.get_conn()
+ head = client.head_object(Bucket=self.bucket_name, Key=self.archive_key)
+ etag: str = head.get("ETag", "")
+
+ marker = self.s3_dags_dir / self.archive_etag_marker
+ if etag and marker.is_file() and marker.read_text() == etag:
+ self._log.debug(
+ "Dag bundle archive 's3://%s/%s' is unchanged (ETag %s), skipping staging",
+ self.bucket_name,
+ self.archive_key,
+ etag,
+ )
+ return
+
+ staging_dir = Path(tempfile.mkdtemp(dir=self.s3_dags_dir.parent, prefix=".s3-archive-staging-"))
+ try:
+ archive_path = staging_dir / "_bundle_archive"
+ client.download_file(self.bucket_name, self.archive_key, os.fspath(archive_path))
Review Comment:
Done in [#73015](https://github.com/apache/airflow/pull/73015/commits/ad07c1b2d0162801dcc30fadf028ce95fb8b6ee1). The archive is fetched with get_object and IfMatch set to the ETag from the HEAD. The body is streamed, and the marker stores the ETag from that response, or the HEAD ETag when the response has none. A replace between the HEAD and the GET returns 412, the previous tree stays in place, and the staging directory is removed. An empty ETag omits IfMatch. An unchanged archive is still one HEAD.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "aoelvp94 (via GitHub) via commits" <co...@airflow.apache.org>.
aoelvp94 commented on PR #73015:
URL: https://github.com/apache/airflow/pull/73015#issuecomment-6554437484
> @aoelvp94 Could you please resolve conflicts?
Rebased onto main. The conflict is resolved.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "aoelvp94 (via GitHub)" <gi...@apache.org>.
aoelvp94 commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4053862005
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -136,6 +179,54 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _refresh_from_archive(self) -> None:
+ """Stage the Dag bundle by downloading and unpacking the single archive object."""
+ client = self.s3_hook.get_conn()
+ head = client.head_object(Bucket=self.bucket_name, Key=self.archive_key)
+ etag: str = head.get("ETag", "")
+
+ marker = self.s3_dags_dir / self.archive_etag_marker
+ if etag and marker.is_file() and marker.read_text() == etag:
+ self._log.debug(
+ "Dag bundle archive 's3://%s/%s' is unchanged (ETag %s), skipping staging",
+ self.bucket_name,
+ self.archive_key,
+ etag,
+ )
+ return
+
+ staging_dir = Path(tempfile.mkdtemp(dir=self.s3_dags_dir.parent, prefix=".s3-archive-staging-"))
+ try:
+ archive_path = staging_dir / "_bundle_archive"
+ client.download_file(self.bucket_name, self.archive_key, os.fspath(archive_path))
+
+ unpack_dir = staging_dir / "unpacked"
+ unpack_dir.mkdir()
+ with tarfile.open(archive_path, "r:*") as tar:
+ tar.extractall(unpack_dir, filter="data")
+ archive_path.unlink()
+
+ if etag:
+ (unpack_dir / self.archive_etag_marker).write_text(etag)
+
+ # Swap the freshly unpacked tree into place so a partially staged
+ # bundle is never observable at self.s3_dags_dir.
+ old_dir = self.s3_dags_dir.parent / f".s3-archive-old-{uuid.uuid4().hex}"
Review Comment:
https://github.com/apache/airflow/pull/73015/changes/3bbbd6d17e02ad20b0221de80546dd520a4e1297 pushed here, check please :)
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -126,6 +155,20 @@ def refresh(self) -> None:
raise AirflowException("Refreshing a specific version is not supported")
with self.lock():
+ if self.archive_key:
+ try:
+ self._refresh_from_archive()
+ return
+ except Exception:
Review Comment:
https://github.com/apache/airflow/pull/73015/changes/3bbbd6d17e02ad20b0221de80546dd520a4e1297 pushed here, check please :)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on PR #73015:
URL: https://github.com/apache/airflow/pull/73015#issuecomment-6492347627
A high-level design question: Should archive mode be a separate class? Using `archive_key` as a switch adds a lot of branching, and it’s easier to miss a branch if additional logic is added to the bundle. We could unify the user-facing interface by having _two_ subclasses, one with `archive_key` and one without, and implement `__new__` on S3DagBundle to dispatch based on arguments. It would make the implementation cleaner IMO.
We could obviously just have S3DagBundle and S3ArchiveDagBundle if it’s good enough.
It might also be a good idea to prepare for extracting things like `_extract_archive` out of the provider since those really are not specific to S3. It’s difficult to manage provider compatibility though.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "aoelvp94 (via GitHub) via commits" <co...@airflow.apache.org>.
aoelvp94 commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4478040308
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +179,88 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _archive_state_dir(self) -> Path:
+ return self.s3_dags_dir.parent / f".{self.name}"
+
+ def _archive_etag_path(self) -> Path:
+ return self._archive_state_dir() / "etag"
+
+ def _remove_stale_archive_staging_dirs(self) -> None:
+ state_dir = self._archive_state_dir()
+ if not state_dir.is_dir():
+ return
+ for leftover in state_dir.iterdir():
+ if leftover.is_dir() and leftover.name.startswith("archive-staging-"):
+ shutil.rmtree(leftover, ignore_errors=True)
+
+ def _download_archive(self, client, archive_path: Path, etag: str) -> str:
+ """Download the archive object named by ``etag``."""
+ request = {"Bucket": self.bucket_name, "Key": self.archive_key}
+ if etag:
+ request["IfMatch"] = etag
+ response = client.get_object(**request)
+ body = response["Body"]
+ try:
+ with archive_path.open("wb") as handle:
+ for chunk in body.iter_chunks():
Review Comment:
Agreed — will switch to `shutil.copyfileobj(..., 1024 * 1024)` in the next push.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "vincbeck (via GitHub)" <gi...@apache.org>.
vincbeck commented on PR #73015:
URL: https://github.com/apache/airflow/pull/73015#issuecomment-5819866637
> I really appreciate this PR in theory and love the speed increase the data shows. But I've been reading it and thinking that it might be a little sharp in ways the users don't expect. Specifically the fallback guarantee. It only holds if the prefix files and the tar are kept in sync, and this really seems to me like one of the pieces that will absolutely get out of sync without extreme diligence from the user. I know it's documented, but just seems like it might cause more harm and confusion.
>
> What about optimizing with the use of client.get_object() to avoid the two request per file and then also using some parallelism to download the objects concurrently to reduce it down further. I think that would get almost all of the same speed gains, but without the sharp edges of the duality this PR has. Either that, or have no fallback to the prefix if you use the archive, essentially make them mutually exclusive so that keeping them in sync isn't critical.
>
> CC @ramitkataria @vincbeck thoughts?
I agree, I would actually do both. Optimize the current mechanism and add the archive option without fallback
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4410417172
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -17,14 +17,40 @@
from __future__ import annotations
import os
+import shutil
+import tarfile
+import tempfile
from pathlib import Path
import structlog
from airflow.dag_processing.bundles.base import BaseDagBundle
from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException
+
+
+def _extract_archive(archive_path: Path, unpack_dir: Path) -> None:
+ """Unpack a bundle archive without writing outside ``unpack_dir``."""
+ with tarfile.open(archive_path, "r:*") as tar:
+ # filter="data" arrived in Python 3.11.4. 3.11.0-3.11.3 still have to reject
+ # the same escapes by hand.
+ if hasattr(tarfile, "data_filter"):
Review Comment:
With 3.11 already in security-only mode and being dropped next year, I wonder if it’s worthwhile to implement a fallback at all. It would be acceptable IMO to simply gate this feature if `filter="data"` is not available.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "aoelvp94 (via GitHub) via commits" <co...@airflow.apache.org>.
aoelvp94 commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4434643445
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -17,14 +17,40 @@
from __future__ import annotations
import os
+import shutil
+import tarfile
+import tempfile
from pathlib import Path
import structlog
from airflow.dag_processing.bundles.base import BaseDagBundle
from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException
+
+
+def _extract_archive(archive_path: Path, unpack_dir: Path) -> None:
+ """Unpack a bundle archive without writing outside ``unpack_dir``."""
+ with tarfile.open(archive_path, "r:*") as tar:
+ # filter="data" arrived in Python 3.11.4. 3.11.0-3.11.3 still have to reject
+ # the same escapes by hand.
+ if hasattr(tarfile, "data_filter"):
Review Comment:
Agreed, the fallback is not worth it. The manual check misses cases the stdlib filter covers, including symlinks. Archive mode now requires `tarfile.data_filter`. Constructing an archive bundle on Python older than 3.11.4 raises `RuntimeError`.
https://github.com/apache/airflow/pull/73015/commits/d046c82a87a055628513a3f34bac846b300de78e
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "aoelvp94 (via GitHub)" <gi...@apache.org>.
aoelvp94 commented on PR #73015:
URL: https://github.com/apache/airflow/pull/73015#issuecomment-6386611579
> Thanks for reworking this.
>
> One small thing: `view_url_template()` still builds the URL from `prefix`, so with `archive_key` set the link points at the bucket root. Could it point at the archive object instead?
>
> Also the docs call the swap atomic, but it's two renames with a short gap where the bundle dir doesn't exist. Maybe just "a swap into place".
Done! thanks for the review
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "aoelvp94 (via GitHub)" <gi...@apache.org>.
aoelvp94 commented on PR #73015:
URL: https://github.com/apache/airflow/pull/73015#issuecomment-6417427708
> Thanks for the PR, I like the idea and I think it's not just for optimizing the cost and overhead. It might also improve the deployment side for the non-Python SDK deployments.
>
> Before we going further to adapt this pattern for all the object store based DagBundle. Would it be better to introduce a new class, e.g. `S3ArchiveDagBundle` instead of having on the existing `S3DagBundle`?
>
> From my perspective, the refreshing scope is completely different (objects under prefix map to dir vs single object map to dir) so having a independent class _might_ be more clear.
>
> Additionally, we _could_ even support versioning for the `S3ArchiveDagBundle`. IIUC, we're able to track the versioning id of that object from S3 as reference in Airflow meta database.
>
> cc @uranusjr who involved with non-Python SDK, @ephraimbuddy who experimented PoC for supporting versioning for object store based DagBundle.
Thanks, this is the right question to settle before the same option lands on GCS and Azure.
I'd keep `archive_key` on `S3DagBundle`. The two refresh scopes are different, which is why `archive_key` and `prefix` are mutually exclusive and rejected together at construction. With `archive_key` unset, the class behaves as it does today, so existing bundle config keeps the same classpath. A separate `S3ArchiveDagBundle` would repeat the hook, the lock, and the local path for a mode that is off unless configured.
Versioning is the part I would not fold into this PR. An archive object has a version id, and a prefix does not, so `supports_versioning` on an archive class is a real follow-up. Both bundles still report `supports_versioning = False`. If we add that, a dedicated class is a better home, and I'm happy to split then.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4410396297
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -17,14 +17,40 @@
from __future__ import annotations
import os
+import shutil
+import tarfile
+import tempfile
from pathlib import Path
import structlog
from airflow.dag_processing.bundles.base import BaseDagBundle
from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException
+
+
+def _extract_archive(archive_path: Path, unpack_dir: Path) -> None:
+ """Unpack a bundle archive without writing outside ``unpack_dir``."""
+ with tarfile.open(archive_path, "r:*") as tar:
+ # filter="data" arrived in Python 3.11.4. 3.11.0-3.11.3 still have to reject
+ # the same escapes by hand.
+ if hasattr(tarfile, "data_filter"):
Review Comment:
```suggestion
if hasattr(tarfile, "data_filter"): # Added in Python 3.11.4.
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4410484234
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -37,23 +63,40 @@ class S3DagBundle(BaseDagBundle):
:param bucket_name: The name of the S3 bucket containing the Dag files.
:param prefix: Optional subdirectory within the S3 bucket where the Dags are stored.
A trailing slash is optional.
- If None, Dags are assumed to be at the root of the bucket (Optional).
+ If None, Dags are assumed to be at the root of the bucket. Mutually exclusive
+ with ``archive_key`` (Optional).
+ :param archive_key: Optional S3 key of a ``.tar.gz`` archive holding the Dag files. When set,
+ the bundle is staged by downloading this single object and unpacking it
+ locally, instead of downloading every object under a prefix one at a time.
+ The archive is then the only source of the Dags: ``archive_key`` is mutually
+ exclusive with ``prefix``, so there is no second copy to keep in sync and no
+ fallback to a per-object sync. Members are extracted relative to the bundle
+ root, e.g. created with ``tar -C <dags_dir> -czf dags.tar.gz .`` (Optional).
"""
supports_versioning = False
+ archive_etag_marker = ".airflow_bundle_archive_etag"
Review Comment:
This would be written into the Dags folder itself and visible to the Dag processor and any user listing the bundle. Consider keeping it next to the bundle dir instead (e.g. in the staging root or `_tracking`).
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "aoelvp94 (via GitHub)" <gi...@apache.org>.
aoelvp94 commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4407460196
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +170,60 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _refresh_from_archive(self) -> None:
+ """Stage the Dag bundle by downloading and unpacking the single archive object."""
+ client = self.s3_hook.get_conn()
+ head = client.head_object(Bucket=self.bucket_name, Key=self.archive_key)
+ etag: str = head.get("ETag", "")
+
+ marker = self.s3_dags_dir / self.archive_etag_marker
+ if etag and marker.is_file() and marker.read_text() == etag:
+ self._log.debug(
+ "Dag bundle archive 's3://%s/%s' is unchanged (ETag %s), skipping staging",
+ self.bucket_name,
+ self.archive_key,
+ etag,
+ )
+ return
+
+ staging_dir = Path(tempfile.mkdtemp(dir=self.s3_dags_dir.parent, prefix=".s3-archive-staging-"))
+ try:
+ archive_path = staging_dir / "_bundle_archive"
+ client.download_file(self.bucket_name, self.archive_key, os.fspath(archive_path))
+
+ unpack_dir = staging_dir / "unpacked"
+ unpack_dir.mkdir()
+ with tarfile.open(archive_path, "r:*") as tar:
+ tar.extractall(unpack_dir, filter="data")
Review Comment:
nice catch! Done in [#73015](https://github.com/apache/airflow/pull/73015/commits/2ceb7e238df73888fe5ed2b8f44eb044892d25bd).
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "aoelvp94 (via GitHub) via commits" <co...@airflow.apache.org>.
aoelvp94 commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4478042469
##########
providers/amazon/docs/bundles/index.rst:
##########
@@ -50,3 +50,39 @@ Example of using the S3DagBundle:
}
}
]'
+
+Staging from a single archive object
+------------------------------------
+
+By default the bundle is staged by downloading every object under ``prefix`` one at a time. Each object
+costs a full request round-trip, so bundles made of many small files can take a long time to stage — a cost
+paid by every component that stages the bundle, which with ephemeral workers (e.g. KubernetesExecutor task
+pods) means on every task start.
+
+Setting ``archive_key`` instead of ``prefix`` stages the bundle from a single ``.tar.gz`` object: one
Review Comment:
Agreed — will call out the 3.11.4+ requirement in the docs in the next push.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "shahar1 (via GitHub)" <gi...@apache.org>.
shahar1 commented on PR #73015:
URL: https://github.com/apache/airflow/pull/73015#issuecomment-6490412747
@aoelvp94 Could you please resolve conflicts?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4410439700
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -79,7 +123,12 @@ def _initialize(self):
if not self.s3_hook.check_for_bucket(bucket_name=self.bucket_name):
raise AirflowException(f"S3 bucket '{self.bucket_name}' does not exist.")
- if self.prefix:
+ if self.archive_key:
+ if not self.s3_hook.check_for_key(key=self.archive_key, bucket_name=self.bucket_name):
+ raise AirflowNotFoundException(
+ f"S3 archive 's3://{self.bucket_name}/{self.archive_key}' does not exist."
+ )
Review Comment:
This adds an HEAD on every `__init__`, and there’s another one immediately in `_refresh_from_archive`. Just let the HEAD in refresh raise.
Also, this raises a different exception type from the prefix check below.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4410474677
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +193,95 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _archive_staging_prefix(self) -> str:
+ # The storage root is shared by every bundle. The lock is per bundle name, so
+ # leftover cleanup must not match another bundle's in-progress staging directory.
+ return f".{self.name}-archive-staging-"
+
+ def _remove_stale_archive_staging_dirs(self) -> None:
+ storage_root = self.s3_dags_dir.parent
+ if not storage_root.is_dir():
+ return
+ prefix = self._archive_staging_prefix()
+ for leftover in storage_root.iterdir():
+ if leftover.is_dir() and leftover.name.startswith(prefix):
+ shutil.rmtree(leftover, ignore_errors=True)
+
+ def _download_archive(self, client, archive_path: Path, etag: str) -> str:
+ """Download the archive object named by ``etag``."""
+ # download_file HEADs again and GETs whatever is current, so a replace in between
+ # can store an ETag for a different object than the one unpacked.
+ request = {"Bucket": self.bucket_name, "Key": self.archive_key}
+ if etag:
+ request["IfMatch"] = etag
+ response = client.get_object(**request)
+ body = response["Body"]
+ try:
+ with archive_path.open("wb") as handle:
+ for chunk in body.iter_chunks():
+ handle.write(chunk)
+ finally:
+ body.close()
+ return response.get("ETag", "")
+
+ def _refresh_from_archive(self) -> None:
+ """Stage the Dag bundle by downloading and unpacking the single archive object."""
+ # A killed process never reaches the finally that removes its staging directory.
+ self._remove_stale_archive_staging_dirs()
+ client = self.s3_hook.get_conn()
+ head = client.head_object(Bucket=self.bucket_name, Key=self.archive_key)
+ etag: str = head.get("ETag", "")
+
+ marker = self.s3_dags_dir / self.archive_etag_marker
+ if etag and marker.is_file() and marker.read_text() == etag:
+ self._log.debug(
+ "Dag bundle archive 's3://%s/%s' is unchanged (ETag %s), skipping staging",
+ self.bucket_name,
+ self.archive_key,
+ etag,
+ )
+ return
+
+ staging_dir = Path(
+ tempfile.mkdtemp(dir=self.s3_dags_dir.parent, prefix=self._archive_staging_prefix())
+ )
+ try:
+ archive_path = staging_dir / "_bundle_archive"
+ downloaded_etag = self._download_archive(client, archive_path, etag)
+
+ unpack_dir = staging_dir / "unpacked"
+ unpack_dir.mkdir()
+ _extract_archive(archive_path, unpack_dir)
+ archive_path.unlink()
+
+ marker_etag = downloaded_etag or etag
+ if marker_etag:
+ (unpack_dir / self.archive_etag_marker).write_text(marker_etag)
+
+ # Swap the freshly unpacked tree into place so a partially staged bundle is
+ # never observable at self.s3_dags_dir. The previous tree is kept inside the
+ # staging directory: a failed swap restores it, and the staging cleanup below
+ # removes it once the swap succeeded.
Review Comment:
Not accurate, as Niko already pointed out. Reword or drop.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "aoelvp94 (via GitHub) via commits" <co...@airflow.apache.org>.
aoelvp94 commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4434647613
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +193,95 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _archive_staging_prefix(self) -> str:
+ # The storage root is shared by every bundle. The lock is per bundle name, so
+ # leftover cleanup must not match another bundle's in-progress staging directory.
+ return f".{self.name}-archive-staging-"
+
+ def _remove_stale_archive_staging_dirs(self) -> None:
+ storage_root = self.s3_dags_dir.parent
+ if not storage_root.is_dir():
+ return
+ prefix = self._archive_staging_prefix()
+ for leftover in storage_root.iterdir():
+ if leftover.is_dir() and leftover.name.startswith(prefix):
+ shutil.rmtree(leftover, ignore_errors=True)
+
+ def _download_archive(self, client, archive_path: Path, etag: str) -> str:
+ """Download the archive object named by ``etag``."""
+ # download_file HEADs again and GETs whatever is current, so a replace in between
+ # can store an ETag for a different object than the one unpacked.
+ request = {"Bucket": self.bucket_name, "Key": self.archive_key}
+ if etag:
+ request["IfMatch"] = etag
+ response = client.get_object(**request)
+ body = response["Body"]
+ try:
+ with archive_path.open("wb") as handle:
+ for chunk in body.iter_chunks():
+ handle.write(chunk)
+ finally:
+ body.close()
+ return response.get("ETag", "")
+
+ def _refresh_from_archive(self) -> None:
+ """Stage the Dag bundle by downloading and unpacking the single archive object."""
+ # A killed process never reaches the finally that removes its staging directory.
Review Comment:
Removed
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +193,95 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _archive_staging_prefix(self) -> str:
+ # The storage root is shared by every bundle. The lock is per bundle name, so
+ # leftover cleanup must not match another bundle's in-progress staging directory.
+ return f".{self.name}-archive-staging-"
+
+ def _remove_stale_archive_staging_dirs(self) -> None:
+ storage_root = self.s3_dags_dir.parent
+ if not storage_root.is_dir():
+ return
+ prefix = self._archive_staging_prefix()
+ for leftover in storage_root.iterdir():
+ if leftover.is_dir() and leftover.name.startswith(prefix):
+ shutil.rmtree(leftover, ignore_errors=True)
+
+ def _download_archive(self, client, archive_path: Path, etag: str) -> str:
+ """Download the archive object named by ``etag``."""
+ # download_file HEADs again and GETs whatever is current, so a replace in between
+ # can store an ETag for a different object than the one unpacked.
Review Comment:
Removed the comments that restated the code.
https://github.com/apache/airflow/pull/73015/commits/4cf1c1d4ac0b51fa7d67ccf0604e30df17ad4e10
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "aoelvp94 (via GitHub) via commits" <co...@airflow.apache.org>.
aoelvp94 commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4434644413
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +193,95 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _archive_staging_prefix(self) -> str:
+ # The storage root is shared by every bundle. The lock is per bundle name, so
+ # leftover cleanup must not match another bundle's in-progress staging directory.
Review Comment:
Removed.
https://github.com/apache/airflow/pull/73015/commits/4cf1c1d4ac0b51fa7d67ccf0604e30df17ad4e10
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +193,95 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _archive_staging_prefix(self) -> str:
+ # The storage root is shared by every bundle. The lock is per bundle name, so
+ # leftover cleanup must not match another bundle's in-progress staging directory.
Review Comment:
Removed.
https://github.com/apache/airflow/pull/73015/commits/4cf1c1d4ac0b51fa7d67ccf0604e30df17ad4e10
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "aoelvp94 (via GitHub) via commits" <co...@airflow.apache.org>.
aoelvp94 commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4478038743
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +179,88 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _archive_state_dir(self) -> Path:
+ return self.s3_dags_dir.parent / f".{self.name}"
+
+ def _archive_etag_path(self) -> Path:
+ return self._archive_state_dir() / "etag"
+
+ def _remove_stale_archive_staging_dirs(self) -> None:
+ state_dir = self._archive_state_dir()
+ if not state_dir.is_dir():
+ return
+ for leftover in state_dir.iterdir():
+ if leftover.is_dir() and leftover.name.startswith("archive-staging-"):
+ shutil.rmtree(leftover, ignore_errors=True)
+
+ def _download_archive(self, client, archive_path: Path, etag: str) -> str:
+ """Download the archive object named by ``etag``."""
+ request = {"Bucket": self.bucket_name, "Key": self.archive_key}
+ if etag:
+ request["IfMatch"] = etag
+ response = client.get_object(**request)
+ body = response["Body"]
+ try:
+ with archive_path.open("wb") as handle:
+ for chunk in body.iter_chunks():
+ handle.write(chunk)
+ finally:
+ body.close()
+ return response.get("ETag", "")
+
+ def _refresh_from_archive(self) -> None:
+ """Stage the Dag bundle by downloading and unpacking the single archive object."""
+ self._remove_stale_archive_staging_dirs()
+ client = self.s3_hook.get_conn()
+ head = client.head_object(Bucket=self.bucket_name, Key=self.archive_key)
+ etag: str = head.get("ETag", "")
+
+ marker = self._archive_etag_path()
+ if etag and marker.is_file() and marker.read_text() == etag:
Review Comment:
Good catch — that's a real bug. The marker lives beside the Dag folder, so after a cleared tree `_initialize` recreated an empty directory, the ETag still matched, and we skipped staging.
Done in https://github.com/apache/airflow/pull/73015/commits/e9ed15391cd0e9f198d50c61f0d37aaf50c38d9a. We now restage unless the staged tree actually has files, and the prefix path drops the marker so switching back to an unchanged archive restages instead of keeping the prefix sync's files. Added a test for the delete-and-reinitialize case (bundle dir gone, matching marker still present → GET/unpack, not skip), plus one for prefix-then-archive.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "shahar1 (via GitHub)" <gi...@apache.org>.
shahar1 commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4052465517
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -136,6 +179,54 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _refresh_from_archive(self) -> None:
+ """Stage the Dag bundle by downloading and unpacking the single archive object."""
+ client = self.s3_hook.get_conn()
+ head = client.head_object(Bucket=self.bucket_name, Key=self.archive_key)
+ etag: str = head.get("ETag", "")
+
+ marker = self.s3_dags_dir / self.archive_etag_marker
+ if etag and marker.is_file() and marker.read_text() == etag:
+ self._log.debug(
+ "Dag bundle archive 's3://%s/%s' is unchanged (ETag %s), skipping staging",
+ self.bucket_name,
+ self.archive_key,
+ etag,
+ )
+ return
+
+ staging_dir = Path(tempfile.mkdtemp(dir=self.s3_dags_dir.parent, prefix=".s3-archive-staging-"))
+ try:
+ archive_path = staging_dir / "_bundle_archive"
+ client.download_file(self.bucket_name, self.archive_key, os.fspath(archive_path))
+
+ unpack_dir = staging_dir / "unpacked"
+ unpack_dir.mkdir()
+ with tarfile.open(archive_path, "r:*") as tar:
+ tar.extractall(unpack_dir, filter="data")
+ archive_path.unlink()
+
+ if etag:
+ (unpack_dir / self.archive_etag_marker).write_text(etag)
+
+ # Swap the freshly unpacked tree into place so a partially staged
+ # bundle is never observable at self.s3_dags_dir.
+ old_dir = self.s3_dags_dir.parent / f".s3-archive-old-{uuid.uuid4().hex}"
Review Comment:
A failed replacement can leave the bundle missing and its backup abandoned.
To install the new bundle, the code:
- Moves the current bundle to a temporary backup location.
- Moves the newly unpacked bundle into the normal location.
- Deletes the backup.
If the second step fails, the current bundle has already been moved away, and nothing restores it. The normal bundle location can therefore remain missing if the fallback also fails.
The backup is also placed outside the temporary staging folder, so the guaranteed cleanup does not remove it. Failed updates can leave abandoned copies consuming disk space.
The suggestion puts the backup inside the staging folder and restores it if replacement fails. Both changes matter: cleanup alone would discard the backup without recovering the working bundle. Using a fixed backup name inside that unique staging folder also makes uuid unnecessary.
---
Drafted by Codex, reviewed by me
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -126,6 +155,20 @@ def refresh(self) -> None:
raise AirflowException("Refreshing a specific version is not supported")
with self.lock():
+ if self.archive_key:
+ try:
+ self._refresh_from_archive()
+ return
+ except Exception:
Review Comment:
If downloading or unpacking the archive fails, Airflow falls back to downloading individual files from the configured S3 prefix - essentially an S3 folder.
But the PR allows a bucket to contain only the archive, with no individual files in that folder. The fallback synchronizes this empty folder with the local bundle and deletes anything that is no longer present in S3. That means it can erase the entire working local bundle.
--
Drafted by Codex, reviewed and verified by me
##########
providers/amazon/docs/bundles/index.rst:
##########
@@ -45,3 +45,45 @@ Example of using the S3DagBundle:
}
}
]'
+
+Staging from a single archive object
+------------------------------------
+
+By default the bundle is staged by downloading every object under ``prefix`` one at a time. Each object
+costs a full request round-trip, so bundles made of many small files can take a long time to stage — a cost
+paid by every component that stages the bundle, which with ephemeral workers (e.g. KubernetesExecutor task
+pods) means on every task start.
+
+Setting the optional ``archive_key`` stages the bundle from a single ``.tar.gz`` object instead: one
+``HEAD`` request to detect changes (unchanged archives are not re-downloaded), one ``GET`` to fetch it, then
+a local unpack and an atomic swap into place. If the archive cannot be fetched or unpacked, staging
+automatically falls back to the per-object sync of ``prefix``.
+
+.. code-block:: bash
+
+ export AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST='[
+ {
+ "name": "my-s3-dags",
+ "classpath": "airflow.providers.amazon.aws.bundles.s3.S3DagBundle",
+ "kwargs": {
+ "aws_conn_id": "aws_default",
+ "bucket_name": "my-airflow-bucket",
+ "prefix": "dags/",
+ "archive_key": "bundle-archives/dags.tar.gz",
+ "refresh_interval": 60
+ }
+ }
+ ]'
+
+Publishing the archive is the responsibility of your deployment process. The archive members must be laid out exactly as
+the objects under ``prefix``, so both staging strategies produce the same local tree — for example, in the
+same CI job that syncs the Dags:
+
+.. code-block:: bash
+
+ tar -C ./dags -czf dags.tar.gz .
+ aws s3 cp dags.tar.gz s3://my-airflow-bucket/bundle-archives/dags.tar.gz
+
+.. note::
+ Keep the archive outside the ``prefix`` location, otherwise a ``aws s3 sync --delete`` of the Dag
Review Comment:
nit:
```suggestion
Keep the archive outside the ``prefix`` location, otherwise an ``aws s3 sync --delete`` of the Dag
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "aoelvp94 (via GitHub)" <gi...@apache.org>.
aoelvp94 commented on PR #73015:
URL: https://github.com/apache/airflow/pull/73015#issuecomment-5854371583
Agreed — I'll drop the fallback.
@ramitkataria's point decides it: it fires exactly when S3 is unhappy and turns one request into ~2N
against the same prefix, from every staging pod. The trigger and the aggravating condition are the
same event, so docs can't fix it. Both earlier findings in this PR were in that same path — dropping
it removes the guard, the LIST probe and two doc caveats with it.
I'd go further than just skipping it and make `archive_key` and `prefix` mutually exclusive. No
second source, no drift, and one artifact to publish instead of two kept in sync.
On the comparison: they optimize different axes. `get_object` drops the per-object HEAD (change
detection already uses `size`/`last_modified` from ListObjects), halving requests — 802 → ~402 for
400 objects in my repro. Concurrency then gets wall-clock close to the archive. But neither takes
request count below ~N, whereas the archive is 1 GET + 1 HEAD, and HEAD-only on an unchanged ETag.
+1 to @vincbeck on both, as separate PRs: the archive is opt-in and inert when unset, whereas
`sync_to_local_dir` affects every current user — and `providers/google` has a byte-identical copy of
the same loop, so that one should probably cover both providers.
Happy to take the sync work as a follow-up, or leave it to whoever wants it. Pushing the no-fallback
version here meanwhile.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "ephraimbuddy (via GitHub)" <gi...@apache.org>.
ephraimbuddy commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4405071497
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +170,60 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _refresh_from_archive(self) -> None:
+ """Stage the Dag bundle by downloading and unpacking the single archive object."""
+ client = self.s3_hook.get_conn()
+ head = client.head_object(Bucket=self.bucket_name, Key=self.archive_key)
+ etag: str = head.get("ETag", "")
+
+ marker = self.s3_dags_dir / self.archive_etag_marker
+ if etag and marker.is_file() and marker.read_text() == etag:
+ self._log.debug(
+ "Dag bundle archive 's3://%s/%s' is unchanged (ETag %s), skipping staging",
+ self.bucket_name,
+ self.archive_key,
+ etag,
+ )
+ return
+
+ staging_dir = Path(tempfile.mkdtemp(dir=self.s3_dags_dir.parent, prefix=".s3-archive-staging-"))
+ try:
+ archive_path = staging_dir / "_bundle_archive"
+ client.download_file(self.bucket_name, self.archive_key, os.fspath(archive_path))
Review Comment:
The marker stores the ETag from the HEAD above, but `download_file` fetches whatever is current at GET time, so if the archive is replaced in between, the marker can describe a different object than the one unpacked. Streaming `get_object(..., IfMatch=etag)` to the file closes the gap (`download_file` doesn't accept `IfMatch`). It also drops the HEAD that `download_file` sends by default, so a changed archive costs 2 HEAD + 1 GET today, not the "one HEAD, one GET" in the docs.
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +170,60 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _refresh_from_archive(self) -> None:
+ """Stage the Dag bundle by downloading and unpacking the single archive object."""
+ client = self.s3_hook.get_conn()
+ head = client.head_object(Bucket=self.bucket_name, Key=self.archive_key)
+ etag: str = head.get("ETag", "")
+
+ marker = self.s3_dags_dir / self.archive_etag_marker
+ if etag and marker.is_file() and marker.read_text() == etag:
+ self._log.debug(
+ "Dag bundle archive 's3://%s/%s' is unchanged (ETag %s), skipping staging",
+ self.bucket_name,
+ self.archive_key,
+ etag,
+ )
+ return
+
+ staging_dir = Path(tempfile.mkdtemp(dir=self.s3_dags_dir.parent, prefix=".s3-archive-staging-"))
+ try:
+ archive_path = staging_dir / "_bundle_archive"
+ client.download_file(self.bucket_name, self.archive_key, os.fspath(archive_path))
+
+ unpack_dir = staging_dir / "unpacked"
+ unpack_dir.mkdir()
+ with tarfile.open(archive_path, "r:*") as tar:
+ tar.extractall(unpack_dir, filter="data")
Review Comment:
I don't think `filter` is available for python 3.11.0-3
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +170,60 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _refresh_from_archive(self) -> None:
+ """Stage the Dag bundle by downloading and unpacking the single archive object."""
+ client = self.s3_hook.get_conn()
+ head = client.head_object(Bucket=self.bucket_name, Key=self.archive_key)
+ etag: str = head.get("ETag", "")
+
+ marker = self.s3_dags_dir / self.archive_etag_marker
+ if etag and marker.is_file() and marker.read_text() == etag:
+ self._log.debug(
+ "Dag bundle archive 's3://%s/%s' is unchanged (ETag %s), skipping staging",
+ self.bucket_name,
+ self.archive_key,
+ etag,
+ )
+ return
+
+ staging_dir = Path(tempfile.mkdtemp(dir=self.s3_dags_dir.parent, prefix=".s3-archive-staging-"))
Review Comment:
If the process is killed mid-staging (OOM, SIGKILL), the `finally` never runs and a full copy of the Dags stays in the shared bundle storage root. Can we put the bundle name in the prefix (e.g. `f".{self.name}-archive-staging-"`) and remove leftovers under `self.lock()` before creating a new one? The lock is per bundle, so without the name a cleanup could hit another bundle's in-progress staging.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4410417172
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -17,14 +17,40 @@
from __future__ import annotations
import os
+import shutil
+import tarfile
+import tempfile
from pathlib import Path
import structlog
from airflow.dag_processing.bundles.base import BaseDagBundle
from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException
+
+
+def _extract_archive(archive_path: Path, unpack_dir: Path) -> None:
+ """Unpack a bundle archive without writing outside ``unpack_dir``."""
+ with tarfile.open(archive_path, "r:*") as tar:
+ # filter="data" arrived in Python 3.11.4. 3.11.0-3.11.3 still have to reject
+ # the same escapes by hand.
+ if hasattr(tarfile, "data_filter"):
Review Comment:
With 3.11 already in security-only mode and being dropped next year, I wonder if it’s worthwhile to implement a fallback at all. It would be acceptable IMO to simply gate this feature if `filter="data"` is not available.
The current fallback implementation already has gaps compared to stdlib regarding symlinks etc.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4410382851
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -17,14 +17,40 @@
from __future__ import annotations
import os
+import shutil
+import tarfile
+import tempfile
from pathlib import Path
import structlog
from airflow.dag_processing.bundles.base import BaseDagBundle
from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException
+
+
+def _extract_archive(archive_path: Path, unpack_dir: Path) -> None:
+ """Unpack a bundle archive without writing outside ``unpack_dir``."""
+ with tarfile.open(archive_path, "r:*") as tar:
+ # filter="data" arrived in Python 3.11.4. 3.11.0-3.11.3 still have to reject
+ # the same escapes by hand.
Review Comment:
```suggestion
filter="data" requires Python 3.11.4 or later.
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4410382851
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -17,14 +17,40 @@
from __future__ import annotations
import os
+import shutil
+import tarfile
+import tempfile
from pathlib import Path
import structlog
from airflow.dag_processing.bundles.base import BaseDagBundle
from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException
+
+
+def _extract_archive(archive_path: Path, unpack_dir: Path) -> None:
+ """Unpack a bundle archive without writing outside ``unpack_dir``."""
+ with tarfile.open(archive_path, "r:*") as tar:
+ # filter="data" arrived in Python 3.11.4. 3.11.0-3.11.3 still have to reject
+ # the same escapes by hand.
Review Comment:
```suggestion
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "aoelvp94 (via GitHub)" <gi...@apache.org>.
aoelvp94 commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4407457141
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +170,60 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _refresh_from_archive(self) -> None:
+ """Stage the Dag bundle by downloading and unpacking the single archive object."""
+ client = self.s3_hook.get_conn()
+ head = client.head_object(Bucket=self.bucket_name, Key=self.archive_key)
+ etag: str = head.get("ETag", "")
+
+ marker = self.s3_dags_dir / self.archive_etag_marker
+ if etag and marker.is_file() and marker.read_text() == etag:
+ self._log.debug(
+ "Dag bundle archive 's3://%s/%s' is unchanged (ETag %s), skipping staging",
+ self.bucket_name,
+ self.archive_key,
+ etag,
+ )
+ return
+
+ staging_dir = Path(tempfile.mkdtemp(dir=self.s3_dags_dir.parent, prefix=".s3-archive-staging-"))
Review Comment:
Done in [#73015](https://github.com/apache/airflow/pull/73015/commits/69bacb5c8bee33a487d55ed5e35fae1c5cf9d176). The staging prefix is now .{bundle name}-archive-staging-. refresh() already holds this bundle's lock, and leftover directories with that prefix are removed before a new one is created. That cleanup also runs when the ETag is unchanged, because a killed process can leave a directory while the marker still matches. The match is iterdir plus startswith, so characters in the bundle name stay literal.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "ramitkataria (via GitHub)" <gi...@apache.org>.
ramitkataria commented on PR #73015:
URL: https://github.com/apache/airflow/pull/73015#issuecomment-5820512236
+1 to both. The fallback can also make things worse: a transient error such as throttling triggers it, turning 1 request into ~2N at exactly the wrong time. I'm also curious how the archive approach compares to concurrent `get_object` calls on performance.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "aoelvp94 (via GitHub) via commits" <co...@airflow.apache.org>.
aoelvp94 commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4434645535
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +193,95 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _archive_staging_prefix(self) -> str:
+ # The storage root is shared by every bundle. The lock is per bundle name, so
+ # leftover cleanup must not match another bundle's in-progress staging directory.
+ return f".{self.name}-archive-staging-"
+
+ def _remove_stale_archive_staging_dirs(self) -> None:
+ storage_root = self.s3_dags_dir.parent
+ if not storage_root.is_dir():
+ return
+ prefix = self._archive_staging_prefix()
+ for leftover in storage_root.iterdir():
+ if leftover.is_dir() and leftover.name.startswith(prefix):
+ shutil.rmtree(leftover, ignore_errors=True)
+
+ def _download_archive(self, client, archive_path: Path, etag: str) -> str:
+ """Download the archive object named by ``etag``."""
+ # download_file HEADs again and GETs whatever is current, so a replace in between
+ # can store an ETag for a different object than the one unpacked.
+ request = {"Bucket": self.bucket_name, "Key": self.archive_key}
+ if etag:
+ request["IfMatch"] = etag
+ response = client.get_object(**request)
+ body = response["Body"]
+ try:
+ with archive_path.open("wb") as handle:
+ for chunk in body.iter_chunks():
+ handle.write(chunk)
+ finally:
+ body.close()
+ return response.get("ETag", "")
+
+ def _refresh_from_archive(self) -> None:
+ """Stage the Dag bundle by downloading and unpacking the single archive object."""
+ # A killed process never reaches the finally that removes its staging directory.
+ self._remove_stale_archive_staging_dirs()
+ client = self.s3_hook.get_conn()
+ head = client.head_object(Bucket=self.bucket_name, Key=self.archive_key)
+ etag: str = head.get("ETag", "")
+
+ marker = self.s3_dags_dir / self.archive_etag_marker
+ if etag and marker.is_file() and marker.read_text() == etag:
+ self._log.debug(
+ "Dag bundle archive 's3://%s/%s' is unchanged (ETag %s), skipping staging",
+ self.bucket_name,
+ self.archive_key,
+ etag,
+ )
+ return
+
+ staging_dir = Path(
+ tempfile.mkdtemp(dir=self.s3_dags_dir.parent, prefix=self._archive_staging_prefix())
+ )
+ try:
+ archive_path = staging_dir / "_bundle_archive"
+ downloaded_etag = self._download_archive(client, archive_path, etag)
+
+ unpack_dir = staging_dir / "unpacked"
+ unpack_dir.mkdir()
+ _extract_archive(archive_path, unpack_dir)
+ archive_path.unlink()
+
+ marker_etag = downloaded_etag or etag
+ if marker_etag:
+ (unpack_dir / self.archive_etag_marker).write_text(marker_etag)
+
+ # Swap the freshly unpacked tree into place so a partially staged bundle is
+ # never observable at self.s3_dags_dir. The previous tree is kept inside the
+ # staging directory: a failed swap restores it, and the staging cleanup below
+ # removes it once the swap succeeded.
Review Comment:
Dropped the comment.
https://github.com/apache/airflow/pull/73015/commits/4cf1c1d4ac0b51fa7d67ccf0604e30df17ad4e10
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +193,95 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _archive_staging_prefix(self) -> str:
+ # The storage root is shared by every bundle. The lock is per bundle name, so
+ # leftover cleanup must not match another bundle's in-progress staging directory.
+ return f".{self.name}-archive-staging-"
Review Comment:
Staging directories are now `.{bundle name}/archive-staging-<random>`. Cleanup only looks inside that directory.
https://github.com/apache/airflow/pull/73015/commits/d046c82a87a055628513a3f34bac846b300de78e
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "aoelvp94 (via GitHub) via commits" <co...@airflow.apache.org>.
aoelvp94 commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4434647066
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -37,23 +63,40 @@ class S3DagBundle(BaseDagBundle):
:param bucket_name: The name of the S3 bucket containing the Dag files.
:param prefix: Optional subdirectory within the S3 bucket where the Dags are stored.
A trailing slash is optional.
- If None, Dags are assumed to be at the root of the bucket (Optional).
+ If None, Dags are assumed to be at the root of the bucket. Mutually exclusive
+ with ``archive_key`` (Optional).
+ :param archive_key: Optional S3 key of a ``.tar.gz`` archive holding the Dag files. When set,
+ the bundle is staged by downloading this single object and unpacking it
+ locally, instead of downloading every object under a prefix one at a time.
+ The archive is then the only source of the Dags: ``archive_key`` is mutually
+ exclusive with ``prefix``, so there is no second copy to keep in sync and no
+ fallback to a per-object sync. Members are extracted relative to the bundle
+ root, e.g. created with ``tar -C <dags_dir> -czf dags.tar.gz .`` (Optional).
"""
supports_versioning = False
+ archive_etag_marker = ".airflow_bundle_archive_etag"
Review Comment:
The marker is `.{bundle name}/etag`, beside the bundle directory, and it is written after the swap. It does not go in `_tracking`, because those files are parsed as version timestamps.
https://github.com/apache/airflow/pull/73015/commits/d046c82a87a055628513a3f34bac846b300de78e
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "ephraimbuddy (via GitHub) via commits" <co...@airflow.apache.org>.
ephraimbuddy commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4455825627
##########
providers/amazon/docs/bundles/index.rst:
##########
@@ -50,3 +50,39 @@ Example of using the S3DagBundle:
}
}
]'
+
+Staging from a single archive object
+------------------------------------
+
+By default the bundle is staged by downloading every object under ``prefix`` one at a time. Each object
+costs a full request round-trip, so bundles made of many small files can take a long time to stage — a cost
+paid by every component that stages the bundle, which with ephemeral workers (e.g. KubernetesExecutor task
+pods) means on every task start.
+
+Setting ``archive_key`` instead of ``prefix`` stages the bundle from a single ``.tar.gz`` object: one
Review Comment:
nit: worth saying here that archive mode needs Python 3.11.4+, since the bundle now fails at construction on older versions.
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +179,88 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _archive_state_dir(self) -> Path:
+ return self.s3_dags_dir.parent / f".{self.name}"
+
+ def _archive_etag_path(self) -> Path:
+ return self._archive_state_dir() / "etag"
+
+ def _remove_stale_archive_staging_dirs(self) -> None:
+ state_dir = self._archive_state_dir()
+ if not state_dir.is_dir():
+ return
+ for leftover in state_dir.iterdir():
+ if leftover.is_dir() and leftover.name.startswith("archive-staging-"):
+ shutil.rmtree(leftover, ignore_errors=True)
+
+ def _download_archive(self, client, archive_path: Path, etag: str) -> str:
+ """Download the archive object named by ``etag``."""
+ request = {"Bucket": self.bucket_name, "Key": self.archive_key}
+ if etag:
+ request["IfMatch"] = etag
+ response = client.get_object(**request)
+ body = response["Body"]
+ try:
+ with archive_path.open("wb") as handle:
+ for chunk in body.iter_chunks():
Review Comment:
nit: `iter_chunks()` defaults to 1 KiB chunks. `shutil.copyfileobj(body, handle, 1024 * 1024)` avoids thousands of small reads for a multi-MB archive.
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -140,6 +179,88 @@ def refresh(self) -> None:
delete_stale=True,
)
+ def _archive_state_dir(self) -> Path:
+ return self.s3_dags_dir.parent / f".{self.name}"
+
+ def _archive_etag_path(self) -> Path:
+ return self._archive_state_dir() / "etag"
+
+ def _remove_stale_archive_staging_dirs(self) -> None:
+ state_dir = self._archive_state_dir()
+ if not state_dir.is_dir():
+ return
+ for leftover in state_dir.iterdir():
+ if leftover.is_dir() and leftover.name.startswith("archive-staging-"):
+ shutil.rmtree(leftover, ignore_errors=True)
+
+ def _download_archive(self, client, archive_path: Path, etag: str) -> str:
+ """Download the archive object named by ``etag``."""
+ request = {"Bucket": self.bucket_name, "Key": self.archive_key}
+ if etag:
+ request["IfMatch"] = etag
+ response = client.get_object(**request)
+ body = response["Body"]
+ try:
+ with archive_path.open("wb") as handle:
+ for chunk in body.iter_chunks():
+ handle.write(chunk)
+ finally:
+ body.close()
+ return response.get("ETag", "")
+
+ def _refresh_from_archive(self) -> None:
+ """Stage the Dag bundle by downloading and unpacking the single archive object."""
+ self._remove_stale_archive_staging_dirs()
+ client = self.s3_hook.get_conn()
+ head = client.head_object(Bucket=self.bucket_name, Key=self.archive_key)
+ etag: str = head.get("ETag", "")
+
+ marker = self._archive_etag_path()
+ if etag and marker.is_file() and marker.read_text() == etag:
Review Comment:
With the marker outside the Dag folder, it can outlive the tree it describes. If the bundle folder is cleared (e.g. to force a resync) and the Dag processor restarts, `_initialize` recreates it empty, this check matches the ETag and skips, and every Dag in the bundle gets deactivated. Switching to `prefix` and back with an unchanged archive similarly keeps the prefix sync's files. Can we re-stage when the bundle folder is empty, drop the marker on the prefix path, and add a test for the delete-and-reinitialize case?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "ephraimbuddy (via GitHub) via commits" <co...@airflow.apache.org>.
ephraimbuddy commented on PR #73015:
URL: https://github.com/apache/airflow/pull/73015#issuecomment-6606291000
> > A high-level design question: Should archive mode be a separate class?
>
> Yeah, same question as I pointed out in [#73015 (review)](https://github.com/apache/airflow/pull/73015#pullrequestreview-5642408146).
Agree with a separate class now. The archive path has grown its own state dir, marker, staging cleanup and download, so `S3DagBundle` would be two implementations behind a flag.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add archive staging support to S3DagBundle [airflow]
Posted by "aoelvp94 (via GitHub) via commits" <co...@airflow.apache.org>.
aoelvp94 commented on PR #73015:
URL: https://github.com/apache/airflow/pull/73015#issuecomment-6653087189
> > > A high-level design question: Should archive mode be a separate class?
> >
> >
> > Yeah, same question as I pointed out in [#73015 (review)](https://github.com/apache/airflow/pull/73015#pullrequestreview-5642408146).
>
> Agree with a separate class now. The archive path has grown its own state dir, marker, staging cleanup and download, so `S3DagBundle` would be two implementations behind a flag.
Agreed — I'll split archive mode into a separate class (e.g. `S3ArchiveDagBundle`) instead of branching inside `S3DagBundle`. The archive path has its own marker, staging dir, and download path now, so a flag on the existing class is the wrong shape.
I'll keep `S3DagBundle` as today's prefix path and move the archive behavior into the new class. Happy to use a thin `__new__` dispatch on `S3DagBundle` if we want one config entry point, or two explicit classpaths — whichever you prefer. I'll update GCS (#75566) and Azure (#75657) the same way once S3 lands.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add S3ArchiveDagBundle for single-archive staging [airflow]
Posted by "o-nikolas (via GitHub) via commits" <co...@airflow.apache.org>.
o-nikolas commented on code in PR #73015:
URL: https://github.com/apache/airflow/pull/73015#discussion_r4486040631
##########
providers/amazon/tests/unit/amazon/aws/bundles/test_s3.py:
##########
@@ -262,3 +282,285 @@ def test_refresh_uses_directory_prefix(self, s3_client, prefix, extra_suffix):
assert repr(bundle) == original_repr
assert bundle.view_url_template() == original_url
assert s3_client.get_object(Bucket=S3_BUCKET_NAME, Key=extra_key)["Body"].read() == b"outside"
+
+ def test_rejects_archive_key(self):
+ with pytest.raises(TypeError, match="archive_key"):
+ S3DagBundle(
+ name="test",
+ aws_conn_id=AWS_CONN_ID_DEFAULT,
+ bucket_name=S3_BUCKET_NAME,
+ prefix=S3_BUCKET_PREFIX,
+ archive_key=S3_ARCHIVE_KEY,
Review Comment:
nit: `test_rejects_archive_key` and `test_rejects_prefix` (below) only check Python's own unexpected-keyword `TypeError`. The first one also passes on main. I think both can go.
##########
providers/amazon/src/airflow/providers/amazon/aws/bundles/s3.py:
##########
@@ -27,33 +30,44 @@
from airflow.providers.common.compat.sdk import AirflowException
-class S3DagBundle(BaseDagBundle):
- """
- S3 Dag bundle - exposes a directory in S3 as a Dag bundle.
+def _require_tar_data_filter() -> None:
+ if hasattr(tarfile, "data_filter"):
+ return
+ raise RuntimeError("Archive Dag bundles require Python 3.11.4 or later (tarfile filter='data').")
- This allows Airflow to load Dags directly from an S3 bucket.
- :param aws_conn_id: Airflow connection ID for AWS. Defaults to AwsBaseHook.default_conn_name.
- :param bucket_name: The name of the S3 bucket containing the Dag files.
- :param prefix: Optional subdirectory within the S3 bucket where the Dags are stored.
- A trailing slash is optional.
- If None, Dags are assumed to be at the root of the bucket (Optional).
+def _extract_archive(archive_path: Path, unpack_dir: Path) -> None:
+ """Unpack a bundle archive without writing outside ``unpack_dir``."""
+ _require_tar_data_filter()
+ with tarfile.open(archive_path, "r:*") as tar:
+ tar.extractall(unpack_dir, filter="data")
+
+
+def _archive_state_dir(bundle_dir: Path, bundle_name: str) -> Path:
+ return bundle_dir.parent / f".{bundle_name}"
+
+
+def _archive_etag_path(bundle_dir: Path, bundle_name: str) -> Path:
+ return _archive_state_dir(bundle_dir, bundle_name) / "etag"
Review Comment:
nit: `_archive_state_dir` / `_archive_etag_path` read like attributes. AGENTS.md asks for verb names, e.g. `_get_archive_state_dir` / `_get_archive_etag_path`.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org