flyteplugins-huggingface 2.2.1__py3-none-any.whl

flyteplugins/huggingface/datasets/_transformers.py

@@ -0,0 +1,352 @@
+from __future__ import annotations
+
+import pathlib
+import typing
+from pathlib import Path
+
+import datasets
+import flyte.storage as storage
+import pyarrow as pa
+import pyarrow.parquet as pq
+from flyte._logging import logger
+from flyte.io import PARQUET, DataFrame
+from flyte.io.extend import DataFrameDecoder, DataFrameEncoder
+from flyteidl2.core import literals_pb2, types_pb2
+from fsspec.core import strip_protocol
+
+from ._io import ensure_hf_cached, join_uri_path, list_parquet_files, run_sync_io
+from ._source import HFSource
+
+_ROWS_PER_SHARD = 100_000
+
+
+def _write_dataset(df: datasets.Dataset, path: str, filesystem) -> None:
+    table = df.data.table
+    writer = pq.ParquetWriter(strip_protocol(path), table.schema, filesystem=filesystem)
+
+    try:
+        for batch in table.to_batches(max_chunksize=10_000):
+            writer.write_batch(batch)
+    finally:
+        writer.close()
+
+
+def _read_parquet_files(
+    parquet_files: list[str],
+    columns: list[str] | None,
+    filesystem=None,
+) -> pa.Table:
+    tables = [
+        pq.read_table(
+            strip_protocol(file_path),
+            columns=columns,
+            filesystem=filesystem if storage.is_remote(file_path) else None,
+        )
+        for file_path in parquet_files
+    ]
+    return pa.concat_tables(tables) if len(tables) > 1 else tables[0]
+
+
+def _batch_to_rows(batch: pa.RecordBatch) -> list[dict]:
+    col_lists = {name: col.to_pylist() for name, col in zip(batch.schema.names, batch.columns)}
+    return [{name: col_lists[name][i] for name in batch.schema.names} for i in range(batch.num_rows)]
+
+
+def _iter_parquet_rows(parquet_files: list[str], columns: list[str] | None) -> typing.Iterator[dict]:
+    for file_path in parquet_files:
+        filesystem = storage.get_underlying_filesystem(path=file_path) if storage.is_remote(file_path) else None
+        pf = pq.ParquetFile(
+            strip_protocol(file_path),
+            filesystem=filesystem,
+        )
+        for batch in pf.iter_batches(batch_size=10_000, columns=columns):
+            yield from _batch_to_rows(batch)
+
+
+def _write_iterable_dataset(ds: datasets.IterableDataset, uri: str, filesystem) -> None:
+    file_idx = 0
+    rows_in_shard = 0
+    writer: pq.ParquetWriter | None = None
+
+    try:
+        for batch in ds.iter(batch_size=10_000):
+            table = pa.table(batch)
+
+            if writer is None:
+                shard_path = join_uri_path(uri, f"{file_idx:05}.parquet")
+                writer = pq.ParquetWriter(
+                    strip_protocol(shard_path),
+                    table.schema,
+                    filesystem=filesystem,
+                )
+
+            for arrow_batch in table.to_batches():
+                writer.write_batch(arrow_batch)
+
+            rows_in_shard += len(table)
+
+            if rows_in_shard >= _ROWS_PER_SHARD:
+                writer.close()
+                writer = None
+
+                file_idx += 1
+                rows_in_shard = 0
+    finally:
+        if writer is not None:
+            writer.close()
+
+
+def _requested_columns(
+    current_task_metadata: literals_pb2.StructuredDatasetMetadata,
+) -> list[str] | None:
+    if current_task_metadata.structured_dataset_type and current_task_metadata.structured_dataset_type.columns:
+        return [c.name for c in current_task_metadata.structured_dataset_type.columns]
+    return None
+
+
+async def _localize_parquet_files(parquet_files: list[str]) -> list[str]:
+    """Download remote parquet files before handing them to PyArrow."""
+    local_files: list[str] = []
+
+    for file_path in parquet_files:
+        if storage.is_remote(file_path):
+            local_dir = storage.get_random_local_directory()
+            local_name = pathlib.PurePosixPath(file_path.rstrip("/")).name
+            local_path = str(local_dir / local_name)
+            logger.info(f"Downloading remote parquet shard {file_path} to {local_path}")
+            local_files.append(await storage.get(file_path, local_path))
+        else:
+            local_files.append(strip_protocol(file_path))
+    return local_files
+
+
+def _probe_remote_parquet_files(parquet_files: list[str], filesystem) -> None:
+    for file_path in parquet_files:
+        if not storage.is_remote(file_path):
+            continue
+        pq.ParquetFile(strip_protocol(file_path), filesystem=filesystem).metadata
+
+
+class HuggingFaceDatasetToParquetEncodingHandler(DataFrameEncoder):
+    def __init__(self):
+        super().__init__(datasets.Dataset, None, PARQUET)
+
+    async def encode(
+        self,
+        dataframe: DataFrame,
+        structured_dataset_type: types_pb2.StructuredDatasetType,
+    ) -> literals_pb2.StructuredDataset:
+        val = dataframe.val
+
+        if val is None and dataframe.uri:
+            structured_dataset_type.format = PARQUET
+            return literals_pb2.StructuredDataset(
+                uri=typing.cast(str, dataframe.uri),
+                metadata=literals_pb2.StructuredDatasetMetadata(structured_dataset_type=structured_dataset_type),
+            )
+
+        if not dataframe.uri:
+            from flyte._context import internal_ctx
+
+            uri = str(internal_ctx().raw_data.get_random_remote_path())
+        else:
+            uri = typing.cast(str, dataframe.uri)
+
+        if not storage.is_remote(uri):
+            Path(uri).mkdir(parents=True, exist_ok=True)
+
+        path = join_uri_path(uri, f"{0:05}.parquet")
+        df = typing.cast(datasets.Dataset, val)
+
+        filesystem = storage.get_underlying_filesystem(path=path)
+        logger.info(
+            f"Writing Hugging Face Dataset output to "
+            f"{'remote' if storage.is_remote(uri) else 'local'} parquet directory {uri}"
+        )
+        await run_sync_io("write HuggingFace dataset", _write_dataset, df, path, filesystem)
+
+        structured_dataset_type.format = PARQUET
+        return literals_pb2.StructuredDataset(
+            uri=uri,
+            metadata=literals_pb2.StructuredDatasetMetadata(structured_dataset_type=structured_dataset_type),
+        )
+
+
+class ParquetToHuggingFaceDatasetDecodingHandler(DataFrameDecoder):
+    def __init__(self, protocol: str | None = None):
+        super().__init__(datasets.Dataset, protocol, PARQUET)
+
+    async def decode(
+        self,
+        flyte_value: literals_pb2.StructuredDataset,
+        current_task_metadata: literals_pb2.StructuredDatasetMetadata,
+    ) -> datasets.Dataset:
+        uri = flyte_value.uri
+
+        if uri.startswith("hf://"):
+            source_uri = uri
+            try:
+                uri = await ensure_hf_cached(HFSource.from_hf_uri(uri))
+            except Exception as e:
+                raise RuntimeError(
+                    f"Failed to materialize Hugging Face dataset from {uri}: {type(e).__name__}: {e!r}"
+                ) from e
+            logger.info(
+                f"Resolved Hugging Face source {source_uri} to "
+                f"{'remote' if storage.is_remote(uri) else 'local'} parquet directory {uri}"
+            )
+
+        filesystem = storage.get_underlying_filesystem(path=uri)
+        logger.info(
+            f"Reading Hugging Face Dataset parquet from "
+            f"{'remote' if storage.is_remote(uri) else 'local'} directory {uri}"
+        )
+
+        parquet_files = await list_parquet_files(uri, filesystem)
+        columns = _requested_columns(current_task_metadata)
+
+        try:
+            if any(storage.is_remote(file_path) for file_path in parquet_files):
+                logger.info(f"Using direct remote parquet reads for {uri} via Flyte storage filesystem")
+            table = await run_sync_io(
+                "read parquet files",
+                _read_parquet_files,
+                parquet_files,
+                columns,
+                filesystem,
+            )
+        except Exception as exc:
+            if any(storage.is_remote(file_path) for file_path in parquet_files):
+                logger.warning(
+                    f"Direct parquet read failed for {uri}: {type(exc).__name__}: {exc}. "
+                    "Falling back to localizing remote parquet shards."
+                )
+                parquet_files = await _localize_parquet_files(parquet_files)
+                logger.info(f"Using localized parquet shard reads for {uri}")
+                table = await run_sync_io(
+                    "read localized parquet files",
+                    _read_parquet_files,
+                    parquet_files,
+                    columns,
+                )
+            else:
+                raise
+
+        return datasets.Dataset(table)
+
+
+class HFToHuggingFaceDatasetDecodingHandler(ParquetToHuggingFaceDatasetDecodingHandler):
+    def __init__(self):
+        super().__init__("hf")
+
+
+class HuggingFaceIterableDatasetToParquetEncodingHandler(DataFrameEncoder):
+    def __init__(self):
+        super().__init__(datasets.IterableDataset, None, PARQUET)
+
+    async def encode(
+        self,
+        dataframe: DataFrame,
+        structured_dataset_type: types_pb2.StructuredDatasetType,
+    ) -> literals_pb2.StructuredDataset:
+        val = dataframe.val
+
+        if val is None and dataframe.uri:
+            structured_dataset_type.format = PARQUET
+            return literals_pb2.StructuredDataset(
+                uri=typing.cast(str, dataframe.uri),
+                metadata=literals_pb2.StructuredDatasetMetadata(structured_dataset_type=structured_dataset_type),
+            )
+
+        if not dataframe.uri:
+            from flyte._context import internal_ctx
+
+            uri = str(internal_ctx().raw_data.get_random_remote_path())
+        else:
+            uri = typing.cast(str, dataframe.uri)
+
+        if not storage.is_remote(uri):
+            Path(uri).mkdir(parents=True, exist_ok=True)
+
+        ds = typing.cast(datasets.IterableDataset, val)
+        filesystem = storage.get_underlying_filesystem(path=uri)
+        logger.info(
+            f"Writing Hugging Face IterableDataset output to "
+            f"{'remote' if storage.is_remote(uri) else 'local'} parquet directory {uri}"
+        )
+        await run_sync_io(
+            "write HuggingFace iterable dataset",
+            _write_iterable_dataset,
+            ds,
+            uri,
+            filesystem,
+        )
+
+        structured_dataset_type.format = PARQUET
+        return literals_pb2.StructuredDataset(
+            uri=uri,
+            metadata=literals_pb2.StructuredDatasetMetadata(structured_dataset_type=structured_dataset_type),
+        )
+
+
+class ParquetToHuggingFaceIterableDatasetDecodingHandler(DataFrameDecoder):
+    def __init__(self, protocol: str | None = None):
+        super().__init__(datasets.IterableDataset, protocol, PARQUET)
+
+    async def decode(
+        self,
+        flyte_value: literals_pb2.StructuredDataset,
+        current_task_metadata: literals_pb2.StructuredDatasetMetadata,
+    ) -> datasets.IterableDataset:
+        uri = flyte_value.uri
+
+        if uri.startswith("hf://"):
+            source_uri = uri
+            try:
+                uri = await ensure_hf_cached(HFSource.from_hf_uri(uri))
+            except Exception as e:
+                raise RuntimeError(
+                    f"Failed to materialize Hugging Face dataset from {uri}: {type(e).__name__}: {e!r}"
+                ) from e
+            logger.info(
+                f"Resolved Hugging Face source {source_uri} to "
+                f"{'remote' if storage.is_remote(uri) else 'local'} parquet directory {uri}"
+            )
+
+        filesystem = storage.get_underlying_filesystem(path=uri)
+        logger.info(
+            f"Reading Hugging Face IterableDataset parquet from "
+            f"{'remote' if storage.is_remote(uri) else 'local'} directory {uri}"
+        )
+        parquet_files = await list_parquet_files(uri, filesystem)
+        columns = _requested_columns(current_task_metadata)
+
+        try:
+            if any(storage.is_remote(file_path) for file_path in parquet_files):
+                logger.info(f"Using direct remote iterable parquet reads for {uri} via Flyte storage filesystem")
+            await run_sync_io(
+                "probe parquet files",
+                _probe_remote_parquet_files,
+                parquet_files,
+                filesystem,
+            )
+        except Exception as exc:
+            if any(storage.is_remote(file_path) for file_path in parquet_files):
+                logger.warning(
+                    f"Direct iterable parquet access failed for {uri}: {type(exc).__name__}: {exc}. "
+                    "Falling back to localizing remote parquet shards."
+                )
+                parquet_files = await _localize_parquet_files(parquet_files)
+                logger.info(f"Using localized iterable parquet shard reads for {uri}")
+            else:
+                raise
+
+        return datasets.IterableDataset.from_generator(
+            _iter_parquet_rows,
+            gen_kwargs={"parquet_files": parquet_files, "columns": columns},
+        )
+
+
+class HFToHuggingFaceIterableDatasetDecodingHandler(ParquetToHuggingFaceIterableDatasetDecodingHandler):
+    def __init__(self):
+        super().__init__("hf")

flyteplugins_huggingface-2.2.1.dist-info/METADATA

@@ -0,0 +1,345 @@
+Metadata-Version: 2.4
+Name: flyteplugins-huggingface
+Version: 2.2.1
+Summary: Hugging Face Plugin for Flyte
+Author-email: André Ahlert <andre@aex.partners>, Samhita Alla <samhita@union.ai>
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: datasets
+Requires-Dist: huggingface-hub
+Requires-Dist: pyarrow
+Requires-Dist: flyte
+
+# Hugging Face Plugin
+
+Native Flyte support for Hugging Face integrations in Flyte.
+
+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.
+2. Pass Hugging Face dataset objects between Flyte tasks with automatic Parquet
+   serialization.
+
+The plugin works by treating Hub datasets as Parquet-backed structured data. For
+Hub sources, it first resolves the dataset's converted Parquet shards, then
+materializes them either into a generated path for the current run or into a
+shared artifact registry rooted at `cache_root`.
+
+## Installation
+
+```bash
+pip install flyteplugins-huggingface
+```
+
+## Quick start
+
+```python
+import datasets
+import flyte
+from flyteplugins.huggingface.datasets import from_hf
+
+env = flyte.TaskEnvironment(name="hf-example")
+
+@env.task
+async def count_reviews(
+    ds: datasets.Dataset = from_hf(
+        "stanfordnlp/imdb",
+        name="plain_text",
+        split="train",
+    ),
+) -> int:
+    return len(ds)
+```
+
+At the Flyte literal level this source is represented as an `hf://` URI, for
+example:
+
+```text
+hf://stanfordnlp/imdb?name=plain_text&split=train
+```
+
+The task receives a hydrated `datasets.Dataset`. The `hf://` URI is only the
+reference used between Flyte and the plugin.
+
+## `from_hf(...)`
+
+`from_hf(...)` is the entry point for Hub-backed task defaults:
+
+```python
+from flyteplugins.huggingface.datasets import from_hf
+
+from_hf(
+    repo: str,
+    *,
+    name: str | None = None,
+    split: str | None = None,
+    revision: str | None = None,
+    cache_root: str | None = None,
+)
+```
+
+Arguments:
+
+- `repo`: Hugging Face dataset repo, such as `"stanfordnlp/imdb"` or `"glue"`.
+- `name`: Optional dataset config/subset.
+- `split`: Optional split such as `"train"` or `"validation"`.
+- `revision`: Optional Hub revision. Defaults to `refs/convert/parquet`.
+- `cache_root`: Optional shared remote cache root for cross-run reuse.
+
+`from_hf(...)` returns a Flyte `DataFrame` reference, not an eagerly loaded
+dataset object. When the task input is typed as `datasets.Dataset` or
+`datasets.IterableDataset`, the plugin decoder materializes that reference into
+the requested Hugging Face type.
+
+## Config resolution
+
+If you specify `name`, the plugin uses that config directly.
+
+If you omit `name`, the plugin resolves the config as follows:
+
+1. Try actual converted-parquet config `default`.
+2. If `default` does not exist and there is exactly one config, use that one.
+3. If there are multiple configs, raise an error and ask for `name=...`.
+
+Examples:
+
+```python
+# Works: imdb has a single converted-parquet config, plain_text.
+from_hf("stanfordnlp/imdb", split="train")
+
+# Required: glue has multiple configs such as mrpc, sst2, qnli, ...
+from_hf("glue", name="mrpc", split="train")
+```
+
+Using `name=` explicitly is recommended in examples and production code because
+it makes the UI literal and task signature more obvious.
+
+## Split behavior
+
+If you specify `split`, only that split is materialized:
+
+```python
+@env.task
+async def train_split(
+    ds: datasets.Dataset = from_hf(
+        "stanfordnlp/imdb",
+        name="plain_text",
+        split="train",
+    ),
+) -> int:
+    return len(ds)
+```
+
+If you omit `split`, the plugin reads every converted Parquet split under the
+resolved config and presents them as one dataset stream/table:
+
+```python
+@env.task
+async def all_splits(
+    ds: datasets.Dataset = from_hf(
+        "stanfordnlp/imdb",
+        name="plain_text",
+    ),
+) -> list[str]:
+    return ds.column_names
+```
+
+That means the result is a combined dataset, not a mapping of split name to
+dataset.
+
+## Cross-run reuse with `cache_root`
+
+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 cache registry so later runs can
+skip the Hub download entirely:
+
+```python
+@env.task
+async def train_cached(
+    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 shared cache layout is:
+
+```text
+{cache_root}/huggingface/datasets/
+  by-key/{source-cache-key}.json
+  blobs/{source-cache-key}/...
+```
+
+The cache key is derived from:
+
+- repo
+- config name
+- split
+- revision
+- resolved Parquet shard metadata
+
+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:
+
+- checking the shared dataset cache
+- materializing from the Hugging Face Hub
+- using a cached artifact
+- reading Parquet from a local or remote directory
+
+This is the easiest way to confirm whether a run is reading from the Hub or
+from your shared cache artifact.
+
+## `datasets.Dataset` between tasks
+
+You can return and pass real `datasets.Dataset` objects between tasks:
+
+```python
+import datasets
+import flyte
+
+env = flyte.TaskEnvironment(name="hf-transform")
+
+
+@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)
+```
+
+Task-produced in-memory datasets are serialized to Parquet automatically. This
+is separate from `from_hf(...)`, which is a source reference rather than a
+materialized dataset object.
+
+## `datasets.IterableDataset`
+
+Use `datasets.IterableDataset` when you want row streaming behavior instead of a
+fully materialized table:
+
+```python
+@env.task
+async def stream_reviews(
+    ds: datasets.IterableDataset = from_hf(
+        "stanfordnlp/imdb",
+        name="plain_text",
+        split="train",
+        cache_root="s3://my-bucket/flyte-hf-cache",
+    ),
+) -> datasets.IterableDataset:
+    def add_length(batch):
+        batch["length"] = [len(text) for text in batch["text"]]
+        return batch
+
+    return ds.map(add_length, batched=True)
+```
+
+Notes:
+
+- The returned Hugging Face `IterableDataset` is consumed with normal synchronous
+  iteration.
+- Internally the plugin streams row batches from Parquet files.
+- Iterable outputs are written back as sharded Parquet directories.
+
+## Column projection
+
+Use a Flyte structured-dataset column annotation when you only want selected
+columns:
+
+```python
+from collections import OrderedDict
+from typing import Annotated
+
+
+@env.task
+async def load_text_only(
+    ds: Annotated[datasets.Dataset, OrderedDict(text=str)] = from_hf(
+        "stanfordnlp/imdb",
+        name="plain_text",
+        split="train",
+    ),
+) -> list[str]:
+    return ds["text"][:10]
+```
+
+The plugin uses the annotation to request only those columns when reading
+Parquet.
+
+## Revision selection
+
+Use `revision=` if you want to pin a specific converted-Parquet revision:
+
+```python
+@env.task
+async def pinned_revision(
+    ds: datasets.Dataset = from_hf(
+        "stanfordnlp/imdb",
+        name="plain_text",
+        split="train",
+        revision="refs/convert/parquet",
+        cache_root="s3://my-bucket/flyte-hf-cache",
+    ),
+) -> int:
+    return len(ds)
+```
+
+If you do not specify a revision, the plugin uses `refs/convert/parquet`.
+
+## Local vs remote behavior
+
+There are two distinct layers to keep in mind:
+
+1. Task inputs and outputs inside Flyte tasks.
+2. What your launcher process sees when a run completes.
+
+Inside a task, a parameter typed as `datasets.Dataset` or
+`datasets.IterableDataset` is hydrated by the plugin into a Hugging Face object.
+
+Outside the task, especially for remote runs, outputs are often represented to
+the launcher as Flyte `DataFrame` references rather than already-opened Hugging
+Face dataset objects. That is expected: the structured dataset literal remains
+the transport format.
+
+## Private datasets
+
+Set `HF_TOKEN` in the task environment to access private Hugging Face datasets.
+Without it, the plugin uses anonymous Hub access.
+
+## Failure modes
+
+Common issues:
+
+- Missing `name` for a dataset with multiple configs:
+  the plugin raises and asks for `name=...`.
+- No converted Parquet shards available:
+  the dataset may not have an auto-converted Parquet representation yet.
+- Remote cache path credentials:
+  your Flyte runtime must be able to read and write the chosen `cache_root`.
+
+## Example
+
+See the example workflow in `plugins/huggingface/examples/hf_dataset_workflow.py`
+for end-to-end local and remote scenarios.

flyteplugins_huggingface-2.2.1.dist-info/RECORD

@@ -0,0 +1,10 @@
+flyteplugins/huggingface/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+flyteplugins/huggingface/datasets/__init__.py,sha256=8wuGxh0ayadFTxcr58VsrOTYrSz5FbZrpIUtjZ_6iqA,1363
+flyteplugins/huggingface/datasets/_io.py,sha256=lZYX6Uc_6ocrlwmQ_GXbf7VUzk6hWZK_wPo5EGDWe5Q,16186
+flyteplugins/huggingface/datasets/_source.py,sha256=YfeCPyN_arxfJ1y8eGm4Jwwu_FmcHsSo_BO4Zh7PuTc,4825
+flyteplugins/huggingface/datasets/_transformers.py,sha256=XjX5zx__PLfZgn3fObbhlWuXxDGWABZAm9ZsqFtLDUk,13008
+flyteplugins_huggingface-2.2.1.dist-info/METADATA,sha256=up4CpJpRR3P8qvzjVS6RlSuVKLKudJtzeItflgENBjQ,9230
+flyteplugins_huggingface-2.2.1.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+flyteplugins_huggingface-2.2.1.dist-info/entry_points.txt,sha256=ixxMpwmr3wlKwLK5GNlxnWeS4i4Pub3O7aNsEQFkjE4,112
+flyteplugins_huggingface-2.2.1.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
+flyteplugins_huggingface-2.2.1.dist-info/RECORD,,