diff --git a/pathwaysutils/experimental/gke/jobset.py b/pathwaysutils/experimental/gke/jobset.py index dfdec0c..6a69af1 100644 --- a/pathwaysutils/experimental/gke/jobset.py +++ b/pathwaysutils/experimental/gke/jobset.py @@ -103,6 +103,7 @@ def __init__( elastic_slices: int = 0, labels: Mapping[str, str] | None = None, annotations: Mapping[str, str] | None = None, + shared_pathways_service: bool = False, ): """Initializes the instance. @@ -124,6 +125,13 @@ def __init__( labels: Optional labels for the JobSet. annotations: Optional annotations for the JobSet. """ + 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 self._namespace = namespace self._jobset_api_version = jobset_api_version @@ -161,6 +169,7 @@ def __init__( user_pod_template=user_pod_template, main_container_name=main_container_name, elastic_slices=elastic_slices, + shared_pathways_service=shared_pathways_service, ) # Build worker template. @@ -176,7 +185,7 @@ def __init__( ) self._success_policy = None - if user_pod_template: + if user_pod_template or shared_pathways_service: self._success_policy = { "operator": "All", "targetReplicatedJobs": [PATHWAYS_HEAD_JOB_NAME], @@ -199,6 +208,7 @@ def _build_head_job_template( user_pod_template: Mapping[str, Any] | None, main_container_name: str, elastic_slices: int, + shared_pathways_service: bool, ) -> client.V1JobTemplateSpec: """Builds the head job template for the JobSet. @@ -357,10 +367,13 @@ def _build_head_job_template( 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=[rm_container, proxy_container], + containers=containers, ) annotations = {} labels = {} diff --git a/pathwaysutils/test/experimental/gke/jobset_test.py b/pathwaysutils/test/experimental/gke/jobset_test.py index dc9d290..6401bdd 100644 --- a/pathwaysutils/test/experimental/gke/jobset_test.py +++ b/pathwaysutils/test/experimental/gke/jobset_test.py @@ -921,6 +921,57 @@ def test_failure_policy(self): config = pw_jobset.to_dict() self.assertEqual(config["spec"]["failurePolicy"]["maxRestarts"], 5) + def test_shared_pathways_service(self): + pw_jobset = self._create_jobset( + name="test-sps", + shared_pathways_service=True, + ) + + config = pw_jobset.to_dict() + helper = JobSetManifestHelper(config) + + self.assertEqual( + config["spec"]["successPolicy"], + { + "operator": "All", + "targetReplicatedJobs": ["pathways-head"], + }, + ) + + self.assertTrue(config["spec"]["network"]["enableDNSHostnames"]) + self.assertTrue(config["spec"]["network"]["publishNotReadyAddresses"]) + + self.assertIn("pathways-head", helper.jobs) + pod_spec = helper.pod_specs["pathways-head"] + # Head job should only have pathways-rm container, no pathways-proxy. + self.assertIn("pathways-rm", helper.containers["pathways-head"]) + 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()