diff --git a/checkpoint/orbax/checkpoint/_src/multihost/multihost.py b/checkpoint/orbax/checkpoint/_src/multihost/multihost.py index 8d4b6f139..775a3df9f 100644 --- a/checkpoint/orbax/checkpoint/_src/multihost/multihost.py +++ b/checkpoint/orbax/checkpoint/_src/multihost/multihost.py @@ -423,9 +423,14 @@ def reached_preemption(step: int) -> bool: if is_proxy_pathways_backend(): return False - preemption_sync_point_reached = multihost_utils.reached_preemption_sync_point( - step - ) + try: + preemption_sync_point_reached = ( + multihost_utils.reached_preemption_sync_point(step) + ) + except RuntimeError: + # JAX raises RuntimeError when the preemption service is disabled + # via jax_enable_preemption_service config. Treat as no preemption. + return False _maybe_log_reached_preemption(step, preemption_sync_point_reached) return preemption_sync_point_reached diff --git a/checkpoint/orbax/checkpoint/_src/multihost/multihost_test.py b/checkpoint/orbax/checkpoint/_src/multihost/multihost_test.py index e4edccf50..dd745c574 100644 --- a/checkpoint/orbax/checkpoint/_src/multihost/multihost_test.py +++ b/checkpoint/orbax/checkpoint/_src/multihost/multihost_test.py @@ -244,5 +244,18 @@ def setUp(self): super().setUp() +class ReachedPreemptionTest(parameterized.TestCase): + + @mock.patch.object(multihost, 'is_proxy_pathways_backend', return_value=False) + @mock.patch.object(multihost.multihost_utils, 'reached_preemption_sync_point') + def test_returns_false_when_service_disabled( + self, mock_sync_point, mock_is_proxy + ): + mock_sync_point.side_effect = RuntimeError( + 'Preemption sync manager has not been initialized.' + ) + self.assertFalse(multihost.reached_preemption(step=5)) + + if __name__ == '__main__': multiprocess_test.main()