From 8e944f757b6eec3cbbba5ade6749bae8ad065444 Mon Sep 17 00:00:00 2001 From: Kewei Li Date: Thu, 16 Jul 2026 15:08:50 -0700 Subject: [PATCH 1/2] Make --dns-endpoint optional in fetch_cluster_credentials for fully-private GKE clusters --- .../shared_pathways_service/gke_utils.py | 12 ++++++++---- .../shared_pathways_service/isc_pathways.py | 10 +++++++++- .../start_vscode_on_cpu_np.py | 9 +++++++++ 3 files changed, 26 insertions(+), 5 deletions(-) diff --git a/pathwaysutils/experimental/shared_pathways_service/gke_utils.py b/pathwaysutils/experimental/shared_pathways_service/gke_utils.py index ac17acf..8436c99 100644 --- a/pathwaysutils/experimental/shared_pathways_service/gke_utils.py +++ b/pathwaysutils/experimental/shared_pathways_service/gke_utils.py @@ -34,7 +34,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 +50,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, diff --git a/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py b/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py index 339b1f3..fc0127f 100644 --- a/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py +++ b/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py @@ -437,6 +437,7 @@ 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, ) -> Iterator["_ISCPathways"]: """Connects to a Pathways server if the cluster exists. If not, creates it. @@ -456,6 +457,10 @@ 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. Yields: The Pathways manager. @@ -466,7 +471,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) diff --git a/pathwaysutils/experimental/shared_pathways_service/start_vscode_on_cpu_np.py b/pathwaysutils/experimental/shared_pathways_service/start_vscode_on_cpu_np.py index 872261b..dee4680 100644 --- a/pathwaysutils/experimental/shared_pathways_service/start_vscode_on_cpu_np.py +++ b/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) From 09e73d0a56b156d72a74bddebb233aa3fc810f69 Mon Sep 17 00:00:00 2001 From: Kewei Li Date: Thu, 16 Jul 2026 15:17:59 -0700 Subject: [PATCH 2/2] Make proxy pod readiness timeout configurable to ride through node cold-start --- .../shared_pathways_service/gke_utils.py | 16 ++++++++++++++-- .../shared_pathways_service/isc_pathways.py | 15 ++++++++++++++- 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/pathwaysutils/experimental/shared_pathways_service/gke_utils.py b/pathwaysutils/experimental/shared_pathways_service/gke_utils.py index 8436c99..cd6dd35 100644 --- a/pathwaysutils/experimental/shared_pathways_service/gke_utils.py +++ b/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. @@ -257,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: @@ -275,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: diff --git a/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py b/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py index fc0127f..70862c9 100644 --- a/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py +++ b/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 @@ -438,6 +444,7 @@ def connect( 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. @@ -461,6 +468,11 @@ def connect( 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. @@ -504,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())