From 6f0c249240aea682386adc611322ff448928ab4c Mon Sep 17 00:00:00 2001 From: Pathways-on-Cloud Team Date: Tue, 13 Oct 2026 14:02:25 -0700 Subject: [PATCH] Discover and validate SPS colocated Python sidecar package versions. PiperOrigin-RevId: 998876736 --- .../deploy_pathways_service.py | 2 + .../shared_pathways_service/gke_utils.py | 122 +++++++++++++++ .../shared_pathways_service/isc_pathways.py | 18 ++- .../shared_pathways_service/validators.py | 142 +++++++++++------- .../shared_pathways_service/gke_utils_test.py | 124 ++++++++++++++- .../isc_pathways_test.py | 141 +++++++++++++++-- .../validators_test.py | 124 ++++++++++++--- 7 files changed, 578 insertions(+), 95 deletions(-) diff --git a/pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py b/pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py index d6e296b..44db48c 100644 --- a/pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py +++ b/pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py @@ -311,6 +311,8 @@ def main(argv: Sequence[str]) -> None: if len(argv) > 1: raise app.UsageError("Too many command-line arguments.") + logging.getLogger().setLevel(logging.INFO) + try: if ( flags.FLAGS["jax_version"].present diff --git a/pathwaysutils/experimental/shared_pathways_service/gke_utils.py b/pathwaysutils/experimental/shared_pathways_service/gke_utils.py index e7a9b04..b80e85e 100644 --- a/pathwaysutils/experimental/shared_pathways_service/gke_utils.py +++ b/pathwaysutils/experimental/shared_pathways_service/gke_utils.py @@ -15,6 +15,7 @@ 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 +from pathwaysutils.experimental.shared_pathways_service import validators import portpicker _logger = logging.getLogger(__name__) @@ -1044,3 +1045,124 @@ def verify_pathways_service_is_up( ready_slices, ready_worker_pods_count, ) + + +_PACKAGES_VERSION_SNIPPET = ( + "import jax, jaxlib, sys; " + "print(f'{sys.version_info.major}.{sys.version_info.minor}," + "{jax.__version__},{jaxlib.__version__}')" +) + + +def get_sidecar_versions( + jobset_name: str, + namespace: str = "default", + sidecar_image: str | None = None, +) -> tuple[str | None, validators.SidecarVersions]: + """Gets sidecar image and JAX/JAXLib/Python versions from the SPS instance. + + Queries a running `colocated-python-sidecar` container in the JobSet for the + installed Python, JAX, and JAXLib versions. + + Args: + jobset_name: The name of the Pathways JobSet. + namespace: The Kubernetes namespace of the JobSet. + sidecar_image: Optional pre-fetched sidecar image string. If None, fetched + from the JobSet via `get_pathways_service_images`. + + Returns: + A tuple of (sidecar_image, SidecarVersions). If colocated python sidecar + is not enabled, sidecar_image is None and SidecarVersions is empty. + + Raises: + RuntimeError: If sidecar versions cannot be queried from a running worker + pod in the JobSet. + """ + _validate_k8s_name(namespace) + _validate_k8s_name(jobset_name) + if sidecar_image is None: + _, sidecar_image = get_pathways_service_images( + f"{jobset_name}-pathways-head:29001", namespace=namespace + ) + if not sidecar_image: + _logger.warning( + "JobSet '%s' does not define a sidecar image for the worker pods.", + jobset_name, + ) + return (None, validators.SidecarVersions()) + + cmd = [ + "kubectl", + "get", + "pods", + "-n", + namespace, + "-l", + f"jobset.sigs.k8s.io/jobset-name={jobset_name}", + "--field-selector=status.phase=Running", + "-o", + "jsonpath={.items[*].metadata.name}", + ] + try: + pod_result = subprocess.run( + cmd, capture_output=True, text=True, check=True, timeout=10 + ) + pod_names = [ + p + for p in pod_result.stdout.strip().split() + if not p.startswith(f"{jobset_name}-pathways-head-") + ] + except Exception as e: + raise RuntimeError( + f"Failed to list running worker pods for JobSet '{jobset_name}' in" + f" namespace '{namespace}': {e}" + ) from e + + if not pod_names: + raise RuntimeError( + f"No running worker pods found for JobSet '{jobset_name}' in" + f" namespace '{namespace}'." + ) + + last_error = None + for pod_name in pod_names[:3]: + try: + _validate_k8s_name(pod_name) + exec_cmd = [ + "kubectl", + "exec", + "-n", + namespace, + f"pod/{pod_name}", + "-c", + "colocated-python-sidecar", + "--", + "python3", + "-c", + _PACKAGES_VERSION_SNIPPET, + ] + exec_result = subprocess.run( + exec_cmd, capture_output=True, text=True, check=True, timeout=10 + ) + py_ver, jax_ver, jaxlib_ver = [ + part.strip() for part in exec_result.stdout.strip().split(",") + ] + if py_ver and jax_ver and jaxlib_ver: + return ( + sidecar_image, + validators.SidecarVersions( + python_version=py_ver, + jax_version=jax_ver, + jaxlib_version=jaxlib_ver, + ), + ) + except Exception as e: # pylint: disable=broad-exception-caught + last_error = e + _logger.debug( + "Could not query live sidecar pod %s for versions: %r", pod_name, e + ) + + raise RuntimeError( + f"Failed to query sidecar versions from running worker pods for JobSet" + f" '{jobset_name}' in namespace '{namespace}'." + ) from last_error diff --git a/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py b/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py index 395a183..f6b8085 100644 --- a/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py +++ b/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py @@ -608,8 +608,21 @@ def connect( proxy_server_image = compatible_proxy_image proxy_options_obj = ProxyOptions.from_list(proxy_options) - if proxy_options_obj.sidecar and sidecar_image: - validators.validate_sidecar_image_versions(sidecar_image) + if sidecar_image: + jobset_name = gke_utils.extract_jobset_name(pathways_service) + _, sidecar_versions = gke_utils.get_sidecar_versions( + jobset_name, sidecar_image=sidecar_image + ) + _logger.info( + "\n%s", + validators.format_sidecar_versions( + pathways_service, sidecar_image, sidecar_versions + ), + ) + if proxy_options_obj.sidecar: + validators.validate_sidecar_image_versions( + sidecar_image, sidecar_versions=sidecar_versions + ) _logger.info("Validation complete.") if not proxy_job_name: @@ -653,3 +666,4 @@ def connect( " _wait_for_placement." ) yield t + diff --git a/pathwaysutils/experimental/shared_pathways_service/validators.py b/pathwaysutils/experimental/shared_pathways_service/validators.py index bcba4b3..2885b94 100644 --- a/pathwaysutils/experimental/shared_pathways_service/validators.py +++ b/pathwaysutils/experimental/shared_pathways_service/validators.py @@ -1,17 +1,24 @@ """Validation functions for Shared Pathways Service.""" from collections.abc import Iterable, Mapping +import dataclasses +import importlib import logging import re import sys from typing import Any from absl import flags -import jax _logger = logging.getLogger(__name__) -_PYTHON_VERSION_REGEX = r"python[-_]?(\d+\.\d+(?:\.\d+)*)" -_JAX_VERSION_REGEX = r"jax[-_]?(\d+\.\d+(?:\.\d+)*)" + +@dataclasses.dataclass(frozen=True) +class SidecarVersions: + """Holds Python, JAX, and JAXLib versions for a colocated python sidecar.""" + + python_version: str | None = None + jax_version: str | None = None + jaxlib_version: str | None = None def validate_proxy_options(proxy_options: Iterable[str] | None) -> None: @@ -124,51 +131,48 @@ def validate_xla_flags(xla_flags: Iterable[str] | None) -> None: ) -def validate_sidecar_image_versions(sidecar_image: str) -> None: - """Checks compatibility of sidecar image versions with user environment. - - Compares the Python and JAX versions in the sidecar image tag with the user - environment's Python and JAX versions. +def format_sidecar_versions( + pathways_service: str, + sidecar_image: str, + sidecar_versions: SidecarVersions, +) -> str: + """Formats the sidecar versions into a human-readable summary string.""" + lines = [ + ( + "Colocated Python sidecar found for Pathways service" + f" '{pathways_service}':" + ), + f" Sidecar Image: {sidecar_image}", + ] + if sidecar_versions.python_version: + lines.append(f" Python: {sidecar_versions.python_version}") + if sidecar_versions.jax_version: + lines.append(f" JAX: {sidecar_versions.jax_version}") + if sidecar_versions.jaxlib_version: + lines.append(f" JAXLib: {sidecar_versions.jaxlib_version}") + return "\n".join(lines) + + +def validate_sidecar_image_versions( + sidecar_image: str, sidecar_versions: SidecarVersions +) -> None: + """Checks compatibility of sidecar package versions with user environment. + + Compares the Python, JAX, and JAXLib versions in the sidecar container with + the user environment's Python, JAX, and JAXLib versions. Args: - sidecar_image: The sidecar image string, e.g., - "us-docker.pkg.dev/.../sidecar:20260423-python_3.12-jax_0.10.0". + sidecar_image: The sidecar image string. + sidecar_versions: The SidecarVersions queried from the sidecar container. Raises: - ValueError: If the sidecar image Python or JAX versions do not match the + ValueError: If the sidecar Python, JAX, or JAXLib versions do not match the user environment. """ _logger.info( "Checking sidecar image version compatibility: %s", sidecar_image ) - parts = sidecar_image.rsplit(":", 1) - if len(parts) < 2: - _logger.warning( - "No tag found in sidecar image: %s. Skipping version validation.", - sidecar_image, - ) - return - tag = parts[1] - - sidecar_python_match = re.search( - _PYTHON_VERSION_REGEX, tag, re.IGNORECASE - ) - sidecar_jax_match = re.search( - _JAX_VERSION_REGEX, tag, re.IGNORECASE - ) - if not sidecar_python_match and not sidecar_jax_match: - _logger.warning( - "No Python or JAX versions found in sidecar image tag: %s. Skipping " - "version validation.", - tag, - ) - return - - def clean_version(version_str: str) -> str: - match = re.match(r"^(\d+(?:\.\d+)*)", version_str) - return match.group(1) if match else version_str - def versions_match(sidecar_ver: str, env_ver: str) -> bool: sidecar_parts = sidecar_ver.split(".") env_parts = env_ver.split(".") @@ -177,8 +181,15 @@ def versions_match(sidecar_ver: str, env_ver: str) -> bool: return False return sidecar_parts[:compare_len] == env_parts[:compare_len] - if sidecar_python_match: - sidecar_python = clean_version(sidecar_python_match.group(1)) + install_hint = "" + if sidecar_versions.jax_version and sidecar_versions.jaxlib_version: + install_hint = ( + f" by running: pip install jax=={sidecar_versions.jax_version}" + f" jaxlib=={sidecar_versions.jaxlib_version}" + ) + + if sidecar_versions.python_version: + sidecar_python = sidecar_versions.python_version.strip() env_python = ( f"{sys.version_info.major}.{sys.version_info.minor}." f"{sys.version_info.micro}" @@ -198,21 +209,38 @@ def versions_match(sidecar_ver: str, env_ver: str) -> bool: env_python, ) - if sidecar_jax_match: - sidecar_jax = clean_version(sidecar_jax_match.group(1)) - env_jax = clean_version(jax.__version__) - if not versions_match(sidecar_jax, env_jax): - raise ValueError( - f"JAX version mismatch: sidecar image matches JAX version " - f"{sidecar_jax}, but the user environment is running JAX " - f"{env_jax}. Either rebuild the sidecar image with a matching " - "JAX version or update the user environment to match the sidecar " - "image." + for pkg_name, display_name, raw_sidecar_ver in ( + ("jax", "JAX", sidecar_versions.jax_version), + ("jaxlib", "JAXLib", sidecar_versions.jaxlib_version), + ): + if not raw_sidecar_ver: + continue + sidecar_ver = raw_sidecar_ver.strip() + env_ver = None + try: + mod = sys.modules.get(pkg_name) + if mod is None: + mod = importlib.import_module(pkg_name) + if hasattr(mod, "__version__"): + env_ver = mod.__version__.strip() + except (ImportError, AttributeError): + env_ver = None + + if env_ver is not None: + if sidecar_ver != env_ver: + raise ValueError( + f"{display_name} version mismatch: sidecar image matches" + f" {display_name} version {sidecar_ver}, but the user environment" + f" is running {display_name} {env_ver}. Either rebuild the sidecar" + f" image with a matching {display_name} version or update the user" + f" environment to match the sidecar image{install_hint}." + ) + _logger.info( + "%s version match: sidecar image matches %s version %s, and the user" + " environment is running %s %s.", + display_name, + display_name, + sidecar_ver, + display_name, + env_ver, ) - _logger.info( - "JAX version match: sidecar image matches JAX version %s, and the user" - " environment is running JAX %s.", - sidecar_jax, - env_jax, - ) - 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 c556fe1..438211a 100644 --- a/pathwaysutils/test/experimental/shared_pathways_service/gke_utils_test.py +++ b/pathwaysutils/test/experimental/shared_pathways_service/gke_utils_test.py @@ -2123,7 +2123,129 @@ def test_verify_pathways_service_is_up_invalid_namespace(self): namespace="invalid namespace!", ) + def test_get_sidecar_versions_no_sidecar(self): + with mock.patch.object( + gke_utils, + "get_pathways_service_images", + return_value=("server_img", None), + ), self.assertLogs(gke_utils._logger, level="WARNING") as log_cm: + sidecar_img, versions = gke_utils.get_sidecar_versions("my-jobset") + self.assertIsNone(sidecar_img) + self.assertIsNone(versions.python_version) + self.assertIsNone(versions.jax_version) + self.assertIsNone(versions.jaxlib_version) + self.assertTrue( + any( + "does not define a sidecar image" in msg + for msg in log_cm.output + ) + ) + + def test_get_sidecar_versions_live_query_success(self): + with mock.patch.object( + gke_utils, + "get_pathways_service_images", + return_value=( + "server_img", + "us-docker.pkg.dev/repo/sidecar:20260423-python_3.12-jax_0.10.0", + ), + ), mock.patch.object(subprocess, "run") as mock_run: + mock_run.side_effect = [ + mock.Mock( + stdout="my-jobset-pathways-head-0-0 my-jobset-worker-0-0\n", + returncode=0, + ), + mock.Mock(stdout="3.12,0.10.0,0.10.0\n", returncode=0), + ] + sidecar_img, versions = gke_utils.get_sidecar_versions("my-jobset") + self.assertEqual( + sidecar_img, + "us-docker.pkg.dev/repo/sidecar:20260423-python_3.12-jax_0.10.0", + ) + self.assertEqual(versions.python_version, "3.12") + self.assertEqual(versions.jax_version, "0.10.0") + self.assertEqual(versions.jaxlib_version, "0.10.0") + + def test_get_sidecar_versions_jobset_name_containing_pathways_head(self): + jobset_name = "test-pathways-head-eval" + with mock.patch.object( + gke_utils, + "get_pathways_service_images", + return_value=("server_img", "us-docker.pkg.dev/repo/sidecar:latest"), + ), mock.patch.object(subprocess, "run") as mock_run: + mock_run.side_effect = [ + mock.Mock( + stdout=( + f"{jobset_name}-pathways-head-0-0" + f" {jobset_name}-pathways-worker-0-0\n" + ), + returncode=0, + ), + mock.Mock(stdout="3.12,0.11.1,0.11.1\n", returncode=0), + ] + _, versions = gke_utils.get_sidecar_versions(jobset_name) + self.assertEqual(versions.python_version, "3.12") + self.assertEqual(versions.jax_version, "0.11.1") + self.assertEqual(versions.jaxlib_version, "0.11.1") + + def test_get_sidecar_versions_first_pod_timeout_retries_second_pod(self): + with mock.patch.object( + gke_utils, + "get_pathways_service_images", + return_value=( + "server_img", + "us-docker.pkg.dev/repo/sidecar:20260423-python_3.12-jax_0.10.0", + ), + ), mock.patch.object(subprocess, "run") as mock_run: + mock_run.side_effect = [ + mock.Mock(stdout="pod-0 pod-1\n", returncode=0), + subprocess.TimeoutExpired(cmd=["kubectl", "exec"], timeout=10), + mock.Mock(stdout="3.12,0.10.0,0.10.0\n", returncode=0), + ] + sidecar_img, versions = gke_utils.get_sidecar_versions("my-jobset") + self.assertEqual( + sidecar_img, + "us-docker.pkg.dev/repo/sidecar:20260423-python_3.12-jax_0.10.0", + ) + self.assertEqual(versions.python_version, "3.12") + self.assertEqual(versions.jax_version, "0.10.0") + self.assertEqual(versions.jaxlib_version, "0.10.0") + self.assertEqual(mock_run.call_count, 3) + + def test_get_sidecar_versions_no_running_worker_pods_raises(self): + with mock.patch.object( + gke_utils, + "get_pathways_service_images", + return_value=("server_img", "us-docker.pkg.dev/repo/sidecar:latest"), + ), mock.patch.object(subprocess, "run") as mock_run: + mock_run.return_value = mock.Mock( + stdout="my-jobset-pathways-head-0-0\n", returncode=0 + ) + with self.assertRaisesRegex( + RuntimeError, "No running worker pods found for JobSet 'my-jobset'" + ): + gke_utils.get_sidecar_versions("my-jobset") + + def test_get_sidecar_versions_bounds_pod_attempts_to_three(self): + with mock.patch.object( + gke_utils, + "get_pathways_service_images", + return_value=("server_img", "us-docker.pkg.dev/repo/sidecar:latest"), + ), mock.patch.object(subprocess, "run") as mock_run: + mock_run.side_effect = [ + mock.Mock(stdout="pod-0 pod-1 pod-2 pod-3 pod-4\n", returncode=0), + subprocess.TimeoutExpired(cmd=["kubectl", "exec"], timeout=10), + subprocess.TimeoutExpired(cmd=["kubectl", "exec"], timeout=10), + subprocess.TimeoutExpired(cmd=["kubectl", "exec"], timeout=10), + ] + with self.assertRaisesRegex( + RuntimeError, + "Failed to query sidecar versions from running worker pods", + ): + gke_utils.get_sidecar_versions("my-jobset") + # 1 call for kubectl get pods + 3 calls for kubectl exec + self.assertEqual(mock_run.call_count, 4) + 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 2fd5fa4..36456ce 100644 --- a/pathwaysutils/test/experimental/shared_pathways_service/isc_pathways_test.py +++ b/pathwaysutils/test/experimental/shared_pathways_service/isc_pathways_test.py @@ -10,6 +10,7 @@ from absl.testing import absltest from absl.testing import parameterized from pathwaysutils.experimental.shared_pathways_service import isc_pathways +from pathwaysutils.experimental.shared_pathways_service import validators class ISCPathwaysTest(parameterized.TestCase): @@ -825,6 +826,7 @@ def test_connect_with_sidecar_validation_success(self): isc_pathways.gke_utils, "fetch_cluster_credentials", autospec=True ) ) + sidecar_img = "us-docker.pkg.dev/.../sidecar:20260423-python_3.12-jax_0.10.0" mock_get_images = self.enter_context( mock.patch.object( isc_pathways.gke_utils, "get_pathways_service_images", autospec=True @@ -832,7 +834,18 @@ def test_connect_with_sidecar_validation_success(self): ) mock_get_images.return_value = ( "us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest", - "us-docker.pkg.dev/.../sidecar:20260423-python_3.12-jax_0.10.0", + sidecar_img, + ) + expected_versions = validators.SidecarVersions( + python_version="3.12", jax_version="0.10.0", jaxlib_version="0.10.0" + ) + mock_get_sidecar_versions = self.enter_context( + mock.patch.object( + isc_pathways.gke_utils, + "get_sidecar_versions", + autospec=True, + return_value=(sidecar_img, expected_versions), + ) ) mock_validate_versions = self.enter_context( mock.patch.object( @@ -852,24 +865,109 @@ def test_connect_with_sidecar_validation_success(self): 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_options=["sidecar:true"], - ): - pass + with self.assertLogs(isc_pathways._logger, level="INFO") as log_cm: + 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_options=["sidecar:true"], + ): + pass + self.assertTrue( + any( + "Colocated Python sidecar found for Pathways service" in msg + and "JAX: 0.10.0" in msg + for msg in log_cm.output + ) + ) mock_get_images.assert_called_once_with( "test-service-pathways-head:1234" ) + mock_get_sidecar_versions.assert_called_once_with( + "test-service", sidecar_image=sidecar_img + ) mock_validate_versions.assert_called_once_with( - "us-docker.pkg.dev/.../sidecar:20260423-python_3.12-jax_0.10.0" + sidecar_img, sidecar_versions=expected_versions ) + def test_connect_logs_sidecar_versions_by_default_when_sidecar_image_present( + self, + ): + self.enter_context(mock.patch.dict(os.environ, {"USER": "testuser"})) + self.enter_context( + mock.patch.object( + isc_pathways.gke_utils, "fetch_cluster_credentials", autospec=True + ) + ) + sidecar_img = "us-docker.pkg.dev/.../sidecar:20260423-python_3.12-jax_0.10.0" + 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", + sidecar_img, + ) + expected_versions = validators.SidecarVersions( + python_version="3.12", jax_version="0.10.0", jaxlib_version="0.10.0" + ) + mock_get_sidecar_versions = self.enter_context( + mock.patch.object( + isc_pathways.gke_utils, + "get_sidecar_versions", + autospec=True, + return_value=(sidecar_img, expected_versions), + ) + ) + mock_validate_versions = self.enter_context( + mock.patch.object( + isc_pathways.validators, + "validate_sidecar_image_versions", + autospec=True, + ) + ) + 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 self.assertLogs(isc_pathways._logger, level="INFO") as log_cm: + 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}, + ): + pass + + self.assertTrue( + any( + "Colocated Python sidecar found for Pathways service" in msg + and "JAX: 0.10.0" in msg + for msg in log_cm.output + ) + ) + mock_get_images.assert_called_once_with( + "test-service-pathways-head:1234" + ) + mock_get_sidecar_versions.assert_called_once_with( + "test-service", sidecar_image=sidecar_img + ) + mock_validate_versions.assert_not_called() + def test_connect_with_sidecar_validation_mismatch_raises_error(self): self.enter_context(mock.patch.dict(os.environ, {"USER": "testuser"})) self.enter_context( @@ -877,6 +975,7 @@ def test_connect_with_sidecar_validation_mismatch_raises_error(self): isc_pathways.gke_utils, "fetch_cluster_credentials", autospec=True ) ) + sidecar_img = "us-docker.pkg.dev/.../sidecar:20260423-python_3.12-jax_0.10.0" mock_get_images = self.enter_context( mock.patch.object( isc_pathways.gke_utils, "get_pathways_service_images", autospec=True @@ -884,7 +983,18 @@ def test_connect_with_sidecar_validation_mismatch_raises_error(self): ) mock_get_images.return_value = ( "us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest", - "us-docker.pkg.dev/.../sidecar:20260423-python_3.12-jax_0.10.0", + sidecar_img, + ) + expected_versions = validators.SidecarVersions( + python_version="3.12", jax_version="0.10.0", jaxlib_version="0.10.0" + ) + mock_get_sidecar_versions = self.enter_context( + mock.patch.object( + isc_pathways.gke_utils, + "get_sidecar_versions", + autospec=True, + return_value=(sidecar_img, expected_versions), + ) ) mock_validate_versions = self.enter_context( mock.patch.object( @@ -910,8 +1020,11 @@ def test_connect_with_sidecar_validation_mismatch_raises_error(self): mock_get_images.assert_called_once_with( "test-service-pathways-head:1234" ) + mock_get_sidecar_versions.assert_called_once_with( + "test-service", sidecar_image=sidecar_img + ) mock_validate_versions.assert_called_once_with( - "us-docker.pkg.dev/.../sidecar:20260423-python_3.12-jax_0.10.0" + sidecar_img, sidecar_versions=expected_versions ) @parameterized.named_parameters( diff --git a/pathwaysutils/test/experimental/shared_pathways_service/validators_test.py b/pathwaysutils/test/experimental/shared_pathways_service/validators_test.py index 484be59..1a2efac 100644 --- a/pathwaysutils/test/experimental/shared_pathways_service/validators_test.py +++ b/pathwaysutils/test/experimental/shared_pathways_service/validators_test.py @@ -1,5 +1,6 @@ """Tests for validation functions for the Shared Pathways service.""" +import sys from unittest import mock from absl import flags @@ -265,22 +266,18 @@ def test_validate_sidecar_image_versions_success(self): mock_sys_info.major = 3 mock_sys_info.minor = 12 mock_sys_info.micro = 8 - with mock.patch("sys.version_info", mock_sys_info), mock.patch( - "jax.__version__", "0.10.0" + mock_jax = mock.Mock(__version__="0.10.0") + mock_jaxlib = mock.Mock(__version__="0.10.0") + with mock.patch("sys.version_info", mock_sys_info), mock.patch.dict( + sys.modules, {"jax": mock_jax, "jaxlib": mock_jaxlib} ): - # Python 3.12 (matches 3.12.8), JAX 0.10.0 (matches 0.10.0) validators.validate_sidecar_image_versions( - "us-docker.pkg.dev/.../sidecar:20260423-python_3.12-jax_0.10.0" - ) - - # Python omitted, JAX 0.10 (matches 0.10.0) - validators.validate_sidecar_image_versions( - "us-docker.pkg.dev/.../sidecar:20260423-jax_0.10" - ) - - # No version info - validators.validate_sidecar_image_versions( - "us-docker.pkg.dev/.../sidecar:latest" + "us-docker.pkg.dev/.../sidecar:20260423-python_3.12-jax_0.10.0", + sidecar_versions=validators.SidecarVersions( + python_version="3.12", + jax_version="0.10.0", + jaxlib_version="0.10.0", + ), ) def test_validate_sidecar_image_versions_python_mismatch(self): @@ -288,12 +285,19 @@ def test_validate_sidecar_image_versions_python_mismatch(self): mock_sys_info.major = 3 mock_sys_info.minor = 11 mock_sys_info.micro = 5 - with mock.patch("sys.version_info", mock_sys_info), mock.patch( - "jax.__version__", "0.10.0" + mock_jax = mock.Mock(__version__="0.10.0") + mock_jaxlib = mock.Mock(__version__="0.10.0") + with mock.patch("sys.version_info", mock_sys_info), mock.patch.dict( + sys.modules, {"jax": mock_jax, "jaxlib": mock_jaxlib} ): with self.assertRaisesRegex(ValueError, "Python version mismatch"): validators.validate_sidecar_image_versions( - "us-docker.pkg.dev/.../sidecar:20260423-python_3.12-jax_0.10.0" + "us-docker.pkg.dev/.../sidecar:20260423-python_3.12-jax_0.10.0", + sidecar_versions=validators.SidecarVersions( + python_version="3.12", + jax_version="0.10.0", + jaxlib_version="0.10.0", + ), ) def test_validate_sidecar_image_versions_jax_mismatch(self): @@ -301,14 +305,92 @@ def test_validate_sidecar_image_versions_jax_mismatch(self): mock_sys_info.major = 3 mock_sys_info.minor = 12 mock_sys_info.micro = 8 - with mock.patch("sys.version_info", mock_sys_info), mock.patch( - "jax.__version__", "0.9.0" + mock_jax = mock.Mock(__version__="0.9.0") + mock_jaxlib = mock.Mock(__version__="0.9.0") + with mock.patch("sys.version_info", mock_sys_info), mock.patch.dict( + sys.modules, {"jax": mock_jax, "jaxlib": mock_jaxlib} + ): + with self.assertRaisesRegex( + ValueError, + r"JAX version mismatch.*pip install jax==0\.10\.0 jaxlib==0\.10\.0", + ): + validators.validate_sidecar_image_versions( + "us-docker.pkg.dev/.../sidecar:20260423-python_3.12-jax_0.10.0", + sidecar_versions=validators.SidecarVersions( + python_version="3.12", + jax_version="0.10.0", + jaxlib_version="0.10.0", + ), + ) + + def test_validate_sidecar_image_versions_jaxlib_mismatch(self): + mock_sys_info = mock.Mock() + mock_sys_info.major = 3 + mock_sys_info.minor = 12 + mock_sys_info.micro = 8 + mock_jax = mock.Mock(__version__="0.10.0") + mock_jaxlib = mock.Mock(__version__="0.9.0") + with mock.patch("sys.version_info", mock_sys_info), mock.patch.dict( + sys.modules, {"jax": mock_jax, "jaxlib": mock_jaxlib} + ): + with self.assertRaisesRegex( + ValueError, + r"JAXLib version mismatch.*pip install jax==0\.10\.0" + r" jaxlib==0\.10\.0", + ): + validators.validate_sidecar_image_versions( + "us-docker.pkg.dev/.../sidecar@sha256:80b6718", + sidecar_versions=validators.SidecarVersions( + python_version="3.12", + jax_version="0.10.0", + jaxlib_version="0.10.0", + ), + ) + + def test_validate_sidecar_image_versions_suffix_mismatch(self): + mock_sys_info = mock.Mock() + mock_sys_info.major = 3 + mock_sys_info.minor = 12 + mock_sys_info.micro = 8 + mock_jax = mock.Mock(__version__="0.11.0") + mock_jaxlib = mock.Mock(__version__="0.11.0rc2") + with mock.patch("sys.version_info", mock_sys_info), mock.patch.dict( + sys.modules, {"jax": mock_jax, "jaxlib": mock_jaxlib} ): - with self.assertRaisesRegex(ValueError, "JAX version mismatch"): + with self.assertRaisesRegex( + ValueError, + r"JAX version mismatch.*0\.11\.0rc2.*0\.11\.0", + ): validators.validate_sidecar_image_versions( - "us-docker.pkg.dev/.../sidecar:20260423-python_3.12-jax_0.10.0" + "us-docker.pkg.dev/.../sidecar:latest", + sidecar_versions=validators.SidecarVersions( + python_version="3.12", + jax_version="0.11.0rc2", + jaxlib_version="0.11.0rc2", + ), ) + def test_format_sidecar_versions(self): + versions = validators.SidecarVersions( + python_version="3.12", + jax_version="0.11.1", + jaxlib_version="0.11.1", + ) + text_out = validators.format_sidecar_versions( + pathways_service="test-service:29001", + sidecar_image="repo/sidecar:tag", + sidecar_versions=versions, + ) + self.assertEqual( + text_out, + "Colocated Python sidecar found for Pathways service" + " 'test-service:29001':\n" + " Sidecar Image: repo/sidecar:tag\n" + " Python: 3.12\n" + " JAX: 0.11.1\n" + " JAXLib: 0.11.1", + ) + if __name__ == "__main__": absltest.main()