Print Cloud Logging link when deploying Shared Pathways Service

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,109 @@ 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_cluster_and_project() -> tuple[str | None, str | None]:
+  """Extracts cluster name and project ID from current kubeconfig or environment."""
+  cluster = None
+  project = None
+  try:
+    _, active_context = k8s_config.list_kube_config_contexts()
+    if active_context:
+      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:
+          project = parts[1]
+          cluster = parts[3]
+      else:
+        cluster = cluster_context or None
+  except Exception as e:  # pylint: disable=broad-except
+    _logger.debug(
+        "Could not determine cluster and project from kubeconfig: %s", e
+    )
+
+  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/test/experimental/shared_pathways_service/deploy_pathways_service_test.py

@@ -5,6 +5,7 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 from pathwaysutils.experimental.shared_pathways_service import deploy_pathways_service
+from pathwaysutils.experimental.shared_pathways_service import gke_utils
 
 
 class DeployPathwaysServiceTest(parameterized.TestCase):
@@ -69,8 +70,10 @@ def test_calculate_vms_per_slice_not_divisible(self):
     with self.assertRaises(ValueError):
       deploy_pathways_service.calculate_vms_per_slice("4x8", 5)
 
+  @mock.patch.object(gke_utils, "get_current_cluster_and_project")
   @mock.patch("pathwaysutils.experimental.shared_pathways_service.deploy_pathways_service.jobset.PathwaysJobSet")
-  def test_run_deployment(self, mock_jobset_cls):
+  def test_run_deployment(self, mock_jobset_cls, mock_detect):
+    mock_detect.return_value = ("test-cluster", "test-project")
     mock_jobset = mock_jobset_cls.return_value
     mock_jobset.to_dict.return_value = {"metadata": {"name": "test-jobset"}}
 
@@ -145,7 +148,9 @@ def test_run_deployment(self, mock_jobset_cls):
     # Verify deploy_func was called with the dict
     mock_deploy.assert_called_once_with({"metadata": {"name": "test-jobset"}})
 
-  def test_run_deployment_worker_backoff_limit(self):
+  @mock.patch.object(gke_utils, "get_current_cluster_and_project")
+  def test_run_deployment_worker_backoff_limit(self, mock_detect):
+    mock_detect.return_value = ("test-cluster", "test-project")
     captured_config = {}
 
     def capture_deploy(config):
@@ -177,6 +182,87 @@ def capture_deploy(config):
     # Verify worker backoff limit is set to a large value
     self.assertGreaterEqual(worker_backoff, 1000000)
 
+  @mock.patch.object(gke_utils, "get_log_link")
+  @mock.patch.object(gke_utils, "get_current_cluster_and_project")
+  def test_run_deployment_cloud_logging_link(
+      self, mock_detect, mock_get_log_link
+  ):
+    mock_detect.return_value = ("test-cluster", "test-project")
+    mock_get_log_link.return_value = (
+        "https://console.cloud.google.com/logs/query;dummy"
+    )
+    mock_deploy = mock.MagicMock()
+
+    deploy_pathways_service.run_deployment(
+        tpu_type="v5e",
+        topology="4x8",
+        num_slices=2,
+        jobset_name="test-jobset",
+        gcs_bucket="gs://test-bucket",
+        server_image="server-image",
+        sidecar_image="sidecar-image",
+        dry_run=False,
+        deploy_func=mock_deploy,
+    )
+
+    mock_detect.assert_called_once()
+    mock_get_log_link.assert_called_once()
+    _, kwargs = mock_get_log_link.call_args
+    self.assertEqual(kwargs["cluster"], "test-cluster")
+    self.assertEqual(kwargs["project"], "test-project")
+    self.assertEqual(kwargs["job_name"], "test-jobset")
+    self.assertIsNotNone(kwargs["start_time"])
+    self.assertIsNotNone(kwargs["end_time"])
+
+  @parameterized.named_parameters(
+      dict(testcase_name="neither", cluster=None, project=None),
+      dict(testcase_name="no_cluster", cluster=None, project="test-project"),
+      dict(testcase_name="no_project", cluster="test-cluster", project=None),
+  )
+  @mock.patch.object(gke_utils, "get_log_link")
+  @mock.patch.object(gke_utils, "get_current_cluster_and_project")
+  def test_run_deployment_missing_cluster_or_project_raises(
+      self, mock_detect, mock_get_log_link, cluster, project
+  ):
+    mock_detect.return_value = (cluster, project)
+    mock_deploy = mock.MagicMock()
+
+    with self.assertRaises(ValueError):
+      deploy_pathways_service.run_deployment(
+          tpu_type="v5e",
+          topology="4x8",
+          num_slices=2,
+          jobset_name="test-jobset",
+          gcs_bucket="gs://test-bucket",
+          server_image="server-image",
+          sidecar_image="sidecar-image",
+          dry_run=False,
+          deploy_func=mock_deploy,
+      )
+
+    mock_detect.assert_called_once()
+    mock_get_log_link.assert_not_called()
+    mock_deploy.assert_not_called()
+
+  @mock.patch.object(gke_utils, "get_log_link")
+  def test_run_deployment_dry_run_no_log_link(self, mock_get_log_link):
+    mock_deploy = mock.MagicMock()
+
+    deploy_pathways_service.run_deployment(
+        tpu_type="v5e",
+        topology="4x8",
+        num_slices=2,
+        jobset_name="test-jobset",
+        gcs_bucket="gs://test-bucket",
+        server_image="server-image",
+        sidecar_image="sidecar-image",
+        dry_run=True,
+        deploy_func=mock_deploy,
+    )
+
+    mock_get_log_link.assert_not_called()
+    mock_deploy.assert_not_called()
+
 
 if __name__ == "__main__":
   FLAGS = flags.FLAGS