diff --git a/pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py b/pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py index d6e296b..25ea3cf 100644 --- a/pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py +++ b/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) diff --git a/pathwaysutils/test/experimental/shared_pathways_service/deploy_pathways_service_test.py b/pathwaysutils/test/experimental/shared_pathways_service/deploy_pathways_service_test.py index 62853cd..ec118f6 100644 --- a/pathwaysutils/test/experimental/shared_pathways_service/deploy_pathways_service_test.py +++ b/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(