From a861e68e56c52353639043df026aa2be84a6f204 Mon Sep 17 00:00:00 2001 From: Pathways-on-Cloud Team Date: Sat, 18 Jul 2026 01:24:28 -0700 Subject: [PATCH] Batch Pathways persistence writes in Orbax handler. PiperOrigin-RevId: 949982430 --- pathwaysutils/persistence/orbax_handler.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/pathwaysutils/persistence/orbax_handler.py b/pathwaysutils/persistence/orbax_handler.py index 12a43e3..fedb22b 100644 --- a/pathwaysutils/persistence/orbax_handler.py +++ b/pathwaysutils/persistence/orbax_handler.py @@ -144,8 +144,21 @@ 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 by location (parent_dir) to batch writes. + by_location = collections.defaultdict(list) + for loc, name, arr in zip(locations, names, arrays): + by_location[loc].append((name, arr)) + + futures_results = [] + for loc, items in by_location.items(): + grouped_names = [item[0] for item in items] + grouped_arrays = [item[1] for item in items] + futures_results.append( + helper.write_arrays( + loc, grouped_names, grouped_arrays, timeout=self.timeout + ) + ) return [ future.CommitFutureAwaitingContractedSignals(