From a5823cdb3f58d9156e2e08a13e8d8395b4c36ca7 Mon Sep 17 00:00:00 2001
From: Sadi Kneipp <ksadi@google.com>
Date: Fri, 18 Sep 2026 17:45:48 -0700
Subject: [PATCH] Batch array persistence writes in
 CloudPathwaysArrayHandler.serialize.

PiperOrigin-RevId: 984158716
---
 pathwaysutils/persistence/helper.py        | 13 --------
 pathwaysutils/persistence/orbax_handler.py | 37 +++++++++++++++-------
 2 files changed, 26 insertions(+), 24 deletions(-)

diff --git a/pathwaysutils/persistence/helper.py b/pathwaysutils/persistence/helper.py
index 915e669..d2bd745 100644
--- a/pathwaysutils/persistence/helper.py
+++ b/pathwaysutils/persistence/helper.py
@@ -225,19 +225,6 @@ def get_bulk_read_request(
   )
 
 
-def write_one_array(
-    location: str,
-    name: str,
-    value: jax.Array,
-    timeout: datetime.timedelta,
-):
-  """Creates the write array plugin program string, compiles it to an executable, calls it and returns an awaitable future."""
-  write_request = get_write_request(location, name, value, timeout)
-  write_executable = plugin_executable.PluginExecutable(write_request)  # pyrefly: ignore[bad-argument-type]
-  _, write_future = write_executable.call([value])
-  return write_future
-
-
 def write_arrays(
     location: str,
     names: Sequence[str],
diff --git a/pathwaysutils/persistence/orbax_handler.py b/pathwaysutils/persistence/orbax_handler.py
index 8b44f30..7d4deb4 100644
--- a/pathwaysutils/persistence/orbax_handler.py
+++ b/pathwaysutils/persistence/orbax_handler.py
@@ -17,7 +17,6 @@
 from collections.abc import Coroutine, Sequence
 import concurrent.futures
 import datetime
-import functools
 import logging
 from typing import Any, cast
 
@@ -144,8 +143,24 @@ async def serialize(
 
     self._wait_for_directory_creation_signals()
     locations, names = extract_parent_dir_and_name(infos)
-    f = functools.partial(helper.write_one_array, timeout=self.timeout)
-    futures_results = list(map(f, locations, names, arrays))
+    # Group arrays by parent directory and device assignment so each batch
+    # satisfies SideChannelLoadedExecutable bulk persistence constraints.
+    grouped_writes: dict[
+        tuple[str, tuple[Any, ...]], tuple[list[str], list[jax.Array]]
+    ] = collections.defaultdict(lambda: ([], []))
+    for loc, name, arr in zip(locations, names, arrays):
+      # pylint:disable=protected-access
+      key = (loc, tuple(arr.sharding._device_assignment))
+      # pylint:enable=protected-access
+      grouped_writes[key][0].append(name)
+      grouped_writes[key][1].append(arr)
+
+    futures_results = [
+        helper.write_arrays(
+            loc, group_names, group_arrays, timeout=self.timeout
+        )
+        for (loc, _), (group_names, group_arrays) in grouped_writes.items()
+    ]
 
     return [
         future.CommitFutureAwaitingContractedSignals(
@@ -229,20 +244,20 @@ async def deserialize(
             for array_metadata in array_metadatas
         }
 
-    # Group inputs by global_mesh so that we can perform batched Array
-    # construction for each global_mesh.
-    inputs_by_global_mesh = collections.defaultdict(list)
-    for i, global_mesh in enumerate(global_meshes):
-      inputs_by_global_mesh[global_mesh].append(i)
+    # Group inputs by parent_dir and global_mesh so that we can perform batched
+    # Array construction for each group.
+    inputs_by_location_and_mesh = collections.defaultdict(list)
+    for i, (info, global_mesh) in enumerate(zip(infos, global_meshes)):
+      inputs_by_location_and_mesh[(str(info.parent_dir), global_mesh)].append(i)
 
     results = cast(list[jax.Array], [None] * len(infos))
 
-    for global_mesh, idxs in inputs_by_global_mesh.items():
+    for (location, global_mesh), idxs in inputs_by_location_and_mesh.items():
       grouped_infos = [infos[idx] for idx in idxs]
       grouped_global_shapes = [global_shapes[idx] for idx in idxs]
       grouped_dtypes = [dtypes[idx] for idx in idxs]
       grouped_shardings = [shardings[idx] for idx in idxs]
-      locations, names = extract_parent_dir_and_name(grouped_infos)
+      _, names = extract_parent_dir_and_name(grouped_infos)
 
       grouped_read_dtypes = []
       grouped_read_shapes = []
@@ -282,7 +297,7 @@ async def deserialize(
         grouped_read_shardings.append(read_sharding)
 
       grouped_arrays, read_future = helper.read_arrays(
-          locations[0],
+          location,
           names,
           grouped_read_dtypes,
           grouped_read_shapes,