Skip fetching GKE cluster credentials if the active kube context matches for the proxy job

pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py

@@ -2,6 +2,7 @@
 
 from collections.abc import Callable, Sequence
 import dataclasses
+import datetime
 import logging
 import math
 from typing import Any
@@ -10,6 +11,7 @@
 from kubernetes import client
 from kubernetes import config
 from pathwaysutils.experimental.gke import jobset
+from pathwaysutils.experimental.shared_pathways_service import gke_utils
 import yaml
 
 _logger = logging.getLogger(__name__)
@@ -195,18 +197,24 @@ def run_deployment(
       if container.name == "pathways-rm":
         container.image = server_image
     # Mutate worker job.
-    for container in pw_jobset.worker_job_template.spec.template.spec.containers:
+    for (
+        container
+    ) in pw_jobset.worker_job_template.spec.template.spec.containers:
       if container.name == "pathways-worker":
         container.image = server_image
 
   # Add colocated python sidecar.
-  pw_jobset.add_colocated_python(image=sidecar_image, shm_mount_path=_SIDECAR_SHM_DIR)
+  pw_jobset.add_colocated_python(
+      image=sidecar_image, shm_mount_path=_SIDECAR_SHM_DIR
+  )
 
   # Mutate the sidecar configuration to match what HEAD expects.
   worker_spec = pw_jobset.worker_job_template.spec.template.spec
 
   # 1. Add extra logging env vars to sidecar.
-  for container in ((worker_spec.containers or []) + (worker_spec.init_containers or [])):
+  for container in (worker_spec.containers or []) + (
+      worker_spec.init_containers or []
+  ):
     if container.name == "colocated-python-sidecar":
       container.env.extend([
           client.V1EnvVar(name="PYTHONUNBUFFERED", value="1"),
@@ -216,15 +224,24 @@ def run_deployment(
           client.V1EnvVar(name="TF_CPP_MIN_LOG_LEVEL", value="0"),
           client.V1EnvVar(name="TF_CPP_MIN_VLOG_LEVEL", value="5"),
           client.V1EnvVar(name="TPU_MIN_LOG_LEVEL", value="0"),
-          client.V1EnvVar(name="GLOG_vmodule", value="jax_array_handlers=5,type_handlers=5,tensorstore_utils=5"),
+          client.V1EnvVar(
+              name="GLOG_vmodule",
+              value=(
+                  "jax_array_handlers=5,type_handlers=5,tensorstore_utils=5"
+              ),
+          ),
       ])
 
   # 2. Add arg to pathways-worker container (in addition to env var set by builder).
   for container in worker_spec.containers:
     if container.name == "pathways-worker":
       args = container.args or []
-      if not any(a.startswith("--cloud_pathways_sidecar_shm_directory=") for a in args):
-        args.append(f"--cloud_pathways_sidecar_shm_directory={_SIDECAR_SHM_DIR}")
+      if not any(
+          a.startswith("--cloud_pathways_sidecar_shm_directory=") for a in args
+      ):
+        args.append(
+            f"--cloud_pathways_sidecar_shm_directory={_SIDECAR_SHM_DIR}"
+        )
       container.args = args
 
   jobset_config = pw_jobset.to_dict()
@@ -234,6 +251,31 @@ def run_deployment(
 
   if not dry_run:
     _logger.info("Deploying JobSet...")
+    cluster, project = gke_utils.get_current_cluster_and_project()
+    if not cluster or not project:
+      raise ValueError(
+          "Cluster or project could not be determined from kubeconfig. Run"
+          " 'gcloud container clusters get-credentials ... && kubectl config"
+          " set-context --current --namespace=default' OR 'kubectl config"
+          " set-context --current --user=... --cluster=...'"
+          " first."
+      )
+    now = datetime.datetime.now(datetime.timezone.utc)
+    start_time = now.isoformat(timespec="milliseconds").replace("+00:00", "Z")
+    end_time = (now + datetime.timedelta(minutes=10)).isoformat(
+        timespec="milliseconds"
+    ).replace("+00:00", "Z")
+    cloud_logging_link = gke_utils.get_log_link(
+        cluster=cluster,
+        project=project,
+        job_name=jobset_name,
+        start_time=start_time,
+        end_time=end_time,
+    )
+    _logger.info(
+        "View SPS deployment logs in Cloud Logging: %s", cloud_logging_link
+    )
+
     deploy_func(jobset_config)
   else:
     _logger.info("Dry run mode, not deploying.")

pathwaysutils/experimental/shared_pathways_service/gke_utils.py

@@ -1,7 +1,9 @@
 """GKE utils for deploying and managing the Pathways proxy."""
 
+import datetime
 import functools
 import logging
+import os
 import re
 import socket
 import subprocess
@@ -239,23 +241,127 @@ def check_pod_ready(pod_name: str, timeout: int = 30) -> str:
   return pod_name
 
 
-def get_log_link(*, cluster: str, project: str, job_name: str) -> str:
-  """Returns a link to Cloud Logging for the given cluster and job name."""
+def _format_log_timestamp(timestamp: str | datetime.datetime) -> str:
+  """Formats a timestamp for use in a Cloud Logging query URL.
+
+  Args:
+    timestamp: An ISO-8601 string or a datetime. Naive datetimes are assumed to
+      be UTC.
+
+  Returns:
+    The timestamp as an ISO-8601 string, or the unmodified string input.
+  """
+  if not isinstance(timestamp, datetime.datetime):
+    return str(timestamp)
+
+  if timestamp.tzinfo is None:
+    timestamp = timestamp.replace(tzinfo=datetime.timezone.utc)
+  timestamp = timestamp.astimezone(datetime.timezone.utc)
+  return timestamp.isoformat(timespec="milliseconds").replace("+00:00", "Z")
+
+
+def get_log_link(
+    *,
+    cluster: str,
+    project: str,
+    job_name: str,
+    namespace: str = "default",
+    start_time: str | datetime.datetime | None = None,
+    end_time: str | datetime.datetime | None = None,
+    duration: str | None = "PT1H",
+) -> str:
+  """Returns a link to Cloud Logging for the given cluster and job name.
+
+  Args:
+    cluster: The name of the GKE cluster.
+    project: The GCP project ID.
+    job_name: The name of the job or jobset.
+    namespace: The Kubernetes namespace. Defaults to "default".
+    start_time: The start time for the time window (ISO-8601 string or
+      datetime).
+    end_time: The end time for the time window (ISO-8601 string or datetime).
+    duration: The duration string (e.g. "PT1H") used when start_time and
+      end_time are not provided.
+
+  Returns:
+    The Cloud Logging query URL.
+  """
   log_filter = (
       'resource.type="k8s_container"\n'
       f'resource.labels.cluster_name="{cluster}"\n'
-      'resource.labels.namespace_name="default"\n'
+      f'resource.labels.namespace_name="{namespace}"\n'
       f'labels.k8s-pod/job-name:"{job_name}"'
   )
   encoded_filter = urllib.parse.quote(log_filter, safe="")
 
+  if start_time is not None and end_time is not None:
+    start_time_str = _format_log_timestamp(start_time)
+    end_time_str = _format_log_timestamp(end_time)
+    time_param = f"startTime={start_time_str};endTime={end_time_str}"
+  elif duration is not None:
+    time_param = f"duration={duration}"
+  else:
+    time_param = ""
+
+  time_part = f";{time_param}" if time_param else ""
   return (
       "https://console.cloud.google.com/logs/query;"
-      f"query={encoded_filter};duration=PT1H"
+      f"query={encoded_filter}{time_part}"
       f"?project={project}"
   )
 
 
+def get_current_kube_context() -> tuple[str | None, str | None, str | None]:
+  """Reads the cluster targeted by the active kube config context.
+
+  Only the kube config is consulted; no environment variable fallbacks are
+  applied. Contexts written by `gcloud container clusters get-credentials` are
+  named `gke_<project>_<location>_<cluster>`.
+
+  Returns:
+    A (cluster, project, location) tuple. All three are populated for a GKE
+    context. For any other context, only the context name is returned as the
+    cluster and the project and location are None. All three are None if the
+    active context cannot be read or is a malformed GKE context.
+  """
+  try:
+    _, active_context = k8s_config.list_kube_config_contexts()
+  except Exception as e:  # pylint: disable=broad-except
+    _logger.debug("Could not read the current kube config context: %s", e)
+    return None, None, None
+
+  if not active_context:
+    return None, None, None
+
+  context_data = active_context.get("context", {})
+  cluster_context = (
+      context_data.get("cluster") or active_context.get("name", "")
+  )
+
+  if cluster_context.startswith("gke_"):
+    parts = cluster_context.split("_", 3)
+    if len(parts) != 4:
+      return None, None, None
+    _, project, location, cluster = parts
+    return cluster, project, location
+
+  return cluster_context or None, None, None
+
+
+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()
+
+  if not cluster:
+    cluster = os.environ.get("GKE_CLUSTER") or os.environ.get("CLUSTER")
+  if not project:
+    project = os.environ.get("PROJECT") or os.environ.get(
+        "GOOGLE_CLOUD_PROJECT"
+    )
+
+  return cluster, project
+
+
 def wait_for_pod(job_name: str) -> str:
   """Waits for the given job's pod to be ready.
 

pathwaysutils/experimental/shared_pathways_service/isc_pathways.py

@@ -534,6 +534,43 @@ 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(
     *,
@@ -584,8 +621,8 @@ def connect(
   validators.validate_pathways_service(pathways_service)
   validators.validate_tpu_instances(expected_tpu_instances)
   validators.validate_proxy_options(proxy_options)
-  gke_utils.fetch_cluster_credentials(
-      cluster_name=cluster, project_id=project, location=region
+  _ensure_cluster_credentials(
+      cluster=cluster, project=project, location=region
   )
 
   server_image, sidecar_image = gke_utils.get_pathways_service_images(