diff --git a/pathwaysutils/experimental/shared_pathways_service/gke_utils.py b/pathwaysutils/experimental/shared_pathways_service/gke_utils.py index 144c999..e7a9b04 100644 --- a/pathwaysutils/experimental/shared_pathways_service/gke_utils.py +++ b/pathwaysutils/experimental/shared_pathways_service/gke_utils.py @@ -13,6 +13,8 @@ from kubernetes import client from kubernetes import config as k8s_config +from pathwaysutils.experimental.gke import jobset as gke_jobset +from pathwaysutils.experimental.shared_pathways_service import tpu_specs import portpicker _logger = logging.getLogger(__name__) @@ -150,7 +152,6 @@ def delete_gke_resource( raise - def get_pod_from_job(job_name: str) -> str: """Returns the pod name for the given job. @@ -391,6 +392,43 @@ def get_current_kube_context() -> tuple[str | None, str | None, str | None]: return cluster_context or None, None, None +def _is_current_kube_context( + *, cluster: str, project: str, location: str +) -> bool: + """Checks whether the active kube config context points at the cluster. + + Args: + cluster: The name of the GKE cluster. + project: The GCP project ID. + location: The GCP region or zone of the cluster. + + Returns: + True if the current kube config context already targets the given cluster. + """ + return get_current_kube_context() == (cluster, project, location) + + +def ensure_cluster_credentials( + *, cluster: str, project: str, location: str +) -> None: + """Fetches the GKE cluster credentials unless kube config already has them.""" + if _is_current_kube_context( + cluster=cluster, project=project, location=location + ): + _logger.info( + "The current kube config context already points to cluster '%s' in" + " project '%s' and location '%s'. Skipping credential fetch.", + cluster, + project, + location, + ) + return + + fetch_cluster_credentials( + cluster_name=cluster, project_id=project, location=location + ) + + def get_current_cluster_and_project() -> tuple[str | None, str | None]: """Extracts cluster name and project ID from current kubeconfig or environment.""" cluster, project, _ = get_current_kube_context() @@ -672,7 +710,6 @@ def stream_pod_logs(pod_name: str) -> subprocess.Popen[str]: raise - def wait_for_deployment( name: str, namespace: str = "default", timeout: int = 300 ) -> None: @@ -764,21 +801,36 @@ def _get_k8s_custom_objects_api() -> client.CustomObjectsApi: return client.CustomObjectsApi() -def get_pathways_service_images( - pathways_service: str, namespace: str = "default" -) -> tuple[str, str | None]: - """Gets the server image and optional worker sidecar image from the JobSet.""" - pathways_head_hostname = pathways_service.split(":")[0] - _validate_k8s_name(namespace) +def extract_jobset_name(pathways_service: str) -> str: + """Extracts the JobSet name from the Pathways service address. + + Args: + pathways_service: The Pathways service address (e.g. + "my-jobset-pathways-head-0-0.my-jobset:29001" or + "my-jobset-pathways-head:29001"). + + Returns: + The JobSet name. - # Try to extract the jobset name from the Pathways service hostname. + Raises: + ValueError: If the JobSet name cannot be extracted from the address. + """ + pathways_head_hostname = pathways_service.split(":")[0] if "-pathways-head" not in pathways_head_hostname: raise ValueError( "Failed to extract jobset name from Pathways service hostname:" f" {pathways_head_hostname}. Expected prefix format:" " -pathways-head" ) - jobset_name = pathways_head_hostname.split("-pathways-head")[0] + return pathways_head_hostname.split("-pathways-head")[0] + + +def get_pathways_service_images( + pathways_service: str, namespace: str = "default" +) -> tuple[str, str | None]: + """Gets the server image and optional worker sidecar image from the JobSet.""" + _validate_k8s_name(namespace) + jobset_name = extract_jobset_name(pathways_service) try: custom_api = _get_k8s_custom_objects_api() @@ -850,4 +902,145 @@ def get_compatible_proxy_server_image(server_image: str) -> str: return new_repo +def _get_pod_labels(pod: client.V1Pod) -> dict[str, str]: + if pod.metadata and pod.metadata.labels: + return pod.metadata.labels + return {} + +def _is_head_pod(pod: client.V1Pod) -> bool: + return ( + _get_pod_labels(pod).get("jobset.sigs.k8s.io/replicatedjob-name") + == gke_jobset.PATHWAYS_HEAD_JOB_NAME + ) + + +def _is_worker_pod(pod: client.V1Pod) -> bool: + return ( + _get_pod_labels(pod).get("jobset.sigs.k8s.io/replicatedjob-name") + == gke_jobset.PATHWAYS_WORKER_JOB_NAME + ) + + +def _is_pod_ready(pod: client.V1Pod) -> bool: + """Returns True if the pod phase is Running and Ready condition is True.""" + status = pod.status + if not status or status.phase != "Running": + return False + for cond in status.conditions or []: + if cond.type == "Ready" and cond.status == "True": + return True + return False + + +def verify_pathways_service_is_up( + *, + cluster: str, + project: str, + region: str, + pathways_service: str, + tpu_type: str, + tpu_count: int, + namespace: str = "default", +) -> None: + """Verifies that the Shared Pathways Service pods are up and ready. + + Args: + cluster: The name of the GKE cluster. + project: The GCP project ID. + region: The GCP region. + pathways_service: The Pathways service address. + tpu_type: The TPU machine type and topology (e.g. 'tpuv6e:4x8'). + tpu_count: Expected number of TPU slices. + namespace: The Kubernetes namespace. + + Raises: + ValueError: If the service address, TPU type, or namespace is invalid. + RuntimeError: If no pods are found for the JobSet or if the head pod or + requested worker slices are not ready. + """ + _validate_k8s_name(namespace) + jobset_name = extract_jobset_name(pathways_service) + _logger.info( + "Verifying Shared Pathways Service '%s' in namespace '%s' on cluster" + " '%s'...", + jobset_name, + namespace, + cluster, + ) + ensure_cluster_credentials( + cluster=cluster, project=project, location=region + ) + + core_api = _get_k8s_core_api() + try: + pod_list = core_api.list_namespaced_pod( + namespace=namespace, + label_selector=f"jobset.sigs.k8s.io/jobset-name={jobset_name}", + ) + except Exception as e: + raise RuntimeError( + "Failed to list pods for Shared Pathways Service JobSet" + f" '{jobset_name}': {e}" + ) from e + + pods = pod_list.items or [] + if not pods: + raise RuntimeError( + f"No pods found for Shared Pathways Service JobSet '{jobset_name}' in" + f" namespace '{namespace}'." + ) + + head_pod_ready = False + ready_worker_pod_count_by_slice: dict[str, int] = {} + ready_worker_pods_count = 0 + for pod in pods: + if not _is_pod_ready(pod): + continue + if _is_head_pod(pod): + head_pod_ready = True + elif _is_worker_pod(pod): + slice_idx = _get_pod_labels(pod).get("jobset.sigs.k8s.io/job-index") + if slice_idx is None: + continue + ready_worker_pods_count += 1 + ready_worker_pod_count_by_slice[slice_idx] = ( + ready_worker_pod_count_by_slice.get(slice_idx, 0) + 1 + ) + + if not head_pod_ready: + raise RuntimeError( + "Shared Pathways Service head pod is not ready for JobSet" + f" '{jobset_name}' in namespace '{namespace}'. Please check the jobset" + " or deploy it again by running `python3 -m" + " pathwaysutils.experimental.shared_pathways_service.deploy_pathways_service`." + ) + + parsed_tpu_type, topology = tpu_specs.parse_tpu_type_string(tpu_type) + tpu_config = tpu_specs.get_tpu_config(parsed_tpu_type) + vms_per_slice = tpu_specs.calculate_vms_per_slice( + topology, tpu_config.chips_per_vm + ) + + ready_slices = sum( + 1 + for count in ready_worker_pod_count_by_slice.values() + if count == vms_per_slice + ) + if ready_slices < tpu_count: + raise RuntimeError( + f"Shared Pathways Service JobSet '{jobset_name}' is not ready:" + f" requested {tpu_count} slice(s) with {vms_per_slice} worker" + f" pod(s) per slice, but only {ready_slices} slice(s)" + f" ({ready_worker_pods_count} worker pod(s)) are ready. Please check" + " the jobset or re-deploy it again by running `python3 -m" + " pathwaysutils.experimental.shared_pathways_service.deploy_pathways_service`." + ) + + _logger.info( + "Shared Pathways Service '%s' is up and ready (%d ready slice(s), %d" + " ready worker pod(s)).", + jobset_name, + ready_slices, + ready_worker_pods_count, + ) diff --git a/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py b/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py index 6c8411d..395a183 100644 --- a/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py +++ b/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py @@ -534,43 +534,6 @@ def _get_username() -> str: return username or "user" -def _is_current_kube_context( - *, cluster: str, project: str, location: str -) -> bool: - """Checks whether the active kube config context points at the cluster. - - Args: - cluster: The name of the GKE cluster. - project: The GCP project ID. - location: The GCP region or zone of the cluster. - - Returns: - True if the current kube config context already targets the given cluster. - """ - return gke_utils.get_current_kube_context() == (cluster, project, location) - - -def _ensure_cluster_credentials( - *, cluster: str, project: str, location: str -) -> None: - """Fetches the GKE cluster credentials unless kube config already has them.""" - if _is_current_kube_context( - cluster=cluster, project=project, location=location - ): - _logger.info( - "The current kube config context already points to cluster '%s' in" - " project '%s' and location '%s'. Skipping credential fetch.", - cluster, - project, - location, - ) - return - - gke_utils.fetch_cluster_credentials( - cluster_name=cluster, project_id=project, location=location - ) - - @contextlib.contextmanager def connect( *, @@ -621,7 +584,7 @@ def connect( validators.validate_pathways_service(pathways_service) validators.validate_tpu_instances(expected_tpu_instances) validators.validate_proxy_options(proxy_options) - _ensure_cluster_credentials( + gke_utils.ensure_cluster_credentials( cluster=cluster, project=project, location=region ) diff --git a/pathwaysutils/experimental/shared_pathways_service/run_workload.py b/pathwaysutils/experimental/shared_pathways_service/run_workload.py index 3ba940a..3c76779 100644 --- a/pathwaysutils/experimental/shared_pathways_service/run_workload.py +++ b/pathwaysutils/experimental/shared_pathways_service/run_workload.py @@ -115,6 +115,7 @@ def run_command( expected to be a callable that returns a context manager. Raises: + RuntimeError: If the Shared Pathways Service is not up or ready. subprocess.CalledProcessError: If the workload command fails. """ if proxy_server_image: @@ -125,6 +126,15 @@ def run_command( DeprecationWarning, stacklevel=2, ) + logging.info("Verifying Shared Pathways Service is up and ready...") + gke_utils.verify_pathways_service_is_up( + cluster=cluster, + project=project, + region=region, + pathways_service=pathways_service, + tpu_type=tpu_type, + tpu_count=tpu_count, + ) logging.info("Connecting to Shared Pathways Service...") with connect_fn( cluster=cluster, 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 dbc9340..c556fe1 100644 --- a/pathwaysutils/test/experimental/shared_pathways_service/gke_utils_test.py +++ b/pathwaysutils/test/experimental/shared_pathways_service/gke_utils_test.py @@ -1593,6 +1593,536 @@ def test_terminate_process_custom_timeout(self): mock_proc.terminate.assert_called_once() mock_proc.wait.assert_called_once_with(timeout=10) + def _make_pod( + self, + name: str, + jobset_name: str, + replicated_job_name: str, + phase: str = "Running", + ready: bool = True, + job_index: str | None = "0", + ) -> client.V1Pod: + labels: dict[str, str] = { + "jobset.sigs.k8s.io/jobset-name": jobset_name, + "jobset.sigs.k8s.io/replicatedjob-name": replicated_job_name, + } + if job_index is not None: + labels["jobset.sigs.k8s.io/job-index"] = job_index + return client.V1Pod( + metadata=client.V1ObjectMeta( + name=name, + labels=labels, + ), + status=client.V1PodStatus( + phase=phase, + conditions=[ + client.V1PodCondition( + type="Ready", + status="True" if ready else "False", + ) + ], + ), + ) + + def test_is_current_kube_context_matching(self): + with mock.patch.object( + gke_utils, + "get_current_kube_context", + return_value=("test-cluster", "test-project", "test-region"), + autospec=True, + ): + self.assertTrue( + gke_utils._is_current_kube_context( + cluster="test-cluster", + project="test-project", + location="test-region", + ) + ) + + def test_is_current_kube_context_not_matching(self): + test_cases = [ + ("different_cluster", ("other-cluster", "test-project", "test-region")), + ("different_project", ("test-cluster", "other-project", "test-region")), + ( + "different_location", + ("test-cluster", "test-project", "other-region"), + ), + ("non_gke_context", ("minikube", None, None)), + ("no_context", (None, None, None)), + ] + for case_name, current_context in test_cases: + with self.subTest(case=case_name): + with mock.patch.object( + gke_utils, + "get_current_kube_context", + return_value=current_context, + autospec=True, + ): + self.assertFalse( + gke_utils._is_current_kube_context( + cluster="test-cluster", + project="test-project", + location="test-region", + ) + ) + + def test_ensure_cluster_credentials_skips_fetch_when_matching(self): + mock_fetch_creds = self.enter_context( + mock.patch.object( + gke_utils, "fetch_cluster_credentials", autospec=True + ) + ) + self.enter_context( + mock.patch.object( + gke_utils, + "get_current_kube_context", + return_value=("test-cluster", "test-project", "test-region"), + autospec=True, + ) + ) + + gke_utils.ensure_cluster_credentials( + cluster="test-cluster", project="test-project", location="test-region" + ) + + mock_fetch_creds.assert_not_called() + + def test_ensure_cluster_credentials_fetches_when_not_matching(self): + mock_fetch_creds = self.enter_context( + mock.patch.object( + gke_utils, "fetch_cluster_credentials", autospec=True + ) + ) + self.enter_context( + mock.patch.object( + gke_utils, + "get_current_kube_context", + return_value=("test-cluster", "other-project", "test-region"), + autospec=True, + ) + ) + + gke_utils.ensure_cluster_credentials( + cluster="test-cluster", project="test-project", location="test-region" + ) + + mock_fetch_creds.assert_called_once_with( + cluster_name="test-cluster", + project_id="test-project", + location="test-region", + ) + + def test_extract_jobset_name_valid(self): + self.assertEqual( + gke_utils.extract_jobset_name( + "my-jobset-pathways-head-0-0.my-jobset:29001" + ), + "my-jobset", + ) + self.assertEqual( + gke_utils.extract_jobset_name("my-jobset-pathways-head:29001"), + "my-jobset", + ) + self.assertEqual( + gke_utils.extract_jobset_name("sps-cluster-pathways-head-0-0:8000"), + "sps-cluster", + ) + + def test_extract_jobset_name_invalid_raises(self): + with self.assertRaises(ValueError): + gke_utils.extract_jobset_name("my-jobset:29001") + + def test_verify_pathways_service_is_up_success(self): + mock_creds = self.enter_context( + mock.patch.object( + gke_utils, "ensure_cluster_credentials", autospec=True + ) + ) + mock_core_api = mock.MagicMock(spec=client.CoreV1Api) + self.enter_context( + mock.patch.object( + gke_utils, "_get_k8s_core_api", return_value=mock_core_api + ) + ) + head_pod = self._make_pod( + name="my-jobset-pathways-head-0-0", + jobset_name="my-jobset", + replicated_job_name="pathways-head", + ready=True, + ) + worker_pods = [] + for slice_idx in range(2): + for vm_idx in range(2): + worker_pods.append( + self._make_pod( + name=f"my-jobset-pathways-worker-{slice_idx}-{vm_idx}", + jobset_name="my-jobset", + replicated_job_name="pathways-worker", + ready=True, + job_index=str(slice_idx), + ) + ) + mock_core_api.list_namespaced_pod.return_value = client.V1PodList( + items=[head_pod] + worker_pods + ) + + gke_utils.verify_pathways_service_is_up( + cluster="my-cluster", + project="my-project", + region="us-central1", + pathways_service="my-jobset-pathways-head-0-0:29001", + tpu_type="tpuv6e:2x4", + tpu_count=2, + ) + + mock_creds.assert_called_once_with( + cluster="my-cluster", + project="my-project", + location="us-central1", + ) + mock_core_api.list_namespaced_pod.assert_called_once_with( + namespace="default", + label_selector="jobset.sigs.k8s.io/jobset-name=my-jobset", + ) + + def test_verify_pathways_service_is_up_list_pods_error(self): + self.enter_context( + mock.patch.object( + gke_utils, "ensure_cluster_credentials", autospec=True + ) + ) + mock_core_api = mock.MagicMock(spec=client.CoreV1Api) + self.enter_context( + mock.patch.object( + gke_utils, "_get_k8s_core_api", return_value=mock_core_api + ) + ) + mock_core_api.list_namespaced_pod.side_effect = client.rest.ApiException( + status=500, reason="Error" + ) + with self.assertRaisesRegex( + RuntimeError, + "Failed to list pods for Shared Pathways Service JobSet 'my-jobset'", + ): + gke_utils.verify_pathways_service_is_up( + cluster="my-cluster", + project="my-project", + region="us-central1", + pathways_service="my-jobset-pathways-head-0-0:29001", + tpu_type="tpuv6e:2x2", + tpu_count=1, + ) + + def test_verify_pathways_service_is_up_no_pods_found(self): + self.enter_context( + mock.patch.object( + gke_utils, "ensure_cluster_credentials", autospec=True + ) + ) + mock_core_api = mock.MagicMock(spec=client.CoreV1Api) + self.enter_context( + mock.patch.object( + gke_utils, "_get_k8s_core_api", return_value=mock_core_api + ) + ) + mock_core_api.list_namespaced_pod.return_value = client.V1PodList(items=[]) + with self.assertRaisesRegex( + RuntimeError, + "No pods found for Shared Pathways Service JobSet 'my-jobset'", + ): + gke_utils.verify_pathways_service_is_up( + cluster="my-cluster", + project="my-project", + region="us-central1", + pathways_service="my-jobset-pathways-head-0-0:29001", + tpu_type="tpuv6e:2x2", + tpu_count=1, + ) + + def test_verify_pathways_service_is_up_no_head_pod(self): + self.enter_context( + mock.patch.object( + gke_utils, "ensure_cluster_credentials", autospec=True + ) + ) + mock_core_api = mock.MagicMock(spec=client.CoreV1Api) + self.enter_context( + mock.patch.object( + gke_utils, "_get_k8s_core_api", return_value=mock_core_api + ) + ) + worker_pod = self._make_pod( + name="my-jobset-pathways-worker-0-0", + jobset_name="my-jobset", + replicated_job_name="pathways-worker", + ready=True, + ) + mock_core_api.list_namespaced_pod.return_value = client.V1PodList( + items=[worker_pod] + ) + with self.assertRaisesRegex( + RuntimeError, "head pod is not ready for JobSet 'my-jobset'" + ): + gke_utils.verify_pathways_service_is_up( + cluster="my-cluster", + project="my-project", + region="us-central1", + pathways_service="my-jobset-pathways-head-0-0:29001", + tpu_type="tpuv6e:2x2", + tpu_count=1, + ) + + def test_verify_pathways_service_is_up_head_pod_not_ready(self): + self.enter_context( + mock.patch.object( + gke_utils, "ensure_cluster_credentials", autospec=True + ) + ) + mock_core_api = mock.MagicMock(spec=client.CoreV1Api) + self.enter_context( + mock.patch.object( + gke_utils, "_get_k8s_core_api", return_value=mock_core_api + ) + ) + head_pod = self._make_pod( + name="my-jobset-pathways-head-0-0", + jobset_name="my-jobset", + replicated_job_name="pathways-head", + phase="Running", + ready=False, + ) + worker_pod = self._make_pod( + name="my-jobset-pathways-worker-0-0", + jobset_name="my-jobset", + replicated_job_name="pathways-worker", + ready=True, + ) + mock_core_api.list_namespaced_pod.return_value = client.V1PodList( + items=[head_pod, worker_pod] + ) + with self.assertRaisesRegex( + RuntimeError, + "head pod is not ready for JobSet 'my-jobset'", + ): + gke_utils.verify_pathways_service_is_up( + cluster="my-cluster", + project="my-project", + region="us-central1", + pathways_service="my-jobset-pathways-head-0-0:29001", + tpu_type="tpuv6e:2x2", + tpu_count=1, + ) + + def test_verify_pathways_service_is_up_no_worker_pods(self): + self.enter_context( + mock.patch.object( + gke_utils, "ensure_cluster_credentials", autospec=True + ) + ) + mock_core_api = mock.MagicMock(spec=client.CoreV1Api) + self.enter_context( + mock.patch.object( + gke_utils, "_get_k8s_core_api", return_value=mock_core_api + ) + ) + head_pod = self._make_pod( + name="my-jobset-pathways-head-0-0", + jobset_name="my-jobset", + replicated_job_name="pathways-head", + ready=True, + ) + mock_core_api.list_namespaced_pod.return_value = client.V1PodList( + items=[head_pod] + ) + with self.assertRaisesRegex( + RuntimeError, + r"Shared Pathways Service JobSet 'my-jobset' is not ready: requested 1" + r" slice\(s\) with 1 worker pod\(s\) per slice, but only 0 slice\(s\)", + ): + gke_utils.verify_pathways_service_is_up( + cluster="my-cluster", + project="my-project", + region="us-central1", + pathways_service="my-jobset-pathways-head-0-0:29001", + tpu_type="tpuv6e:2x2", + tpu_count=1, + ) + + def test_verify_pathways_service_is_up_zero_ready_worker_pods(self): + self.enter_context( + mock.patch.object( + gke_utils, "ensure_cluster_credentials", autospec=True + ) + ) + mock_core_api = mock.MagicMock(spec=client.CoreV1Api) + self.enter_context( + mock.patch.object( + gke_utils, "_get_k8s_core_api", return_value=mock_core_api + ) + ) + head_pod = self._make_pod( + name="my-jobset-pathways-head-0-0", + jobset_name="my-jobset", + replicated_job_name="pathways-head", + ready=True, + ) + worker_pod = self._make_pod( + name="my-jobset-pathways-worker-0-0", + jobset_name="my-jobset", + replicated_job_name="pathways-worker", + phase="Pending", + ready=False, + ) + mock_core_api.list_namespaced_pod.return_value = client.V1PodList( + items=[head_pod, worker_pod] + ) + with self.assertRaisesRegex( + RuntimeError, + r"Shared Pathways Service JobSet 'my-jobset' is not ready: requested 1" + r" slice\(s\) with 1 worker pod\(s\) per slice, but only 0 slice\(s\)" + r" \(0 worker pod\(s\)\) are ready", + ): + gke_utils.verify_pathways_service_is_up( + cluster="my-cluster", + project="my-project", + region="us-central1", + pathways_service="my-jobset-pathways-head-0-0:29001", + tpu_type="tpuv6e:2x2", + tpu_count=1, + ) + + def test_verify_pathways_service_is_up_insufficient_ready_slices(self): + self.enter_context( + mock.patch.object( + gke_utils, "ensure_cluster_credentials", autospec=True + ) + ) + mock_core_api = mock.MagicMock(spec=client.CoreV1Api) + self.enter_context( + mock.patch.object( + gke_utils, "_get_k8s_core_api", return_value=mock_core_api + ) + ) + head_pod = self._make_pod( + name="my-jobset-pathways-head-0-0", + jobset_name="my-jobset", + replicated_job_name="pathways-head", + ready=True, + ) + worker_pod_0 = self._make_pod( + name="my-jobset-pathways-worker-0-0", + jobset_name="my-jobset", + replicated_job_name="pathways-worker", + ready=True, + job_index="0", + ) + worker_pod_1 = self._make_pod( + name="my-jobset-pathways-worker-0-1", + jobset_name="my-jobset", + replicated_job_name="pathways-worker", + ready=True, + job_index="0", + ) + worker_pod_2 = self._make_pod( + name="my-jobset-pathways-worker-1-0", + jobset_name="my-jobset", + replicated_job_name="pathways-worker", + phase="Failed", + ready=False, + job_index="1", + ) + worker_pod_3 = self._make_pod( + name="my-jobset-pathways-worker-1-1", + jobset_name="my-jobset", + replicated_job_name="pathways-worker", + phase="Failed", + ready=False, + job_index="1", + ) + mock_core_api.list_namespaced_pod.return_value = client.V1PodList( + items=[ + head_pod, + worker_pod_0, + worker_pod_1, + worker_pod_2, + worker_pod_3, + ] + ) + with self.assertRaisesRegex( + RuntimeError, + r"requested 2 slice\(s\) with 2 worker pod\(s\) per slice, but only 1" + r" slice\(s\) \(2 worker pod\(s\)\) are ready", + ): + gke_utils.verify_pathways_service_is_up( + cluster="my-cluster", + project="my-project", + region="us-central1", + pathways_service="my-jobset-pathways-head-0-0:29001", + tpu_type="tpuv6e:2x4", + tpu_count=2, + ) + + def test_verify_pathways_service_is_up_partial_slice_not_ready(self): + self.enter_context( + mock.patch.object( + gke_utils, "ensure_cluster_credentials", autospec=True + ) + ) + mock_core_api = mock.MagicMock(spec=client.CoreV1Api) + self.enter_context( + mock.patch.object( + gke_utils, "_get_k8s_core_api", return_value=mock_core_api + ) + ) + head_pod = self._make_pod( + name="my-jobset-pathways-head-0-0", + jobset_name="my-jobset", + replicated_job_name="pathways-head", + ready=True, + ) + worker_pod_0 = self._make_pod( + name="my-jobset-pathways-worker-0-0", + jobset_name="my-jobset", + replicated_job_name="pathways-worker", + ready=True, + job_index="0", + ) + worker_pod_1 = self._make_pod( + name="my-jobset-pathways-worker-1-0", + jobset_name="my-jobset", + replicated_job_name="pathways-worker", + ready=True, + job_index="1", + ) + mock_core_api.list_namespaced_pod.return_value = client.V1PodList( + items=[head_pod, worker_pod_0, worker_pod_1] + ) + with self.assertRaisesRegex( + RuntimeError, + r"requested 1 slice\(s\) with 2 worker pod\(s\) per slice, but only 0" + r" slice\(s\) \(2 worker pod\(s\)\) are ready", + ): + gke_utils.verify_pathways_service_is_up( + cluster="my-cluster", + project="my-project", + region="us-central1", + pathways_service="my-jobset-pathways-head-0-0:29001", + tpu_type="tpuv6e:2x4", + tpu_count=1, + ) + + def test_verify_pathways_service_is_up_invalid_namespace(self): + with self.assertRaises(ValueError): + gke_utils.verify_pathways_service_is_up( + cluster="my-cluster", + project="my-project", + region="us-central1", + pathways_service="my-jobset-pathways-head-0-0:29001", + tpu_type="tpuv6e:2x2", + tpu_count=1, + namespace="invalid namespace!", + ) + if __name__ == "__main__": absltest.main() 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 9eb7620..2fd5fa4 100644 --- a/pathwaysutils/test/experimental/shared_pathways_service/isc_pathways_test.py +++ b/pathwaysutils/test/experimental/shared_pathways_service/isc_pathways_test.py @@ -26,7 +26,7 @@ def setUp(self): # cluster so that credentials are fetched. self.mock_is_current_kube_context = self.enter_context( mock.patch.object( - isc_pathways, "_is_current_kube_context", autospec=True + isc_pathways.gke_utils, "_is_current_kube_context", autospec=True ) ) self.mock_is_current_kube_context.return_value = False @@ -1419,89 +1419,5 @@ def test_connect_proxy_server_image_deprecation_warning(self): pass -class KubeConfigCredentialsTest(parameterized.TestCase): - """Tests for the kube config context helpers.""" - - def setUp(self): - super().setUp() - self.mock_get_current_kube_context = self.enter_context( - mock.patch.object( - isc_pathways.gke_utils, "get_current_kube_context", autospec=True - ) - ) - self.mock_fetch_creds = self.enter_context( - mock.patch.object( - isc_pathways.gke_utils, "fetch_cluster_credentials", autospec=True - ) - ) - - def test_is_current_kube_context_matching(self): - """Tests that a matching active context is detected.""" - self.mock_get_current_kube_context.return_value = ( - "test-cluster", - "test-project", - "test-region", - ) - - self.assertTrue( - isc_pathways._is_current_kube_context( - cluster="test-cluster", - project="test-project", - location="test-region", - ) - ) - - @parameterized.named_parameters( - ("different_cluster", ("other-cluster", "test-project", "test-region")), - ("different_project", ("test-cluster", "other-project", "test-region")), - ("different_location", ("test-cluster", "test-project", "other-region")), - ("non_gke_context", ("minikube", None, None)), - ("no_context", (None, None, None)), - ) - def test_is_current_kube_context_not_matching(self, current_context): - """Tests that a non-matching active context is detected.""" - self.mock_get_current_kube_context.return_value = current_context - - self.assertFalse( - isc_pathways._is_current_kube_context( - cluster="test-cluster", - project="test-project", - location="test-region", - ) - ) - - def test_ensure_cluster_credentials_skips_fetch_when_matching(self): - """Tests that credentials are not fetched for the active context.""" - self.mock_get_current_kube_context.return_value = ( - "test-cluster", - "test-project", - "test-region", - ) - - isc_pathways._ensure_cluster_credentials( - cluster="test-cluster", project="test-project", location="test-region" - ) - - self.mock_fetch_creds.assert_not_called() - - def test_ensure_cluster_credentials_fetches_when_not_matching(self): - """Tests that credentials are fetched for a different context.""" - self.mock_get_current_kube_context.return_value = ( - "test-cluster", - "other-project", - "test-region", - ) - - isc_pathways._ensure_cluster_credentials( - cluster="test-cluster", project="test-project", location="test-region" - ) - - self.mock_fetch_creds.assert_called_once_with( - cluster_name="test-cluster", - project_id="test-project", - location="test-region", - ) - - if __name__ == "__main__": absltest.main() diff --git a/pathwaysutils/test/experimental/shared_pathways_service/run_workload_test.py b/pathwaysutils/test/experimental/shared_pathways_service/run_workload_test.py index 7323000..71232cc 100644 --- a/pathwaysutils/test/experimental/shared_pathways_service/run_workload_test.py +++ b/pathwaysutils/test/experimental/shared_pathways_service/run_workload_test.py @@ -7,6 +7,7 @@ from absl.testing import absltest from absl.testing import flagsaver +from pathwaysutils.experimental.shared_pathways_service import gke_utils from pathwaysutils.experimental.shared_pathways_service import run_workload @@ -33,6 +34,14 @@ def __exit__( class RunTpuWorkloadTest(absltest.TestCase): + def setUp(self): + super().setUp() + self.mock_verify_service = self.enter_context( + mock.patch.object( + gke_utils, "verify_pathways_service_is_up", autospec=True + ) + ) + def test_run_workload_success(self): fake_instances = [] @@ -78,6 +87,44 @@ def fake_connect_fn(**kwargs: Any) -> Generator[FakeConnect, None, None]: ) mock_proc.wait.assert_called_once() + with self.subTest("Service verified before connecting"): + self.mock_verify_service.assert_called_once_with( + cluster="test-cluster", + project="test-project", + region="test-region", + pathways_service="test-service:1234", + tpu_type="tpuv6e:4x8", + tpu_count=1, + ) + + def test_run_command_fails_when_service_not_up(self): + self.mock_verify_service.side_effect = RuntimeError( + "No pods found for Shared Pathways Service JobSet 'test-service'" + ) + mock_connect_fn = mock.MagicMock() + mock_popen = self.enter_context( + mock.patch.object(subprocess, "Popen", autospec=True) + ) + + with self.assertRaisesRegex( + RuntimeError, + "No pods found for Shared Pathways Service JobSet 'test-service'", + ): + run_workload.run_command( + cluster="test-cluster", + project="test-project", + region="test-region", + gcs_bucket="test-bucket", + pathways_service="test-service-pathways-head-0-0:1234", + tpu_type="tpuv6e:4x8", + tpu_count=1, + command="echo hello", + connect_fn=mock_connect_fn, + ) + + mock_connect_fn.assert_not_called() + mock_popen.assert_not_called() + def test_run_command_runs_command_inside_context(self): """Verifies that the command is executed while the connection is active.""" connection_active_during_run = False