diff --git a/pathwaysutils/experimental/shared_pathways_service/gke_utils.py b/pathwaysutils/experimental/shared_pathways_service/gke_utils.py index ac17acf..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. @@ -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: diff --git a/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py b/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py index 339b1f3..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 @@ -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()) 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)