jinchuntan · Oct 18, 2026
diff --git a/src/datasets/utils/tf_utils.py b/src/datasets/utils/tf_utils.py
@@ -18,6 +18,7 @@
 import warnings
 from functools import partial
 from math import ceil
+from unittest.mock import patch
 from uuid import uuid4
 
 import numpy as np
@@ -344,7 +345,9 @@ def __iter__(self):
                     **base_args,
                 }
                 worker = ctx.Process(target=self.worker_loop, kwargs=worker_kwargs, daemon=True)
-                worker.start()
+                # Spawn imports user code before worker_loop can hide GPUs through TensorFlow.
+                with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}):
+                    worker.start()
                 workers.append(worker)
 
             end_signal_received = False

diff --git a/tests/fixtures/tf_initialized_collate.py b/tests/fixtures/tf_initialized_collate.py
@@ -0,0 +1,16 @@
+import os
+
+import numpy as np
+import tensorflow as tf
+
+
+# Importing a user's collator can initialize TensorFlow before the worker target runs.
+IMPORT_TIME_TENSOR = tf.constant(1, dtype=tf.int64)
+GPUS_HIDDEN_AT_IMPORT = os.environ.get("CUDA_VISIBLE_DEVICES") == "" and not tf.config.get_visible_devices("GPU")
+
+
+def collate_with_initialized_tensorflow(examples):
+    return {
+        "x": tf.convert_to_tensor([example["x"] for example in examples], dtype=tf.int64) * IMPORT_TIME_TENSOR,
+        "gpus_hidden_at_import": np.full(len(examples), GPUS_HIDDEN_AT_IMPORT),
+    }
diff --git a/tests/test_tf_utils.py b/tests/test_tf_utils.py
@@ -0,0 +1,64 @@
+import os
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import numpy as np
+import pytest
+
+from datasets import Dataset
+from datasets.utils import tf_utils
+
+from .utils import require_tf
+
+
+@require_tf
+@pytest.mark.parametrize("cuda_devices", [None, "0"])
+@pytest.mark.parametrize("num_workers", [1, 2])
+def test_tf_workers_hide_gpus_before_importing_collator(monkeypatch, cuda_devices, num_workers):
+    import tensorflow as tf
+
+    from .fixtures.tf_initialized_collate import collate_with_initialized_tensorflow
+
+    parent_gpus = tf.config.get_visible_devices("GPU")
+    if cuda_devices is None:
+        monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
+    else:
+        monkeypatch.setenv("CUDA_VISIBLE_DEVICES", cuda_devices)
+    dataset = Dataset.from_dict({"x": list(range(6))})
+    tf_dataset = dataset.to_tf_dataset(
+        collate_fn=collate_with_initialized_tensorflow, batch_size=2, num_workers=num_workers, prefetch=False
+    )
+    batches = list(tf_dataset.as_numpy_iterator())
+    np.testing.assert_array_equal(np.concatenate([batch["x"] for batch in batches]), np.arange(6))
+    assert np.concatenate([batch["gpus_hidden_at_import"] for batch in batches]).all()
+    assert os.environ.get("CUDA_VISIBLE_DEVICES") == cuda_devices
+    assert tf.config.get_visible_devices("GPU") == parent_gpus
+
+
+@pytest.mark.parametrize("cuda_devices", [None, "0"])
+def test_tf_worker_start_failure_restores_cuda_environment(monkeypatch, cuda_devices):
+    if cuda_devices is None:
+        monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
+    else:
+        monkeypatch.setenv("CUDA_VISIBLE_DEVICES", cuda_devices)
+    context = tf_utils.get_context("spawn")
+    failure = RuntimeError("worker start failed")
+    worker = Mock()
+    worker.start.side_effect = failure
+    monkeypatch.setattr(context, "Process", lambda **kwargs: worker)
+    generator = tf_utils.NumpyMultiprocessingGenerator(
+        dataset=Dataset.from_dict({"x": [0, 1]}),
+        cols_to_retain=["x"],
+        collate_fn=None,
+        collate_fn_args={},
+        columns_to_np_types={"x": np.int64},
+        output_signature={"x": SimpleNamespace(shape=SimpleNamespace(rank=1))},
+        shuffle=False,
+        batch_size=1,
+        drop_remainder=False,
+        num_workers=1,
+    )
+    with pytest.raises(RuntimeError) as exc_info:
+        next(iter(generator))
+    assert exc_info.value is failure
+    assert os.environ.get("CUDA_VISIBLE_DEVICES") == cuda_devices