diff --git a/pathwaysutils/experimental/gke/jobset.py b/pathwaysutils/experimental/gke/jobset.py index 37bdf90..62f3693 100644 --- a/pathwaysutils/experimental/gke/jobset.py +++ b/pathwaysutils/experimental/gke/jobset.py @@ -478,6 +478,10 @@ def __init__( self._labels = dict(labels) if labels else {} self._annotations = dict(annotations) if annotations else {} + if num_slices <= 0: + raise ValueError(f"num_slices ({num_slices}) must be positive") + if elastic_slices < 0: + raise ValueError(f"elastic_slices ({elastic_slices}) cannot be negative") if elastic_slices > num_slices: raise ValueError( f"elastic_slices ({elastic_slices}) cannot be greater than " @@ -492,11 +496,25 @@ def __init__( ) # Calculate VMs. - dims = [int(x) for x in topology.split("x")] + if not topology.strip(): + raise ValueError(f"Invalid topology: {topology!r}") + try: + dims = [int(x) for x in topology.split("x")] + except ValueError as e: + raise ValueError(f"Invalid topology: {topology!r}") from e + if any(d <= 0 for d in dims): + raise ValueError( + f"Topology dimensions must be positive integers, got {topology!r}" + ) total_chips = math.prod(dims) chips_per_vm = 8 if tpu_type.lower().endswith("8t") else 4 if total_chips < chips_per_vm: num_vms = 1 + elif total_chips % chips_per_vm != 0: + raise ValueError( + f"Total chips ({total_chips}) for topology {topology!r} must be" + f" divisible by chips_per_vm ({chips_per_vm}) for {tpu_type!r}" + ) else: num_vms = total_chips // chips_per_vm @@ -1271,7 +1289,7 @@ def import_yaml(cls, filepath: str) -> "PathwaysJobSet": head_job_template = _deserialize_dict( api_client, job["template"], client.V1JobTemplateSpec ) - elif job["name"] in ("worker", PATHWAYS_WORKER_JOB_NAME): + elif job["name"] == PATHWAYS_WORKER_JOB_NAME: worker_job_template = _deserialize_dict( api_client, job["template"], client.V1JobTemplateSpec ) @@ -1289,6 +1307,7 @@ def import_yaml(cls, filepath: str) -> "PathwaysJobSet": rm_image = DEFAULT_PATHWAYS_RM_AND_WORKER_IMAGE proxy_image = DEFAULT_PATHWAYS_PROXY_IMAGE + has_proxy = False if ( head_job_template.spec and head_job_template.spec.template @@ -1301,8 +1320,11 @@ def import_yaml(cls, filepath: str) -> "PathwaysJobSet": for c in all_containers: if c.name == "pathways-rm" and c.image: rm_image = c.image - elif c.name == "pathways-proxy" and c.image: - proxy_image = c.image + elif c.name == "pathways-proxy": + has_proxy = True + if c.image: + proxy_image = c.image + instance._shared_pathways_service = not has_proxy instance._pathways_rm_and_worker_image = rm_image instance._pathways_proxy_image = proxy_image @@ -1317,14 +1339,13 @@ def _validate_config(cls, config: dict[str, Any]) -> None: jobs = { j["name"]: j for j in config.get("spec", {}).get("replicatedJobs", []) } - if "head" not in jobs and PATHWAYS_HEAD_JOB_NAME not in jobs: + if PATHWAYS_HEAD_JOB_NAME not in jobs: raise ValueError( - f"Missing head replicated job ('head' or '{PATHWAYS_HEAD_JOB_NAME}')" + f"Missing head replicated job ('{PATHWAYS_HEAD_JOB_NAME}')" ) - if "worker" not in jobs and PATHWAYS_WORKER_JOB_NAME not in jobs: + if PATHWAYS_WORKER_JOB_NAME not in jobs: raise ValueError( - "Missing worker replicated job ('worker' or" - f" '{PATHWAYS_WORKER_JOB_NAME}')" + f"Missing worker replicated job ('{PATHWAYS_WORKER_JOB_NAME}')" ) def apply( diff --git a/pathwaysutils/test/experimental/gke/jobset_test.py b/pathwaysutils/test/experimental/gke/jobset_test.py index a11e1ef..596ee75 100644 --- a/pathwaysutils/test/experimental/gke/jobset_test.py +++ b/pathwaysutils/test/experimental/gke/jobset_test.py @@ -941,7 +941,7 @@ def test_import_validation_failures(self): "metadata": {"name": "test"}, "spec": { "replicatedJobs": [ - {"name": "worker", "replicas": 1, "template": {}} + {"name": "pathways-worker", "replicas": 1, "template": {}} ] }, } @@ -1771,6 +1771,128 @@ def test_with_custom_libtpu_rejects_invalid_paths( with self.assertRaisesRegex(ValueError, expected_error): pw_jobset.with_custom_libtpu(path, mount_host_dir=mount_host_dir) + @parameterized.named_parameters( + dict( + testcase_name="zero_slices", + topology="2x2", + num_slices=0, + elastic_slices=0, + expected_error="num_slices .* must be positive", + ), + dict( + testcase_name="negative_slices", + topology="2x2", + num_slices=-1, + elastic_slices=0, + expected_error="num_slices .* must be positive", + ), + dict( + testcase_name="negative_elastic_slices", + topology="2x2", + num_slices=2, + elastic_slices=-1, + expected_error="elastic_slices .* cannot be negative", + ), + dict( + testcase_name="empty_topology", + topology="", + num_slices=1, + elastic_slices=0, + expected_error="Invalid topology", + ), + dict( + testcase_name="non_numeric_topology", + topology="2xabc", + num_slices=1, + elastic_slices=0, + expected_error="Invalid topology", + ), + dict( + testcase_name="zero_dim_topology", + topology="0x4", + num_slices=1, + elastic_slices=0, + expected_error="Topology dimensions must be positive", + ), + dict( + testcase_name="negative_dim_topology", + topology="-2x4", + num_slices=1, + elastic_slices=0, + expected_error="Topology dimensions must be positive", + ), + dict( + testcase_name="non_divisible_chips_per_vm", + topology="2x3", + num_slices=1, + elastic_slices=0, + expected_error="must be divisible by chips_per_vm", + ), + ) + def test_invalid_slices_and_topology( + self, topology, num_slices, elastic_slices, expected_error + ): + with self.assertRaisesRegex(ValueError, expected_error): + self._create_jobset( + topology=topology, + num_slices=num_slices, + elastic_slices=elastic_slices, + ) + + def test_import_yaml_sets_shared_pathways_service_and_supports_add_user_workload( + self, + ): + temp_dir = self.create_tempdir().full_path + # Normal JobSet roundtrip + add_user_workload + pw_jobset = self._create_jobset(topology="2x2", num_slices=1) + normal_path = os.path.join(temp_dir, "normal.yaml") + pw_jobset.export_yaml(normal_path) + imported_normal = jobset.PathwaysJobSet.import_yaml(normal_path) + self.assertFalse(imported_normal._shared_pathways_service) + imported_normal.add_user_workload( + image="us-docker.pkg.dev/my-project/test:v1", + command="python3 test.py", + ) + self.assertIn("successPolicy", imported_normal.to_dict()["spec"]) + + # Shared Pathways Service JobSet roundtrip + add_user_workload + sps_jobset = self._create_jobset( + topology="2x2", num_slices=1, shared_pathways_service=True + ) + sps_path = os.path.join(temp_dir, "sps.yaml") + sps_jobset.export_yaml(sps_path) + imported_sps = jobset.PathwaysJobSet.import_yaml(sps_path) + self.assertTrue(imported_sps._shared_pathways_service) + imported_sps.add_user_workload( + image="us-docker.pkg.dev/my-project/test:v1", + command="python3 test.py", + ) + self.assertNotIn("successPolicy", imported_sps.to_dict()["spec"]) + + def test_import_yaml_rejects_legacy_head_and_worker_job_names(self): + temp_dir = self.create_tempdir().full_path + pw_jobset = self._create_jobset(topology="2x2", num_slices=1) + + head_config = pw_jobset.to_dict() + for job in head_config["spec"]["replicatedJobs"]: + if job["name"] == jobset.PATHWAYS_HEAD_JOB_NAME: + job["name"] = "head" + head_path = os.path.join(temp_dir, "head_legacy.yaml") + with open(head_path, "w") as f: + yaml.dump(head_config, f) + with self.assertRaisesRegex(ValueError, "Missing head replicated job"): + jobset.PathwaysJobSet.import_yaml(head_path) + + worker_config = pw_jobset.to_dict() + for job in worker_config["spec"]["replicatedJobs"]: + if job["name"] == jobset.PATHWAYS_WORKER_JOB_NAME: + job["name"] = "worker" + worker_path = os.path.join(temp_dir, "worker_legacy.yaml") + with open(worker_path, "w") as f: + yaml.dump(worker_config, f) + with self.assertRaisesRegex(ValueError, "Missing worker replicated job"): + jobset.PathwaysJobSet.import_yaml(worker_path) + if __name__ == "__main__": absltest.main()