Discover and validate SPS colocated Python sidecar package versions.

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

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

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
+

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,
-    )
-