Introduces multiprocess_prefetch as a simpler, more efficient version of the existing MultiprocessPrefetchIterDataset that supports setting state without restarting workers.

PiperOrigin-RevId: 832523331

Author: Grain Team

CHANGELOG.md

@@ -13,6 +13,9 @@ 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.
+  * Introduces `grain.experimental.multiprocess_prefetch` as an simpler version
+    of `IterDataset.mp_prefetch` that supports setting state without restarting
+    workers.
 
 * Breaking changes:
 

grain/BUILD

@@ -64,6 +64,7 @@ py_library(
         "//grain/_src/python/dataset/transformations:cache",
         "//grain/_src/python/dataset/transformations:interleave",
         "//grain/_src/python/dataset/transformations:prefetch_autotune",
+        "//grain/_src/python/dataset/transformations:process_prefetch",
         "//grain/_src/python/dataset/transformations:limit",
         "//grain/_src/python/dataset/transformations:packing",
         "//grain/_src/python/dataset/transformations:packing_concat_then_split",

grain/_src/python/BUILD

@@ -338,6 +338,7 @@ py_library(
     srcs = ["shared_memory_array.py"],
     srcs_version = "PY3",
     deps = [
+        "//grain/_src/core:tree_lib",
         "@pypi//numpy:pkg",
     ],
 )

grain/_src/python/dataset/transformations/BUILD

@@ -360,3 +360,40 @@ py_test(
         "@pypi//numpy:pkg",
     ],
 )
+
+py_library(
+    name = "process_prefetch",
+    srcs = ["process_prefetch.py"],
+    srcs_version = "PY3",
+    deps = [
+        "//grain/_src/core:config",
+        "//grain/_src/core:tree_lib",
+        "//grain/_src/python:shared_memory_array",
+        "//grain/_src/python/dataset",
+        "//grain/_src/python/dataset:base",
+        "//grain/_src/python/dataset:stats",
+        "//grain/_src/python/dataset/transformations:interleave",
+        "@abseil-py//absl/flags",
+        "@pypi//cloudpickle:pkg",
+        "@pypi//numpy:pkg",
+    ],
+)
+
+py_test(
+    name = "process_prefetch_test",
+    timeout = "long",
+    srcs = ["process_prefetch_test.py"],
+    shard_count = 50,
+    srcs_version = "PY3",
+    deps = [
+        ":process_prefetch",
+        "//grain/_src/core:transforms",
+        "//grain/_src/python:options",
+        "//grain/_src/python/dataset",
+        "//grain/_src/python/dataset:base",
+        "@abseil-py//absl/logging",
+        "@abseil-py//absl/testing:absltest",
+        "@abseil-py//absl/testing:parameterized",
+        "@pypi//numpy:pkg",
+    ],
+)