feat: Add Hugging Face Datasets plugin

[andreahlert]

Summary

Port of the flytekit-huggingface plugin to flyte-sdk v2, enabling native support for datasets.Dataset as a Flyte DataFrame type.

• DataFrameEncoder/Decoder for datasets.Dataset with Parquet serialization
• Cloud storage support (S3, GCS, Azure) via fsspec-compatible storage options
• Anonymous S3 fallback for public datasets (mirrors Polars plugin pattern)
• Column filtering on both encode and decode via type annotations
• Auto-registration via flyte.plugins.types entry point

Follows the same patterns as the existing Polars plugin.

Usage Example

import flyte
import datasets

env = flyte.TaskEnvironment("hf-example")

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

@env.task
async def process(ds: datasets.Dataset) -> int:
    return len(ds)

Test plan

• Type recognition tests (Dataset, with columns, with format annotations)
• Handler registration and property tests
• Encode/decode roundtrip tests
• DataFrame wrapper integration test
• Raw task I/O test via flyte.TaskEnvironment
• Column subsetting on decode
• Various data types (int, float, str, bool)
• Empty dataset roundtrip
• Storage options tests (S3, GCS, Azure, anonymous, unknown protocol)

[andreahlert]

@pingsutw @cosmicBboy @kumare3 could you take a look? This ports the huggingface plugin from flytekit to v2, same approach as the polars plugin.

[andreahlert]

Hi Ketan! With 2.0 out I’ve rebased and addressed several of your comments: using the public storage API instead of get_storage, and public flyte.io imports where possible.

[andreahlert]

Addressed all comments. Switched both encode and decode to use pq.write_table/pq.read_table with the Flyte filesystem directly. This avoids HuggingFace's DownloadManager cache entirely and keeps credentials handled by Flyte's storage layer.

As a side effect, _get_storage_options is gone and the empty dataset test is no longer skipped since that was a from_parquet limitation. Also dropped column filtering from encode to match the Polars plugin pattern.

[kumare3]

Can you add an example?

I am thinking either hf:// file system do we need it
Cc @samhita-alla

[andreahlert]

Can you add an example?

I am thinking either hf:// file system do we need it Cc @samhita-alla

I think we do, at least for the read side. HF Hub auto-converts all datasets to Parquet now, so hf:// would let Flyte track what was actually consumed instead of users doing opaque load_dataset calls. The fsspec backend already exists.

Immutability is a concern since Hub data can change, but for pinned dataset revisions (hf://datasets/squad@v1.0) it works.

Write side is trickier and probably not needed for a first pass.

Open to what is best in your opinion.

[samhita-alla]

hey @andreahlert, thank you for creating this PR! sharing my thoughts on how we can improve the plugin from a UX perspective:

---

the two primary use cases i see are:

• loading a dataset from HF Hub: to use in training or processing tasks
• creating a dataset in a task: to pass between tasks for training and optionally publish to HF Hub

for loading, i think this should follow the same pattern as our HF model prefetch utility i.e. stream from the hub rather than fully materializing the dataset into memory. the output should be a reference (e.g. Dir) stored in cloud storage, which downstream tasks can then stream from.

for creating, when a task outputs a datasets.Dataset, we should:

• stream-write it to remote storage
• optionally support pushing to HF Hub

i don’t think a type transformer alone will be useful here.

---

i know the PR mirrors what’s in the v1 plugin, but i think this approach would make it genuinely useful for users. not sure if you have the bandwidth to work on this. if not, totally understandable. we can leave the PR open and i can take a pass at it later!

[andreahlert]

hey @andreahlert, thank you for creating this PR! sharing my thoughts on how we can improve the plugin from a UX perspective:

---

the two primary use cases i see are:

• loading a dataset from HF Hub: to use in training or processing tasks

• creating a dataset in a task: to pass between tasks for training and optionally publish to HF Hub

for loading, i think this should follow the same pattern as our HF model prefetch utility i.e. stream from the hub rather than fully materializing the dataset into memory. the output should be a reference (e.g. Dir) stored in cloud storage, which downstream tasks can then stream from.

for creating, when a task outputs a datasets.Dataset, we should:

• stream-write it to remote storage

• optionally support pushing to HF Hub

i don’t think a type transformer alone will be useful here.

---

i know the PR mirrors what’s in the v1 plugin, but i think this approach would make it genuinely useful for users. not sure if you have the bandwidth to work on this. if not, totally understandable. we can leave the PR open and i can take a pass at it later!

Thanks for reviewing this. I still have sometime to work on this. Will address later on this week.

[andreahlert]

Hey everybody. Updated based on @samhita-alla feedback.

Took a step back and rethought the approach. The plugin now has a hf_dataset() prefetch that streams parquet from HF Hub to Flyte storage (same pattern as flyte.prefetch.hf_model), batched writes in the encoder, and an example showing the full flow.

Not sure if I got the prefetch design exactly right so would appreciate your take on it.

Hub push is left out for now, figured that's better as a follow-up. WDYT?

[samhita-alla]

closing this in favor of #992