Set max_slice_restarts to 1000000 in Shared Pathways Service deployment

This change configures the max_slice_restarts parameter to 1000000 when deploying the shared Pathways service. Using a very high backOffLimit for the worker pods keeps the jobset from failing on multiple client reconnections.

FUTURE_COPYBARA_INTEGRATE_REVIEW=#288 from AI-Hypercomputer:add-gke-test-extras-dependencies 672c861
PiperOrigin-RevId: 952331464

Author: guptaaka

pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py

@@ -4,7 +4,6 @@
 import dataclasses
 import logging
 import math
-import os
 from typing import Any
 from absl import app
 from absl import flags
@@ -142,7 +141,6 @@ def calculate_vms_per_slice(topology: str, chips_per_vm: int) -> int:
     ) from e
 
 
-
 def deploy_jobset(jobset_yaml: dict[str, Any]) -> None:
   """Deploys the JobSet to the current Kubernetes cluster."""
   try:
@@ -185,6 +183,9 @@ def run_deployment(
       topology=topology,
       num_slices=num_slices,
       shared_pathways_service=True,
+      # TODO(b/496958026): Remove this once go/sps-worker-pod-stability is
+      # implemented
+      max_slice_restarts=1000000,
   )
 
   # If custom server_image is provided, mutate the templates to use it.

pathwaysutils/proxy_backend.py

@@ -13,17 +13,27 @@
 # limitations under the License.
 """Register the IFRT Proxy as a backend for JAX."""
 
+import os
 import jax
 from jax.extend import backend
 from jax.extend.backend import ifrt_proxy
 
 
 def register_backend_factory() -> None:
+  """Registers the IFRT Proxy backend factory with JAX."""
+
+  def make_client():
+    options = ifrt_proxy.ClientConnectionOptions()
+    timeout_secs = os.environ.get("PATHWAYS_PROXY_CONNECTION_TIMEOUT_SECS")
+    if timeout_secs:
+      options.connection_timeout_in_seconds = int(timeout_secs)
+    return ifrt_proxy.get_client(
+        jax.config.read("jax_backend_target"),
+        options,
+    )
+
   backend.register_backend_factory(
       "proxy",
-      lambda: ifrt_proxy.get_client(
-          jax.config.read("jax_backend_target"),
-          ifrt_proxy.ClientConnectionOptions(),
-      ),
+      make_client,
       priority=-1,
   )

pathwaysutils/test/proxy_backend_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Tests for the proxy backend module."""
 
+import os
 from unittest import mock
 
 from absl.testing import absltest
@@ -54,6 +55,46 @@ def test_proxy_backend_registration(self):
     proxy_backend.register_backend_factory()
     self.assertIn("proxy", backend.backends())
 
+  def test_proxy_backend_registration_with_timeout(self):
+    mock_get_client = self.enter_context(
+        mock.patch.object(
+            ifrt_proxy,
+            "get_client",
+            return_value=mock.MagicMock(),
+        )
+    )
+    self.enter_context(
+        mock.patch.dict(
+            os.environ, {"PATHWAYS_PROXY_CONNECTION_TIMEOUT_SECS": "42"}
+        )
+    )
+    proxy_backend.register_backend_factory()
+    self.assertIn("proxy", backend.backends())
+    mock_get_client.assert_called_once()
+    args, _ = mock_get_client.call_args
+    self.assertEqual(args[0], "grpc://localhost:12345")
+    options = args[1]
+    self.assertEqual(options.connection_timeout_in_seconds, 42)
+
+  def test_proxy_backend_registration_without_timeout(self):
+    mock_get_client = self.enter_context(
+        mock.patch.object(
+            ifrt_proxy,
+            "get_client",
+            return_value=mock.MagicMock(),
+        )
+    )
+    self.enter_context(mock.patch.dict(os.environ))
+    os.environ.pop("PATHWAYS_PROXY_CONNECTION_TIMEOUT_SECS", None)
+
+    proxy_backend.register_backend_factory()
+    self.assertIn("proxy", backend.backends())
+    mock_get_client.assert_called_once()
+    args, _ = mock_get_client.call_args
+    self.assertEqual(args[0], "grpc://localhost:12345")
+    options = args[1]
+    self.assertNotEqual(options.connection_timeout_in_seconds, 42)
+
 
 if __name__ == "__main__":
   absltest.main()