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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
