add HF dataset plugin

[samhita-alla]

This PR adds a Hugging Face datasets plugin for Flyte. It lets tasks consume Hub datasets directly via from_hf(...) and transparently materializes them as datasets.Dataset or datasets.IterableDataset with optional shared caching through cache_root.

from flyteplugins.huggingface.datasets import from_hf

@env.task
async def train(
    ds: datasets.Dataset = from_hf(
        "stanfordnlp/imdb",
        name="plain_text",
        split="train",
        cache_root="s3://my-bucket/flyte-hf-cache",
    ),
) -> int:
    return len(ds)

The plugin streams converted parquet shards from the Hub, stores them as sharded parquet artifacts, and reuses cached materializations across runs when available.

You can also return and pass real datasets.Dataset objects between tasks:

@env.task
async def create_dataset() -> datasets.Dataset:
    return datasets.Dataset.from_dict(
        {
            "text": ["hello", "world", "flyte"],
            "label": [0, 1, 0],
        }
    )

@env.task
async def filter_positive(ds: datasets.Dataset) -> datasets.Dataset:
    return ds.filter(lambda row: row["label"] == 1)

[andreahlert]

a few questions. But after all, it has an overall better DX than my version. tks!

[andreahlert]

dropping hf-transfer is fine (HfFileSystem doesn't use it anyway), but the hf_hub_download fallback went with it. was killing the snapshot-download escapehatch intentional?

[samhita-alla]

since we're using HFS, we don't need it anymore.

[andreahlert]

triple registration, entry point + both init.pys. lru_cache saves it but any of the three can go, right?

Suggested change

ParquetToHuggingFaceIterableDatasetDecodingHandler,

[samhita-alla]

thanks for the catch! we have init registration as a fallback to support backward compatibility, but i think we can remove it in this case.

[andreahlert]

with cache_root=None the shards/sha256/manifest work runs and gets thrown away, every decode re-downloads to a new local dir. was cache-root meant to be strictly opt-in, or would defaulting to a raw_data remote path make sense?

[samhita-alla]

opt-in for now as we don't support artifacts in v2 yet. even if we default to a raw_data remote path, it won't persist across runs. so it isn't really helpful. hope that makes sense.

[andreahlert]

artifact_uri only matters on hit, the miss branch resets to default_remote_path at 504. migration case in mind, but i really dont know.. maybe can the indirection come out

[samhita-alla]

you're right. we don't really need artifact_uri. removed it.

[andreahlert]

imho back to lazy_module("datasets")? entry-point discovery triggers this and drags pyarrow/numpy/fsspec in for users who never touch the plugin. Is there any other reason i'm not seeing here?

[samhita-alla]

moved imports inside register_huggingface_dataset_transformers() to avoid eager imports. think that should do.

[andreahlert]

localizing full shards before pyarrow kills projection/row-group pushdown (pq.read_table(..., filesystem=fs, columns=...) range-requests only what it needs). was there a case where fsspec+pyarrow broke down for you?

this one is not a problem. Just questioning to for learning purposes :)

[samhita-alla]

actually, i could use fsspec with pyarrow. i initially ran into remote storage auth-related errors, so i thought this would make things easier. but you’re right, i could use flyte’s storage abstraction to get the filesystem and read remote parquet files directly.

[andreahlert]

understood! thank you

[cosmicBboy]

is there ever a world where we also want to support download models? In that case would be future-proof to rename this dataset_from_hf or expose a module flyteplugins.huggingface.datasets.from_hf so that we can also do flyteplugins.huggingface.models.from_hf

[samhita-alla]

makes sense! it’s now exposed in the module so we avoid the risk of collisions, and each namespace owns its own registration.