Internal

PiperOrigin-RevId: 838933101

Author: Grain Team

grain/_src/python/BUILD

@@ -208,49 +208,6 @@ py_test(
     ],
 )
 
-py_library(
-    name = "grain_pool",
-    srcs = ["grain_pool.py"],
-    srcs_version = "PY3",
-    target_compatible_with = select({
-        "@platforms//os:windows": ["@platforms//:incompatible"],
-        "//conditions:default": [],
-    }),
-    deps = [
-        ":grain_logging",
-        ":multiprocessing_common",
-        ":options",
-        ":record",
-        ":shared_memory_array",
-        "//grain/_src/core:config",
-        "//grain/_src/core:monitoring",
-        "//grain/_src/core:parallel",
-        "//grain/_src/core:tree_lib",
-        "@abseil-py//absl/flags",
-        "@abseil-py//absl/logging",
-        "@pypi//cloudpickle:pkg",
-    ],
-)
-
-py_test(
-    name = "grain_pool_test",
-    srcs = ["grain_pool_test.py"],
-    shard_count = 20,
-    srcs_version = "PY3",
-    tags = ["not_run:arm"],
-    deps = [
-        ":data_sources",
-        ":grain_pool",
-        ":options",
-        ":record",
-        "//grain/_src/core:config",
-        "//grain/_src/core:monitoring",
-        "@abseil-py//absl/flags",
-        "@abseil-py//absl/testing:absltest",
-        "@abseil-py//absl/testing:parameterized",
-    ],
-)
-
 py_library(
     name = "checkpoint_handlers",
     srcs = ["checkpoint_handlers.py"],

grain/_src/python/data_loader.py

@@ -39,7 +39,6 @@
 from grain._src.python.dataset import dataset
 from grain._src.python.dataset.transformations import batch as batch_ds
 from grain._src.python.dataset.transformations import flatmap
-from grain._src.python.dataset.transformations import prefetch
 from grain._src.python.operations import BatchOperation
 from grain._src.python.operations import Operation
 from grain._src.python.samplers import Sampler
@@ -462,10 +461,8 @@ def _create_dataset(self) -> dataset.IterDataset:
       ds = _apply_transform_to_dataset(operation, ds)
     ds = ds.map(lambda r: r.data)
     if self.multiprocessing_options.num_workers > 0:
-      ds = prefetch.MultiprocessPrefetchIterDataset(
-          ds,
+      ds = ds.mp_prefetch(
           self.multiprocessing_options,
-          always_report_worker_state=True,
       )
     if not self._use_native_dataset_checkpointing:
       ds = _DataLoaderStateIterDataset(

grain/_src/python/dataset/BUILD

@@ -53,7 +53,6 @@ py_library(
         "//grain/_src/core:tree_lib",
         "//grain/_src/python:checkpointing",
         "//grain/_src/python:grain_logging",
-        "//grain/_src/python:grain_pool",
         "//grain/_src/python:options",
         "//grain/_src/python:shared_memory_array",
         "//grain/proto:execution_summary_py_pb2",

grain/_src/python/dataset/dataset.py

@@ -1324,13 +1324,14 @@ def mp_prefetch(
       A dataset prefetching input elements in separate processes.
     """
     options = options or grain_options.MultiprocessingOptions(num_workers=10)
-    # Loaded lazily due to a circular dependency (dataset <-> prefetch).
+    # Loaded lazily due to a circular dependency (dataset <-> process_prefetch).
     # pylint: disable=g-import-not-at-top
-    from grain._src.python.dataset.transformations import prefetch
+    from grain._src.python.dataset.transformations import process_prefetch
     # pylint: enable=g-import-not-at-top
-    return prefetch.MultiprocessPrefetchIterDataset(
+    return process_prefetch.multiprocess_prefetch(
         self,
-        multiprocessing_options=options,
+        num_workers=options.num_workers,
+        buffer_size=options.per_worker_buffer_size,
         worker_init_fn=worker_init_fn,
         sequential_slice=sequential_slice,
     )

grain/_src/python/dataset/transformations/BUILD

@@ -53,7 +53,7 @@ py_test(
 
 py_test(
     name = "prefetch_test",
-    timeout = "long",
+    timeout = "eternal",
     srcs = ["prefetch_test.py"],
     shard_count = 50,
     srcs_version = "PY3",