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,6 +479,19 @@ def is_local_port_free(port: int) -> bool:
   return portpicker.is_port_free(port)
 
 
+@functools.lru_cache(maxsize=1)
+def _get_k8s_core_api() -> client.CoreV1Api:
+  """Initializes and returns the Kubernetes CoreV1Api client."""
+  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
+  return client.CoreV1Api()
+
+
 def get_worker_sidecar_image(
     pathways_service: str, namespace: str = "default"
 ) -> str | None:
@@ -485,49 +500,109 @@ def get_worker_sidecar_image(
   _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]
+  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,
+    v1 = _get_k8s_core_api()
+    pod_list = v1.list_namespaced_pod(
+        namespace=namespace,
+        label_selector=f"jobset.sigs.k8s.io/jobset-name={jobset_name}",
     )
-  except subprocess.CalledProcessError as e:
-    _logger.exception("Failed to get pods. kubectl output:\n%r", e.stderr)
+  except Exception as e:  # pylint: disable=broad-except
+    _logger.exception("Failed to list pods for sidecar image: %r", e)
     return None
 
+  for pod in pod_list.items:
+    spec = pod.spec
+    if not spec:
+      continue
+    containers = (spec.init_containers or []) + (spec.containers or [])
+    for container in containers:
+      if container.name == "colocated-python-sidecar":
+        if container.image:
+          return container.image
+
+  return None
+
+
+def get_server_image(
+    pathways_service: str, namespace: str = "default"
+) -> str:
+  """Gets the server image used by the Pathways service."""
+  pathways_head_hostname = pathways_service.split(":")[0]
+  _validate_k8s_name(namespace)
+
+  # Try to extract the jobset name from the Pathways service hostname.
+  if "-pathways-head" not in pathways_head_hostname:
+    raise ValueError(
+        "Failed to extract jobset name from Pathways service hostname:"
+        f" {pathways_head_hostname}"
+    )
+  jobset_name = pathways_head_hostname.split("-pathways-head")[0]
+
   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
+    v1 = _get_k8s_core_api()
+    pod_list = v1.list_namespaced_pod(
+        namespace=namespace,
+        label_selector=f"jobset.sigs.k8s.io/jobset-name={jobset_name}",
+    )
+  except client.rest.ApiException as e:
+    _logger.exception("Failed to list pods: %r", e)
+    raise
+  except Exception as e:
+    _logger.exception("Failed to get pods: %r", e)
+    raise
 
-  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
+  for pod in pod_list.items:
+    spec = pod.spec
+    if not spec:
+      continue
+    containers = (spec.containers or []) + (spec.init_containers or [])
+    for container in containers:
+      if container.name in ("pathways-rm", "pathways-worker"):
+        if container.image:
+          return container.image
+
+  raise RuntimeError(
+      "Failed to get server image for Pathways service:"
+      f" {pathways_service} in namespace: {namespace}"
+  )
+
+
+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
+
+
+def get_proxy_server_image(
+    pathways_service: str, namespace: str = "default"
+) -> str:
+  """Gets the compatible proxy server image for the given Pathways service."""
+  server_image = get_server_image(pathways_service, namespace=namespace)
+  return get_compatible_proxy_server_image(server_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,15 +484,36 @@ 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
   )
 
+  compatible_proxy_image = gke_utils.get_proxy_server_image(pathways_service)
+  _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,
+    )
+  proxy_server_image = compatible_proxy_image
+
   proxy_options_obj = ProxyOptions.from_list(proxy_options)
   if proxy_options_obj.sidecar:
     sidecar_image = gke_utils.get_worker_sidecar_image(

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: