No public description

PiperOrigin-RevId: 889041970

Author: SujeethJinesh

checkpoint/orbax/checkpoint/_src/futures/signaling_client.py

@@ -17,13 +17,45 @@
 import functools
 import threading
 import time
-from typing import Sequence
+from typing import Optional, Sequence
 from absl import logging
 import jax
 from orbax.checkpoint._src.multihost import multihost
 from typing_extensions import Protocol
 
 
+_FORCE_THREADSAFE_SIGNALING_CLIENT = False
+_THREADSAFE_SIGNALING_OVERRIDE_REFCOUNT = 0
+_OVERRIDE_LOCK = threading.RLock()
+
+
+def set_force_threadsafe_signaling_client(force: bool) -> None:
+  """Overrides signaling-client selection for the current process."""
+  global _FORCE_THREADSAFE_SIGNALING_CLIENT
+  with _OVERRIDE_LOCK:
+    _FORCE_THREADSAFE_SIGNALING_CLIENT = force
+    get_signaling_client.cache_clear()
+
+
+def acquire_threadsafe_signaling_client_override() -> None:
+  """Enables the threadsafe signaling override for the current process."""
+  global _THREADSAFE_SIGNALING_OVERRIDE_REFCOUNT
+  with _OVERRIDE_LOCK:
+    _THREADSAFE_SIGNALING_OVERRIDE_REFCOUNT += 1
+    set_force_threadsafe_signaling_client(True)
+
+
+def release_threadsafe_signaling_client_override() -> None:
+  """Releases one threadsafe signaling override reference for this process."""
+  global _THREADSAFE_SIGNALING_OVERRIDE_REFCOUNT
+  with _OVERRIDE_LOCK:
+    _THREADSAFE_SIGNALING_OVERRIDE_REFCOUNT = max(
+        0, _THREADSAFE_SIGNALING_OVERRIDE_REFCOUNT - 1
+    )
+    if _THREADSAFE_SIGNALING_OVERRIDE_REFCOUNT == 0:
+      set_force_threadsafe_signaling_client(False)
+
+
 class SignalingClient(Protocol):
   """Client that supports signaling between threads and processes."""
 
@@ -50,7 +82,7 @@ def blocking_key_value_get(self, key: str, timeout_secs: int) -> str:
     """
     ...
 
-  def key_value_try_get(self, key: str) -> str | None:
+  def key_value_try_get(self, key: str) -> Optional[str]:
     """Tries to get the value for a given key in the client without blocking.
 
     Args:
@@ -82,7 +114,7 @@ def wait_at_barrier(
       key: str,
       *,
       timeout_secs: int,
-      process_ids: Sequence[int] | None = None,
+      process_ids: Optional[Sequence[int]] = None,
   ):
     """Waits at a barrier identified by key.
 
@@ -136,7 +168,7 @@ def blocking_key_value_get(self, key: str, timeout_secs: int) -> str:
     """
     return str(self._client.blocking_key_value_get(key, timeout_secs * 1000))
 
-  def key_value_try_get(self, key: str) -> str | None:
+  def key_value_try_get(self, key: str) -> Optional[str]:
     """Tries to get the value for a given key in the client without blocking.
 
     Args:
@@ -187,7 +219,7 @@ def wait_at_barrier(
       key: str,
       *,
       timeout_secs: int,
-      process_ids: Sequence[int] | None = None,
+      process_ids: Optional[Sequence[int]] = None,
   ):
     """Waits at a barrier identified by key.
 
@@ -284,7 +316,7 @@ def blocking_key_value_get(self, key: str, timeout_secs: int) -> str:
 
       return self._data[key]
 
-  def key_value_try_get(self, key: str) -> str | None:
+  def key_value_try_get(self, key: str) -> Optional[str]:
     """Tries to get the value for a key without blocking.
 
     Args:
@@ -342,7 +374,7 @@ def wait_at_barrier(
       key: str,
       *,
       timeout_secs: int,
-      process_ids: Sequence[int] | None = None,
+      process_ids: Optional[Sequence[int]] = None,
   ):
     """Waits at a barrier identified by key.
 
@@ -363,12 +395,24 @@ def get_signaling_client() -> SignalingClient:
   if multihost.is_jax_distributed_client_initialized():
     logging.info("Using JaxDistributedSignalingClient")
     return JaxDistributedSignalingClient()
-  else:
-    process_count = multihost.process_count()
-    if process_count > 1:
-      raise RuntimeError(
-          "ThreadSafeKeyValueSignalingClient should only be used in a single"
-          f" controller setup, process count: {process_count}."
-      )
+
+  process_count = multihost.process_count()
+  if process_count == 1:
     logging.info("Using ThreadSafeKeyValueSignalingClient")
     return ThreadSafeKeyValueSignalingClient()
+
+  # Pathways sidecars run with process_count > 1 but do not initialize the JAX
+  # distributed client. Allow the in-process client only when that sidecar path
+  # explicitly opts in via `set_force_threadsafe_signaling_client(True)`.
+  if _FORCE_THREADSAFE_SIGNALING_CLIENT:
+    logging.warning(
+        "Using ThreadSafeKeyValueSignalingClient with process_count=%s because"
+        " force override is enabled and JAX distributed client is unavailable.",
+        process_count,
+    )
+    return ThreadSafeKeyValueSignalingClient()
+
+  raise RuntimeError(
+      "ThreadSafeKeyValueSignalingClient should only be used in a single"
+      f" controller setup, process count: {process_count}."
+  )

checkpoint/orbax/checkpoint/_src/futures/signaling_client_test.py

@@ -285,9 +285,15 @@ class TestGetSignalingClient(absltest.TestCase):
 
   def setUp(self):
     super().setUp()
+    signaling_client.set_force_threadsafe_signaling_client(False)
     signaling_client.get_signaling_client.cache_clear()
 
   def tearDown(self):
+    signaling_client.set_force_threadsafe_signaling_client(False)
+    while getattr(
+        signaling_client, "_THREADSAFE_SIGNALING_OVERRIDE_REFCOUNT", 0
+    ):
+      signaling_client.release_threadsafe_signaling_client_override()
     super().tearDown()
     signaling_client.get_signaling_client.cache_clear()
 
@@ -344,6 +350,57 @@ def test_raises_error_when_multiprocess_and_not_initialized(
     mock_is_init.assert_called_once()
     mock_process_count.assert_called_once()
 
+  @mock.patch.object(
+      multihost, "is_jax_distributed_client_initialized", return_value=False
+  )
+  @mock.patch.object(multihost, "process_count", return_value=2)
+  def test_force_override_allows_threadsafe_client_in_multiprocess(
+      self, mock_is_init, mock_process_count
+  ):
+    signaling_client.set_force_threadsafe_signaling_client(True)
+
+    client = signaling_client.get_signaling_client()
+
+    self.assertIsInstance(
+        client, signaling_client.ThreadSafeKeyValueSignalingClient
+    )
+    mock_is_init.assert_called_once()
+    mock_process_count.assert_called_once()
+
+  @mock.patch.object(
+      multihost, "is_jax_distributed_client_initialized", return_value=False
+  )
+  @mock.patch.object(multihost, "process_count", return_value=2)
+  def test_acquire_release_override_is_reference_counted(
+      self, mock_is_init, mock_process_count
+  ):
+    signaling_client.acquire_threadsafe_signaling_client_override()
+    signaling_client.acquire_threadsafe_signaling_client_override()
+    self.addCleanup(
+        lambda: setattr(
+            signaling_client, "_THREADSAFE_SIGNALING_OVERRIDE_REFCOUNT", 0
+        )
+    )
+
+    client = signaling_client.get_signaling_client()
+    self.assertIsInstance(
+        client, signaling_client.ThreadSafeKeyValueSignalingClient
+    )
+
+    signaling_client.release_threadsafe_signaling_client_override()
+    signaling_client.get_signaling_client.cache_clear()
+    client = signaling_client.get_signaling_client()
+    self.assertIsInstance(
+        client, signaling_client.ThreadSafeKeyValueSignalingClient
+    )
+
+    signaling_client.release_threadsafe_signaling_client_override()
+    signaling_client.get_signaling_client.cache_clear()
+    with self.assertRaisesRegex(RuntimeError, "process count: 2"):
+      signaling_client.get_signaling_client()
+    mock_is_init.assert_called()
+    mock_process_count.assert_called()
+
 
 if __name__ == "__main__":
   absltest.main()