Add Pathways support and proper testing.

PiperOrigin-RevId: 914897467

Author: angel-core

tests/sft/checkpoint_manager_pathways_test.py

@@ -0,0 +1,88 @@
+"""Pathways tests for CheckpointManager."""
+
+import os
+from unittest import mock
+
+import jax
+import orbax.checkpoint as ocp
+from tunix.sft import checkpoint_manager
+from tunix.sft import checkpoint_manager_test
+
+from GOOGLE_INTERNAL_PACKAGE_PATH.pyglib.contrib.g3_multiprocessing import g3_multiprocessing
+from GOOGLE_INTERNAL_PACKAGE_PATH.testing.pybase import googletest
+
+
+_ORIG_REGISTER = ocp.pathways.register_type_handlers
+
+
+# Required for Pathways for testing.
+def wrapped_register(*args, **kwargs):
+  kwargs['thinmint_testing'] = True
+  return _ORIG_REGISTER(*args, **kwargs)
+
+
+class DirectPathwaysPersistenceTest(
+    checkpoint_manager_test.BaseCheckpointManagerTest
+):
+
+  def setUp(self):
+    super().setUp()
+    self.enterContext(
+        mock.patch.dict(os.environ, {'ENABLE_PATHWAYS_PERSISTENCE': '1'})
+    )
+    # Intercept calls in production code to ensure thinmint_testing stays True
+    self.mock_register = self.enterContext(
+        mock.patch.object(
+            ocp.pathways,
+            'register_type_handlers',
+            side_effect=wrapped_register,
+        )
+    )
+
+  def test_restore_different_sharding(self):
+    self.skipTest('Pathways only supports jax.sharding.NamedSharding.')
+
+  def test_register_type_handlers(self):
+    cp_path = f'{self.temp_path}/{self.id()}'
+    _ = checkpoint_manager.CheckpointManager(cp_path)
+
+    self.mock_register.assert_called_once()
+    _, kwargs = self.mock_register.call_args
+    impl = kwargs['checkpointing_impl']
+    self.assertEqual(
+        impl,
+        ocp.pathways.CheckpointingImpl.from_options(
+            use_remote_python=True
+        ),
+    )
+
+
+class DirectPathwaysNoPersistenceTest(
+    checkpoint_manager_test.BaseCheckpointManagerTest
+):
+
+  def setUp(self):
+    super().setUp()
+    self.enterContext(
+        mock.patch.dict(os.environ, {'ENABLE_PATHWAYS_PERSISTENCE': ''})
+    )
+    self.mock_register = self.enterContext(
+        mock.patch.object(
+            ocp.pathways,
+            'register_type_handlers',
+            side_effect=wrapped_register,
+        )
+    )
+
+  def test_restore_different_sharding(self):
+    self.skipTest('Pathways only supports jax.sharding.NamedSharding.')
+
+  def test_register_type_handlers(self):
+    cp_path = f'{self.temp_path}/{self.id()}'
+    _ = checkpoint_manager.CheckpointManager(cp_path)
+
+    self.mock_register.assert_not_called()
+
+if __name__ == '__main__':
+  jax.config.parse_flags_with_absl()
+  g3_multiprocessing.handle_test_main(googletest.main)

tests/sft/checkpoint_manager_test.py

@@ -101,6 +101,34 @@ def setUp(self):
         axis_names=('fsdp', 'tp'),
     )
 
+  def test_handlers_options(self):
+    """Verifies OCDBT/Zarr3 options match active platform configuration."""
+    cp_path = f'{self.temp_path}/{self.id()}'
+    cp_manager = checkpoint_manager.CheckpointManager(cp_path)
+
+    platforms = jax.config.jax_platforms or ''
+    is_pathways_or_proxy = 'proxy' in platforms or 'pathways' in platforms
+
+    handler = cp_manager._checkpoint_manager._checkpointer._handler  # pytype: disable=attribute-error
+    registry_entries = handler._handler_registry.get_all_entries()  # pytype: disable=attribute-error
+
+    handlers = {}
+    for (item_name, _), h in registry_entries.items():
+      handlers[item_name] = h
+
+    self.assertIn('model_params', handlers)
+    self.assertIn('optimizer_state', handlers)
+
+    if is_pathways_or_proxy:
+      self.assertFalse(handlers['model_params']._use_ocdbt)  # pytype: disable=attribute-error
+      self.assertFalse(handlers['optimizer_state']._use_ocdbt)  # pytype: disable=attribute-error
+    else:
+      self.assertTrue(handlers['model_params']._use_ocdbt)  # pytype: disable=attribute-error
+      self.assertTrue(handlers['optimizer_state']._use_ocdbt)  # pytype: disable=attribute-error
+
+    self.assertFalse(handlers['model_params']._use_zarr3)  # pytype: disable=attribute-error
+    self.assertFalse(handlers['optimizer_state']._use_zarr3)  # pytype: disable=attribute-error
+
   def test_empty_root_directory(self):
     cp_manager = checkpoint_manager.CheckpointManager(root_directory=None)
     self.assertIsNone(cp_manager.latest_step())
@@ -299,6 +327,12 @@ def test_restore_with_backward_compatibility(self, ckpt_path):
     # The checkpoints in test_data is saved with StandardSave. The test is to
     # verify the checkpoint manager with PyTreeRestore can still restore the
     # checkpoints saved with StandardSave.
+    if os.getenv('ENABLE_PATHWAYS_PERSISTENCE', '') == '1':
+      self.skipTest(
+          'Pathways persistence cannot read standard backwards-compatible'
+          ' checkpoints.'
+      )
+
     ckpt_manager = checkpoint_manager.CheckpointManager(
         os.path.join(os.path.dirname(__file__), ckpt_path)
     )

tunix/sft/checkpoint_manager.py

@@ -50,7 +50,11 @@ def __init__(
     if root_directory is not None:
       # When using Pathways, the checkpoint manager only supports persistence
       # APIs now.
-      if 'proxy' in os.getenv('JAX_PLATFORMS', ''):
+      platforms = jax.config.jax_platforms or ''
+      if (
+          'proxy' in platforms
+          or 'pathways' in platforms
+      ):
         item_handlers = {
             'model_params': ocp.PyTreeCheckpointHandler(
                 use_ocdbt=False,
@@ -65,6 +69,11 @@ def __init__(
           logging.info(
               'Using persistence API for checkpointing with Pathways.'
           )
+          ocp.pathways.register_type_handlers(
+              checkpointing_impl=ocp.pathways.CheckpointingImpl.from_options(
+                  use_remote_python=True,
+              )
+          )
         else:
           logging.warning(
               'Checkpointing without the persistence API, be aware of potential'