Deprecate proxy_server_image and auto-detect compatible proxy server image from Pathways service

PiperOrigin-RevId: 974690013

Author: guptaaka

pathwaysutils/experimental/shared_pathways_service/gke_utils.py

@@ -1,13 +1,15 @@
 """GKE utils for deploying and managing the Pathways proxy."""
 
-import json
+import functools
 import logging
 import re
 import socket
 import subprocess
 import time
 import urllib.parse
 
+from kubernetes import client
+from kubernetes import config as k8s_config
 import portpicker
 
 _logger = logging.getLogger(__name__)
@@ -477,57 +479,116 @@ def is_local_port_free(port: int) -> bool:
   return portpicker.is_port_free(port)
 
 
-def get_worker_sidecar_image(
+@functools.lru_cache(maxsize=1)
+def _init_k8s_config() -> None:
+  """Initializes the Kubernetes configuration."""
+  try:
+    k8s_config.load_kube_config()
+  except Exception:  # pylint: disable=broad-except
+    try:
+      k8s_config.load_incluster_config()
+    except Exception as e:
+      raise RuntimeError("Failed to load Kubernetes configuration") from e
+
+
+@functools.lru_cache(maxsize=1)
+def _get_k8s_core_api() -> client.CoreV1Api:
+  """Initializes and returns the Kubernetes CoreV1Api client."""
+  _init_k8s_config()
+  return client.CoreV1Api()
+
+
+@functools.lru_cache(maxsize=1)
+def _get_k8s_custom_objects_api() -> client.CustomObjectsApi:
+  """Initializes and returns the Kubernetes CustomObjectsApi client."""
+  _init_k8s_config()
+  return client.CustomObjectsApi()
+
+
+def get_pathways_service_images(
     pathways_service: str, namespace: str = "default"
-) -> str | None:
-  """Gets the image of the sidecar container used by the workers."""
+) -> 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)
 
   # Try to extract the jobset name from the Pathways service hostname.
-  jobset_name = None
-  if "-pathways-head" in pathways_head_hostname:
-    jobset_name = pathways_head_hostname.split("-pathways-head")[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:"
+        " <jobset_name>-pathways-head"
+    )
+  jobset_name = pathways_head_hostname.split("-pathways-head")[0]
 
-  command = ["kubectl", "get", "pods", "-n", namespace, "-o", "json"]
   try:
-    result = subprocess.run(
-        command,
-        check=True,
-        capture_output=True,
-        text=True,
+    custom_api = _get_k8s_custom_objects_api()
+    jobset = custom_api.get_namespaced_custom_object(
+        group="jobset.x-k8s.io",
+        version="v1alpha2",
+        namespace=namespace,
+        plural="jobsets",
+        name=jobset_name,
     )
-  except subprocess.CalledProcessError as e:
-    _logger.exception("Failed to get pods. kubectl output:\n%r", e.stderr)
-    return None
+  except Exception as e:
+    _logger.exception("Failed to get JobSet: %r", e)
+    raise
+
+  server_image = None
+  sidecar_image = None
+
+  # Find the worker job and extract both images
+  for job in jobset.get("spec", {}).get("replicatedJobs", []):
+    if job.get("name") in ("pathways-worker", "worker"):
+      spec = (
+          job.get("template", {})
+          .get("spec", {})
+          .get("template", {})
+          .get("spec", {})
+      )
+      containers = spec.get("containers", []) + spec.get("initContainers", [])
+      for c in containers:
+        if c.get("name") == "pathways-worker" and c.get("image"):
+          server_image = c["image"]
+        elif c.get("name") == "colocated-python-sidecar" and c.get("image"):
+          sidecar_image = c["image"]
+      break
+
+  if not server_image:
+    raise RuntimeError(
+        "Failed to get server image of the worker job for Pathways service:"
+        f" {pathways_service} in namespace: {namespace}"
+    )
+
+  return (server_image, sidecar_image)
+
+
+def get_compatible_proxy_server_image(server_image: str) -> str:
+  """Converts a Pathways server image to its compatible proxy server image."""
+  if not server_image:
+    return server_image
+
+  # Extract tag or digest if present.
+  if "@" in server_image:
+    repo, tag_or_digest = server_image.split("@", 1)
+    sep = "@"
+  else:
+    last_slash = server_image.rfind("/")
+    if ":" in server_image[last_slash + 1:]:
+      repo, tag_or_digest = server_image.rsplit(":", 1)
+      sep = ":"
+    else:
+      repo = server_image
+      tag_or_digest = None
+      sep = ""
+
+  prefix, sep_slash, last_component = repo.rpartition("/")
+  new_last_component = last_component.replace("server", "proxy_server")
+
+  new_repo = f"{prefix}{sep_slash}{new_last_component}"
+  if tag_or_digest is not None:
+    return f"{new_repo}{sep}{tag_or_digest}"
+  return new_repo
 
-  try:
-    pods_data = json.loads(result.stdout)
-  except json.JSONDecodeError as e:
-    _logger.exception("Failed to parse kubectl get pods output: %r", e)
-    return None
-
-  items = pods_data.get("items", [])
-
-  # Look for pods belonging to the jobset and having the sidecar
-  # container/initContainer.
-  if jobset_name:
-    for pod in items:
-      metadata = pod.get("metadata", {})
-      labels = metadata.get("labels", {})
-      pod_jobset_name = labels.get("jobset.sigs.k8s.io/jobset-name")
-      pod_name = metadata.get("name", "")
-
-      if pod_jobset_name == jobset_name or pod_name.startswith(jobset_name):
-        spec = pod.get("spec", {})
-        for container in spec.get("initContainers", []) + spec.get(
-            "containers", []
-        ):
-          if container.get("name") == "colocated-python-sidecar":
-            image = container.get("image")
-            if image:
-              return image
-
-  return None
 
 

pathwaysutils/experimental/shared_pathways_service/isc_pathways.py

@@ -12,6 +12,7 @@
 import threading
 import time
 from typing import Any
+import warnings
 
 import jax
 import jax.extend.backend as jax_backend
@@ -455,7 +456,7 @@ def connect(
     pathways_service: str,
     expected_tpu_instances: Mapping[str, int],
     proxy_job_name: str | None = None,
-    proxy_server_image: str = DEFAULT_PROXY_IMAGE,
+    proxy_server_image: str | None = None,
     proxy_options: Sequence[str] | None = None,
     collect_service_metrics: bool = False,
 ) -> Iterator["_ISCPathways"]:
@@ -471,8 +472,10 @@ def connect(
       of instances. For example: {"tpuv6e:2x2": 2}
     proxy_job_name: The name to use for the deployed proxy. If not provided, a
       random name will be generated.
-    proxy_server_image: The proxy server image to use. If not provided, a
-      default will be used.
+    proxy_server_image: (Deprecated) The proxy server image to use. If not
+      provided, it will be auto-detected from the Pathways service. If the given
+      proxy image is incompatible with the Pathways service, it will be
+      replaced with the compatible proxy image.
     proxy_options: Configuration options for the Pathways proxy. If not
       provided, no extra options will be used.
     collect_service_metrics: Whether to collect usage metrics for Shared
@@ -481,22 +484,44 @@ def connect(
   Yields:
     The Pathways manager.
   """
+  if proxy_server_image is not None:
+    warnings.warn(
+        "`proxy_server_image` is deprecated and will be removed in a future"
+        " release. The proxy server image is automatically detected from the"
+        " Pathways service.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
   _logger.info("Validating Pathways service and TPU instances...")
   validators.validate_pathways_service(pathways_service)
   validators.validate_tpu_instances(expected_tpu_instances)
-  validators.validate_proxy_server_image(proxy_server_image)
   validators.validate_proxy_options(proxy_options)
   gke_utils.fetch_cluster_credentials(
       cluster_name=cluster, project_id=project, location=region
   )
 
-  proxy_options_obj = ProxyOptions.from_list(proxy_options)
-  if proxy_options_obj.sidecar:
-    sidecar_image = gke_utils.get_worker_sidecar_image(
-        pathways_service=pathways_service
+  server_image, sidecar_image = gke_utils.get_pathways_service_images(
+      pathways_service
+  )
+  compatible_proxy_image = gke_utils.get_compatible_proxy_server_image(
+      server_image
+  )
+  _logger.info(
+      "Auto-detected compatible proxy server image: %s", compatible_proxy_image
+  )
+  if proxy_server_image and proxy_server_image != compatible_proxy_image:
+    _logger.warning(
+        "The provided proxy image '%s' is incompatible with the service"
+        " '%s'. Replacing it with the compatible proxy image '%s'.",
+        proxy_server_image,
+        pathways_service,
+        compatible_proxy_image,
     )
-    if sidecar_image:
-      validators.validate_sidecar_image_versions(sidecar_image)
+  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)
   _logger.info("Validation complete.")
 
   if not proxy_job_name:

pathwaysutils/experimental/shared_pathways_service/run_connect_example.py

@@ -36,7 +36,8 @@
 flags.DEFINE_string(
     "proxy_server_image",
     None,
-    "The proxy server image to use. If not provided, a default will be used.",
+    "Deprecated: The proxy server image to use. If not provided, it will be"
+    " auto-detected from the Pathways service.",
 )
 flags.DEFINE_list(
     "proxy_options",
@@ -73,8 +74,7 @@ def main(argv: Sequence[str]) -> None:
       pathways_service=FLAGS.pathways_service,
       expected_tpu_instances={FLAGS.tpu_type: FLAGS.tpu_count},
       proxy_job_name=FLAGS.proxy_job_name,
-      proxy_server_image=FLAGS.proxy_server_image
-      or isc_pathways.DEFAULT_PROXY_IMAGE,
+      proxy_server_image=FLAGS.proxy_server_image,
       proxy_options=FLAGS.proxy_options,
       collect_service_metrics=FLAGS.collect_service_metrics,
   ):

pathwaysutils/experimental/shared_pathways_service/run_workload.py

@@ -22,6 +22,7 @@
 import shlex
 import subprocess
 from typing import Any, ContextManager
+import warnings
 
 from absl import app
 from absl import flags
@@ -55,7 +56,8 @@
 _PROXY_SERVER_IMAGE = flags.DEFINE_string(
     "proxy_server_image",
     "",
-    "The proxy server image to use. If not provided, a default will be used.",
+    "Deprecated: The proxy server image to use. If not provided, it will be"
+    " auto-detected from the Pathways service.",
 )
 _PROXY_OPTIONS = flags.DEFINE_list(
     "proxy_options",
@@ -103,7 +105,8 @@ def run_command(
     tpu_type: The TPU machine type and topology.
     tpu_count: The number of TPU slices.
     command: The command to run on TPUs.
-    proxy_server_image: The proxy server image to use.
+    proxy_server_image: (Deprecated) The proxy server image to use. If not
+      provided, it will be auto-detected from the Pathways service.
     proxy_options: Configuration options for the Pathways proxy.
     collect_service_metrics: Whether to collect usage metrics for Shared
       Pathways Service. Defaults to False.
@@ -113,6 +116,14 @@ def run_command(
   Raises:
     subprocess.CalledProcessError: If the workload command fails.
   """
+  if proxy_server_image:
+    warnings.warn(
+        "`proxy_server_image` is deprecated and will be removed in a future"
+        " release. The proxy server image is automatically detected from the"
+        " Pathways service.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
   logging.info("Connecting to Shared Pathways Service...")
   with connect_fn(
       cluster=cluster,
@@ -122,9 +133,7 @@ def run_command(
       pathways_service=pathways_service,
       expected_tpu_instances={tpu_type: tpu_count},
       proxy_server_image=(
-          proxy_server_image
-          if proxy_server_image
-          else isc_pathways.DEFAULT_PROXY_IMAGE
+          proxy_server_image if proxy_server_image else None
       ),
       proxy_options=proxy_options,
       collect_service_metrics=collect_service_metrics,

pathwaysutils/experimental/shared_pathways_service/validators.py

@@ -113,22 +113,6 @@ def validate_tpu_instances(expected_tpu_instances: Mapping[Any, Any]) -> None:
   _validate_tpu_supported(inst)
 
 
-def validate_proxy_server_image(proxy_server_image: str) -> None:
-  """Validates the proxy server image format."""
-  if not proxy_server_image or not proxy_server_image.strip():
-    raise ValueError("Proxy server image cannot be empty.")
-  if "/" not in proxy_server_image:
-    raise ValueError(
-        f"Proxy server image '{proxy_server_image}' must contain '/', "
-        "separating the registry or namespace from the final image name."
-    )
-  if ":" not in proxy_server_image and "@" not in proxy_server_image:
-    raise ValueError(
-        f"Proxy server image '{proxy_server_image}' must contain a tag with ':'"
-        " or a digest with '@'."
-    )
-
-
 def validate_xla_flags(xla_flags: Iterable[str] | None) -> None:
   """Validates that all XLA flags start with '--xla_'."""
   if not xla_flags: