diff --git a/pathwaysutils/experimental/gke/jobset.py b/pathwaysutils/experimental/gke/jobset.py index 28bee85..88dd451 100644 --- a/pathwaysutils/experimental/gke/jobset.py +++ b/pathwaysutils/experimental/gke/jobset.py @@ -18,7 +18,8 @@ import logging import math import time -from typing import TYPE_CHECKING, Any, Mapping, Sequence +from typing import Any, Mapping, Sequence, TYPE_CHECKING + import yaml try: @@ -86,7 +87,7 @@ def _format_image(image: str, default_tag: str) -> str: if "@" in image: return image last_slash = image.rfind("/") - if ":" in image[last_slash + 1:]: + if ":" in image[last_slash + 1 :]: return image return f"{image}:{default_tag}" @@ -135,14 +136,16 @@ def __init__( topology: TPU topology (e.g., "2x2"). num_slices: Number of slices. max_restarts: Maximum number of restarts for the JobSet. - max_slice_restarts: Maximum number of slice restarts (defaults to 1_000_000 in headless and SPS mode). + max_slice_restarts: Maximum number of slice restarts (defaults to + 1_000_000 in headless and SPS mode). termination_grace_period_seconds: Optional termination grace period. pathways_version: Version tag for Pathways images. jobset_api_version: API version of JobSet. 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. + shared_pathways_service: Whether to run only RM for Shared Pathways + Service. pathways_rm_and_worker_image: Base Docker image for Resource Manager and Worker containers. pathways_proxy_image: Base Docker image for Proxy container. @@ -251,7 +254,8 @@ def _build_head_job_template( instance_type: TPU instance type (e.g., "tpuv5:2x2"). image_tag: Version tag for Pathways images. elastic_slices: Number of elastic slices. - shared_pathways_service: Whether to run only RM for Shared Pathways Service. + shared_pathways_service: Whether to run only RM for Shared Pathways + Service. pathways_rm_and_worker_image: Base Docker image for Resource Manager. pathways_proxy_image: Base Docker image for Proxy container. @@ -359,23 +363,33 @@ def _build_head_job_template( head_pod_spec = client.V1PodSpec( containers=containers, restart_policy="Never", + host_network=True, + dns_policy="ClusterFirstWithHostNet", + node_selector={"cloud.google.com/gke-nodepool": "cpu-np"}, + priority_class_name="high", + volumes=[ + client.V1Volume( + name="shared-tmp", + host_path=client.V1HostPathVolumeSource( + path="/tmp", type="DirectoryOrCreate" + ), + ) + ], ) - job_annotations = { - "alpha.jobset.sigs.k8s.io/exclusive-topology": "kubernetes.io/hostname" + job_annos = { + "kueue.x-k8s.io/safe-to-forcefully-delete": "true", } head_job_template = client.V1JobTemplateSpec( - metadata=client.V1ObjectMeta(annotations=job_annotations), + metadata=client.V1ObjectMeta(annotations=job_annos), spec=client.V1JobSpec( backoff_limit=0, completion_mode="Indexed", completions=1, parallelism=1, template=client.V1PodTemplateSpec( - metadata=client.V1ObjectMeta( - annotations=job_annotations, labels={} - ), + metadata=client.V1ObjectMeta(labels={}), spec=head_pod_spec, ), ), @@ -504,11 +518,24 @@ def _build_worker_job_template( ) ], restart_policy="OnFailure", + host_network=True, + dns_policy="ClusterFirstWithHostNet", + priority_class_name="high", ) - if termination_grace_period_seconds is not None: - worker_pod_spec.termination_grace_period_seconds = ( - termination_grace_period_seconds - ) + worker_pod_spec.termination_grace_period_seconds = ( + termination_grace_period_seconds + if termination_grace_period_seconds is not None + else 60 + ) + + pod_annos = { + "alpha.jobset.sigs.k8s.io/exclusive-topology": ( + "cloud.google.com/gke-nodepool" + ) + } + pod_lbls = { + "kueue.x-k8s.io/podset": "worker", + } worker_job_template = client.V1JobTemplateSpec( metadata=client.V1ObjectMeta(), @@ -519,11 +546,7 @@ def _build_worker_job_template( parallelism=num_vms, template=client.V1PodTemplateSpec( metadata=client.V1ObjectMeta( - annotations={ - "alpha.jobset.sigs.k8s.io/exclusive-topology": ( - "cloud.google.com/gke-nodepool" - ) - } + annotations=pod_annos, labels=pod_lbls ), spec=worker_pod_spec, ), @@ -574,6 +597,94 @@ def _add_volume_to_pod_spec( volumes.append(volume) pod_spec.volumes = volumes + def add_user_workload( + self, + image: str, + command: Sequence[str] | str, + ) -> "PathwaysJobSet": + """Adds a user workload container to the head pod and converts RM/Proxy to sidecars. + + Args: + image: Docker image for the user workload. + command: Command to execute in the container (string or sequence of + strings). + + Returns: + The PathwaysJobSet instance for chaining. + """ + pod_spec = self._head_job_template.spec.template.spec + + # 1. Convert existing head containers to sidecar initContainers. + init_containers = pod_spec.init_containers or [] + for c in pod_spec.containers or []: + c.restart_policy = "Always" + if not any(ic.name == c.name for ic in init_containers): + init_containers.append(c) + pod_spec.init_containers = init_containers + + # 2. Ensure shared-tmp volume is on head pod. + self._add_volume_to_pod_spec( + pod_spec, + client.V1Volume( + name="shared-tmp", + host_path=client.V1HostPathVolumeSource( + path="/tmp", type="DirectoryOrCreate" + ), + ), + ) + + # 3. Handle command. + cmd = ["sh", "-c", command] if isinstance(command, str) else list(command) + + # 4. Hardcoded env vars. + user_env_list = [ + client.V1EnvVar(name="JAX_PLATFORMS", value="proxy"), + client.V1EnvVar( + name="JAX_BACKEND_TARGET", + value=f"grpc://localhost:{PATHWAYS_PROXY_PORT}", + ), + client.V1EnvVar( + name="MEGASCALE_NUM_SLICES", + value=str(self._worker_replicas), + ), + client.V1EnvVar( + name="JOBSET_NAME", + value_from=client.V1EnvVarSource( + field_ref=client.V1ObjectFieldSelector( + field_path=( + "metadata.annotations['jobset.sigs.k8s.io/jobset-name']" + ) + ) + ), + ), + ] + + # 5. Hardcoded resources. + resources = client.V1ResourceRequirements( + limits={"cpu": "24", "memory": "100G"} + ) + + # 6. Hardcoded volume mounts. + volume_mounts = [client.V1VolumeMount(name="shared-tmp", mount_path="/tmp")] + + user_container = client.V1Container( + name="user-workload", + image=image, + image_pull_policy="Always", + command=cmd, + env=user_env_list, + resources=resources, + volume_mounts=volume_mounts, + ) + + pod_spec.containers = [user_container] + self._success_policy = { + "operator": "All", + "targetReplicatedJobs": [PATHWAYS_HEAD_JOB_NAME], + } + + return self + def add_colocated_python( self, image: str, @@ -699,6 +810,27 @@ def _compile_config(self) -> dict[str, Any]: self._worker_job_template ) + # Preserve restartPolicy on initContainers (e.g. native sidecars) if set. + for job_tmpl, ser_tmpl in ( + (self._head_job_template, serialized_head), + (self._worker_job_template, serialized_worker), + ): + if ( + job_tmpl.spec + and job_tmpl.spec.template + and job_tmpl.spec.template.spec + ): + orig_init = job_tmpl.spec.template.spec.init_containers or [] + ser_init = ( + ser_tmpl.get("spec", {}) + .get("template", {}) + .get("spec", {}) + .get("initContainers", []) + ) + for orig_c, ser_c in zip(orig_init, ser_init): + if getattr(orig_c, "restart_policy", None): + ser_c["restartPolicy"] = orig_c.restart_policy + head_job = { "name": PATHWAYS_HEAD_JOB_NAME, "replicas": 1, @@ -793,6 +925,32 @@ def import_yaml(cls, filepath: str) -> "PathwaysJobSet": ) instance._worker_replicas = job["replicas"] + # Preserve restartPolicy on deserialized init_containers. + for job in config["spec"]["replicatedJobs"]: + target_template = None + if job["name"] == PATHWAYS_HEAD_JOB_NAME: + target_template = head_job_template + elif job["name"] in ("worker", PATHWAYS_WORKER_JOB_NAME): + target_template = worker_job_template + + if ( + target_template + and target_template.spec + and target_template.spec.template + and target_template.spec.template.spec + ): + raw_init = ( + job.get("template", {}) + .get("spec", {}) + .get("template", {}) + .get("spec", {}) + .get("initContainers", []) + ) + des_init = target_template.spec.template.spec.init_containers or [] + for raw_c, des_c in zip(raw_init, des_init): + if "restartPolicy" in raw_c: + des_c.restart_policy = raw_c["restartPolicy"] + if head_job_template is None: raise ValueError(f"Missing head job ({PATHWAYS_HEAD_JOB_NAME}) in config") if worker_job_template is None: diff --git a/pathwaysutils/test/experimental/gke/jobset_test.py b/pathwaysutils/test/experimental/gke/jobset_test.py index 41fd382..07c5d5d 100644 --- a/pathwaysutils/test/experimental/gke/jobset_test.py +++ b/pathwaysutils/test/experimental/gke/jobset_test.py @@ -106,7 +106,6 @@ def get_all_containers_by_name( return matches - class PathwaysJobSetTest(parameterized.TestCase): def _create_jobset( @@ -149,8 +148,12 @@ 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.assertNotIn("hostNetwork", pod_spec) - self.assertNotIn("dnsPolicy", pod_spec) + self.assertTrue(pod_spec["hostNetwork"]) + self.assertEqual(pod_spec["dnsPolicy"], "ClusterFirstWithHostNet") + self.assertEqual(pod_spec["priorityClassName"], "high") + self.assertEqual( + pod_spec["nodeSelector"]["cloud.google.com/gke-nodepool"], "cpu-np" + ) self.assertEqual(pod_spec["restartPolicy"], "Never") def test_headless_head_job_containers(self): @@ -207,8 +210,9 @@ def test_worker_job_pod_spec(self): helper = JobSetManifestHelper(config) pod_spec = helper.pod_specs["pathways-worker"] - self.assertNotIn("hostNetwork", pod_spec) - self.assertNotIn("dnsPolicy", pod_spec) + self.assertTrue(pod_spec["hostNetwork"]) + self.assertEqual(pod_spec["dnsPolicy"], "ClusterFirstWithHostNet") + self.assertEqual(pod_spec["priorityClassName"], "high") self.assertEqual(pod_spec["restartPolicy"], "OnFailure") self.assertEqual(pod_spec["terminationGracePeriodSeconds"], 60) @@ -328,7 +332,7 @@ def test_add_gcsfuse_volumes_and_annotations(self): def test_add_gcsfuse_handles_none_metadata(self): pw_jobset = self._create_jobset(topology="2x2", num_slices=1) - + # Force metadata to be None to simulate imported templates or raw specs without metadata pw_jobset._head_job_template.metadata = None pw_jobset._head_job_template.spec.template.metadata = None @@ -342,7 +346,7 @@ def test_add_gcsfuse_handles_none_metadata(self): bucket="my-bucket", ) helper = JobSetManifestHelper(pw_jobset.to_dict()) - + self.assertEqual(helper.job_metadatas["pathways-head"].get("annotations", {}).get("gke-gcsfuse/volumes"), "true") self.assertEqual(helper.pod_metadatas["pathways-head"].get("annotations", {}).get("gke-gcsfuse/volumes"), "true") self.assertEqual(helper.job_metadatas["pathways-worker"].get("annotations", {}).get("gke-gcsfuse/volumes"), "true") @@ -350,7 +354,7 @@ def test_add_gcsfuse_handles_none_metadata(self): def test_add_gcsfuse_preserves_existing_metadata(self): pw_jobset = self._create_jobset(topology="2x2", num_slices=1) - + # Pre-populate metadata, annotations, and labels pw_jobset._head_job_template.metadata = client.V1ObjectMeta( labels={"existing-job-label": "value"}, @@ -403,14 +407,14 @@ def test_add_gcsfuse_preserves_existing_volumes(self): def test_add_colocated_python_handles_none_volumes(self): pw_jobset = self._create_jobset(topology="2x2", num_slices=1) - + # Force volumes to be None pw_jobset._worker_job_template.spec.template.spec.volumes = None # Should not crash and should correctly add volume pw_jobset.add_colocated_python(image="gcr.io/my-project/colocated-python:custom") helper = JobSetManifestHelper(pw_jobset.to_dict()) - + self.assertIn("shared-memory", helper.volumes["pathways-worker"]) def test_add_colocated_python_sidecar(self): @@ -438,7 +442,7 @@ def test_add_colocated_python_sidecar(self): def test_add_colocated_python_preserves_init_containers(self): pw_jobset = self._create_jobset(topology="2x2", num_slices=1) - + # Pre-populate init container on worker pod worker_spec = pw_jobset._worker_job_template.spec.template.spec existing_init = client.V1Container(name="existing-init-container", image="ubuntu:latest") @@ -817,6 +821,118 @@ def test_failure_policy(self): config = pw_jobset.to_dict() self.assertEqual(config["spec"]["failurePolicy"]["maxRestarts"], 5) + def test_add_user_workload(self): + pw_jobset = self._create_jobset(topology="2x2", num_slices=1) + pw_jobset.add_user_workload( + image="us-docker.pkg.dev/my-project/test:v1", + command="python3 -m test_module", + ) + + config = pw_jobset.to_dict() + helper = JobSetManifestHelper(config) + + # Verify success policy targets pathways-head + self.assertEqual( + config["spec"]["successPolicy"], + { + "operator": "All", + "targetReplicatedJobs": ["pathways-head"], + }, + ) + + # Verify head pod has user-workload container in containers + self.assertIn("user-workload", helper.containers["pathways-head"]) + user_c = helper.containers["pathways-head"]["user-workload"] + self.assertEqual(user_c["image"], "us-docker.pkg.dev/my-project/test:v1") + self.assertEqual(user_c["command"], ["sh", "-c", "python3 -m test_module"]) + self.assertTrue( + any( + e["name"] == "MEGASCALE_NUM_SLICES" and e["value"] == "1" + for e in user_c["env"] + ) + ) + self.assertTrue( + any( + e["name"] == "JAX_PLATFORMS" and e["value"] == "proxy" + for e in user_c["env"] + ) + ) + + # Verify RM and Proxy were moved to initContainers with restartPolicy Always + self.assertIn("pathways-rm", helper.init_containers["pathways-head"]) + self.assertIn("pathways-proxy", helper.init_containers["pathways-head"]) + self.assertEqual( + helper.init_containers["pathways-head"]["pathways-rm"]["restartPolicy"], + "Always", + ) + self.assertEqual( + helper.init_containers["pathways-head"]["pathways-proxy"][ + "restartPolicy" + ], + "Always", + ) + + # Verify shared-tmp volume is mounted + self.assertIn("shared-tmp", helper.volumes["pathways-head"]) + self.assertTrue( + any( + m["name"] == "shared-tmp" and m["mountPath"] == "/tmp" + for m in user_c["volumeMounts"] + ) + ) + + def test_add_user_workload_roundtrip(self): + pw_jobset = self._create_jobset(topology="2x2", num_slices=1) + pw_jobset.add_user_workload( + image="us-docker.pkg.dev/my-project/test:v1", + command=["python3", "test.py"], + ) + + temp_filepath = os.path.join( + self.create_tempdir().full_path, "jobset_user_workload.yaml" + ) + pw_jobset.export_yaml(temp_filepath) + imported = jobset.PathwaysJobSet.import_yaml(temp_filepath) + + self.assertEqual( + normalize_k8s_spec(pw_jobset.to_dict()), + normalize_k8s_spec(imported.to_dict()), + ) + + def test_pod_template_metadata(self): + pw_jobset = jobset.PathwaysJobSet( + name="test-workload", + namespace="default", + pathways_dir="gs://bucket/scratch", + tpu_type="v5e", + topology="2x2", + num_slices=1, + ) + config = pw_jobset.to_dict() + helper = JobSetManifestHelper(config) + + # Head job and pod annotations/labels + self.assertEqual( + helper.job_metadatas["pathways-head"]["annotations"], + {"kueue.x-k8s.io/safe-to-forcefully-delete": "true"}, + ) + self.assertNotIn("annotations", helper.pod_metadatas["pathways-head"]) + self.assertEqual(helper.pod_metadatas["pathways-head"]["labels"], {}) + + # Worker job and pod annotations/labels + self.assertEqual( + helper.pod_metadatas["pathways-worker"]["annotations"][ + "alpha.jobset.sigs.k8s.io/exclusive-topology" + ], + "cloud.google.com/gke-nodepool", + ) + self.assertEqual( + helper.pod_metadatas["pathways-worker"]["labels"][ + "kueue.x-k8s.io/podset" + ], + "worker", + ) + def test_shared_pathways_service(self): pw_jobset = self._create_jobset( name="test-sps", @@ -845,6 +961,76 @@ def test_shared_pathways_service(self): self.assertNotIn("pathways-proxy", helper.containers["pathways-head"]) self.assertLen(pod_spec["containers"], 1) + def test_kokoro_pretraining_workload_generation(self): + """Verifies that PathwaysJobSet can generate a JobSet matching Kokoro pretraining test workloads.""" + pw_jobset = jobset.PathwaysJobSet( + name="maxtext-pretraining-test", + namespace="default", + pathways_dir="gs://my-bucket/scratch", + tpu_type="v5e", + topology="4x8", + num_slices=1, + labels={"kueue.x-k8s.io/queue-name": "multislice-queue"}, + ) + pw_jobset.add_user_workload( + image="us-docker.pkg.dev/my-project/maxtext:latest", + command="python3 MaxText/train.py MaxText/configs/base.yml", + ) + + config = pw_jobset.to_dict() + helper = JobSetManifestHelper(config) + + # Verify queue label + self.assertEqual( + config["metadata"]["labels"]["kueue.x-k8s.io/queue-name"], + "multislice-queue", + ) + + # Verify success policy + self.assertEqual( + config["spec"]["successPolicy"]["targetReplicatedJobs"], + ["pathways-head"], + ) + + # Verify head pod spec: hostNetwork, dnsPolicy, priorityClassName, nodeSelector + head_pod_spec = helper.pod_specs["pathways-head"] + self.assertTrue(head_pod_spec["hostNetwork"]) + self.assertEqual(head_pod_spec["dnsPolicy"], "ClusterFirstWithHostNet") + self.assertEqual(head_pod_spec["priorityClassName"], "high") + self.assertEqual( + head_pod_spec["nodeSelector"]["cloud.google.com/gke-nodepool"], + "cpu-np", + ) + + # Verify head init containers (RM and Proxy sidecars) + self.assertIn("pathways-rm", helper.init_containers["pathways-head"]) + self.assertIn("pathways-proxy", helper.init_containers["pathways-head"]) + self.assertEqual( + helper.init_containers["pathways-head"]["pathways-rm"]["restartPolicy"], + "Always", + ) + self.assertEqual( + helper.init_containers["pathways-head"]["pathways-proxy"][ + "restartPolicy" + ], + "Always", + ) + + # Verify user workload container + self.assertIn("user-workload", helper.containers["pathways-head"]) + user_container = helper.containers["pathways-head"]["user-workload"] + self.assertEqual( + user_container["command"], + ["sh", "-c", "python3 MaxText/train.py MaxText/configs/base.yml"], + ) + + # Verify worker pod spec: hostNetwork, dnsPolicy, priorityClassName, terminationGracePeriodSeconds + worker_pod_spec = helper.pod_specs["pathways-worker"] + self.assertTrue(worker_pod_spec["hostNetwork"]) + self.assertEqual(worker_pod_spec["dnsPolicy"], "ClusterFirstWithHostNet") + self.assertEqual(worker_pod_spec["priorityClassName"], "high") + self.assertEqual(worker_pod_spec["terminationGracePeriodSeconds"], 60) + if __name__ == "__main__": absltest.main()