Fix MultiprocessingOptions.enable_profiling being ignored.

CHANGELOG.md

@@ -19,6 +19,7 @@ changes. Best viewed [here](https://google-grain.readthedocs.io/en/latest/change
 * Bug fixes:
   * Fixed bug in DataLoader where sharding remainder was dropped even when ShardOptions.drop_remainder=False.
   * Fixes reference cycle in BatchMapDataset.
+  * Fixed `MultiprocessingOptions.enable_profiling` being ignored; it can now be used to enable worker profiling programmatically when the global flag is False.
 
 ## Grain 0.2.18 (June 17, 2026)
 

grain/_src/core/profiler.py

@@ -107,9 +107,14 @@ def register_subprocess(pid: int, port: int) -> Callable[[], None]:
   )
 
 
-def get_worker_init_fn(port: int) -> Callable[[], None]:
+def get_worker_init_fn(
+    port: int, enable_profiling: bool = False
+) -> Callable[[], None]:
   """Start the profiler server in a worker process."""
-  if not is_worker_profiling_enabled():
+  profiling_enabled = (
+      enable_profiling and is_worker_profiling_supported()
+  ) or is_worker_profiling_enabled()
+  if not profiling_enabled:
     return lambda: None
 
   def _worker_init_fn() -> None:

grain/_src/python/dataset/dataset.py

@@ -1394,6 +1394,7 @@ def mp_prefetch(
         buffer_size=options.per_worker_buffer_size,
         worker_init_fn=worker_init_fn,
         sequential_slice=sequential_slice,
+        enable_profiling=options.enable_profiling,
     )
 
   def __init_subclass__(cls, **kwargs):

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

@@ -147,6 +147,7 @@ def __init__(
       parent: dataset.IterDataset[T],
       buffer_size: int,
       worker_init_fn: Callable[[], None] | None = None,
+      enable_profiling: bool = False,
   ):
     """Initializes the ProcessPrefetchIterDataset.
 
@@ -155,6 +156,8 @@ def __init__(
       buffer_size: The size of the buffer used for prefetching.
       worker_init_fn: An optional function to run in the worker process at
         startup.
+      enable_profiling: Whether to enable profiling in worker processes if not
+        enabled via the flag.
     """
     if buffer_size <= 0:
       raise ValueError(
@@ -163,16 +166,22 @@ def __init__(
     super().__init__(parent)
     self._buffer_size = buffer_size
     self._worker_init_fn = worker_init_fn
+    self._enable_profiling = enable_profiling
 
   def __str__(self) -> str:
     return f"ProcessPrefetchIterDataset(buffer_size={self._buffer_size})"
 
   def __iter__(self) -> dataset.DatasetIterator[T]:
     worker_init_fn = self._worker_init_fn
     worker_profiler_port = None
-    if profiler.is_worker_profiling_enabled() and profiler.is_loaded():
+    enable_worker_profiling = (
+        self._enable_profiling and profiler.is_worker_profiling_supported()
+    ) or profiler.is_worker_profiling_enabled()
+    if enable_worker_profiling and profiler.is_loaded():
       worker_profiler_port = portpicker.pick_unused_port()
-      profiler_init_fn = profiler.get_worker_init_fn(worker_profiler_port)
+      profiler_init_fn = profiler.get_worker_init_fn(
+          worker_profiler_port, enable_profiling=self._enable_profiling
+      )
       if worker_init_fn is None:
         worker_init_fn = profiler_init_fn
       else:
@@ -679,6 +688,7 @@ def multiprocess_prefetch(
     buffer_size: int = 1,
     worker_init_fn: Callable[[int, int], None] | None = None,
     sequential_slice: bool = False,
+    enable_profiling: bool = False,
 ) -> dataset.IterDataset[T]:
   """Uses a multiple processes to prefetch elements ahead of time.
 
@@ -693,6 +703,9 @@ def multiprocess_prefetch(
     buffer_size: The size of the prefetch buffer for each process.
     worker_init_fn: A function that is called in each worker process.
     sequential_slice: Whether to use sequential slicing.
+    enable_profiling: Whether to enable profiling in worker processes. If False,
+      it defaults to the global --grain_enable_multiprocess_worker_profiling
+      flag.
 
   Returns:
     `IterDataset` that prefetches elements from `ds` using multiple processes.
@@ -737,6 +750,7 @@ def multiprocess_prefetch(
         worker_ds,
         buffer_size=buffer_size,
         worker_init_fn=functools.partial(_run_all, worker_init_fns),
+        enable_profiling=enable_profiling,
     )
     shards.append(worker_ds)
 

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

@@ -1132,6 +1132,52 @@ def test_element_spec(self):
     self.assertEqual(spec.dtype, np.int64)
     self.assertEqual(spec.shape, ())
 
+  @mock.patch.object(
+      process_prefetch.profiler,
+      'is_worker_profiling_supported',
+      return_value=True,
+  )
+  @mock.patch.object(
+      process_prefetch.profiler,
+      'is_worker_profiling_enabled',
+      return_value=False,
+  )
+  @mock.patch.object(process_prefetch.profiler, 'is_loaded', return_value=True)
+  @mock.patch.object(
+      process_prefetch.portpicker, 'pick_unused_port', return_value=12345
+  )
+  @mock.patch.object(process_prefetch.profiler, 'start_server')
+  @mock.patch.object(process_prefetch, 'ProcessPrefetchDatasetIterator')
+  def test_enable_profiling_options(
+      self,
+      mock_iterator,
+      mock_start_server,
+      mock_pick_port,
+      mock_is_loaded,
+      mock_enabled,
+      mock_supported,
+  ):
+    del mock_pick_port, mock_is_loaded, mock_enabled, mock_supported
+    mock_iter_instance = mock.create_autospec(dataset.DatasetIterator)
+    mock_iter_instance._ctx = base.IteratorContext()
+    mock_iterator.return_value = mock_iter_instance
+    mock_iter_instance.__next__.side_effect = StopIteration
+
+    ds = dataset.MapDataset.range(10).to_iter_dataset()
+    ds = ds.mp_prefetch(
+        options.MultiprocessingOptions(num_workers=1, enable_profiling=True)
+    )
+
+    results = list(ds)
+    self.assertEqual(results, [])
+
+    mock_iterator.assert_called_once()
+    called_worker_init_fn = mock_iterator.call_args[0][2]
+
+    self.assertIsNotNone(called_worker_init_fn)
+    called_worker_init_fn()
+    mock_start_server.assert_called_once_with(12345)
+
 
 if __name__ == '__main__':
   absltest.main()

grain/_src/python/options.py

@@ -95,8 +95,9 @@ class MultiprocessingOptions:
       each worker maintains. These are elements after all transformations. If
       your transformations include batching this means a single element is a
       batch.
-    enable_profiling: If True, profiling info is logged. This is only available
-      when num_workers >= 1.
+    enable_profiling: Whether to enable profiling in worker processes. If False,
+      it defaults to the global --grain_enable_multiprocess_worker_profiling
+      flag. This is only available when num_workers >= 1.
   """
 
   num_workers: int = 0