mp_prefetch uses multithreading when the GIL is disabled.

PiperOrigin-RevId: 833406009

Author: Grain Team

CHANGELOG.md

@@ -13,6 +13,8 @@ changes. Best viewed [here](https://google-grain.readthedocs.io/en/latest/change
     hiding first batch processing behind model checkpoint recovery.
   * Introduces `grain.experimental.multithread_prefetch` as an
     alternative to multiprocessing prefetch in free-threading Python.
+  * Switches to multithreading instead of multiprocessing in
+    `IterDataset.mp_prefetch` when free-threaded Python is detected.
 
 * Breaking changes:
 

grain/_src/python/dataset/dataset.py

@@ -57,6 +57,7 @@
 from grain._src.core import monitoring as grain_monitoring
 from grain._src.core import transforms
 from grain._src.core import usage_logging
+import multiprocessing as mp
 from grain._src.python import checkpointing
 from grain._src.python import options as grain_options
 from grain._src.python.dataset import base
@@ -1312,6 +1313,18 @@ def mp_prefetch(
     # pylint: disable=g-import-not-at-top
     from grain._src.python.dataset.transformations import prefetch
     # pylint: enable=g-import-not-at-top
+    if mp.is_in_free_threaded_python():
+      if worker_init_fn is not None:
+        warnings.warn(
+            "Free-threaded Python is used: `mp_prefetch` falls back to"
+            " thread-based implementation and `worker_init_fn` is ignored."
+        )
+      return prefetch.multithread_prefetch(
+          self,
+          num_threads=options.num_workers,
+          buffer_size=options.per_worker_buffer_size,
+          sequential_slice=sequential_slice,
+      )
     return prefetch.MultiprocessPrefetchIterDataset(
         self,
         multiprocessing_options=options,

grain/_src/python/dataset/dataset_test.py

@@ -31,6 +31,7 @@
 from grain._src.python.dataset import base
 from grain._src.python.dataset import dataset
 from grain._src.python.dataset import stats as dataset_stats
+from grain._src.python.dataset.transformations import prefetch
 import grain._src.python.testing.experimental as test_util
 from grain.proto import execution_summary_pb2
 import numpy as np
@@ -1182,6 +1183,21 @@ def test_apply(self, ds):
         ],
     )
 
+  def test_mp_prefetch_switches_to_threads_for_free_threaded_python(self):
+    ds = dataset.MapDataset.range(15).to_iter_dataset()
+    prefetched_ds = ds.mp_prefetch()
+    is_free_threaded = (
+        hasattr(sys, "_is_gil_enabled") and not sys._is_gil_enabled()
+    )
+    if is_free_threaded:
+      self.assertNotIsInstance(
+          prefetched_ds, prefetch.MultiprocessPrefetchIterDataset
+      )
+    else:
+      self.assertIsInstance(
+          prefetched_ds, prefetch.MultiprocessPrefetchIterDataset
+      )
+
 
 class TfRandomMapAlwaysAddingOne(transforms.TfRandomMapTransform):
 

grain/_src/python/dataset/transformations/prefetch_test.py

@@ -679,6 +679,13 @@ def map(self, x):
     )
     iter_ds = iter_ds.mp_prefetch()
 
+    print(
+        'nithintati:'
+        ' test_reports_first_unpicklable_dataset_when_with_multiple_parents'
+    )
+    print('nithintati: ', iter_ds)
+    print('nithintati: ', mp.is_in_free_threaded_python())
+    print('nithintati: ', sys.version_info)
     with self.assertRaisesRegex(
         ValueError,
         r'UnpicklableObject is not picklable',
@@ -691,6 +698,7 @@ def map(self, x):
           r'Dataset: MapMapDataset\(transform=LeftTransform\) cannot be'
           r' pickled!',
       )
+    raise ValueError('nithintati: ' + str(sys.version_info))
 
   def test_reports_unpicklable_issue_when_only_one_parent_unpicklable(self):
     class UnpicklableObject: