diff --git a/CHANGELOG.md b/CHANGELOG.md index 63df64230..af8406262 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/grain/_src/core/profiler.py b/grain/_src/core/profiler.py index cd347e7d6..bbec43a77 100644 --- a/grain/_src/core/profiler.py +++ b/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: diff --git a/grain/_src/python/dataset/dataset.py b/grain/_src/python/dataset/dataset.py index 200ccfc54..3e7bcdb74 100644 --- a/grain/_src/python/dataset/dataset.py +++ b/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): diff --git a/grain/_src/python/dataset/transformations/process_prefetch.py b/grain/_src/python/dataset/transformations/process_prefetch.py index bce8a48b0..689ffccc9 100644 --- a/grain/_src/python/dataset/transformations/process_prefetch.py +++ b/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,6 +166,7 @@ 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})" @@ -170,9 +174,14 @@ def __str__(self) -> str: 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) diff --git a/grain/_src/python/dataset/transformations/process_prefetch_test.py b/grain/_src/python/dataset/transformations/process_prefetch_test.py index 610cb7eb2..2b72524fa 100644 --- a/grain/_src/python/dataset/transformations/process_prefetch_test.py +++ b/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() diff --git a/grain/_src/python/options.py b/grain/_src/python/options.py index f855aa932..9ee5702a1 100644 --- a/grain/_src/python/options.py +++ b/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