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/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",