Support fully-private GKE clusters and configurable proxy pod readiness timeout

pathwaysutils/experimental/shared_pathways_service/gke_utils.py

@@ -12,6 +12,13 @@
 
 _logger = logging.getLogger(__name__)
 
+# Default readiness timeout for the proxy pod. Wider than kubectl's short waits
+# so it can ride through a cold-start node provisioning from zero (cluster
+# autoscaling plus the image pull), which can take longer than a very short
+# timeout even though scheduling and the image pull complete quickly once a node
+# is warm. Callers that expect long cold starts can pass a larger value.
+DEFAULT_POD_READY_TIMEOUT_S = 60
+
 # TODO(b/456189271): Evaluate and replace the subprocess calls with Kubernetes
 # Python API for kubectl calls.
 
@@ -34,7 +41,11 @@ def _validate_k8s_name(name: str) -> None:
 
 
 def fetch_cluster_credentials(
-    *, cluster_name: str, project_id: str, location: str
+    *,
+    cluster_name: str,
+    project_id: str,
+    location: str,
+    use_dns_endpoint: bool = True,
 ) -> None:
   """Fetches credentials for the GKE cluster."""
   _validate_k8s_name(cluster_name)
@@ -46,10 +57,10 @@ def fetch_cluster_credentials(
       "get-credentials",
       f"--location={location}",
       f"--project={project_id}",
-      "--dns-endpoint",
-      "--",
-      cluster_name,
   ]
+  if use_dns_endpoint:
+    get_credentials_command.append("--dns-endpoint")
+  get_credentials_command += ["--", cluster_name]
   try:
     subprocess.run(
         get_credentials_command,
@@ -253,11 +264,16 @@ def get_log_link(*, cluster: str, project: str, job_name: str) -> str:
   )
 
 
-def wait_for_pod(job_name: str) -> str:
+def wait_for_pod(
+    job_name: str, timeout: int = DEFAULT_POD_READY_TIMEOUT_S
+) -> str:
   """Waits for the given job's pod to be ready.
 
   Args:
     job_name: The name of the job.
+    timeout: The maximum time in seconds to wait for the pod to be ready.
+      Defaults to a cold-start-tolerant value so the wait rides through node
+      provisioning from zero.
   Returns:
     The name of the pod.
   Raises:
@@ -271,7 +287,7 @@ def wait_for_pod(job_name: str) -> str:
       "Pod created: %s. Waiting for it to be ready...", pod_name
   )
 
-  return check_pod_ready(pod_name)
+  return check_pod_ready(pod_name, timeout=timeout)
 
 
 def _test_remote_connection(port: int) -> None:

pathwaysutils/experimental/shared_pathways_service/isc_pathways.py

@@ -248,6 +248,8 @@ class _ISCPathways:
     metrics_collector: The metrics collector instance if enabled.
     start_time: The start time of the TPU assignment.
     total_chips: The total number of TPU chips expected across all instances.
+    proxy_pod_ready_timeout_s: The maximum time in seconds to wait for the proxy
+      pod to become ready.
   """
 
   def __init__(
@@ -263,6 +265,7 @@ def __init__(
       proxy_server_image: str,
       proxy_options: ProxyOptions | None = None,
       collect_service_metrics: bool = False,
+      proxy_pod_ready_timeout_s: int = gke_utils.DEFAULT_POD_READY_TIMEOUT_S,
   ):
     """Initializes the TPU manager."""
     self.cluster = cluster
@@ -277,6 +280,7 @@ def __init__(
     self._proxy_port = None
     self.proxy_server_image = proxy_server_image
     self.proxy_options = proxy_options or ProxyOptions()
+    self._proxy_pod_ready_timeout_s = proxy_pod_ready_timeout_s
     self._old_jax_platforms = None
     if collect_service_metrics:
       raw_collector = metrics_collector.MetricsCollector(
@@ -350,7 +354,9 @@ def __enter__(self):
       )
       _logger.info("View proxy logs in Cloud Logging: %s", cloud_logging_link)
 
-      self.proxy_pod_name = gke_utils.wait_for_pod(self._proxy_job_name)
+      self.proxy_pod_name = gke_utils.wait_for_pod(
+          self._proxy_job_name, timeout=self._proxy_pod_ready_timeout_s
+      )
       self._proxy_port, self._port_forward_process = (
           gke_utils.enable_port_forwarding(
               f"pod/{self.proxy_pod_name}", PROXY_SERVER_PORT
@@ -437,6 +443,8 @@ def connect(
     proxy_server_image: str = DEFAULT_PROXY_IMAGE,
     proxy_options: Sequence[str] | None = None,
     collect_service_metrics: bool = False,
+    use_dns_endpoint: bool = True,
+    proxy_pod_ready_timeout_s: int = gke_utils.DEFAULT_POD_READY_TIMEOUT_S,
 ) -> Iterator["_ISCPathways"]:
   """Connects to a Pathways server if the cluster exists. If not, creates it.
 
@@ -456,6 +464,15 @@ def connect(
       provided, no extra options will be used.
     collect_service_metrics: Whether to collect usage metrics for Shared
       Pathways Service.
+    use_dns_endpoint: If True (default), fetch cluster credentials via the
+      cluster's DNS endpoint. Set to False for fully-private clusters whose
+      DNS and public endpoints are both disabled, so credentials are fetched
+      via the reachable private IP endpoint instead.
+    proxy_pod_ready_timeout_s: The maximum time in seconds to wait for the proxy
+      pod to become ready. Defaults to a cold-start-tolerant value so the wait
+      rides through node provisioning from zero. Increase this when the proxy
+      pod schedules onto a pool that autoscales from zero and cold starts can
+      take several minutes.
 
   Yields:
     The Pathways manager.
@@ -466,7 +483,10 @@ def connect(
   validators.validate_proxy_server_image(proxy_server_image)
   validators.validate_proxy_options(proxy_options)
   gke_utils.fetch_cluster_credentials(
-      cluster_name=cluster, project_id=project, location=region
+      cluster_name=cluster,
+      project_id=project,
+      location=region,
+      use_dns_endpoint=use_dns_endpoint,
   )
 
   proxy_options_obj = ProxyOptions.from_list(proxy_options)
@@ -496,6 +516,7 @@ def connect(
       proxy_server_image=proxy_server_image,
       proxy_options=proxy_options_obj,
       collect_service_metrics=collect_service_metrics,
+      proxy_pod_ready_timeout_s=proxy_pod_ready_timeout_s,
   ) as t:
     if t.proxy_pod_name:
       num_slices = sum(t.expected_tpu_instances.values())

pathwaysutils/experimental/shared_pathways_service/start_vscode_on_cpu_np.py

@@ -37,6 +37,14 @@
     False,
     "If true, only print the generated YAML without deploying.",
 )
+_USE_DNS_ENDPOINT = flags.DEFINE_boolean(
+    "use_dns_endpoint",
+    True,
+    "If true (default), fetch cluster credentials via the cluster's DNS "
+    "endpoint. Set to false for fully-private clusters whose DNS and public "
+    "endpoints are both disabled, so credentials are fetched via the reachable "
+    "private IP endpoint instead.",
+)
 
 _TEMPLATE_FILE = os.path.join(
     os.path.dirname(__file__), "yamls/code-server.yaml"
@@ -150,6 +158,7 @@ def main(argv):
       cluster_name=_CLUSTER.value,
       project_id=_PROJECT.value,
       location=_REGION.value,
+      use_dns_endpoint=_USE_DNS_ENDPOINT.value,
   )
   try:
     _deploy_vscode(service_name, deployment_yaml)