Fix Pathways colocated typed-key transport and post-restore reconstruction.

checkpoint/orbax/checkpoint/_src/multihost/colocated_transport.py

@@ -14,7 +14,7 @@
 
 """Helpers for transporting values through colocated Python."""
 
-from collections.abc import Mapping, Sequence
+from collections.abc import Callable, Mapping, Sequence
 import dataclasses
 import functools
 import re
@@ -24,16 +24,33 @@
 from absl import logging
 import jax
 import jax.experimental.colocated_python as cp
+from jax.experimental.colocated_python import func as cp_func
 from jax.experimental.colocated_python import serialization as cp_serialization
 import numpy as np
 from orbax.checkpoint._src.arrays import abstract_arrays
 from orbax.checkpoint._src.metadata import sharding as sharding_metadata
 from orbax.checkpoint._src.serialization import jax_array_restore_args
 
-
 PyTree = Any
 _PATHWAYS_SERIALIZATION_PATCH_INSTALLED = False
 _PJRT_IFRT_DEVICE_ID_RE = re.compile(r'PjRtIFRTDeviceId=(\d+)')
+_original_make_prng_wrapped_fun = getattr(
+    cp_func, '_make_prng_wrapped_fun', None
+)
+
+
+def _make_prng_wrapped_fun_without_copied_lock(
+    fun: Callable[..., Any],
+    *args: Any,
+    **kwargs: Any,
+) -> Callable[..., Any]:
+  """Calls JAX's PRNG wrapper without retaining its redundant lock copy."""
+  assert _original_make_prng_wrapped_fun is not None
+  wrapped_fun = _original_make_prng_wrapped_fun(fun, *args, **kwargs)
+  lock = getattr(fun, '_lock', None)
+  if lock is not None and wrapped_fun.__dict__.get('_lock') is lock:
+    del wrapped_fun.__dict__['_lock']
+  return wrapped_fun
 
 
 def _to_serializable_cpu_device(device: jax.Device) -> jax.Device:
@@ -271,6 +288,13 @@ def _orbax_reduce_single_device_sharding(
   cp_serialization._reduce_device_list = _orbax_reduce_device_list  # pylint: disable=protected-access
   cp_serialization._reduce_single_device_sharding = _orbax_reduce_single_device_sharding  # pylint: disable=protected-access
   cp_serialization._get_cpu_device_map = _get_cpu_device_map  # pylint: disable=protected-access
+  if _original_make_prng_wrapped_fun is not None:
+    # JAX's PRNG wrapper copies MethodCallerAtBackend._lock onto an internal
+    # function even though the callable already defines its own pickle state.
+    cp_func._make_prng_wrapped_fun = (  # pylint: disable=protected-access
+        _make_prng_wrapped_fun_without_copied_lock
+    )
+
   _PATHWAYS_SERIALIZATION_PATCH_INSTALLED = True
 
 
@@ -459,7 +483,17 @@ def _to_final_spec(leaf: Any, tpu_or_cpu_spec: Any) -> Any:
 
 
 def shape_dtype_struct_for_array(array: jax.Array) -> jax.ShapeDtypeStruct:
-  """Builds a ShapeDtypeStruct from a jax.Array."""
+  """Builds a ShapeDtypeStruct compatible with colocated Python."""
+  # JAX versions with typed-PRNG transport need the logical key dtype so they
+  # can convert it to the physical uint32 representation before IFRT dispatch.
+  if _original_make_prng_wrapped_fun is not None and jax.dtypes.issubdtype(
+      array.dtype, jax.dtypes.prng_key
+  ):
+    return jax.ShapeDtypeStruct(
+        shape=array.shape,
+        dtype=array.dtype,
+        sharding=array.sharding,
+    )
   return cast(
       jax.ShapeDtypeStruct, abstract_arrays.to_shape_dtype_struct(array)  # pyrefly: ignore[bad-argument-type]
   )

checkpoint/orbax/checkpoint/_src/multihost/colocated_transport_test.py

@@ -14,6 +14,8 @@
 
 # pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring,too-few-public-methods
 
+import pickle
+import threading
 from unittest import mock
 
 from absl.testing import absltest
@@ -27,6 +29,18 @@
 from orbax.checkpoint._src.serialization import type_handlers
 
 
+class _LockingCallable:
+
+  def __init__(self):
+    self._lock = threading.Lock()
+
+  def __call__(self, value):
+    return value
+
+  def __reduce__(self):
+    return type(self), ()
+
+
 class ColocatedTransportTest(absltest.TestCase):
 
   def setUp(self):
@@ -124,6 +138,74 @@ def test_shape_dtype_struct_for_array_preserves_array_metadata(self):
     self.assertEqual(spec.dtype, self.arr.dtype)
     self.assertIs(spec.sharding, self.arr.sharding)
 
+  def test_shape_dtype_struct_preserves_prng_dtype_when_supported(self):
+    key = jax.random.key(0)
+    with mock.patch.object(
+        colocated_transport,
+        '_original_make_prng_wrapped_fun',
+        mock.sentinel.make_prng_wrapped_fun,
+    ):
+      spec = colocated_transport.shape_dtype_struct_for_array(key)
+
+    self.assertEqual(spec.shape, key.shape)
+    self.assertEqual(spec.dtype, key.dtype)
+    self.assertIs(spec.sharding, key.sharding)
+
+  def test_shape_dtype_struct_uses_legacy_prng_representation(self):
+    key = jax.random.key(0)
+    key_data = jax.random.key_data(key)
+    with mock.patch.object(
+        colocated_transport, '_original_make_prng_wrapped_fun', None
+    ):
+      spec = colocated_transport.shape_dtype_struct_for_array(key)
+
+    self.assertEqual(spec.shape, key_data.shape)
+    self.assertEqual(spec.dtype, key_data.dtype)
+
+  def test_prng_wrapper_drops_only_redundant_lock_copy(self):
+
+    def make_wrapper(fun, unused_in_info, unused_out_info):
+      def wrapped(*args, **kwargs):
+        return fun(*args, **kwargs)
+
+      wrapped.__dict__.update(fun.__dict__)
+      setattr(wrapped, '__wrapped__', fun)
+      return wrapped
+
+    fun = _LockingCallable()
+    with mock.patch.object(
+        colocated_transport,
+        '_original_make_prng_wrapped_fun',
+        make_wrapper,
+    ):
+      make_prng_wrapper = getattr(
+          colocated_transport,
+          '_make_prng_wrapped_fun_without_copied_lock',
+      )
+      wrapped = make_prng_wrapper(fun, {}, {})
+
+    self.assertNotIn('_lock', wrapped.__dict__)
+    self.assertIs(getattr(wrapped, '__wrapped__'), fun)
+    serialize = getattr(cp_serialization, '_serialize')
+    deserialize = getattr(cp_serialization, '_deserialize')
+    restored = deserialize(serialize(wrapped))
+    self.assertEqual(restored(7), 7)
+
+  def test_jax_prng_wrapper_serializes_locking_callable(self):
+    original_wrapper = getattr(
+        colocated_transport, '_original_make_prng_wrapped_fun'
+    )
+    if original_wrapper is None:
+      self.skipTest('Installed JAX has no colocated PRNG wrapper.')
+
+    make_prng_wrapper = getattr(
+        colocated_transport,
+        '_make_prng_wrapped_fun_without_copied_lock',
+    )
+    wrapped = make_prng_wrapper(_LockingCallable(), {}, {})
+
+    getattr(cp_serialization, '_serialize')(wrapped)
+
   def test_zeros_like_spec(self):
     spec = jax.ShapeDtypeStruct(
         (4,),
@@ -312,6 +394,8 @@ def test_install_pathways_colocated_serialization_patch_is_idempotent(self):
         patched_get_cpu_device_map,
         cp_serialization._get_cpu_device_map,  # pylint: disable=protected-access
     )
+    with self.assertRaises(TypeError):
+      pickle.dumps(threading.Lock())
 
 
 if __name__ == '__main__':

checkpoint/orbax/checkpoint/experimental/emergency/multi_tier_checkpointing/colocated_controller.py

@@ -584,6 +584,40 @@ def _validate_restored_shape_dtype(
       )
 
 
+def _restore_typed_prng_keys(
+    restored_state: PyTree,
+    template_state: PyTree,
+) -> PyTree:
+  """Rewraps physical PRNG key data according to the caller template."""
+  restored_leaves, restored_treedef = jax.tree.flatten(restored_state)
+  template_leaves = jax.tree.leaves(
+      template_state, is_leaf=_is_restore_spec_leaf
+  )
+  if len(restored_leaves) != len(template_leaves):
+    return restored_state
+
+  for index, template_leaf in enumerate(template_leaves):
+    restored_leaf = restored_leaves[index]
+    expected = _expected_leaf_from_template(template_leaf)
+    if (
+        expected is None
+        or expected.dtype is None
+        or not isinstance(restored_leaf, jax.Array)
+        or not jax.dtypes.issubdtype(expected.dtype, jax.dtypes.prng_key)
+        or jax.dtypes.issubdtype(restored_leaf.dtype, jax.dtypes.prng_key)
+    ):
+      continue
+
+    # `wrap_key_data(dtype=...)` is newer than the JAX versions supported by
+    # this experimental path. KeyTy._impl is accepted by both old and new JAX.
+    restored_leaves[index] = jax.random.wrap_key_data(
+        restored_leaf,
+        impl=getattr(expected.dtype, '_impl'),
+    )
+
+  return jax.tree.unflatten(restored_treedef, restored_leaves)
+
+
 def _validate_restore_template_compatibility(
     *,
     item: PyTree,
@@ -1073,6 +1107,7 @@ def restore(
         else state_restore_args.restore_args
     )
     state = self._rebuild_restored_state(result, rebuild_template)
+    state = _restore_typed_prng_keys(state, rebuild_template)
     _validate_restored_shape_dtype(
         restored_state=state, template_state=rebuild_template
     )

checkpoint/orbax/checkpoint/experimental/emergency/multi_tier_checkpointing/colocated_controller_test.py

@@ -1853,6 +1853,43 @@ def test_restore_rejects_protocol_no_checkpoint_sentinel(self):
 
     controller._worker_manager.restore_infer.specialize.assert_not_called()
 
+  def test_restore_rewraps_physical_prng_key_data(self):
+    controller, sharding = self._make_controller_for_restore()
+    key = jax.random.key(7, impl='unsafe_rbg')
+    key_data = jax.random.key_data(key)
+    restored_cpu_state = {
+        'key': jax.device_put(key_data, sharding),
+    }
+    template_state = {
+        'key': jax.ShapeDtypeStruct((), key.dtype, sharding=sharding),
+    }
+    restore_args = {
+        'key': type_handlers.ArrayRestoreArgs(
+            sharding=sharding,
+            global_shape=(),
+            dtype=key.dtype,
+        ),
+    }
+    self._set_worker_restore_result(controller, return_value=restored_cpu_state)
+
+    result = controller.restore(
+        7,
+        args_lib.Composite(
+            state=args_lib.PyTreeRestore(
+                item=template_state,
+                restore_args=restore_args,
+            )
+        ),
+        default_item_mode=True,
+    )
+
+    self.assertEqual(result['key'].shape, key.shape)
+    self.assertEqual(result['key'].dtype, key.dtype)
+    np.testing.assert_array_equal(
+        np.asarray(jax.random.key_data(result['key'])),
+        np.asarray(key_data),
+    )
+
   def test_restore_rejects_shape_mismatch_before_final_remap(self):
     controller, sharding = self._make_controller_for_restore()
     mesh = sharding.mesh