Preserve shared service mode in import_yaml and validate topologies in PathwaysJobSet.

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(

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()