diff --git a/pathwaysutils/experimental/shared_pathways_service/gke_utils.py b/pathwaysutils/experimental/shared_pathways_service/gke_utils.py index 5ba47f8..9a1501e 100644 --- a/pathwaysutils/experimental/shared_pathways_service/gke_utils.py +++ b/pathwaysutils/experimental/shared_pathways_service/gke_utils.py @@ -9,7 +9,7 @@ import socket import subprocess import time -from typing import Any +from typing import Any, NamedTuple import urllib.parse from kubernetes import client @@ -24,6 +24,22 @@ # Default size of the time window covered by Cloud Logging links. LOG_LINK_WINDOW = datetime.timedelta(minutes=10) +# Default time in seconds to wait for a pod to become ready. This has to cover +# scheduling the pod, which can include a cluster autoscaler scale up when no +# node has room for it, as well as pulling its image and starting it. +DEFAULT_POD_READY_TIMEOUT_S = 60 + +# Time in seconds between checks of a pod's status while waiting for it. +_POD_READY_POLL_INTERVAL_S = 5 + +# Container waiting reasons that do not resolve without user intervention. +_FATAL_CONTAINER_WAITING_REASONS = frozenset({ + "CrashLoopBackOff", + "CreateContainerConfigError", + "ImagePullBackOff", + "InvalidImageName", +}) + # TODO(b/456189271): Evaluate and replace the subprocess calls with Kubernetes # Python API for kubectl calls. @@ -312,42 +328,123 @@ def get_pod_from_job(job_name: str) -> str: return pod_name -def check_pod_ready(pod_name: str, timeout: int = 30) -> str: - """Checks if the given pod is ready. +class _PodNotReadyReason(NamedTuple): + """Describes why a pod is not ready. + + Attributes: + message: A human-readable description of why the pod is not ready. + is_fatal: Whether the pod cannot become ready without user intervention. + is_unscheduled: Whether the pod is waiting to be scheduled onto a node, for + example while the cluster autoscaler adds a node for it. + """ + + message: str + is_fatal: bool = False + is_unscheduled: bool = False + + +def _join_reason(reason: str | None, message: str | None) -> str: + """Joins a Kubernetes status reason and message into one string.""" + return ": ".join(part for part in (reason, message) if part) or "no details" + + +def _get_pod_not_ready_reason(pod: client.V1Pod) -> _PodNotReadyReason: + """Returns why the given pod is not ready.""" + status = pod.status + if not status: + return _PodNotReadyReason("Pod has no status yet.") + + if status.phase in ("Failed", "Succeeded"): + return _PodNotReadyReason( + f"Pod has terminated with phase {status.phase}" + f" ({_join_reason(status.reason, status.message)}).", + is_fatal=True, + ) + + for condition in status.conditions or []: + if condition.type == "PodScheduled" and condition.status == "False": + return _PodNotReadyReason( + "Pod is not scheduled onto a node yet" + f" ({_join_reason(condition.reason, condition.message)}). If the" + " cluster autoscaler is adding a node for the pod, this can take a" + " few minutes.", + is_unscheduled=True, + ) + + for container_status in status.container_statuses or []: + waiting = container_status.state.waiting if container_status.state else None + if waiting and waiting.reason: + return _PodNotReadyReason( + f"Container '{container_status.name}' is waiting" + f" ({_join_reason(waiting.reason, waiting.message)}).", + is_fatal=waiting.reason in _FATAL_CONTAINER_WAITING_REASONS, + ) + + return _PodNotReadyReason(f"Pod phase is {status.phase}.") + + +def check_pod_ready( + pod_name: str, + timeout: int = DEFAULT_POD_READY_TIMEOUT_S, + namespace: str = "default", +) -> str: + """Waits for the given pod to become ready. + + Polls the pod until it is ready, the timeout expires, or the pod reaches a + state that it cannot recover from on its own (e.g. an image pull failure). + Changes in why the pod is not ready yet are logged, e.g. that the pod waits + for the cluster autoscaler to add a node. Args: pod_name: The name of the pod. timeout: The maximum time in seconds to wait for the pod to be ready. + namespace: The namespace of the pod. Returns: The name of the pod. Raises: - RuntimeError: If the pod fails to become ready within the timeout. + RuntimeError: If the pod fails to become ready within the timeout, cannot + become ready, or its status cannot be read. """ _validate_k8s_name(pod_name) - wait_command = [ - "kubectl", - "wait", - "--for=condition=Ready", - f"--timeout={timeout}s", - "--", - f"pod/{pod_name}", - ] - try: - subprocess.run(wait_command, check=True, capture_output=True, text=True) - except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e: - _logger.exception("Pod failed to become ready: %r", e) - - raise RuntimeError( - f"Pod did not become ready: {e.stderr}." - ) from e - except Exception as e: - _logger.exception("Error setting up the pod: %r", e) - raise + _validate_k8s_name(namespace) + core_api = _get_k8s_core_api() + deadline = time.monotonic() + timeout + last_message = None + while True: + try: + pod = core_api.read_namespaced_pod(name=pod_name, namespace=namespace) + except client.rest.ApiException as e: + _logger.exception("Failed to get the status of pod %s: %r", pod_name, e) + raise RuntimeError( + f"Failed to get the status of pod {pod_name}: {e}" + ) from e + + if _is_pod_ready(pod): + _logger.info("Pod is ready: %s.", pod_name) + return pod_name + + reason = _get_pod_not_ready_reason(pod) + if reason.is_fatal: + raise RuntimeError( + f"Pod did not become ready: pod/{pod_name}. {reason.message}" + ) - _logger.info("Pod is ready: %s.", pod_name) - return pod_name + if reason.message != last_message: + # Use a warning for a pod waiting for a node: this is the slow path, and + # INFO logs are not shown by default when using the Python API. + log = _logger.warning if reason.is_unscheduled else _logger.info + log("Waiting for pod %s to become ready: %s", pod_name, reason.message) + last_message = reason.message + + remaining = deadline - time.monotonic() + if remaining <= 0: + raise RuntimeError( + f"Pod did not become ready within {timeout} seconds: pod/{pod_name}." + f" Last status: {reason.message}" + ) + time.sleep(min(_POD_READY_POLL_INTERVAL_S, remaining)) def _format_log_timestamp(timestamp: str | datetime.datetime) -> str: @@ -548,11 +645,14 @@ def get_current_cluster_and_project() -> tuple[str | None, str | None]: return cluster, project -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. Returns: The name of the pod. Raises: @@ -566,7 +666,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 5fbb69f..a628928 100644 --- a/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py +++ b/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py @@ -257,6 +257,8 @@ class _ISCPathways: proxy_pod_name: The name of the proxy pod, assigned during deployment. proxy_server_image: The image to use for the proxy server. proxy_options: Configuration options for the Pathways proxy. + proxy_ready_timeout: The maximum time in seconds to wait for the proxy pod + to become ready. 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. @@ -275,6 +277,7 @@ def __init__( proxy_server_image: str, proxy_options: ProxyOptions | None = None, collect_service_metrics: bool = False, + proxy_ready_timeout: int = gke_utils.DEFAULT_POD_READY_TIMEOUT_S, ): """Initializes the TPU manager.""" self.cluster = cluster @@ -290,6 +293,7 @@ def __init__( self._proxy_port = None self.proxy_server_image = proxy_server_image self.proxy_options = proxy_options or ProxyOptions() + self.proxy_ready_timeout = proxy_ready_timeout self._old_jax_platforms = None if collect_service_metrics: raw_collector = metrics_collector.MetricsCollector( @@ -413,7 +417,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_ready_timeout + ) self._proxy_port, self._port_forward_process = ( gke_utils.start_port_forwarding( @@ -549,6 +555,7 @@ def connect( proxy_server_image: str | None = None, proxy_options: Sequence[str] | None = None, collect_service_metrics: bool = False, + proxy_ready_timeout: int = gke_utils.DEFAULT_POD_READY_TIMEOUT_S, ) -> Iterator["_ISCPathways"]: """Connects to a Pathways server if the cluster exists. If not, creates it. @@ -570,6 +577,9 @@ def connect( provided, no extra options will be used. collect_service_metrics: Whether to collect usage metrics for Shared Pathways Service. + proxy_ready_timeout: The maximum time in seconds to wait for the proxy pod + to become ready. This includes the time the cluster autoscaler needs to + add a node when no existing node has room for the proxy pod. Yields: The Pathways manager. @@ -646,6 +656,7 @@ def connect( proxy_server_image=proxy_server_image, proxy_options=proxy_options_obj, collect_service_metrics=collect_service_metrics, + proxy_ready_timeout=proxy_ready_timeout, ) as t: if t.proxy_pod_name: num_slices = sum(t.expected_tpu_instances.values()) diff --git a/pathwaysutils/test/experimental/shared_pathways_service/gke_utils_test.py b/pathwaysutils/test/experimental/shared_pathways_service/gke_utils_test.py index ce4a577..8184bdc 100644 --- a/pathwaysutils/test/experimental/shared_pathways_service/gke_utils_test.py +++ b/pathwaysutils/test/experimental/shared_pathways_service/gke_utils_test.py @@ -3,8 +3,10 @@ import datetime import io +import itertools import socket import subprocess +import time from typing import Any from unittest import mock import urllib.parse @@ -515,53 +517,239 @@ def test_get_pod_from_job_failure(self): with self.assertRaises(subprocess.CalledProcessError): gke_utils.get_pod_from_job("test-proxy-job") - def test_check_pod_ready_success(self): - """Tests that check_pod_ready returns the pod name on success.""" - mock_run = self.enter_context( + def _mock_core_api(self) -> mock.MagicMock: + """Patches the Kubernetes CoreV1Api client and returns the mock.""" + mock_core_api = mock.MagicMock(spec=client.CoreV1Api) + self.enter_context( mock.patch.object( - subprocess, - "run", - autospec=True, + gke_utils, "_get_k8s_core_api", return_value=mock_core_api ) ) - mock_wait_success_result = subprocess.CompletedProcess( - args=["kubectl", "wait"], - returncode=0, + return mock_core_api + + def _make_proxy_pod( + self, + *, + phase: str = "Running", + ready: bool = True, + unschedulable_message: str | None = None, + waiting_reason: str | None = None, + ) -> client.V1Pod: + """Returns a proxy pod with the given status.""" + conditions = [ + client.V1PodCondition( + type="Ready", status="True" if ready else "False" + ) + ] + if unschedulable_message is not None: + conditions.append( + client.V1PodCondition( + type="PodScheduled", + status="False", + reason="Unschedulable", + message=unschedulable_message, + ) + ) + container_statuses = None + if waiting_reason is not None: + container_statuses = [ + client.V1ContainerStatus( + name="pathways-proxy", + image="test-image", + image_id="", + ready=False, + restart_count=0, + state=client.V1ContainerState( + waiting=client.V1ContainerStateWaiting(reason=waiting_reason) + ), + ) + ] + return client.V1Pod( + metadata=client.V1ObjectMeta(name="test-pod-123"), + status=client.V1PodStatus( + phase=phase, + conditions=conditions, + container_statuses=container_statuses, + ), ) - mock_run.return_value = mock_wait_success_result + + def test_check_pod_ready_success(self): + """Tests that check_pod_ready returns the pod name on success.""" + mock_core_api = self._mock_core_api() + mock_core_api.read_namespaced_pod.return_value = self._make_proxy_pod() pod_name = gke_utils.check_pod_ready("test-pod-123") self.assertEqual(pod_name, "test-pod-123") - mock_run.assert_called_once_with( + mock_core_api.read_namespaced_pod.assert_called_once_with( + name="test-pod-123", namespace="default" + ) + + def test_check_pod_ready_waits_for_node_scale_up(self): + """Tests that check_pod_ready waits for an unschedulable pod.""" + mock_sleep = self.enter_context( + mock.patch.object(time, "sleep", autospec=True) + ) + mock_core_api = self._mock_core_api() + mock_core_api.read_namespaced_pod.side_effect = [ + self._make_proxy_pod( + phase="Pending", + ready=False, + unschedulable_message=( + "0/2 nodes are available: 1 node(s) didn't have free ports for" + " the requested pod ports." + ), + ), + self._make_proxy_pod( + phase="Pending", ready=False, waiting_reason="ContainerCreating" + ), + self._make_proxy_pod(), + ] + + with self.assertLogs(gke_utils._logger, level="WARNING") as logs: + pod_name = gke_utils.check_pod_ready("test-pod-123") + + self.assertEqual(pod_name, "test-pod-123") + self.assertEqual(mock_core_api.read_namespaced_pod.call_count, 3) + self.assertEqual(mock_sleep.call_count, 2) + # Only the wait for a node is logged as a warning. + self.assertLen(logs.output, 1) + self.assertIn("didn't have free ports", logs.output[0]) + self.assertIn("cluster autoscaler", logs.output[0]) + + def test_check_pod_ready_logs_status_while_waiting(self): + """Tests the status logged while waiting for a pod without details.""" + self.enter_context(mock.patch.object(time, "sleep", autospec=True)) + mock_core_api = self._mock_core_api() + pod_without_status = client.V1Pod( + metadata=client.V1ObjectMeta(name="test-pod-123") + ) + # A waiting container without a reason does not explain the wait, so the + # pod phase is reported instead. + pod_with_unexplained_wait = self._make_proxy_pod( + phase="Pending", ready=False + ) + pod_with_unexplained_wait.status.container_statuses = [ + client.V1ContainerStatus( + name="pathways-proxy", + image="test-image", + image_id="", + ready=False, + restart_count=0, + state=client.V1ContainerState( + waiting=client.V1ContainerStateWaiting() + ), + ) + ] + mock_core_api.read_namespaced_pod.side_effect = [ + pod_without_status, + pod_with_unexplained_wait, + self._make_proxy_pod(), + ] + + with self.assertLogs(gke_utils._logger, level="INFO") as logs: + pod_name = gke_utils.check_pod_ready("test-pod-123") + + self.assertEqual(pod_name, "test-pod-123") + waiting_messages = [ + record.getMessage() + for record in logs.records + if record.getMessage().startswith("Waiting for pod") + ] + self.assertEqual( + waiting_messages, [ - "kubectl", - "wait", - "--for=condition=Ready", - "--timeout=30s", - "--", - "pod/test-pod-123", + "Waiting for pod test-pod-123 to become ready: Pod has no status" + " yet.", + "Waiting for pod test-pod-123 to become ready: Pod phase is" + " Pending.", ], - check=True, - capture_output=True, - text=True, ) - def test_check_pod_ready_failure(self): - """Tests that check_pod_ready raises a RuntimeError if kubectl wait fails.""" - mock_run = self.enter_context( + def test_check_pod_ready_invalid_name(self): + """Tests that check_pod_ready rejects invalid pod and namespace names.""" + mock_core_api = self._mock_core_api() + mock_core_api.read_namespaced_pod.return_value = self._make_proxy_pod() + test_cases = [ + ("invalid_pod_name", "Invalid_Pod", "default"), + ("invalid_namespace", "test-pod-123", "Invalid_Namespace"), + ] + for case_name, pod_name, namespace in test_cases: + with self.subTest(case=case_name): + with self.assertRaisesRegex( + ValueError, "Invalid Kubernetes resource name" + ): + gke_utils.check_pod_ready(pod_name, namespace=namespace) + + mock_core_api.read_namespaced_pod.assert_not_called() + + def test_check_pod_ready_fails_fast_if_pod_cannot_become_ready(self): + """Tests that check_pod_ready fails fast for a pod that cannot recover.""" + mock_sleep = self.enter_context( + mock.patch.object(time, "sleep", autospec=True) + ) + mock_core_api = self._mock_core_api() + test_cases = [ + ( + "image_pull_backoff", + self._make_proxy_pod( + phase="Pending", ready=False, waiting_reason="ImagePullBackOff" + ), + "ImagePullBackOff", + ), + ( + "failed_phase", + self._make_proxy_pod(phase="Failed", ready=False), + "terminated with phase Failed", + ), + ] + for case_name, pod, expected_message in test_cases: + with self.subTest(case=case_name): + mock_core_api.read_namespaced_pod.return_value = pod + with self.assertRaisesRegex(RuntimeError, expected_message): + gke_utils.check_pod_ready("test-pod-123") + + mock_sleep.assert_not_called() + + def test_check_pod_ready_timeout(self): + """Tests that check_pod_ready raises if the pod is not ready in time.""" + mock_sleep = self.enter_context( + mock.patch.object(time, "sleep", autospec=True) + ) + # Each call advances the clock by 5 seconds. + self.enter_context( mock.patch.object( - subprocess, - "run", - autospec=True, + time, "monotonic", side_effect=itertools.count(step=5) ) ) - mock_run.side_effect = subprocess.CalledProcessError( - returncode=1, cmd="kubectl wait", stderr="error" + mock_core_api = self._mock_core_api() + mock_core_api.read_namespaced_pod.return_value = self._make_proxy_pod( + phase="Pending", + ready=False, + unschedulable_message="0/1 nodes are available.", + ) + + with self.assertRaisesRegex( + RuntimeError, + r"Pod did not become ready within 8 seconds: pod/test-pod-123\..*" + r"Unschedulable: 0/1 nodes are available\.", + ): + gke_utils.check_pod_ready("test-pod-123", timeout=8) + + self.assertEqual(mock_core_api.read_namespaced_pod.call_count, 2) + # The deadline is at 8s and the first poll happens at 5s, so the sleep is + # capped to the remaining 3 seconds instead of the full poll interval. + mock_sleep.assert_called_once_with(3) + + def test_check_pod_ready_api_error(self): + """Tests that check_pod_ready raises if the pod status cannot be read.""" + mock_core_api = self._mock_core_api() + mock_core_api.read_namespaced_pod.side_effect = client.rest.ApiException( + status=404, reason="Not Found" ) with self.assertRaisesRegex( - RuntimeError, "Pod did not become ready: error." + RuntimeError, "Failed to get the status of pod test-pod-123" ): gke_utils.check_pod_ready("test-pod-123") @@ -951,23 +1139,22 @@ def test_wait_for_pod_success(self): autospec=True, ) ) - mock_get_pods_result = subprocess.CompletedProcess( + mock_run.return_value = subprocess.CompletedProcess( args=["kubectl", "get", "pods"], returncode=0, stdout="pod/test-pod-123\n", ) - mock_wait_success_result = subprocess.CompletedProcess( - args=["kubectl", "wait"], - returncode=0, - ) - mock_run.side_effect = [mock_get_pods_result, mock_wait_success_result] + mock_core_api = self._mock_core_api() + mock_core_api.read_namespaced_pod.return_value = self._make_proxy_pod() pod_name = gke_utils.wait_for_pod("test-proxy-job") self.assertEqual(pod_name, "test-pod-123") - self.assertEqual(mock_run.call_count, 2) - self.assertIn("get", mock_run.call_args_list[0].args[0]) - self.assertIn("wait", mock_run.call_args_list[1].args[0]) + mock_run.assert_called_once() + self.assertIn("get", mock_run.call_args.args[0]) + mock_core_api.read_namespaced_pod.assert_called_once_with( + name="test-pod-123", namespace="default" + ) def test_wait_for_pod_get_pods_fails(self): """Tests that wait_for_pod raises an error if 'get pods' fails.""" @@ -987,7 +1174,7 @@ def test_wait_for_pod_get_pods_fails(self): self.assertEqual(mock_run.call_count, 1) def test_wait_for_pod_wait_fails(self): - """Tests that wait_for_pod raises a RuntimeError if 'wait' times out.""" + """Tests that wait_for_pod raises if the pod does not become ready.""" mock_run = self.enter_context( mock.patch.object( subprocess, @@ -995,19 +1182,23 @@ def test_wait_for_pod_wait_fails(self): autospec=True, ) ) - mock_get_pods_result = subprocess.CompletedProcess( + mock_run.return_value = subprocess.CompletedProcess( args=["kubectl", "get", "pods"], returncode=0, stdout="pod/test-pod-123\n", ) - mock_run.side_effect = [ - mock_get_pods_result, - subprocess.TimeoutExpired(cmd="kubectl wait", timeout=30), - ] + mock_check_pod_ready = self.enter_context( + mock.patch.object( + gke_utils, + "check_pod_ready", + autospec=True, + side_effect=RuntimeError("Pod did not become ready"), + ) + ) with self.assertRaises(RuntimeError): - gke_utils.wait_for_pod("test-proxy-job") - self.assertEqual(mock_run.call_count, 2) + gke_utils.wait_for_pod("test-proxy-job", timeout=42) + mock_check_pod_ready.assert_called_once_with("test-pod-123", timeout=42) def test_enable_port_forwarding_success(self): """Tests successful port forwarding.""" diff --git a/pathwaysutils/test/experimental/shared_pathways_service/isc_pathways_test.py b/pathwaysutils/test/experimental/shared_pathways_service/isc_pathways_test.py index ab5cc43..9d34137 100644 --- a/pathwaysutils/test/experimental/shared_pathways_service/isc_pathways_test.py +++ b/pathwaysutils/test/experimental/shared_pathways_service/isc_pathways_test.py @@ -359,8 +359,18 @@ def test_isc_pathways_pod_wait_failure_raises_error(self): self.enter_context( mock.patch("urllib.parse.quote", return_value="encoded_filter") ) - # Simulate 'kubectl get' returning a pod, 'kubectl wait' timing out, - # then 'kubectl delete' succeeding. + mock_check_pod_ready = self.enter_context( + mock.patch.object( + isc_pathways.gke_utils, + "check_pod_ready", + autospec=True, + side_effect=RuntimeError( + "Pod did not become ready within 60 seconds: pod/test-pod-123." + ), + ) + ) + # Simulate 'kubectl get' returning a pod, the pod not becoming ready, then + # 'kubectl delete' succeeding. mock_get_pods_result = subprocess.CompletedProcess( args=["kubectl", "get", "pods"], returncode=0, @@ -373,9 +383,7 @@ def test_isc_pathways_pod_wait_failure_raises_error(self): mock_run.side_effect = [ # First call for 'get' mock_get_pods_result, - # Second call for 'wait' - subprocess.TimeoutExpired(cmd="kubectl wait", timeout=30), - # Third call for 'delete' + # Second call for 'delete' mock_delete_result, ] @@ -398,12 +406,15 @@ def test_isc_pathways_pod_wait_failure_raises_error(self): self.assertIn("Pod did not become ready", str(context.exception)) mock_deploy.assert_called_once() + mock_check_pod_ready.assert_called_once_with( + "test-pod-123", + timeout=isc_pathways.gke_utils.DEFAULT_POD_READY_TIMEOUT_S, + ) - # Check that subprocess.run was called for get, wait, and delete. - self.assertEqual(mock_run.call_count, 3) + # Check that subprocess.run was called for get and delete. + self.assertEqual(mock_run.call_count, 2) self.assertIn("get", mock_run.call_args_list[0][0][0]) - self.assertIn("wait", mock_run.call_args_list[1][0][0]) - self.assertIn("delete", mock_run.call_args_list[2][0][0]) + self.assertIn("delete", mock_run.call_args_list[1][0][0]) def test_isc_pathways(self): """Tests the full lifecycle of ISCPathways.""" @@ -481,7 +492,7 @@ def test_isc_pathways(self): mock_random.choices.return_value = list("abcde") mock_random.randint.return_value = 29005 - # Mock for 'kubectl wait' and 'kubectl delete' + # Mock for 'kubectl delete' mock_success_result = subprocess.CompletedProcess( args=["kubectl"], returncode=0, @@ -497,9 +508,16 @@ def test_isc_pathways(self): mock_run.side_effect = [ mock_get_pods_result, # For 'get_pod_from_job' - mock_success_result, # For 'check_pod_ready' mock_success_result, # For 'kubectl delete job' in __exit__ ] + self.enter_context( + mock.patch.object( + isc_pathways.gke_utils, + "check_pod_ready", + autospec=True, + return_value="test-pod-123", + ) + ) mock_pick_unused_port.return_value = 29007 @@ -632,6 +650,47 @@ def test_connect_skips_credentials_when_kube_config_matches(self): ) mock_fetch_creds.assert_not_called() + def test_connect_passes_proxy_ready_timeout(self): + """Tests that connect passes the proxy readiness timeout to the manager.""" + self.enter_context( + mock.patch.object( + isc_pathways.gke_utils, "fetch_cluster_credentials", autospec=True + ) + ) + mock_get_images = self.enter_context( + mock.patch.object( + isc_pathways.gke_utils, "get_pathways_service_images", autospec=True + ) + ) + mock_get_images.return_value = ( + "us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest", + None, + ) + mock_isc_pathways = self.enter_context( + mock.patch.object(isc_pathways, "_ISCPathways", autospec=True) + ) + self.enter_context(mock.patch("threading.Thread", autospec=True)) + mock_manager_instance = ( + mock_isc_pathways.return_value.__enter__.return_value + ) + mock_manager_instance.proxy_pod_name = "test-pod-123" + mock_manager_instance.expected_tpu_instances = {"tpuv6e:2x2": 1} + + with isc_pathways.connect( + cluster="test-cluster", + project="test-project", + region="test-region", + gcs_bucket="test-bucket", + pathways_service="test-service-pathways-head:1234", + expected_tpu_instances={"tpuv6e:2x2": 1}, + proxy_ready_timeout=900, + ): + pass + + self.assertEqual( + mock_isc_pathways.call_args.kwargs["proxy_ready_timeout"], 900 + ) + def test_connect_success(self): """Tests that connect calls the dependencies and yields the manager.""" # Arrange @@ -710,6 +769,9 @@ def test_connect_success(self): proxy_server_image=isc_pathways.DEFAULT_PROXY_IMAGE, proxy_options=isc_pathways.ProxyOptions(), collect_service_metrics=False, + proxy_ready_timeout=( + isc_pathways.gke_utils.DEFAULT_POD_READY_TIMEOUT_S + ), ) self.assertIs(tm, mock_manager_instance)