kaxil commented on code in PR #76317: URL: https://github.com/apache/airflow/pull/76317#discussion_r4514638260
########## providers/huggingface/src/airflow/providers/huggingface/operators/datasets.py: ########## @@ -0,0 +1,647 @@ +# 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. +"""Operators to load, transform, and publish Hugging Face datasets.""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path +import tempfile +from typing import TYPE_CHECKING, Any, Callable + +if TYPE_CHECKING: + from datasets import Dataset, DatasetDict + +from airflow.providers.common.compat.sdk import Asset, BaseOperator, Context, ObjectStoragePath + +from airflow.providers.huggingface.hooks.huggingface import HuggingFaceHook + + +def _path(uri: str, *, conn_id: str | None = None) -> ObjectStoragePath: + """Resolve a URI through Airflow's configured object-storage filesystems.""" + return ObjectStoragePath(uri, conn_id=conn_id) + + +def _reference(reference: str | dict[str, Any]) -> str: + if isinstance(reference, dict): + uri = reference.get("uri") + if not uri: + raise ValueError("Artifact reference must contain a non-empty 'uri'.") + return uri + if not isinstance(reference, str) or not reference: + raise ValueError("Artifact reference must be a URI or contain a 'uri'.") + return reference + + +def _sha256_path(uri: str, *, conn_id: str | None = None) -> str: + digest = hashlib.sha256() + with _path(uri, conn_id=conn_id).open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _validate_feature(feature: Any, field_name: str) -> None: + """Allow Parquet-safe scalar, class-label, and nested sequence features.""" + from datasets import ClassLabel, LargeList, Sequence, Value + + if isinstance(feature, (Value, ClassLabel)): + return + if isinstance(feature, (Sequence, LargeList)): + _validate_feature(feature.feature, field_name) + return + if isinstance(feature, dict): + for key, value in feature.items(): + _validate_feature(value, f"{field_name}.{key}") + return + if isinstance(feature, (list, tuple)): + for value in feature: + _validate_feature(value, field_name) + return + raise ValueError( + f"unsupported Parquet feature type for {field_name!r}: {type(feature).__name__}; " + "supported types are Value, ClassLabel, and nested Sequence or LargeList features" + ) + + +def _dataset_parts(dataset: Dataset | DatasetDict) -> list[tuple[str, Dataset]]: + from datasets import Dataset, DatasetDict + + if isinstance(dataset, DatasetDict): + return list(dataset.items()) + if isinstance(dataset, Dataset): + return [("dataset", dataset)] + raise TypeError("Expected a datasets.Dataset or datasets.DatasetDict.") + + +def _save( + dataset: Dataset | DatasetDict, + uri: str, + metadata: dict[str, Any] | None = None, + *, + storage_conn_id: str | None = None, +) -> dict[str, Any]: + """Write split Parquet objects and a JSON manifest through ObjectStoragePath.""" + import pyarrow.parquet as parquet + + parts = _dataset_parts(dataset) + written: list[ObjectStoragePath] = [] + manifest_path = _path(uri, conn_id=storage_conn_id) + split_manifest: dict[str, Any] = {} + try: + for index, (name, part) in enumerate(parts): + for field_name, feature in part.features.items(): + _validate_feature(feature, field_name) + split_uri = f"{uri}.split-{index:05d}.parquet" + split_path = _path(split_uri, conn_id=storage_conn_id) + split_path.parent.mkdir(parents=True, exist_ok=True) + with split_path.open("wb") as stream: + parquet.write_table(part.data.table, stream) + written.append(split_path) + split_manifest[name] = { + "uri": split_uri, + "rows": len(part), + "features": part.features.to_dict(), + "feature_order": list(part.features), + "split": ( + str(part.split) + if getattr(part, "split", None) is not None + else (name if name != "dataset" else None) + ), + "sha256": _sha256_path(split_uri, conn_id=storage_conn_id), + } + + manifest = { + "format": "parquet", + "version": 1, + "splits": split_manifest, + "metadata": metadata or {}, + } + manifest_bytes = json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode("utf-8") + manifest_path.parent.mkdir(parents=True, exist_ok=True) + with manifest_path.open("wb") as stream: + stream.write(manifest_bytes) + return { + "uri": uri, + "format": "parquet", + "splits": { + name: { + "rows": part_info["rows"], + "features": part_info["features"], + "split": part_info["split"], + } + for name, part_info in split_manifest.items() + }, + **(metadata or {}), + "fingerprint": hashlib.sha256(manifest_bytes).hexdigest(), + } + except Exception: + # Remove partial split outputs. A failed write must not leave an artifact + # reference that a downstream task could mistake for a complete dataset. + for path in written: + try: + path.unlink(missing_ok=True) + except (OSError, NotImplementedError): + pass + try: + manifest_path.unlink(missing_ok=True) + except (OSError, NotImplementedError): + pass + raise + + +def _read_artifact_manifest( + reference: str | dict[str, Any], *, storage_conn_id: str | None = None +) -> tuple[str, dict[str, Any]]: + uri = _reference(reference) + manifest_path = _path(uri, conn_id=storage_conn_id) + with manifest_path.open("rb") as stream: + manifest_bytes = stream.read() + if isinstance(reference, dict) and reference.get("fingerprint"): + actual = hashlib.sha256(manifest_bytes).hexdigest() + if actual != reference["fingerprint"]: + raise ValueError("Dataset artifact fingerprint does not match its reference") + + manifest = json.loads(manifest_bytes) + if manifest.get("format") != "parquet" or manifest.get("version") != 1: + raise ValueError(f"Unsupported dataset artifact manifest: {uri}") + return uri, manifest + + +def _load_artifact( + reference: str | dict[str, Any], + *, + cache_dir: str | None = None, + manifest: dict[str, Any] | None = None, + storage_conn_id: str | None = None, +) -> Dataset | DatasetDict: + from datasets import Dataset, DatasetDict, Features + + uri, manifest = ( + _read_artifact_manifest(reference, storage_conn_id=storage_conn_id) + if manifest is None + else (_reference(reference), manifest) + ) + resolved_cache_dir = cache_dir or str(Path(tempfile.gettempdir()) / "airflow-hf-datasets-cache") + Path(resolved_cache_dir).mkdir(parents=True, exist_ok=True) + parts: dict[str, Dataset] = {} + with tempfile.TemporaryDirectory(prefix="airflow-hf-parquet-") as local_dir: + for index, (name, item) in enumerate(manifest["splits"].items()): + split_uri = item["uri"] + local_path = Path(local_dir) / f"split-{index:05d}.parquet" + digest = hashlib.sha256() + with _path(split_uri, conn_id=storage_conn_id).open("rb") as source, local_path.open("wb") as target: + for chunk in iter(lambda: source.read(1024 * 1024), b""): + digest.update(chunk) + target.write(chunk) + if digest.hexdigest() != item.get("sha256"): + raise ValueError(f"Dataset artifact split fingerprint mismatch: {split_uri}") + + feature_data = item["features"] + feature_order = item.get("feature_order", list(feature_data)) + features = Features.from_dict({key: feature_data[key] for key in feature_order}) + raw_split = item.get("split") or (name if name != "dataset" else None) + split_name = raw_split.split("[")[0] if raw_split and "[" in raw_split else raw_split + restored = Dataset.from_parquet( + str(local_path), + features=features, + split=split_name, + cache_dir=resolved_cache_dir, + ) + # Parquet readers may return fields in physical/file order; restore the + # manifest's declared feature order for stable downstream transforms. + parts[name] = restored.select_columns(list(features)) + return parts["dataset"] if list(parts) == ["dataset"] else DatasetDict(parts) + + +def _state_key(prefix: str, identity: str) -> str: + """Task state stores already scope keys to a DAG run and task instance.""" + return f"{prefix}:{identity}" + + +def _publication_marker(prefix: str, identity: dict[str, Any]) -> str: + payload = json.dumps(identity, sort_keys=True, separators=(",", ":")).encode() + return f"{prefix}:{hashlib.sha256(payload).hexdigest()}" + + +def _commit_details(commit: Any, repo_id: str, hook: HuggingFaceHook) -> tuple[str, str]: + """Normalize CommitInfo and GitCommitInfo returned by Hub APIs.""" + commit_sha = getattr(commit, "oid", None) or getattr(commit, "commit_id", None) + if not commit_sha: + raise ValueError(f"Hub publication for {repo_id!r} returned no commit identifier") + commit_url = getattr(commit, "commit_url", None) + if not commit_url: + endpoint = getattr(hook.get_conn(), "endpoint", "https://huggingface.co";).rstrip("/") + commit_url = f"{endpoint}/datasets/{repo_id}/commit/{commit_sha}" + return commit_sha, commit_url + + +def _is_hub_repo(path: str) -> bool: + """Match a Hub dataset repo ID without treating an existing local path as remote.""" + return path.count("/") == 1 and not path.startswith((".", "/")) and not Path(path).exists() + + +class LoadDatasetOperator(BaseOperator): + """Load a dataset revision and persist it as a Parquet artifact.""" + + template_fields = ("path", "name", "split", "revision", "cache_dir", "output_uri") + + def __init__( + self, + *, + path: str, + output_uri: str, + name: str | None = None, + split: str | None = None, + revision: str | None = None, + cache_dir: str | None = None, + storage_conn_id: str | None = None, + download_mode: str = "reuse_dataset_if_exists", + huggingface_conn_id: str | None = "huggingface_default", + load_kwargs: dict[str, Any] | None = None, + durable: bool = True, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.path, self.output_uri, self.name, self.split = path, output_uri, name, split + self.revision, self.cache_dir, self.download_mode = revision, cache_dir, download_mode + self.storage_conn_id = storage_conn_id + self.huggingface_conn_id, self.load_kwargs = huggingface_conn_id, load_kwargs or {} + self.durable = durable + + def execute(self, context: Context) -> dict[str, Any]: + hook = HuggingFaceHook(self.huggingface_conn_id) + requested_revision = self.revision + is_hub_repo = _is_hub_repo(self.path) + state = context.get("task_state_store") if self.durable else None + state_identity = json.dumps([self.path, requested_revision], separators=(",", ":")) + state_key = _state_key("load-revision", state_identity) + revision = state.get(state_key) if state is not None else None + if revision is None: + revision = requested_revision + if is_hub_repo: + revision = hook.get_dataset_info(self.path, requested_revision).sha + if state is not None and revision is not None: + state.set(state_key, revision) + elif state is not None: + self.log.info("Reusing pinned Hugging Face dataset revision %s for %s", revision, self.path) + + kwargs = dict(self.load_kwargs) + kwargs["download_mode"] = self.download_mode + if self.name is not None: + kwargs["name"] = self.name + if self.split is not None: + kwargs["split"] = self.split + if revision is not None: + kwargs["revision"] = revision + if self.cache_dir is not None: + kwargs["cache_dir"] = self.cache_dir + dataset = hook.load_dataset(self.path, **kwargs) + metadata = { + "repository_id": self.path if is_hub_repo else None, + "revision": revision, + "configuration": self.name, + "split": self.split, + } + return _save(dataset, self.output_uri, metadata, storage_conn_id=self.storage_conn_id) + + +class StreamDatasetOperator(BaseOperator): + """Transform a streamed dataset lazily and publish it directly to the Hub.""" + + template_fields = ("path", "name", "split", "revision", "config", "commit_message") + + def __init__( + self, + *, + path: str, + repo_id: str, + transform: Callable[..., dict[str, Any]] | None = None, + predicate: Callable[[dict[str, Any]], bool] | None = None, + name: str | None = None, + split: str | None = None, + revision: str | None = None, + config: str | None = None, + private: bool | None = None, + commit_message: str | None = None, + transform_kwargs: dict[str, Any] | None = None, + push_kwargs: dict[str, Any] | None = None, + huggingface_conn_id: str | None = "huggingface_default", + durable: bool = True, + **kwargs: Any, + ) -> None: + if "{{" in repo_id: + raise ValueError("repo_id cannot be templated: it defines the Asset that downstream DAGs schedule on.") + if transform is not None and not callable(transform): + raise TypeError("transform must be a callable or None") + if predicate is not None and not callable(predicate): + raise TypeError("predicate must be a callable or None") + self.path, self.repo_id = path, repo_id + self.transform, self.predicate = transform, predicate + self.name, self.split, self.revision = name, split, revision + self.config, self.private, self.commit_message = config, private, commit_message + self.transform_kwargs = transform_kwargs or {} + self.push_kwargs = push_kwargs or {} + self.huggingface_conn_id = huggingface_conn_id + self.durable = durable + self.published_asset = Asset(uri=f"hf://datasets/{repo_id}") + user_outlets = kwargs.pop("outlets", None) + if user_outlets is None: + user_outlets = [] + elif not isinstance(user_outlets, (list, tuple)): + user_outlets = [user_outlets] + super().__init__(outlets=[*user_outlets, self.published_asset], **kwargs) + + def execute(self, context: Context) -> dict[str, Any]: + hook = HuggingFaceHook(self.huggingface_conn_id) + is_hub_repo = _is_hub_repo(self.path) + requested_revision = self.revision + state = context.get("task_state_store") if self.durable else None + revision_identity = json.dumps( + [self.path, requested_revision, self.name, self.split], separators=(",", ":") + ) + revision_key = _state_key("stream-revision", revision_identity) + revision = state.get(revision_key) if state is not None else None + if revision is None: + revision = requested_revision + if is_hub_repo: + revision = hook.get_dataset_info(self.path, requested_revision).sha + if state is not None and revision is not None: + state.set(revision_key, revision) + elif state is not None: + self.log.info("Reusing pinned Hugging Face streaming revision %s for %s", revision, self.path) + + run_id = context.get("run_id", "manual") + push_kwargs = dict(self.push_kwargs) + if self.config is not None: + push_kwargs["config_name"] = self.config + if self.private is not None: + push_kwargs["private"] = self.private + marker = _publication_marker( + "huggingface-stream-publish", + { + "task_id": self.task_id, + "run_id": run_id, + "source": self.path, + "source_revision": revision, + "name": self.name, + "split": self.split, + "repo_id": self.repo_id, + "config_name": push_kwargs.get("config_name", "default"), + "revision": push_kwargs.get("revision", "main"), + "create_pr": bool(push_kwargs.get("create_pr", False)), + }, + ) + state_key = _state_key("stream-publication", marker) + if state is not None: + prior_result = state.get(state_key) + if prior_result: + self.log.info("Reusing stored streaming publication result for %s", self.repo_id) + self._emit_asset_event(context, prior_result) + return prior_result + + target_revision = push_kwargs.get("revision") + include_pull_requests = bool(push_kwargs.get("create_pr", False)) + commit = None + if self.durable: + commit = hook.find_commit_by_marker( + self.repo_id, + marker, + revision=target_revision, + include_pull_requests=include_pull_requests, + ) + if commit is not None: + self.log.info("Reusing existing completed streaming publication to %s", self.repo_id) + + if commit is None: + load_kwargs: dict[str, Any] = {"streaming": True} + if self.name is not None: + load_kwargs["name"] = self.name + if self.split is not None: + load_kwargs["split"] = self.split + if revision is not None: + load_kwargs["revision"] = revision + stream = hook.load_dataset(self.path, **load_kwargs) + + if self.transform is not None: + stream = stream.map(self.transform, fn_kwargs=self.transform_kwargs) + if self.predicate is not None: + stream = stream.filter(self.predicate) + + message = self.commit_message or "Stream dataset from Airflow" + push_kwargs["commit_message"] = f"{message} [{marker}]" + commit = hook.push_to_hub(stream, self.repo_id, **push_kwargs) + # A retry can recover the completed Hub commit if the worker exited + # after upload but before Airflow recorded the task result. + if commit is None: + if self.durable: + commit = hook.find_commit_by_marker( + self.repo_id, + marker, + revision=target_revision, + include_pull_requests=include_pull_requests, + ) + + if commit is None: + raise RuntimeError( + f"Hub push to {self.repo_id!r} completed without returning or exposing a commit for marker {marker}" + ) + commit_sha, commit_url = _commit_details(commit, self.repo_id, hook) + result = { + "repository_id": self.repo_id, + "commit_sha": commit_sha, + "commit_url": commit_url, + "url": f"https://huggingface.co/datasets/{self.repo_id}";, + "source_repository_id": self.path if is_hub_repo else None, + "source_revision": revision, + "configuration": self.name, + "split": self.split, + "streaming": True, + } + if state is not None: + state.set(state_key, result) + self._emit_asset_event(context, result) + return result + + def _emit_asset_event(self, context: Context, publication: dict[str, Any]) -> None: + events = context.get("outlet_events") + if events is not None: + events[self.published_asset].extra = { + "repository_id": self.repo_id, + "commit_sha": publication["commit_sha"], + "commit_url": publication["commit_url"], + "source_revision": publication["source_revision"], + } + + +class MapDatasetOperator(BaseOperator): + """Apply a Python callable and persist the transformed artifact.""" + + template_fields = ("input_uri", "output_uri", "transform_kwargs", "cache_dir") + + def __init__( + self, + *, + input_uri: str | dict[str, Any], + output_uri: str, + transform: Callable[..., Any], + transform_kwargs: dict[str, Any] | None = None, + map_kwargs: dict[str, Any] | None = None, + cache_dir: str | None = None, + storage_conn_id: str | None = None, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + if not callable(transform): + raise TypeError("transform must be a callable") + self.input_uri, self.output_uri, self.transform = input_uri, output_uri, transform + self.transform_kwargs, self.map_kwargs = transform_kwargs or {}, map_kwargs or {} + self.cache_dir = cache_dir + self.storage_conn_id = storage_conn_id + + def execute(self, context: Context) -> dict[str, Any]: + dataset = _load_artifact( + self.input_uri, cache_dir=self.cache_dir, storage_conn_id=self.storage_conn_id + ) + mapped = dataset.map(self.transform, fn_kwargs=self.transform_kwargs, **self.map_kwargs) + source = self.input_uri if isinstance(self.input_uri, dict) else {} + provenance = { + key: source[key] + for key in ("repository_id", "revision", "configuration", "split") + if key in source + } + return _save(mapped, self.output_uri, provenance, storage_conn_id=self.storage_conn_id) + + +class PushToHubOperator(BaseOperator): + """Synchronously publish an artifact and emit its successful asset event.""" + + template_fields = ("input_uri", "config", "commit_message") + + def __init__( + self, + *, + input_uri: str | dict[str, Any], + repo_id: str, + config: str | None = None, + private: bool | None = None, + commit_message: str | None = None, + huggingface_conn_id: str | None = "huggingface_default", + storage_conn_id: str | None = None, + push_kwargs: dict[str, Any] | None = None, + durable: bool = True, + **kwargs: Any, + ) -> None: + if "{{" in repo_id: + raise ValueError("repo_id cannot be templated: it defines the Asset that downstream DAGs schedule on.") + self.input_uri, self.repo_id, self.config = input_uri, repo_id, config + self.private, self.commit_message = private, commit_message + self.huggingface_conn_id, self.push_kwargs = huggingface_conn_id, push_kwargs or {} + self.storage_conn_id = storage_conn_id + self.durable = durable + self.published_asset = Asset(uri=f"hf://datasets/{repo_id}") + user_outlets = kwargs.pop("outlets", None) + if user_outlets is None: + user_outlets = [] + elif not isinstance(user_outlets, (list, tuple)): + user_outlets = [user_outlets] + super().__init__(outlets=[*user_outlets, self.published_asset], **kwargs) + + def execute(self, context: Context) -> dict[str, Any]: + source = self.input_uri if isinstance(self.input_uri, dict) else {} + fingerprint = source.get("fingerprint") + if not fingerprint: + raise ValueError("Publication requires an artifact reference with a fingerprint") + kwargs = dict(self.push_kwargs) + if self.config is not None: + kwargs["config_name"] = self.config + if self.private is not None: + kwargs["private"] = self.private + target_revision = kwargs.get("revision") + create_pr = bool(kwargs.get("create_pr", False)) + target_marker = _publication_marker( + "target", + { + "repo_id": self.repo_id, + "config_name": kwargs.get("config_name", "default"), Review Comment: This target identity leaves out `split` and `data_dir` from `push_kwargs`, and both decide where the data lands on the Hub. Pushing the same artifact with `push_kwargs={"split": "train"}` and then with `{"split": "validation"}` builds the same marker, so the second task finds the first commit and succeeds without uploading anything. Could both go into this dict? ########## providers/huggingface/src/airflow/providers/huggingface/hooks/huggingface.py: ########## @@ -0,0 +1,140 @@ +# 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. +"""Connection and client helpers for Hugging Face services.""" +from __future__ import annotations + +from typing import Any + +from huggingface_hub import HfApi +from huggingface_hub.utils import RepositoryNotFoundError, RevisionNotFoundError + +from airflow.exceptions import AirflowNotFoundException +from airflow.providers.common.compat.sdk import BaseHook + + +class HuggingFaceHook(BaseHook): + """Access Hugging Face Datasets and Hub APIs using an Airflow connection. + + If the connection is absent, huggingface_hub may use ``HF_TOKEN`` or a Review Comment: This now only holds for `huggingface_default` or `None`. A missing named connection raises, and a named connection with an empty password sends no token at all. Can the docstring say that? ########## providers/huggingface/src/airflow/providers/huggingface/asset.py: ########## @@ -0,0 +1,55 @@ +# 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. +"""Airflow asset URI and OpenLineage integration for Hub datasets.""" +from __future__ import annotations + +from urllib.parse import SplitResult, urlsplit + + +def sanitize_uri(uri: str | SplitResult) -> str | SplitResult: Review Comment: `breeze run mypy` fails on this file: `create_asset` passes the `str | SplitResult` this returns into `Asset(uri=...)`, and the `except ImportError` fallback in `create_asset` redefines `Asset`. Core calls the handler as `Callable[[SplitResult], SplitResult]` (the S3 one is typed that way), so could this keep that signature, with a small private helper for the two string callers? ########## providers/huggingface/.gitignore: ########## @@ -0,0 +1,221 @@ +*.iml +# Byte-compiled / optimized / DLL files Review Comment: The other providers' `.gitignore` is just `*.iml`. Everything from here down is the GitHub Python template, which also ignores paths like `build/` and `lib/` inside the provider. Can it be only the first line? ########## providers/huggingface/README.rst: ########## @@ -0,0 +1,80 @@ + +.. 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. + +.. NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE OVERWRITTEN! + +.. IF YOU WANT TO MODIFY TEMPLATE FOR THIS FILE, YOU SHOULD MODIFY THE TEMPLATE + ``PROVIDER_README_TEMPLATE.rst.jinja2`` IN the ``dev/breeze/src/airflow_breeze/templates`` DIRECTORY + +Package ``apache-airflow-providers-huggingface`` + +Release: ``0.1.0`` + + +Load, transform, and publish Hugging Face Datasets. + +Provider package +---------------- + +This is a provider package for ``huggingface`` provider. All classes for this provider package +are in ``airflow.providers.huggingface`` python package. + +You can find package information and changelog for the provider +in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-huggingface/0.1.0/>`_. + +Installation +------------ + +You can install this package on top of an existing Airflow installation (see ``Requirements`` below +for the minimum Airflow version supported) via +``pip install apache-airflow-providers-huggingface`` + +The package supports the following python versions: 3.11,3.12,3.13,3.14 + +Requirements +------------ + +=================== ================== +PIP package Version required +=================== ================== +``apache-airflow`` ``>=3.0,<4`` +``datasets`` ``>=3.4.0,<5`` Review Comment: This is still the old requirements table, and the `[common.compat]` extra further down no longer exists. It's generated from pyproject.toml, so regenerating should fix it. `docs/index.rst:37-38` also still says `<5` and `<2` in the prose. ########## providers/huggingface/provider.yaml: ########## @@ -0,0 +1,59 @@ +# 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. + +package-name: apache-airflow-providers-huggingface +name: Hugging Face +description: Load, transform, and publish Hugging Face Datasets. +state: not-ready +source-date-epoch: 0 +lifecycle: incubation +versions: + - 0.1.0 +integrations: + - integration-name: Hugging Face Datasets + external-doc-url: https://huggingface.co/docs/datasets/ + logo: /docs/integration-logos/HuggingFace.png Review Comment: The PNG is an Airflow plus Datasets banner rather than the Hugging Face logo. The other integration logos are the vendor's mark alone, so could this be the plain Hugging Face logo? ########## providers/huggingface/tests/system/huggingface/example_write.py: ########## @@ -0,0 +1,244 @@ +# 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. +"""Publish to a temporary private Hub repo and validate XCom, asset, and Hub head.""" +import os +import uuid + +import airflow +from packaging.version import Version +from airflow.sdk import Asset, DAG, get_current_context, task +from airflow.providers.common.compat.sdk import TriggerRule + +from airflow.providers.huggingface.hooks.huggingface import HuggingFaceHook +from airflow.providers.huggingface.operators.datasets import ( + LoadDatasetOperator, + MapDatasetOperator, + PushToHubOperator, +) + +HUB_ORG = os.environ.get("HF_TEST_HUB_ORG", "airflow-hf") +SYSTEM_TESTS_ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID", uuid.uuid4().hex) Review Comment: Breeze sets `SYSTEM_TESTS_ENV_ID` to an empty string, so `.get()` returns `""` here and `REPO_ID` ends in `-`, which `validate_repo_id` rejects as soon as `create_repo` runs. `os.environ.get("SYSTEM_TESTS_ENV_ID") or uuid.uuid4().hex`? ########## providers/huggingface/src/airflow/providers/huggingface/operators/datasets.py: ########## @@ -0,0 +1,647 @@ +# 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. +"""Operators to load, transform, and publish Hugging Face datasets.""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path +import tempfile +from typing import TYPE_CHECKING, Any, Callable + +if TYPE_CHECKING: + from datasets import Dataset, DatasetDict + +from airflow.providers.common.compat.sdk import Asset, BaseOperator, Context, ObjectStoragePath + +from airflow.providers.huggingface.hooks.huggingface import HuggingFaceHook + + +def _path(uri: str, *, conn_id: str | None = None) -> ObjectStoragePath: + """Resolve a URI through Airflow's configured object-storage filesystems.""" + return ObjectStoragePath(uri, conn_id=conn_id) + + +def _reference(reference: str | dict[str, Any]) -> str: + if isinstance(reference, dict): + uri = reference.get("uri") + if not uri: + raise ValueError("Artifact reference must contain a non-empty 'uri'.") + return uri + if not isinstance(reference, str) or not reference: + raise ValueError("Artifact reference must be a URI or contain a 'uri'.") + return reference + + +def _sha256_path(uri: str, *, conn_id: str | None = None) -> str: + digest = hashlib.sha256() + with _path(uri, conn_id=conn_id).open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _validate_feature(feature: Any, field_name: str) -> None: + """Allow Parquet-safe scalar, class-label, and nested sequence features.""" + from datasets import ClassLabel, LargeList, Sequence, Value + + if isinstance(feature, (Value, ClassLabel)): + return + if isinstance(feature, (Sequence, LargeList)): + _validate_feature(feature.feature, field_name) + return + if isinstance(feature, dict): + for key, value in feature.items(): + _validate_feature(value, f"{field_name}.{key}") + return + if isinstance(feature, (list, tuple)): + for value in feature: + _validate_feature(value, field_name) + return + raise ValueError( + f"unsupported Parquet feature type for {field_name!r}: {type(feature).__name__}; " + "supported types are Value, ClassLabel, and nested Sequence or LargeList features" + ) + + +def _dataset_parts(dataset: Dataset | DatasetDict) -> list[tuple[str, Dataset]]: + from datasets import Dataset, DatasetDict + + if isinstance(dataset, DatasetDict): + return list(dataset.items()) + if isinstance(dataset, Dataset): + return [("dataset", dataset)] + raise TypeError("Expected a datasets.Dataset or datasets.DatasetDict.") + + +def _save( + dataset: Dataset | DatasetDict, + uri: str, + metadata: dict[str, Any] | None = None, + *, + storage_conn_id: str | None = None, +) -> dict[str, Any]: + """Write split Parquet objects and a JSON manifest through ObjectStoragePath.""" + import pyarrow.parquet as parquet + + parts = _dataset_parts(dataset) + written: list[ObjectStoragePath] = [] + manifest_path = _path(uri, conn_id=storage_conn_id) + split_manifest: dict[str, Any] = {} + try: + for index, (name, part) in enumerate(parts): + for field_name, feature in part.features.items(): + _validate_feature(feature, field_name) + split_uri = f"{uri}.split-{index:05d}.parquet" + split_path = _path(split_uri, conn_id=storage_conn_id) + split_path.parent.mkdir(parents=True, exist_ok=True) + with split_path.open("wb") as stream: + parquet.write_table(part.data.table, stream) + written.append(split_path) + split_manifest[name] = { + "uri": split_uri, + "rows": len(part), + "features": part.features.to_dict(), + "feature_order": list(part.features), + "split": ( + str(part.split) + if getattr(part, "split", None) is not None + else (name if name != "dataset" else None) + ), + "sha256": _sha256_path(split_uri, conn_id=storage_conn_id), Review Comment: `_sha256_path` reads the split back right after writing it, so every Load and Map task sends each split to S3 or GCS and then downloads it again. Could the digest be computed while writing, with a small wrapper around `stream` that hashes what `write_table` writes? This can wait for a follow-up since it only costs transfer time. ########## providers/huggingface/src/airflow/providers/huggingface/operators/datasets.py: ########## @@ -0,0 +1,647 @@ +# 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. +"""Operators to load, transform, and publish Hugging Face datasets.""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path +import tempfile +from typing import TYPE_CHECKING, Any, Callable + +if TYPE_CHECKING: + from datasets import Dataset, DatasetDict + +from airflow.providers.common.compat.sdk import Asset, BaseOperator, Context, ObjectStoragePath + +from airflow.providers.huggingface.hooks.huggingface import HuggingFaceHook + + +def _path(uri: str, *, conn_id: str | None = None) -> ObjectStoragePath: + """Resolve a URI through Airflow's configured object-storage filesystems.""" + return ObjectStoragePath(uri, conn_id=conn_id) + + +def _reference(reference: str | dict[str, Any]) -> str: + if isinstance(reference, dict): + uri = reference.get("uri") + if not uri: + raise ValueError("Artifact reference must contain a non-empty 'uri'.") + return uri + if not isinstance(reference, str) or not reference: + raise ValueError("Artifact reference must be a URI or contain a 'uri'.") + return reference + + +def _sha256_path(uri: str, *, conn_id: str | None = None) -> str: + digest = hashlib.sha256() + with _path(uri, conn_id=conn_id).open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _validate_feature(feature: Any, field_name: str) -> None: + """Allow Parquet-safe scalar, class-label, and nested sequence features.""" + from datasets import ClassLabel, LargeList, Sequence, Value + + if isinstance(feature, (Value, ClassLabel)): + return + if isinstance(feature, (Sequence, LargeList)): + _validate_feature(feature.feature, field_name) + return + if isinstance(feature, dict): + for key, value in feature.items(): + _validate_feature(value, f"{field_name}.{key}") + return + if isinstance(feature, (list, tuple)): + for value in feature: + _validate_feature(value, field_name) + return + raise ValueError( + f"unsupported Parquet feature type for {field_name!r}: {type(feature).__name__}; " + "supported types are Value, ClassLabel, and nested Sequence or LargeList features" + ) + + +def _dataset_parts(dataset: Dataset | DatasetDict) -> list[tuple[str, Dataset]]: + from datasets import Dataset, DatasetDict + + if isinstance(dataset, DatasetDict): + return list(dataset.items()) + if isinstance(dataset, Dataset): + return [("dataset", dataset)] + raise TypeError("Expected a datasets.Dataset or datasets.DatasetDict.") + + +def _save( + dataset: Dataset | DatasetDict, + uri: str, + metadata: dict[str, Any] | None = None, + *, + storage_conn_id: str | None = None, +) -> dict[str, Any]: + """Write split Parquet objects and a JSON manifest through ObjectStoragePath.""" + import pyarrow.parquet as parquet + + parts = _dataset_parts(dataset) + written: list[ObjectStoragePath] = [] + manifest_path = _path(uri, conn_id=storage_conn_id) + split_manifest: dict[str, Any] = {} + try: + for index, (name, part) in enumerate(parts): + for field_name, feature in part.features.items(): + _validate_feature(feature, field_name) + split_uri = f"{uri}.split-{index:05d}.parquet" + split_path = _path(split_uri, conn_id=storage_conn_id) + split_path.parent.mkdir(parents=True, exist_ok=True) + with split_path.open("wb") as stream: + parquet.write_table(part.data.table, stream) + written.append(split_path) + split_manifest[name] = { + "uri": split_uri, + "rows": len(part), + "features": part.features.to_dict(), + "feature_order": list(part.features), + "split": ( + str(part.split) + if getattr(part, "split", None) is not None + else (name if name != "dataset" else None) + ), + "sha256": _sha256_path(split_uri, conn_id=storage_conn_id), + } + + manifest = { + "format": "parquet", + "version": 1, + "splits": split_manifest, + "metadata": metadata or {}, + } + manifest_bytes = json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode("utf-8") + manifest_path.parent.mkdir(parents=True, exist_ok=True) + with manifest_path.open("wb") as stream: + stream.write(manifest_bytes) + return { + "uri": uri, + "format": "parquet", + "splits": { + name: { + "rows": part_info["rows"], + "features": part_info["features"], + "split": part_info["split"], + } + for name, part_info in split_manifest.items() + }, + **(metadata or {}), + "fingerprint": hashlib.sha256(manifest_bytes).hexdigest(), + } + except Exception: + # Remove partial split outputs. A failed write must not leave an artifact + # reference that a downstream task could mistake for a complete dataset. + for path in written: + try: + path.unlink(missing_ok=True) + except (OSError, NotImplementedError): + pass + try: + manifest_path.unlink(missing_ok=True) + except (OSError, NotImplementedError): + pass + raise + + +def _read_artifact_manifest( + reference: str | dict[str, Any], *, storage_conn_id: str | None = None +) -> tuple[str, dict[str, Any]]: + uri = _reference(reference) + manifest_path = _path(uri, conn_id=storage_conn_id) + with manifest_path.open("rb") as stream: + manifest_bytes = stream.read() + if isinstance(reference, dict) and reference.get("fingerprint"): + actual = hashlib.sha256(manifest_bytes).hexdigest() + if actual != reference["fingerprint"]: + raise ValueError("Dataset artifact fingerprint does not match its reference") + + manifest = json.loads(manifest_bytes) + if manifest.get("format") != "parquet" or manifest.get("version") != 1: + raise ValueError(f"Unsupported dataset artifact manifest: {uri}") + return uri, manifest + + +def _load_artifact( + reference: str | dict[str, Any], + *, + cache_dir: str | None = None, + manifest: dict[str, Any] | None = None, + storage_conn_id: str | None = None, +) -> Dataset | DatasetDict: + from datasets import Dataset, DatasetDict, Features + + uri, manifest = ( + _read_artifact_manifest(reference, storage_conn_id=storage_conn_id) + if manifest is None + else (_reference(reference), manifest) + ) + resolved_cache_dir = cache_dir or str(Path(tempfile.gettempdir()) / "airflow-hf-datasets-cache") + Path(resolved_cache_dir).mkdir(parents=True, exist_ok=True) + parts: dict[str, Dataset] = {} + with tempfile.TemporaryDirectory(prefix="airflow-hf-parquet-") as local_dir: + for index, (name, item) in enumerate(manifest["splits"].items()): + split_uri = item["uri"] + local_path = Path(local_dir) / f"split-{index:05d}.parquet" + digest = hashlib.sha256() + with _path(split_uri, conn_id=storage_conn_id).open("rb") as source, local_path.open("wb") as target: + for chunk in iter(lambda: source.read(1024 * 1024), b""): + digest.update(chunk) + target.write(chunk) + if digest.hexdigest() != item.get("sha256"): + raise ValueError(f"Dataset artifact split fingerprint mismatch: {split_uri}") + + feature_data = item["features"] + feature_order = item.get("feature_order", list(feature_data)) + features = Features.from_dict({key: feature_data[key] for key in feature_order}) + raw_split = item.get("split") or (name if name != "dataset" else None) + split_name = raw_split.split("[")[0] if raw_split and "[" in raw_split else raw_split + restored = Dataset.from_parquet( + str(local_path), + features=features, + split=split_name, + cache_dir=resolved_cache_dir, + ) + # Parquet readers may return fields in physical/file order; restore the + # manifest's declared feature order for stable downstream transforms. + parts[name] = restored.select_columns(list(features)) + return parts["dataset"] if list(parts) == ["dataset"] else DatasetDict(parts) + + +def _state_key(prefix: str, identity: str) -> str: + """Task state stores already scope keys to a DAG run and task instance.""" + return f"{prefix}:{identity}" + + +def _publication_marker(prefix: str, identity: dict[str, Any]) -> str: + payload = json.dumps(identity, sort_keys=True, separators=(",", ":")).encode() + return f"{prefix}:{hashlib.sha256(payload).hexdigest()}" + + +def _commit_details(commit: Any, repo_id: str, hook: HuggingFaceHook) -> tuple[str, str]: + """Normalize CommitInfo and GitCommitInfo returned by Hub APIs.""" + commit_sha = getattr(commit, "oid", None) or getattr(commit, "commit_id", None) + if not commit_sha: + raise ValueError(f"Hub publication for {repo_id!r} returned no commit identifier") + commit_url = getattr(commit, "commit_url", None) + if not commit_url: + endpoint = getattr(hook.get_conn(), "endpoint", "https://huggingface.co";).rstrip("/") + commit_url = f"{endpoint}/datasets/{repo_id}/commit/{commit_sha}" + return commit_sha, commit_url + + +def _is_hub_repo(path: str) -> bool: + """Match a Hub dataset repo ID without treating an existing local path as remote.""" + return path.count("/") == 1 and not path.startswith((".", "/")) and not Path(path).exists() + + +class LoadDatasetOperator(BaseOperator): + """Load a dataset revision and persist it as a Parquet artifact.""" + + template_fields = ("path", "name", "split", "revision", "cache_dir", "output_uri") + + def __init__( + self, + *, + path: str, + output_uri: str, + name: str | None = None, + split: str | None = None, + revision: str | None = None, + cache_dir: str | None = None, + storage_conn_id: str | None = None, + download_mode: str = "reuse_dataset_if_exists", + huggingface_conn_id: str | None = "huggingface_default", + load_kwargs: dict[str, Any] | None = None, + durable: bool = True, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.path, self.output_uri, self.name, self.split = path, output_uri, name, split + self.revision, self.cache_dir, self.download_mode = revision, cache_dir, download_mode + self.storage_conn_id = storage_conn_id + self.huggingface_conn_id, self.load_kwargs = huggingface_conn_id, load_kwargs or {} + self.durable = durable + + def execute(self, context: Context) -> dict[str, Any]: + hook = HuggingFaceHook(self.huggingface_conn_id) + requested_revision = self.revision + is_hub_repo = _is_hub_repo(self.path) + state = context.get("task_state_store") if self.durable else None + state_identity = json.dumps([self.path, requested_revision], separators=(",", ":")) + state_key = _state_key("load-revision", state_identity) + revision = state.get(state_key) if state is not None else None + if revision is None: + revision = requested_revision + if is_hub_repo: + revision = hook.get_dataset_info(self.path, requested_revision).sha + if state is not None and revision is not None: + state.set(state_key, revision) + elif state is not None: + self.log.info("Reusing pinned Hugging Face dataset revision %s for %s", revision, self.path) + + kwargs = dict(self.load_kwargs) + kwargs["download_mode"] = self.download_mode + if self.name is not None: + kwargs["name"] = self.name + if self.split is not None: + kwargs["split"] = self.split + if revision is not None: + kwargs["revision"] = revision + if self.cache_dir is not None: + kwargs["cache_dir"] = self.cache_dir + dataset = hook.load_dataset(self.path, **kwargs) + metadata = { + "repository_id": self.path if is_hub_repo else None, + "revision": revision, + "configuration": self.name, + "split": self.split, + } + return _save(dataset, self.output_uri, metadata, storage_conn_id=self.storage_conn_id) + + +class StreamDatasetOperator(BaseOperator): + """Transform a streamed dataset lazily and publish it directly to the Hub.""" + + template_fields = ("path", "name", "split", "revision", "config", "commit_message") + + def __init__( + self, + *, + path: str, + repo_id: str, + transform: Callable[..., dict[str, Any]] | None = None, + predicate: Callable[[dict[str, Any]], bool] | None = None, + name: str | None = None, + split: str | None = None, + revision: str | None = None, + config: str | None = None, + private: bool | None = None, + commit_message: str | None = None, + transform_kwargs: dict[str, Any] | None = None, + push_kwargs: dict[str, Any] | None = None, + huggingface_conn_id: str | None = "huggingface_default", + durable: bool = True, + **kwargs: Any, + ) -> None: + if "{{" in repo_id: + raise ValueError("repo_id cannot be templated: it defines the Asset that downstream DAGs schedule on.") + if transform is not None and not callable(transform): + raise TypeError("transform must be a callable or None") + if predicate is not None and not callable(predicate): + raise TypeError("predicate must be a callable or None") + self.path, self.repo_id = path, repo_id + self.transform, self.predicate = transform, predicate + self.name, self.split, self.revision = name, split, revision + self.config, self.private, self.commit_message = config, private, commit_message + self.transform_kwargs = transform_kwargs or {} + self.push_kwargs = push_kwargs or {} + self.huggingface_conn_id = huggingface_conn_id + self.durable = durable + self.published_asset = Asset(uri=f"hf://datasets/{repo_id}") + user_outlets = kwargs.pop("outlets", None) + if user_outlets is None: + user_outlets = [] + elif not isinstance(user_outlets, (list, tuple)): + user_outlets = [user_outlets] + super().__init__(outlets=[*user_outlets, self.published_asset], **kwargs) + + def execute(self, context: Context) -> dict[str, Any]: + hook = HuggingFaceHook(self.huggingface_conn_id) + is_hub_repo = _is_hub_repo(self.path) + requested_revision = self.revision + state = context.get("task_state_store") if self.durable else None + revision_identity = json.dumps( + [self.path, requested_revision, self.name, self.split], separators=(",", ":") + ) + revision_key = _state_key("stream-revision", revision_identity) + revision = state.get(revision_key) if state is not None else None + if revision is None: + revision = requested_revision + if is_hub_repo: + revision = hook.get_dataset_info(self.path, requested_revision).sha + if state is not None and revision is not None: + state.set(revision_key, revision) + elif state is not None: + self.log.info("Reusing pinned Hugging Face streaming revision %s for %s", revision, self.path) + + run_id = context.get("run_id", "manual") + push_kwargs = dict(self.push_kwargs) + if self.config is not None: + push_kwargs["config_name"] = self.config + if self.private is not None: + push_kwargs["private"] = self.private + marker = _publication_marker( + "huggingface-stream-publish", + { + "task_id": self.task_id, + "run_id": run_id, + "source": self.path, + "source_revision": revision, + "name": self.name, + "split": self.split, + "repo_id": self.repo_id, + "config_name": push_kwargs.get("config_name", "default"), + "revision": push_kwargs.get("revision", "main"), + "create_pr": bool(push_kwargs.get("create_pr", False)), + }, + ) + state_key = _state_key("stream-publication", marker) + if state is not None: + prior_result = state.get(state_key) + if prior_result: + self.log.info("Reusing stored streaming publication result for %s", self.repo_id) + self._emit_asset_event(context, prior_result) + return prior_result + + target_revision = push_kwargs.get("revision") + include_pull_requests = bool(push_kwargs.get("create_pr", False)) + commit = None + if self.durable: + commit = hook.find_commit_by_marker( + self.repo_id, + marker, + revision=target_revision, + include_pull_requests=include_pull_requests, + ) + if commit is not None: + self.log.info("Reusing existing completed streaming publication to %s", self.repo_id) + + if commit is None: + load_kwargs: dict[str, Any] = {"streaming": True} + if self.name is not None: + load_kwargs["name"] = self.name + if self.split is not None: + load_kwargs["split"] = self.split + if revision is not None: + load_kwargs["revision"] = revision + stream = hook.load_dataset(self.path, **load_kwargs) + + if self.transform is not None: + stream = stream.map(self.transform, fn_kwargs=self.transform_kwargs) + if self.predicate is not None: + stream = stream.filter(self.predicate) + + message = self.commit_message or "Stream dataset from Airflow" + push_kwargs["commit_message"] = f"{message} [{marker}]" + commit = hook.push_to_hub(stream, self.repo_id, **push_kwargs) + # A retry can recover the completed Hub commit if the worker exited + # after upload but before Airflow recorded the task result. + if commit is None: + if self.durable: + commit = hook.find_commit_by_marker( + self.repo_id, + marker, + revision=target_revision, + include_pull_requests=include_pull_requests, + ) + + if commit is None: + raise RuntimeError( + f"Hub push to {self.repo_id!r} completed without returning or exposing a commit for marker {marker}" + ) + commit_sha, commit_url = _commit_details(commit, self.repo_id, hook) + result = { + "repository_id": self.repo_id, + "commit_sha": commit_sha, + "commit_url": commit_url, + "url": f"https://huggingface.co/datasets/{self.repo_id}";, + "source_repository_id": self.path if is_hub_repo else None, + "source_revision": revision, + "configuration": self.name, + "split": self.split, + "streaming": True, + } + if state is not None: + state.set(state_key, result) + self._emit_asset_event(context, result) + return result + + def _emit_asset_event(self, context: Context, publication: dict[str, Any]) -> None: + events = context.get("outlet_events") + if events is not None: + events[self.published_asset].extra = { + "repository_id": self.repo_id, + "commit_sha": publication["commit_sha"], + "commit_url": publication["commit_url"], + "source_revision": publication["source_revision"], + } + + +class MapDatasetOperator(BaseOperator): + """Apply a Python callable and persist the transformed artifact.""" + + template_fields = ("input_uri", "output_uri", "transform_kwargs", "cache_dir") + + def __init__( + self, + *, + input_uri: str | dict[str, Any], + output_uri: str, + transform: Callable[..., Any], + transform_kwargs: dict[str, Any] | None = None, + map_kwargs: dict[str, Any] | None = None, + cache_dir: str | None = None, + storage_conn_id: str | None = None, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + if not callable(transform): + raise TypeError("transform must be a callable") + self.input_uri, self.output_uri, self.transform = input_uri, output_uri, transform + self.transform_kwargs, self.map_kwargs = transform_kwargs or {}, map_kwargs or {} + self.cache_dir = cache_dir + self.storage_conn_id = storage_conn_id + + def execute(self, context: Context) -> dict[str, Any]: + dataset = _load_artifact( + self.input_uri, cache_dir=self.cache_dir, storage_conn_id=self.storage_conn_id + ) + mapped = dataset.map(self.transform, fn_kwargs=self.transform_kwargs, **self.map_kwargs) + source = self.input_uri if isinstance(self.input_uri, dict) else {} + provenance = { + key: source[key] + for key in ("repository_id", "revision", "configuration", "split") + if key in source + } + return _save(mapped, self.output_uri, provenance, storage_conn_id=self.storage_conn_id) + + +class PushToHubOperator(BaseOperator): + """Synchronously publish an artifact and emit its successful asset event.""" + + template_fields = ("input_uri", "config", "commit_message") + + def __init__( + self, + *, + input_uri: str | dict[str, Any], + repo_id: str, + config: str | None = None, + private: bool | None = None, + commit_message: str | None = None, + huggingface_conn_id: str | None = "huggingface_default", + storage_conn_id: str | None = None, + push_kwargs: dict[str, Any] | None = None, + durable: bool = True, + **kwargs: Any, + ) -> None: + if "{{" in repo_id: + raise ValueError("repo_id cannot be templated: it defines the Asset that downstream DAGs schedule on.") + self.input_uri, self.repo_id, self.config = input_uri, repo_id, config + self.private, self.commit_message = private, commit_message + self.huggingface_conn_id, self.push_kwargs = huggingface_conn_id, push_kwargs or {} + self.storage_conn_id = storage_conn_id + self.durable = durable + self.published_asset = Asset(uri=f"hf://datasets/{repo_id}") + user_outlets = kwargs.pop("outlets", None) + if user_outlets is None: + user_outlets = [] + elif not isinstance(user_outlets, (list, tuple)): + user_outlets = [user_outlets] + super().__init__(outlets=[*user_outlets, self.published_asset], **kwargs) + + def execute(self, context: Context) -> dict[str, Any]: + source = self.input_uri if isinstance(self.input_uri, dict) else {} + fingerprint = source.get("fingerprint") + if not fingerprint: + raise ValueError("Publication requires an artifact reference with a fingerprint") + kwargs = dict(self.push_kwargs) + if self.config is not None: + kwargs["config_name"] = self.config + if self.private is not None: + kwargs["private"] = self.private + target_revision = kwargs.get("revision") + create_pr = bool(kwargs.get("create_pr", False)) + target_marker = _publication_marker( + "target", + { + "repo_id": self.repo_id, + "config_name": kwargs.get("config_name", "default"), + "revision": target_revision or "main", + "create_pr": create_pr, + }, + ) + marker = f"airflow-artifact-sha256:{fingerprint}:{target_marker}" Review Comment: This marker still matches a commit of the same artifact anywhere in the branch history. If one run publishes A, the next publishes B, and a later run publishes A again (a rollback or a backfill), the lookup returns the first run's commit, skips the push, and emits an asset event for A while the Hub serves B. Could a match count only when it's the head of the target ref, or carry the task instance identity so only retries of this task match? -- 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]
