Skip to content

datasets/huggingface.py: concurrent workers for TransformAndMapHuggingFaceDatasetJob + HF Hub transient-drop retries - #664

Open
robin-p-schmitt wants to merge 5 commits into
mainfrom
robin-hf-dataset
Open

robin-p-schmitt wants to merge 5 commits into
mainfrom
robin-hf-dataset

Conversation

@robin-p-schmitt

Copy link
Copy Markdown
Contributor

Two changes for downloading/processing very large HF datasets (e.g. People's Speech, ~5000 WebDataset tar shards, several TB):

  1. TransformAndMapHuggingFaceDatasetJob(concurrent=N)
    Splits load_dataset_opts["data_files"] round-robin into N buckets, each run as its own task array index writing to its own out_dirs[i] instead of a single out_dir. Buckets are independently resumable and progress-trackable, and double as shards for training. sis_hash_exclude = {"concurrent": None} keeps existing instances' hashes.
    Along with it:
  • the tmp map-output dir is suffixed by task_id (concurrent tasks share one job dir),
  • load_dataset_opts["features"] may be a plain function returning the Features object — datasets.Features/pyarrow types can't be stored as a job attribute (sisyphus' object-state reflection crashes on them),
  • streaming=True loads are allowed: Iterable* datasets pass the load/transform asserts, with an explicit assert that some transform materialized them before .map()/.save_to_disk().
  1. patch_hf_hub_transient_retries() (called at the start of run())
    huggingface_hub (1.16.1) does not survive a CDN closing the connection mid-body (httpx.RemoteProtocolError), which loses a whole multi-hour shard: _DEFAULT_RETRY_ON_EXCEPTIONS doesn't cover it, and HfFileSystemStreamFile.read retries exactly once. The patch adds the exception to the retry tuple (also in the kwdefaults of _http_backoff_base/http_backoff/http_stream_backoff, since defaults bind at def time) and wraps the stream read in a capped exponential backoff that clears self.response, so _open_connection() re-requests with a Range header and the read resumes at self.loc. Best-effort and idempotent: every step is individually guarded and falls back to the original
    behaviour.

Comment thread datasets/huggingface.py Outdated
Comment on lines +366 to +372
# `features` (a datasets.Features/Audio/Value/... object) can't be stored directly as a
# job attribute: sisyphus's generic object-state reflection (used e.g. to find tk.Path
# inputs, regardless of whether the containing dict is hashed) walks the whole job
# __dict__ and fails on the underlying pyarrow DataType objects. So `features` may instead
# be a plain top-level function that builds and returns the Features object -- exactly
# like `map_func`/`transform`, which are already handled this way -- called here instead
# of being stored anywhere.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please check (or have your LLM do it) if some comments/docstrings can be shortened or omitted. (e.g. I don't think we need 7 lines of explanation on why we would call this if it is callable.)

@robin-p-schmitt robin-p-schmitt self-assigned this Oct 8, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants