Print Cloud Logging link when deploying SPS + Few lints

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__)
@@ -207,36 +209,65 @@ 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 [])):
+  # These make the colocated Python sidecar logs as verbose as possible so that
+  # issues in the sidecar (e.g. array (de)serialization during checkpointing)
+  # can be debugged from the container logs.
+  for container in (worker_spec.containers or []) + (
+      worker_spec.init_containers or []
+  ):
     if container.name == "colocated-python-sidecar":
       container.env.extend([
+          # Disable Python stdout/stderr buffering so logs are emitted
+          # immediately and are not lost if the container crashes.
           client.V1EnvVar(name="PYTHONUNBUFFERED", value="1"),
+          # Python logging level for the sidecar.
           client.V1EnvVar(name="LOGLEVEL", value="DEBUG"),
+          # glog (C++) logging: emit INFO and above (0 = INFO) and enable
+          # VLOG(n) messages for n <= 5.
           client.V1EnvVar(name="GLOG_minloglevel", value="0"),
           client.V1EnvVar(name="GLOG_v", value="5"),
+          # TSL/XLA (C++) logging used by JAX: emit INFO and above and enable
+          # VLOG(n) messages for n <= 5.
           client.V1EnvVar(name="TF_CPP_MIN_LOG_LEVEL", value="0"),
           client.V1EnvVar(name="TF_CPP_MIN_VLOG_LEVEL", value="5"),
+          # TPU runtime (libtpu) logging: emit INFO and above.
           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"),
+          # Per-module verbose logging (level 5) for the array serialization
+          # modules used when transferring/checkpointing arrays via the
+          # sidecar.
+          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).
+  # 2. Add arg to pathways-worker container.
   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()
@@ -246,6 +277,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 + gke_utils.LOG_LINK_WINDOW).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
@@ -15,6 +17,9 @@
 
 _logger = logging.getLogger(__name__)
 
+# Default size of the time window covered by Cloud Logging links.
+LOG_LINK_WINDOW = datetime.timedelta(minutes=10)
+
 # TODO(b/456189271): Evaluate and replace the subprocess calls with Kubernetes
 # Python API for kubectl calls.
 
@@ -239,23 +244,167 @@ 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 _parse_log_timestamp(
+    timestamp: str | datetime.datetime,
+) -> datetime.datetime:
+  """Parses a timestamp into a timezone-aware datetime.
+
+  Args:
+    timestamp: An ISO-8601 string or a datetime. Naive values are assumed to be
+      UTC.
+
+  Returns:
+    A timezone-aware datetime.
+
+  Raises:
+    ValueError: If the string is not a valid ISO-8601 timestamp.
+  """
+  if isinstance(timestamp, datetime.datetime):
+    parsed = timestamp
+  else:
+    # `fromisoformat` does not accept a trailing "Z" before Python 3.11.
+    parsed = datetime.datetime.fromisoformat(
+        re.sub(r"[Zz]$", "+00:00", timestamp)
+    )
+  if parsed.tzinfo is None:
+    parsed = parsed.replace(tzinfo=datetime.timezone.utc)
+  return parsed
+
+
+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). If provided without end_time, the end time is set to
+      `LOG_LINK_WINDOW` after start_time.
+    end_time: The end time for the time window (ISO-8601 string or datetime).
+      If provided without start_time, the start time is set to
+      `LOG_LINK_WINDOW` before end_time.
+    duration: The duration string (e.g. "PT1H") used when neither start_time
+      nor end_time is provided.
+
+  Returns:
+    The Cloud Logging query URL.
+
+  Raises:
+    ValueError: If only one of start_time and end_time is provided and it is a
+      string that is not a valid ISO-8601 timestamp.
+  """
   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}"'
   )
+
+  if start_time is not None and end_time is None:
+    end_time = _parse_log_timestamp(start_time) + LOG_LINK_WINDOW
+  elif end_time is not None and start_time is None:
+    start_time = _parse_log_timestamp(end_time) - LOG_LINK_WINDOW
+
+  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 = ""
+
   encoded_filter = urllib.parse.quote(log_filter, safe="")
 
+  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.