Enable quick restart for Shared Pathways Service workers

pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py

@@ -61,6 +61,15 @@
     " Jobs have an effectively unlimited backoff limit.",
     lower_bound=0,
 )
+_ENABLE_WORKER_QUICK_RESTART = flags.DEFINE_boolean(
+    "enable_worker_quick_restart",
+    True,
+    "If true, Pathways workers restart in place, without a Kubernetes"
+    " container restart, when they request a quick restart (e.g. when a proxy"
+    " using the worker exits). This avoids the Kubernetes back-off on worker"
+    " container restarts. Server images that do not support quick restart"
+    " ignore this and restart the worker container instead.",
+)
 _SIDECAR_SHM_DIR = "/tmp/sidecar_dir"
 
 
@@ -185,6 +194,7 @@ def run_deployment(
     dry_run,
     deploy_func: Callable[[dict[str, Any]], None] = deploy_jobset,
     max_restarts: int = 30,
+    enable_worker_quick_restart: bool = True,
 ) -> None:
   """Executes the deployment logic."""
   # Use PathwaysJobSet builder instead of YAML template.
@@ -197,8 +207,9 @@ def run_deployment(
       num_slices=num_slices,
       shared_pathways_service=True,
       max_restarts=max_restarts,
-      # TODO(b/496958026): Remove this once go/sps-worker-pod-stability is
-      # implemented
+      # TODO(b/496958026): Remove this once all supported server images
+      # support quick restart on Cloud. Until then, older images still restart
+      # the worker container on proxy exit.
       max_slice_restarts=1000000,
   )
 
@@ -258,7 +269,7 @@ def run_deployment(
           ),
       ])
 
-  # 2. Add arg to pathways-worker container.
+  # 2. Add args to pathways-worker container
   for container in worker_spec.containers:
     if container.name == "pathways-worker":
       args = container.args or []
@@ -268,6 +279,13 @@ def run_deployment(
         args.append(
             f"--cloud_pathways_sidecar_shm_directory={_SIDECAR_SHM_DIR}"
         )
+      if enable_worker_quick_restart:
+        # Server images that don't support quick restart don't define
+        # --enable_quick_restart, so --undefok makes them ignore it instead of
+        # failing to start.
+        args.extend(
+            ["--enable_quick_restart=true", "--undefok=enable_quick_restart"]
+        )
       container.args = args
 
   jobset_config = pw_jobset.to_dict()
@@ -333,6 +351,7 @@ def main(argv: Sequence[str]) -> None:
         sidecar_image=_SIDECAR_IMAGE.value,
         dry_run=_DRY_RUN.value,
         max_restarts=_MAX_RESTARTS.value,
+        enable_worker_quick_restart=_ENABLE_WORKER_QUICK_RESTART.value,
     )
   except ValueError as e:
     _logger.exception("Error: %s", e)

pathwaysutils/test/experimental/shared_pathways_service/deploy_pathways_service_test.py

@@ -147,6 +147,12 @@ def test_run_deployment(self, mock_jobset_cls, mock_detect):
         mock_worker_job.spec.template.spec.containers[0].args,
     )
 
+    # Verify quick restart is enabled on pathways-worker by default
+    self.assertContainsSubset(
+        ["--enable_quick_restart=true", "--undefok=enable_quick_restart"],
+        mock_worker_job.spec.template.spec.containers[0].args,
+    )
+
     # Verify deploy_func was called with the dict
     mock_deploy.assert_called_once_with({"metadata": {"name": "test-jobset"}})
 
@@ -215,6 +221,50 @@ def capture_deploy(config):
         {"restartStrategy": "Recreate", "maxRestarts": 3},
     )
 
+  @parameterized.named_parameters(
+      dict(testcase_name="enabled", enable=True),
+      dict(testcase_name="disabled", enable=False),
+  )
+  @mock.patch.object(gke_utils, "get_current_cluster_and_project")
+  def test_run_deployment_worker_quick_restart(self, mock_detect, enable):
+    mock_detect.return_value = ("test-cluster", "test-project")
+    captured_config = {}
+
+    def capture_deploy(config):
+      nonlocal captured_config
+      captured_config = config
+
+    deploy_pathways_service.run_deployment(
+        tpu_type="v5e",
+        topology="4x8",
+        num_slices=2,
+        jobset_name="test-jobset",
+        gcs_bucket="gs://test-bucket",
+        server_image=(
+            "us-docker.pkg.dev/test-project/test-repo/server:test-tag"
+        ),
+        sidecar_image=(
+            "us-docker.pkg.dev/test-project/test-repo/sidecar:test-tag"
+        ),
+        dry_run=False,
+        deploy_func=capture_deploy,
+        enable_worker_quick_restart=enable,
+    )
+
+    containers_with_quick_restart = []
+    for job in captured_config["spec"]["replicatedJobs"]:
+      pod_spec = job["template"]["spec"]["template"]["spec"]
+      for container in pod_spec["containers"] + pod_spec.get(
+          "initContainers", []
+      ):
+        if "--enable_quick_restart=true" in container.get("args", []):
+          containers_with_quick_restart.append(container["name"])
+
+    # Verify quick restart is only enabled on pathways-worker, if requested
+    self.assertEqual(
+        containers_with_quick_restart, ["pathways-worker"] if enable else []
+    )
+
   @mock.patch.object(gke_utils, "get_log_link")
   @mock.patch.object(gke_utils, "get_current_cluster_and_project")
   def test_run_deployment_cloud_logging_link(