incorporate suggestions @andreahlert and @cosmicBboy

Signed-off-by: Samhita Alla <aallasamhita@gmail.com>

Author: samhita-alla

plugins/huggingface/README.md

@@ -1,9 +1,9 @@
-# Hugging Face Datasets Plugin
+# Hugging Face Plugin
 
-Native Flyte support for Hugging Face `datasets.Dataset` and
-`datasets.IterableDataset`.
+Native Flyte support for Hugging Face integrations in Flyte.
 
-This plugin gives you two related capabilities:
+This plugin provides dataset support for Hugging Face `datasets.Dataset`
+and `datasets.IterableDataset` objects. It gives you two related capabilities:
 
 1. Use `from_hf(...)` to reference a dataset on the Hugging Face Hub as a task
    input default.
@@ -26,7 +26,7 @@ pip install flyteplugins-huggingface
 ```python
 import datasets
 import flyte
-from flyteplugins.huggingface import from_hf
+from flyteplugins.huggingface.datasets import from_hf
 
 env = flyte.TaskEnvironment(name="hf-example")
 
@@ -56,7 +56,7 @@ reference used between Flyte and the plugin.
 `from_hf(...)` is the entry point for Hub-backed task defaults:
 
 ```python
-from flyteplugins.huggingface import from_hf
+from flyteplugins.huggingface.datasets import from_hf
 
 from_hf(
     repo: str,
@@ -142,7 +142,7 @@ dataset.
 Without `cache_root`, a Hub source is materialized into a generated path for the
 current execution only.
 
-With `cache_root`, the plugin uses a shared artifact registry so later runs can
+With `cache_root`, the plugin uses a shared cache registry so later runs can
 skip the Hub download entirely:
 
 ```python
@@ -158,7 +158,7 @@ async def train_cached(
     return len(ds)
 ```
 
-The shared registry layout is:
+The shared cache layout is:
 
 ```text
 {cache_root}/huggingface/datasets/
@@ -177,6 +177,10 @@ The cache key is derived from:
 This means the cache is stable across runs as long as the underlying converted
 Parquet source does not change.
 
+The canonical artifact location is always
+`{cache_root}/huggingface/datasets/blobs/{source-cache-key}/...`. The registry
+record under `by-key/` is metadata for that cache key.
+
 ## What the plugin logs
 
 When `LOG_LEVEL` is `INFO` or lower, the plugin logs whether it is:

plugins/huggingface/examples/hf_dataset_workflow.py

@@ -22,7 +22,7 @@
 from flyte._image import PythonWheels
 from flyte.io import DataFrame
 
-from flyteplugins.huggingface import from_hf
+from flyteplugins.huggingface.datasets import from_hf
 
 env = flyte.TaskEnvironment(
     name="hf-dataset-example",

plugins/huggingface/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "flyteplugins-huggingface"
 dynamic = ["version"]
-description = "Hugging Face Datasets plugin for flyte"
+description = "Hugging Face Plugin for Flyte"
 readme = "README.md"
 authors = [{ name = "André Ahlert", email = "andre@aex.partners" }, { name = "Samhita Alla", email = "samhita@union.ai" }]
 requires-python = ">=3.10"
@@ -13,7 +13,7 @@ dependencies = [
 ]
 
 [project.entry-points."flyte.plugins.types"]
-huggingface = "flyteplugins.huggingface:register_huggingface_dataset_transformers"
+huggingface = "flyteplugins.huggingface.datasets:register_huggingface_dataset_transformers"
 
 [build-system]
 requires = ["setuptools", "setuptools_scm"]

plugins/huggingface/src/flyteplugins/huggingface/__init__.py

@@ -1,10 +0,0 @@
-from .dataset import (
-    HFSource,
-    from_hf,
-    register_huggingface_dataset_transformers,
-)
-
-__all__ = ["HFSource", "from_hf"]
-
-
-register_huggingface_dataset_transformers()

…eplugins/huggingface/dataset/__init__.py …plugins/huggingface/datasets/__init__.pyplugins/huggingface/src/flyteplugins/huggingface/dataset/__init__.py renamed to plugins/huggingface/src/flyteplugins/huggingface/datasets/__init__.py plugins/huggingface/src/flyteplugins/huggingface/dataset/__init__.py renamed to plugins/huggingface/src/flyteplugins/huggingface/datasets/__init__.py

@@ -1,14 +1,6 @@
 import functools
 
 from ._source import HFSource, from_hf
-from ._transformers import (
-    HFToHuggingFaceDatasetDecodingHandler,
-    HFToHuggingFaceIterableDatasetDecodingHandler,
-    HuggingFaceDatasetToParquetEncodingHandler,
-    HuggingFaceIterableDatasetToParquetEncodingHandler,
-    ParquetToHuggingFaceDatasetDecodingHandler,
-    ParquetToHuggingFaceIterableDatasetDecodingHandler,
-)
 
 __all__ = ["HFSource", "from_hf"]
 
@@ -18,6 +10,15 @@ def register_huggingface_dataset_transformers():
     """Register Hugging Face Dataset encoders and decoders."""
     from flyte.io.extend import DataFrameTransformerEngine
 
+    from ._transformers import (
+        HFToHuggingFaceDatasetDecodingHandler,
+        HFToHuggingFaceIterableDatasetDecodingHandler,
+        HuggingFaceDatasetToParquetEncodingHandler,
+        HuggingFaceIterableDatasetToParquetEncodingHandler,
+        ParquetToHuggingFaceDatasetDecodingHandler,
+        ParquetToHuggingFaceIterableDatasetDecodingHandler,
+    )
+
     DataFrameTransformerEngine.register(HuggingFaceDatasetToParquetEncodingHandler(), default_format_for_type=True)
     DataFrameTransformerEngine.register(ParquetToHuggingFaceDatasetDecodingHandler(), default_format_for_type=True)
     DataFrameTransformerEngine.register(HFToHuggingFaceDatasetDecodingHandler())
@@ -30,6 +31,3 @@ def register_huggingface_dataset_transformers():
         default_format_for_type=True,
     )
     DataFrameTransformerEngine.register(HFToHuggingFaceIterableDatasetDecodingHandler())
-
-
-register_huggingface_dataset_transformers()

…/flyteplugins/huggingface/dataset/_io.py …flyteplugins/huggingface/datasets/_io.pyplugins/huggingface/src/flyteplugins/huggingface/dataset/_io.py renamed to plugins/huggingface/src/flyteplugins/huggingface/datasets/_io.py plugins/huggingface/src/flyteplugins/huggingface/dataset/_io.py renamed to plugins/huggingface/src/flyteplugins/huggingface/datasets/_io.py

@@ -146,7 +146,7 @@ async def list_parquet_files(uri: str, filesystem) -> list[str]:
             raw = [f"{proto}{f}" for f in raw]
         return raw
     except Exception as e:
-        logger.debug(f"Unable to list parquet files under {uri}: {e}")
+        logger.warning(f"Unable to list parquet files under {uri}: {e}")
         return [join_uri_path(uri, f"{0:05}.parquet")]
 
 
@@ -307,9 +307,7 @@ def hf_cache_manifest(
     }
 
 
-async def read_cache_manifest(
-    remote_path: str,
-) -> dict[str, typing.Any] | None:
+async def read_cache_manifest(remote_path: str) -> dict[str, typing.Any] | None:
     path = manifest_path(remote_path)
     try:
         if not await storage_path_exists(path):
@@ -348,13 +346,9 @@ async def read_registry_record(
 async def write_registry_record(
     source: HFSource,
     cache_key: str,
-    artifact_uri: str,
     manifest: dict[str, typing.Any],
 ) -> None:
-    record = {
-        **manifest,
-        "artifact_uri": artifact_uri,
-    }
+    record = dict(manifest)
     data = json.dumps(record, sort_keys=True, indent=2).encode("utf-8")
     await storage_write_bytes(get_hf_registry_record_path(source, cache_key), data)
 
@@ -481,32 +475,25 @@ async def ensure_hf_cached(source: HFSource) -> str:
         f"under {source.cache_root}"
     )
     registry_record = await read_registry_record(source, cache_key)
-    remote_path = (
-        registry_record.get("artifact_uri", default_remote_path) if registry_record is not None else default_remote_path
-    )
-
-    if await read_cache_manifest(remote_path) == expected_manifest:
+    if await read_cache_manifest(default_remote_path) == expected_manifest:
         if registry_record is None:
             await write_registry_record(
                 source,
                 cache_key,
-                remote_path,
                 expected_manifest,
             )
-        logger.info(f"Using cached Hugging Face dataset at {remote_path}")
-        return remote_path
+        logger.info(f"Using cached Hugging Face dataset at {default_remote_path}")
+        return default_remote_path
 
     logger.info(
         f"Materializing Hugging Face dataset {source.repo} "
         f"({_source_log_description(source)}) "
-        f"to remote cache artifact {remote_path}"
+        f"to remote cache artifact {default_remote_path}"
     )
-    remote_path = default_remote_path
-    await stream_hf_to_remote(source, remote_path, shards, expected_manifest)
+    await stream_hf_to_remote(source, default_remote_path, shards, expected_manifest)
     await write_registry_record(
         source,
         cache_key,
-        remote_path,
         expected_manifest,
     )
-    return remote_path
+    return default_remote_path

…teplugins/huggingface/dataset/_source.py …eplugins/huggingface/datasets/_source.pyplugins/huggingface/src/flyteplugins/huggingface/dataset/_source.py renamed to plugins/huggingface/src/flyteplugins/huggingface/datasets/_source.py plugins/huggingface/src/flyteplugins/huggingface/dataset/_source.py renamed to plugins/huggingface/src/flyteplugins/huggingface/datasets/_source.py

@@ -21,6 +21,29 @@ class HFSource:
     revision: str | None = None
     cache_root: str | None = None
 
+    def __post_init__(self) -> None:
+        self.repo = self._normalize_required_field("repo", self.repo)
+        self.name = self._normalize_optional_field("name", self.name)
+        self.split = self._normalize_optional_field("split", self.split)
+        self.revision = self._normalize_optional_field("revision", self.revision)
+        self.cache_root = self._normalize_optional_field("cache_root", self.cache_root)
+
+    @staticmethod
+    def _normalize_required_field(field_name: str, value: str) -> str:
+        normalized = value.strip()
+        if not normalized:
+            raise ValueError(f"HFSource {field_name} must not be empty")
+        return normalized
+
+    @staticmethod
+    def _normalize_optional_field(field_name: str, value: str | None) -> str | None:
+        if value is None:
+            return None
+        normalized = value.strip()
+        if not normalized:
+            raise ValueError(f"HFSource {field_name} must not be blank")
+        return normalized
+
     def to_hf_uri(self) -> str:
         uri = f"hf://{self.repo}"
         params = {}