diff --git a/pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py b/pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py index 913aa25..51b893c 100644 --- a/pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py +++ b/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.") diff --git a/pathwaysutils/experimental/shared_pathways_service/gke_utils.py b/pathwaysutils/experimental/shared_pathways_service/gke_utils.py index 184f8a2..41c8c40 100644 --- a/pathwaysutils/experimental/shared_pathways_service/gke_utils.py +++ b/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. diff --git a/pathwaysutils/test/experimental/shared_pathways_service/deploy_pathways_service_test.py b/pathwaysutils/test/experimental/shared_pathways_service/deploy_pathways_service_test.py index a541175..4e92bb5 100644 --- a/pathwaysutils/test/experimental/shared_pathways_service/deploy_pathways_service_test.py +++ b/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 diff --git a/pathwaysutils/test/experimental/shared_pathways_service/gke_utils_test.py b/pathwaysutils/test/experimental/shared_pathways_service/gke_utils_test.py index fde1830..d4caac0 100644 --- a/pathwaysutils/test/experimental/shared_pathways_service/gke_utils_test.py +++ b/pathwaysutils/test/experimental/shared_pathways_service/gke_utils_test.py @@ -1,6 +1,7 @@ """Tests for gke_utils.py. """ +import datetime import io import socket import subprocess @@ -353,6 +354,203 @@ def test_get_log_link(self): "duration=PT1H?project=test-project", ) + def test_get_log_link_with_time_window_str(self): + cluster = "test-cluster" + project = "test-project" + job_name = "test-job" + start_time = "2026-09-10T10:00:00.000Z" + end_time = "2026-09-10T10:10:00.000Z" + log_link = gke_utils.get_log_link( + cluster=cluster, + project=project, + job_name=job_name, + start_time=start_time, + end_time=end_time, + ) + self.assertEqual( + log_link, + r"https://console.cloud.google.com/logs/query;query=resource.type%3D" + r"%22k8s_container%22%0Aresource.labels.cluster_name%3D" + "%22test-cluster%22%0Aresource.labels.namespace_name%3D" + "%22default%22%0Alabels.k8s-pod%2Fjob-name%3A%22test-job%22;" + "startTime=2026-09-10T10:00:00.000Z;endTime=2026-09-10T10:10:00.000Z" + "?project=test-project", + ) + + def test_get_log_link_with_time_window_datetime(self): + cluster = "test-cluster" + project = "test-project" + job_name = "test-job" + start_dt = datetime.datetime( + 2026, 9, 10, 10, 0, 0, tzinfo=datetime.timezone.utc + ) + end_dt = datetime.datetime( + 2026, 9, 10, 10, 10, 0, tzinfo=datetime.timezone.utc + ) + log_link = gke_utils.get_log_link( + cluster=cluster, + project=project, + job_name=job_name, + start_time=start_dt, + end_time=end_dt, + ) + self.assertEqual( + log_link, + r"https://console.cloud.google.com/logs/query;query=resource.type%3D" + r"%22k8s_container%22%0Aresource.labels.cluster_name%3D" + "%22test-cluster%22%0Aresource.labels.namespace_name%3D" + "%22default%22%0Alabels.k8s-pod%2Fjob-name%3A%22test-job%22;" + "startTime=2026-09-10T10:00:00.000Z;endTime=2026-09-10T10:10:00.000Z" + "?project=test-project", + ) + + def test_get_log_link_with_naive_datetime_assumes_utc(self): + start_dt = datetime.datetime(2026, 9, 10, 10, 0, 0) + end_dt = datetime.datetime(2026, 9, 10, 10, 10, 0) + log_link = gke_utils.get_log_link( + cluster="test-cluster", + project="test-project", + job_name="test-job", + start_time=start_dt, + end_time=end_dt, + ) + self.assertEqual( + log_link, + r"https://console.cloud.google.com/logs/query;query=resource.type%3D" + r"%22k8s_container%22%0Aresource.labels.cluster_name%3D" + "%22test-cluster%22%0Aresource.labels.namespace_name%3D" + "%22default%22%0Alabels.k8s-pod%2Fjob-name%3A%22test-job%22;" + "startTime=2026-09-10T10:00:00.000Z;endTime=2026-09-10T10:10:00.000Z" + "?project=test-project", + ) + + def test_get_log_link_with_aware_non_utc_datetime_converts_to_utc(self): + tz = datetime.timezone(datetime.timedelta(hours=2)) + start_dt = datetime.datetime(2026, 9, 10, 12, 0, 0, tzinfo=tz) + end_dt = datetime.datetime(2026, 9, 10, 12, 10, 0, tzinfo=tz) + log_link = gke_utils.get_log_link( + cluster="test-cluster", + project="test-project", + job_name="test-job", + start_time=start_dt, + end_time=end_dt, + ) + self.assertIn( + "startTime=2026-09-10T10:00:00.000Z;endTime=2026-09-10T10:10:00.000Z", + log_link, + ) + + def test_get_log_link_without_time_window_omits_time_param(self): + log_link = gke_utils.get_log_link( + cluster="test-cluster", + project="test-project", + job_name="test-job", + duration=None, + ) + self.assertEqual( + log_link, + r"https://console.cloud.google.com/logs/query;query=resource.type%3D" + r"%22k8s_container%22%0Aresource.labels.cluster_name%3D" + "%22test-cluster%22%0Aresource.labels.namespace_name%3D" + "%22default%22%0Alabels.k8s-pod%2Fjob-name%3A%22test-job%22" + "?project=test-project", + ) + + def test_get_log_link_with_custom_namespace(self): + log_link = gke_utils.get_log_link( + cluster="test-cluster", + project="test-project", + job_name="test-job", + namespace="custom-ns", + ) + self.assertIn("%22custom-ns%22", log_link) + + def test_get_current_cluster_and_project_from_kubeconfig(self): + mock_active_context = { + "name": "gke_test-proj_us-central1_test-cl", + "context": {"cluster": "gke_test-proj_us-central1_test-cl"}, + } + with mock.patch.object( + k8s_config, + "list_kube_config_contexts", + return_value=([mock_active_context], mock_active_context), + ): + cluster, project = gke_utils.get_current_cluster_and_project() + self.assertEqual(cluster, "test-cl") + self.assertEqual(project, "test-proj") + + def test_get_current_cluster_and_project_fallback_env(self): + with mock.patch.object( + k8s_config, + "list_kube_config_contexts", + return_value=([], None), + ): + with mock.patch.dict( + "os.environ", + {"GKE_CLUSTER": "env-cl", "PROJECT": "env-proj"}, + clear=False, + ): + cluster, project = gke_utils.get_current_cluster_and_project() + self.assertEqual(cluster, "env-cl") + self.assertEqual(project, "env-proj") + + def test_get_current_cluster_and_project_non_gke_context(self): + mock_active_context = { + "name": "minikube", + "context": {"cluster": "minikube"}, + } + with mock.patch.object( + k8s_config, + "list_kube_config_contexts", + return_value=([mock_active_context], mock_active_context), + ): + with mock.patch.dict("os.environ", {}, clear=True): + cluster, project = gke_utils.get_current_cluster_and_project() + self.assertEqual(cluster, "minikube") + self.assertIsNone(project) + + def test_get_current_cluster_and_project_empty_context_name(self): + mock_active_context = {"name": "", "context": {}} + with mock.patch.object( + k8s_config, + "list_kube_config_contexts", + return_value=([mock_active_context], mock_active_context), + ): + with mock.patch.dict("os.environ", {}, clear=True): + cluster, project = gke_utils.get_current_cluster_and_project() + self.assertIsNone(cluster) + self.assertIsNone(project) + + def test_get_current_cluster_and_project_malformed_gke_context(self): + mock_active_context = { + "name": "gke_test-proj_us-central1", + "context": {"cluster": "gke_test-proj_us-central1"}, + } + with mock.patch.object( + k8s_config, + "list_kube_config_contexts", + return_value=([mock_active_context], mock_active_context), + ): + with mock.patch.dict("os.environ", {}, clear=True): + cluster, project = gke_utils.get_current_cluster_and_project() + self.assertIsNone(cluster) + self.assertIsNone(project) + + def test_get_current_cluster_and_project_kubeconfig_error_falls_back(self): + with mock.patch.object( + k8s_config, + "list_kube_config_contexts", + side_effect=k8s_config.ConfigException("no kubeconfig"), + ): + with mock.patch.dict( + "os.environ", + {"CLUSTER": "env-cl", "GOOGLE_CLOUD_PROJECT": "env-proj"}, + clear=True, + ): + cluster, project = gke_utils.get_current_cluster_and_project() + self.assertEqual(cluster, "env-cl") + self.assertEqual(project, "env-proj") + def test_wait_for_pod_success(self): """Tests that wait_for_pod returns the pod name on success.""" mock_run = self.enter_context(