Use a flag to control whether readiness prob is enabled

GitOrigin-RevId: 26af676

Author: muyangyuapple

axlearn/cloud/gcp/pathways_utils.py

@@ -883,6 +883,7 @@ class Config(BaseLeaderWorkerTemplate.Config):
 
         target_port: Optional[int] = None
         enable_service: bool = None
+        enable_health_probes: Optional[bool] = None
 
         # Health probe configuration for inference pods.
         # - startup_probe_failure_threshold: Max consecutive failures before probe fails.
@@ -928,6 +929,12 @@ def define_flags(cls, fv):
             "Whether to enable creation of service for LWS",
             **common_kwargs,
         )
+        flags.DEFINE_boolean(
+            "enable_health_probes",
+            False,
+            "Whether to enable health probes for LWS service",
+            **common_kwargs,
+        )
         flags.DEFINE_integer(
             "target_port",
             None,
@@ -962,6 +969,7 @@ def set_defaults(cls, fv):
         fv.set_default("pathways_head_mem", fv.pathways_head_mem or "16")
         fv.set_default("target_port", fv.target_port or 9000)
         fv.set_default("enable_service", fv.enable_service or False)
+        fv.set_default("enable_health_probes", fv.enable_health_probes or False)
 
     @classmethod
     def default_config(cls):
@@ -1128,7 +1136,7 @@ def _build_head_container(self) -> dict:
         # Add health probes for inference pods when service is enabled.
         # This ensures K8s services only route traffic to healthy pods that have
         # finished loading models and are ready to serve requests.
-        if cfg.enable_service:
+        if cfg.enable_health_probes:
             # startupProbe: Allows long startup time for model loading.
             # Max startup time = failureThreshold * periodSeconds.
             container["startupProbe"] = {

axlearn/cloud/gcp/pathways_utils_test.py

@@ -614,7 +614,7 @@ def test_health_probes_configuration(
         with (
             self._job_config(
                 CloudBuildBundler,
-                enable_service=True,
+                enable_health_probes=True,
                 health_check_path=health_check_path,
                 health_check_port=health_check_port,
                 startup_probe_failure_threshold=startup_probe_failure_threshold,
@@ -646,12 +646,12 @@ def test_health_probes_configuration(
             # readinessProbe failureThreshold is always 3
             self.assertEqual(readiness_probe["failureThreshold"], 3)
 
-    def test_no_health_probes_when_service_disabled(self):
-        """Tests that health probes are not added when enable_service is False."""
+    def test_no_health_probes_when_readiness_probe_disabled(self):
+        """Tests that health probes are not added when enable_health_probes is False."""
         with (
             self._job_config(
                 CloudBuildBundler,
-                enable_service=False,
+                enable_health_probes=False,
             ) as (cfg, bundler_cfg),
         ):
             cfg.inner.set(