better correspondence between cached and saved datasets created using from_generator

[vttrifonov]

Feature request

At the moment .from_generator can only create a dataset that lives in the cache. The cached dataset cannot be loaded with load_from_disk because the cache folder is missing state.json. So the only way to convert this cached dataset to a regular is to use save_to_disk which needs to create a copy of the cached dataset. For large datasets this can end up wasting a lot of space. In my case the saving operation failed so I am stuck with a large cached dataset and no clear way to convert to a Dataset that I can use. The requested feature is to provide a way to be able to load a cached dataset using .load_from_disk. Alternatively .from_generator can create the dataset at a specified location so that it can be loaded from there with .load_from_disk.

Motivation

I have the following workflow which has exposed some awkwardness about the Datasets saving/caching.

1. I created a cached dataset using .from_generator which was cached in a folder. This dataset is rather large (~600GB) with many shards.
2. I tried to save this dataset using .save_to_disk to another location so that I can use later as a Dataset. This essentially creates another copy (for a total of 1.2TB!) of what is already in the cache... In my case the saving operation keeps dying for some reason and I am stuck with a cached dataset and no copy.
3. Now I am trying to "save" the existing cached dataset but it is not clear how to access the cached files after .from_generator has finished e.g. from a different process. I should not be even looking at the cache but I really do not want to waste another 2hr to generate the set so that if fails agains (I already did this couple of times).

• I tried .load_from_disk but it does not work with cached files and complains that this is not a Dataset (!).
• I looked at .from_file which takes one file but the cached file has many (shards) so I am not sure how to make this work.
• I tried .load_dataset but this seems to either try to "download" a copy (of a file which is already in the local file system!) which I will then need to save or I need to use streaming=False to create an IterableDataset which then I need to convert (using the cache) to Dataset so that I can save it. With both options I will end up with 3 copies of the same dataset for a total of ~2TB! I am hoping here is another way to do this...

Maybe I am missing something here: I looked at docs and forums but no luck. I have a bunch of arrow files cached by Dataset.from_generator and no clean way to make them into a Dataset that I can use.

This all could be so much easer if load_from_disk can recognize the cached files and produce a Dataset: after the cache is created I would not have to "save" it again and I can just load it when I need. At the moment load_from_disk needs state.json which is lacking in the cache folder. So perhaps .from_generator could be made to "finalize" (e.g. create state.json) the dataset once it is done so that it can be loaded easily. Or provide .from_generator with a save_to_dir parameter in addition to cache_dir which can be used for the whole process including creating the state.json at the end.

As a proof of concept I just created state.json by hand and load_from_disk worked using the cache! So it seems to be the missing piece here.

Your contribution

Time permitting I can look into .from_generator to see if adding state.json is feasible.

[Tommigun1980]

Until this lands, in case anyone needs this, I made this helper (based on HF sources) that merely serializes the metadata from a dataset, thus allowing to bypass the completely redundant additional save.

import json
import posixpath
import os
from typing import Optional
from fsspec.core import url_to_fs
from datasets.utils.typing import PathLike

from datasets import Dataset, config
from datasets.info import asdict

def save_dataset_metadata(
    dataset: Dataset,
    dataset_path: PathLike,
    storage_options: Optional[dict] = None
):
    """
    Saves only the dataset metadata (state.json and dataset_info.json) to a dataset directory.

    This does NOT rewrite or move any Arrow shard files.
    It reuses the existing data files referenced by `dataset.cache_files`.

    Args:
        dataset (`datasets.Dataset`):
            The in-memory Dataset instance whose current cache-backed data files
            should be exposed as a load_from_disk()-compatible dataset.
        dataset_path (`path-like`):
            Path (e.g. `dataset/train`) or remote URI (e.g. `s3://my-bucket/dataset/train`)
            of the dataset directory where the metadata will be saved to.
        storage_options (`dict`, *optional*):
            Key/value pairs to be passed on to the file-system backend, if any.
    """
    fs, _ = url_to_fs(dataset_path, **(storage_options or {}))
    fs.makedirs(dataset_path, exist_ok=True)

    # Build json-serializable state (copied from HF save_to_disk, but without shard writing)
    state = {
        key: dataset.__dict__[key]
        for key in [
            "_fingerprint",
            "_format_columns",
            "_format_kwargs",
            "_format_type",
            "_output_all_columns",
        ]
    }
    state["_split"] = str(dataset.split) if dataset.split is not None else dataset.split

    # Relies on `dataset.cache_files` to point at the real Arrow shards.
    state["_data_files"] = [
        {"filename": os.path.relpath(cache_filename["filename"], dataset_path)} for cache_filename in dataset.cache_files
    ]

    # Same json-serializability check as HF
    for k in state["_format_kwargs"].keys():
        json.dumps(state["_format_kwargs"][k])

    # Dataset info (same as HF)
    dataset_info = asdict(dataset._info)

    # Write state.json
    with fs.open(
        posixpath.join(dataset_path, config.DATASET_STATE_JSON_FILENAME), "w", encoding="utf-8"
    ) as state_file:
        json.dump(state, state_file, indent=2, sort_keys=True)

    # Write dataset_info.json
    with fs.open(
        posixpath.join(dataset_path, config.DATASET_INFO_FILENAME), "w", encoding="utf-8"
    ) as dataset_info_file:
        # Sort only the top-level keys, like HF does
        sorted_keys_dataset_info = {key: dataset_info[key] for key in sorted(dataset_info)}
        json.dump(sorted_keys_dataset_info, dataset_info_file, indent=2)

Usage:

from .save_dataset_metadata import save_dataset_metadata

# ...

dataset = datasets.Dataset.from_generator(..., cache_dir="xxx")

# ...

save_dataset_metadata(dataset, Path(directory_path))

[Tommigun1980]

@lhoestq Do you think this issue could get some traction to be resolved? Currently if you have a large dataset created in cache and need to persist it, I don't think it's great that you have to use double disk space and hours on re-serializing the dataset, when in reality only the metadata is missing. Imo Datasets should be able to persist metadata only in this case.

Thanks for considering.

[lhoestq]

Have you tried using pickle ?

ds = Dataset.from_genetator(...)  # generates and caches the dataset
with open("dataset.pickle", "wb") as f:
    pickle.dump(ds, f)  # saves the Dataset obj including the cache files locations (and not the data! unless some of them live in memory because of add_column or add_item for example)

# later
with open("dataset.pickle", "rb") as f:
    ds = pickle.load(ds, f)  # reload the Dataset obj from the cache files

I'm not a big fan of using pickle for this kind of use cases but it should do the job.

Anyways, moving forward it would be cool to change how save_to_disk / load_from_disk work to make them compatible with load_dataset. This implies replacing the state.json file with the usual README.md which can list the data files per config/split. save_to_disk would still materialize the full dataset on disk in the output directory though.

The difference would be load_from_disk would use the dataset directory as cache directory to keep it self-contained and avoid copies, contrary to load_dataset which copies the data to a directory per dataset inside ~/.cache/huggingface/datasets/. Contributions on this are welcome on this subject btw !