From f9af778540938af360cc99b588e0d51898dc6f7e Mon Sep 17 00:00:00 2001 From: Luke Baumann Date: Thu, 23 Jul 2026 15:34:29 -0700 Subject: [PATCH] pathwaysutils: Remove hostNetwork configuration from PathwaysJobSet pod specifications. PiperOrigin-RevId: 952991560 --- pathwaysutils/experimental/gke/jobset.py | 120 +++------- .../test/experimental/gke/jobset_test.py | 222 ++---------------- .../gke/testdata/model_lite_tpuv5e_4x8.yaml | 4 - 3 files changed, 51 insertions(+), 295 deletions(-) diff --git a/pathwaysutils/experimental/gke/jobset.py b/pathwaysutils/experimental/gke/jobset.py index 6a69af1..8f8cdae 100644 --- a/pathwaysutils/experimental/gke/jobset.py +++ b/pathwaysutils/experimental/gke/jobset.py @@ -11,17 +11,27 @@ # limitations under the License. """Pathways JobSet generator and builder (with Worker Job Config).""" +from __future__ import annotations + import hashlib import json import logging import math import time -from typing import Any, Mapping, Sequence -from kubernetes import client -from kubernetes import config as k8s_config +from typing import TYPE_CHECKING, Any, Mapping, Sequence import yaml -# GKE sidecar containers restartPolicy compatibility placeholder. +try: + import kubernetes +except ImportError as e: + raise ImportError( + "GKE utilities require `kubernetes`. " + "Please install pathwaysutils with GKE support:\n\n" + " pip install 'pathwaysutils[gke]'\n" + ) from e + +from kubernetes import client +from kubernetes import config as k8s_config _logger = logging.getLogger(__name__) @@ -72,7 +82,7 @@ def _deserialize_dict( - api_client: client.ApiClient, data_dict: Mapping[str, Any], klass: Any + api_client: Any, data_dict: Mapping[str, Any], klass: Any ) -> Any: class FakeResponse: @@ -93,8 +103,6 @@ def __init__( tpu_type: str, topology: str, num_slices: int, - user_pod_template: Mapping[str, Any] | None = None, - main_container_name: str = "main", max_restarts: int = 0, max_slice_restarts: int = 0, termination_grace_period_seconds: int | None = None, @@ -114,8 +122,6 @@ def __init__( tpu_type: TPU type (e.g., "v5e"). topology: TPU topology (e.g., "2x2"). num_slices: Number of slices. - user_pod_template: Optional user pod template for the head job. - main_container_name: Name of the main container in user_pod_template. max_restarts: Maximum number of restarts for the JobSet. max_slice_restarts: Maximum number of slice restarts. termination_grace_period_seconds: Optional termination grace period. @@ -124,12 +130,8 @@ def __init__( elastic_slices: Number of elastic slices. labels: Optional labels for the JobSet. annotations: Optional annotations for the JobSet. + shared_pathways_service: Whether to run only RM for Shared Pathways Service. """ - if shared_pathways_service and user_pod_template: - raise ValueError( - "Cannot enable shared_pathways_service when user_pod_template is" - " provided." - ) self._shared_pathways_service = shared_pathways_service self._name = name @@ -166,8 +168,6 @@ def __init__( num_slices=num_slices, instance_type=instance_type, image_tag=image_tag, - user_pod_template=user_pod_template, - main_container_name=main_container_name, elastic_slices=elastic_slices, shared_pathways_service=shared_pathways_service, ) @@ -185,7 +185,7 @@ def __init__( ) self._success_policy = None - if user_pod_template or shared_pathways_service: + if shared_pathways_service: self._success_policy = { "operator": "All", "targetReplicatedJobs": [PATHWAYS_HEAD_JOB_NAME], @@ -205,8 +205,6 @@ def _build_head_job_template( num_slices: int, instance_type: str, image_tag: str, - user_pod_template: Mapping[str, Any] | None, - main_container_name: str, elastic_slices: int, shared_pathways_service: bool, ) -> client.V1JobTemplateSpec: @@ -217,9 +215,8 @@ def _build_head_job_template( num_slices: Number of slices. instance_type: TPU instance type (e.g., "tpuv5:2x2"). image_tag: Version tag for Pathways images. - user_pod_template: Optional user pod template for the head job. - main_container_name: Name of the main container in user_pod_template. elastic_slices: Number of elastic slices. + shared_pathways_service: Whether to run only RM for Shared Pathways Service. Returns: The head job template. @@ -318,74 +315,18 @@ def _build_head_job_template( ), ) - api_client = client.ApiClient() + containers = [rm_container] + if not shared_pathways_service: + containers.append(proxy_container) - if user_pod_template: - user_template_obj = _deserialize_dict( - api_client, user_pod_template, client.V1PodTemplateSpec - ) - head_pod_spec = user_template_obj.spec - head_pod_spec.host_network = True - head_pod_spec.dns_policy = "ClusterFirstWithHostNet" - - rm_container.restart_policy = "Always" # pyrefly: ignore[missing-attribute] - proxy_container.restart_policy = "Always" # pyrefly: ignore[missing-attribute] - - init_containers = head_pod_spec.init_containers or [] - init_containers.extend([rm_container, proxy_container]) - head_pod_spec.init_containers = init_containers - - # Inject JAX env vars into main container. - jax_env = [ - client.V1EnvVar( - name="PATHWAYS_HEAD", - value_from=client.V1EnvVarSource( - field_ref=client.V1ObjectFieldSelector( - field_path=( - "metadata.labels['jobset.sigs.k8s.io/coordinator']" - ) - ) - ), - ), - client.V1EnvVar(name="JAX_PLATFORMS", value="proxy"), - client.V1EnvVar(name="XCLOUD_ENVIRONMENT", value="GCP"), - client.V1EnvVar( - name="JAX_BACKEND_TARGET", - value=f"grpc://$(PATHWAYS_HEAD):{PATHWAYS_PROXY_PORT}", - ), - ] - containers = head_pod_spec.containers or [] - for c in containers: - if c.name == main_container_name: - env = c.env or [] - env.extend(jax_env) - c.env = env - break - head_pod_spec.containers = containers - - annotations = user_pod_template.get("metadata", {}).get("annotations", {}) - labels = user_pod_template.get("metadata", {}).get("labels", {}) - else: - # Headless mode. - containers = [rm_container] - if not shared_pathways_service: - containers.append(proxy_container) - head_pod_spec = client.V1PodSpec( - host_network=True, - dns_policy="ClusterFirstWithHostNet", - containers=containers, - ) - annotations = {} - labels = {} - - if not head_pod_spec.restart_policy: - head_pod_spec.restart_policy = "Never" + head_pod_spec = client.V1PodSpec( + containers=containers, + restart_policy="Never", + ) - # Default annotations job_annotations = { "alpha.jobset.sigs.k8s.io/exclusive-topology": "kubernetes.io/hostname" } - job_annotations.update(annotations) head_job_template = client.V1JobTemplateSpec( metadata=client.V1ObjectMeta(annotations=job_annotations), @@ -396,7 +337,7 @@ def _build_head_job_template( parallelism=1, template=client.V1PodTemplateSpec( metadata=client.V1ObjectMeta( - annotations=job_annotations, labels=labels + annotations=job_annotations, labels={} ), spec=head_pod_spec, ), @@ -525,8 +466,6 @@ def _build_worker_job_template( ), ) ], - host_network=True, - dns_policy="ClusterFirstWithHostNet", restart_policy="OnFailure", ) if termination_grace_period_seconds is not None: @@ -640,11 +579,10 @@ def add_colocated_python( client.V1VolumeMount(name=shm_volume_name, mount_path=shm_mount_path), ], ) - colocated_container.restart_policy = "Always" # pyrefly: ignore[missing-attribute] - init_containers = pod_spec.init_containers or [] - init_containers.append(colocated_container) - pod_spec.init_containers = init_containers + containers = pod_spec.containers or [] + containers.append(colocated_container) + pod_spec.containers = containers # Add volume mount to pathways-worker. for container in pod_spec.containers: diff --git a/pathwaysutils/test/experimental/gke/jobset_test.py b/pathwaysutils/test/experimental/gke/jobset_test.py index 6401bdd..5ef1b64 100644 --- a/pathwaysutils/test/experimental/gke/jobset_test.py +++ b/pathwaysutils/test/experimental/gke/jobset_test.py @@ -2,12 +2,22 @@ import itertools import os from typing import Any +import unittest from unittest import mock from absl.testing import absltest from absl.testing import parameterized +import yaml + +try: + import kubernetes +except ImportError: + raise unittest.SkipTest( + "kubernetes is not installed (requires pathwaysutils[test])" + ) + from kubernetes import client +from kubernetes import config as k8s_config from pathwaysutils.experimental.gke import jobset -import yaml def normalize_k8s_spec(spec: Any) -> Any: @@ -139,8 +149,8 @@ def test_headless_head_job_pod_spec(self): self.assertIn("pathways-head", helper.jobs) self.assertEqual(helper.jobs["pathways-head"]["replicas"], 1) pod_spec = helper.pod_specs["pathways-head"] - self.assertTrue(pod_spec["hostNetwork"]) - self.assertEqual(pod_spec["dnsPolicy"], "ClusterFirstWithHostNet") + self.assertNotIn("hostNetwork", pod_spec) + self.assertNotIn("dnsPolicy", pod_spec) self.assertEqual(pod_spec["restartPolicy"], "Never") def test_headless_head_job_containers(self): @@ -167,98 +177,6 @@ def test_headless_head_job_containers(self): ) self.assertIn("--num_elastic_slices=2", proxy_container["args"]) - def test_non_headless_head_job_init_containers(self): - user_pod_template = { - "spec": { - "containers": [{ - "name": "jax-tpu", - "image": "ubuntu:latest", - }] - } - } - js = self._create_jobset( - user_pod_template=user_pod_template, - main_container_name="jax-tpu", - ) - - config = js.to_dict() - helper = JobSetManifestHelper(config) - - pod_spec = helper.pod_specs["pathways-head"] - self.assertLen(pod_spec["initContainers"], 2) - self.assertIn("pathways-rm", helper.init_containers["pathways-head"]) - self.assertIn("pathways-proxy", helper.init_containers["pathways-head"]) - - rm_container = helper.init_containers["pathways-head"]["pathways-rm"] - proxy_container = helper.init_containers["pathways-head"]["pathways-proxy"] - self.assertEqual(rm_container["restartPolicy"], "Always") - self.assertEqual(proxy_container["restartPolicy"], "Always") - - def test_non_headless_head_job_jax_env(self): - user_pod_template = { - "spec": { - "containers": [{ - "name": "jax-tpu", - "image": "ubuntu:latest", - }] - } - } - js = self._create_jobset( - user_pod_template=user_pod_template, - main_container_name="jax-tpu", - ) - - config = js.to_dict() - helper = JobSetManifestHelper(config) - - self.assertIn("jax-tpu", helper.containers["pathways-head"]) - main_container = helper.containers["pathways-head"]["jax-tpu"] - env_names = [e["name"] for e in main_container["env"]] - self.assertIn("PATHWAYS_HEAD", env_names) - self.assertIn("JAX_PLATFORMS", env_names) - self.assertIn("XCLOUD_ENVIRONMENT", env_names) - self.assertIn("JAX_BACKEND_TARGET", env_names) - - def test_non_headless_head_job_annotations(self): - user_pod_template = { - "metadata": {"annotations": {"example.com/annotation": "value"}}, - "spec": { - "containers": [{ - "name": "jax-tpu", - "image": "ubuntu:latest", - }] - } - } - js = self._create_jobset( - user_pod_template=user_pod_template, - main_container_name="jax-tpu", - ) - - config = js.to_dict() - helper = JobSetManifestHelper(config) - - head_job = helper.jobs["pathways-head"] - self.assertEqual( - head_job["template"]["metadata"]["annotations"][ - "example.com/annotation" - ], - "value", - ) - self.assertEqual( - head_job["template"]["spec"]["template"]["metadata"]["annotations"][ - "example.com/annotation" - ], - "value", - ) - - def test_monkeypatch_restart_policy(self): - # Construct V1Container with restart_policy to test monkeypatch. - c = client.V1Container( - name="test", - restart_policy="Always" # pyrefly: ignore[unexpected-keyword] - ) # pytype: disable=wrong-keyword-args - self.assertEqual(getattr(c, "restart_policy"), "Always") - def test_worker_job_replicas(self): js = self._create_jobset(num_slices=2) @@ -289,8 +207,8 @@ def test_worker_job_pod_spec(self): helper = JobSetManifestHelper(config) pod_spec = helper.pod_specs["pathways-worker"] - self.assertTrue(pod_spec["hostNetwork"]) - self.assertEqual(pod_spec["dnsPolicy"], "ClusterFirstWithHostNet") + self.assertNotIn("hostNetwork", pod_spec) + self.assertNotIn("dnsPolicy", pod_spec) self.assertEqual(pod_spec["restartPolicy"], "OnFailure") self.assertEqual(pod_spec["terminationGracePeriodSeconds"], 60) @@ -358,7 +276,7 @@ def test_add_gcsfuse_read_only(self, read_only): ("worker", "pathways-worker", ["pathways-worker"]), ("explicit", ["pathways-worker", "pathways-rm"], ["pathways-worker", "pathways-rm"]), ) - def test_add_gcsfuse_container_filtering_headless( + def test_add_gcsfuse_container_filtering( self, containers_param, expected_containers ): pw_jobset = self._create_jobset(topology="2x2", num_slices=1) @@ -383,49 +301,6 @@ def test_add_gcsfuse_container_filtering_headless( else: self.assertFalse(has_mount, f"Expected {c_name} in {job_name} NOT to have mount") - @parameterized.named_parameters( - ("all", "all", ["pathways-rm", "pathways-proxy", "pathways-worker", "jax-tpu"]), - ("explicit", ["pathways-worker", "jax-tpu"], ["pathways-worker", "jax-tpu"]), - ("explicit_rm", ["pathways-worker", "pathways-rm"], ["pathways-worker", "pathways-rm"]), - ) - def test_add_gcsfuse_container_filtering_non_headless( - self, containers_param, expected_containers - ): - user_pod_template = { - "spec": { - "containers": [{ - "name": "jax-tpu", - "image": "gcr.io/my-project/jax-tpu:latest", - }] - } - } - pw_jobset = self._create_jobset( - topology="2x2", - num_slices=1, - user_pod_template=user_pod_template, - main_container_name="jax-tpu", - ) - bucket_hash = int(hashlib.md5("my-bucket".encode()).hexdigest(), 16) % (10**8) - expected_vol_name = f"gcsfuse-{bucket_hash}" - - pw_jobset.add_gcsfuse( - containers=containers_param, - mount_path="/gcs/data", - bucket="my-bucket", - ) - helper = JobSetManifestHelper(pw_jobset.to_dict()) - - all_possible_containers = ["pathways-rm", "pathways-proxy", "pathways-worker", "jax-tpu"] - for c_name in all_possible_containers: - matches = helper.get_all_containers_by_name(c_name) - self.assertNotEmpty(matches) - for job_name, container in matches: - has_mount = any(m["mountPath"] == "/gcs/data" for m in container.get("volumeMounts", [])) - if c_name in expected_containers: - self.assertTrue(has_mount, f"Expected {c_name} in {job_name} to have mount") - else: - self.assertFalse(has_mount, f"Expected {c_name} in {job_name} NOT to have mount") - def test_add_gcsfuse_volumes_and_annotations(self): pw_jobset = self._create_jobset(topology="2x2", num_slices=1) bucket_hash = int(hashlib.md5("my-bucket".encode()).hexdigest(), 16) % (10**8) @@ -544,9 +419,8 @@ def test_add_colocated_python_sidecar(self): pw_jobset.add_colocated_python(image="gcr.io/my-project/colocated-python:custom") helper = JobSetManifestHelper(pw_jobset.to_dict()) - self.assertIn("colocated-python-sidecar", helper.init_containers["pathways-worker"]) - sidecar = helper.init_containers["pathways-worker"]["colocated-python-sidecar"] - self.assertEqual(sidecar["restartPolicy"], "Always") + self.assertIn("colocated-python-sidecar", helper.containers["pathways-worker"]) + sidecar = helper.containers["pathways-worker"]["colocated-python-sidecar"] self.assertEqual(sidecar["image"], "gcr.io/my-project/colocated-python:custom") self.assertTrue( any( @@ -575,7 +449,7 @@ def test_add_colocated_python_preserves_init_containers(self): # Verify both exist self.assertIn("existing-init-container", helper.init_containers["pathways-worker"]) - self.assertIn("colocated-python-sidecar", helper.init_containers["pathways-worker"]) + self.assertIn("colocated-python-sidecar", helper.containers["pathways-worker"]) def test_add_colocated_python_volume_default(self): pw_jobset = self._create_jobset(topology="2x2", num_slices=1) @@ -619,8 +493,8 @@ def test_add_colocated_python_custom_shm(self): ) helper = JobSetManifestHelper(pw_jobset.to_dict()) - self.assertIn("colocated-python-sidecar", helper.init_containers["pathways-worker"]) - sidecar = helper.init_containers["pathways-worker"]["colocated-python-sidecar"] + self.assertIn("colocated-python-sidecar", helper.containers["pathways-worker"]) + sidecar = helper.containers["pathways-worker"]["colocated-python-sidecar"] self.assertTrue( any( m["name"] == "shared-memory" and m["mountPath"] == "/tmp/custom-shm" @@ -655,35 +529,6 @@ def test_add_colocated_python_custom_shm(self): ) ) - def test_colocated_python_with_jax_command(self): - jax_command = "import jax; print(jax.devices());" - user_pod_template = { - "spec": { - "containers": [{ - "name": "jax-tpu", - "image": "gcr.io/my-project/jax-tpu:latest", - "command": ["python3", "-c", jax_command], - }] - } - } - pw_jobset = self._create_jobset( - topology="2x2", - num_slices=1, - user_pod_template=user_pod_template, - main_container_name="jax-tpu", - ) - - pw_jobset.add_colocated_python(image="gcr.io/my-project/colocated-python:custom") - helper = JobSetManifestHelper(pw_jobset.to_dict()) - - self.assertIn("pathways-head", helper.jobs) - self.assertIn("jax-tpu", helper.containers["pathways-head"]) - jax_container = helper.containers["pathways-head"]["jax-tpu"] - self.assertEqual(jax_container["command"], ["python3", "-c", jax_command]) - - self.assertIn("pathways-worker", helper.jobs) - self.assertIn("colocated-python-sidecar", helper.init_containers["pathways-worker"]) - # Reference similar GKE / JobSet custom object unit test suites: # - Google3 GKE JobSet test suite: //depot/google3/cloud/ai/map/catmint/supervisor/client/python/orchestrators/gke_callbacks_test.py # - Upstream Kubernetes SIGs JobSet unit tests: https://github.com/kubernetes-sigs/jobset/blob/main/pkg/controllers/jobset_controller_test.go @@ -949,29 +794,6 @@ def test_shared_pathways_service(self): self.assertNotIn("pathways-proxy", helper.containers["pathways-head"]) self.assertLen(pod_spec["containers"], 1) - def test_shared_pathways_service_with_user_template_fails(self): - user_pod_template = { - "spec": { - "containers": [{ - "name": "jax-tpu", - "image": "gcr.io/my-project/jax-tpu:latest", - }] - } - } - with self.assertRaisesRegex( - ValueError, - "Cannot enable shared_pathways_service when user_pod_template is" - " provided.", - ): - jobset.PathwaysJobSet( - name="test-sps", - namespace="default", - pathways_dir="gs://bucket/scratch", - tpu_type="v5e", - topology="2x2", - num_slices=2, - shared_pathways_service=True, - user_pod_template=user_pod_template, - ) + if __name__ == "__main__": absltest.main() diff --git a/pathwaysutils/test/experimental/gke/testdata/model_lite_tpuv5e_4x8.yaml b/pathwaysutils/test/experimental/gke/testdata/model_lite_tpuv5e_4x8.yaml index ad72f17..e84e4c8 100644 --- a/pathwaysutils/test/experimental/gke/testdata/model_lite_tpuv5e_4x8.yaml +++ b/pathwaysutils/test/experimental/gke/testdata/model_lite_tpuv5e_4x8.yaml @@ -73,8 +73,6 @@ spec: volumeMounts: - mountPath: /tmp name: shared-tmp - dnsPolicy: ClusterFirstWithHostNet - hostNetwork: true initContainers: - args: - --server_port=29001 @@ -209,8 +207,6 @@ spec: volumeMounts: - mountPath: /tmp name: shared-tmp - dnsPolicy: ClusterFirstWithHostNet - hostNetwork: true nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 4x8