Hide GPUs before spawning TensorFlow data workers by jinchuntan · Pull Request #8847 · huggingface/datasets · GitHub
Skip to content

Hide GPUs before spawning TensorFlow data workers - #8847

Open
jinchuntan wants to merge 1 commit into
huggingface:mainfrom
jinchuntan:fix/tf-worker-gpu-startup
Open

jinchuntan wants to merge 1 commit into
huggingface:mainfrom
jinchuntan:fix/tf-worker-gpu-startup

Conversation

@jinchuntan

@jinchuntan jinchuntan commented Oct 18, 2026 •

Copy link
Copy Markdown
Contributor

to_tf_dataset(num_workers=...) can import a user's collator and initialize TensorFlow before worker_loop hides its GPUs. Hide GPUs in the spawn environment so initialization already uses the CPU, then restore the parent environment. Fixes #6948.

Six new cases cover imported collators with eager TensorFlow initialization, one/two workers, unset/existing CUDA_VISIBLE_DEVICES and failed startup. All 10 selected worker/conversion tests pass, along with lint and formatting. Tested with TensorFlow 2.21 on CPU. GPU hardware was unavailable.

Set worker GPU visibility before spawn imports user modules, then restore the parent environment. Cover eager TensorFlow initialization in imported collators, normal data conversion, and failed process startup.

Fixes huggingface#6948

Assisted-by: Codex
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.

to_tf_dataset: Visible devices cannot be modified after being initialized

1 participant