Add base PathwaysJobSet builder

PiperOrigin-RevId: 919912906

Author: lukebaumann

pathwaysutils/experimental/gke/__init__.py

@@ -0,0 +1 @@
+# Empty __init__.py

pathwaysutils/experimental/gke/jobset.py

@@ -0,0 +1,142 @@
+"""Pathways JobSet generator and builder (Skeleton)."""
+
+import math
+from typing import Any, Optional
+from kubernetes import client
+import yaml
+
+# Core constants
+PATHWAYS_HEAD_JOB_NAME = "pathways-head"
+PATHWAYS_WORKER_JOB_NAME = "pathways-worker"
+
+MACHINE_TYPE_TO_TPU_VERSION_MAP = {
+    "tpu7x-standard-4t": "tpu7x",
+    "tpu7x": "tpu7x",
+    "ct6e-standard-4t": "tpuv6e",
+    "v6e": "tpuv6e",
+    "ct6e-standard-8t": "tpuv6e1t",
+    "ct5p-hightpu-4t": "tpuv5",
+    "v5p": "tpuv5",
+    "ct5lp-hightpu-4t": "tpuv5e",
+    "v5e": "tpuv5e",
+    "ct5lp-hightpu-8t": "tpuv5e1t",
+    "ct4p-hightpu-4t": "tpuv4",
+    "v4": "tpuv4",
+}
+
+MACHINE_TYPE_TO_GKE_ACCELERATOR_TYPE_MAP = {
+    "tpu7x-standard-4t": "tpu7x",
+    "tpu7x": "tpu7x",
+    "ct6e-standard-4t": "tpu-v6e-slice",
+    "v6e": "tpu-v6e-slice",
+    "ct6e-standard-8t": "tpu-v6e-slice",
+    "ct5p-hightpu-4t": "tpu-v5p-slice",
+    "v5p": "tpu-v5p-slice",
+    "ct5lp-hightpu-4t": "tpu-v5-lite-podslice",
+    "v5e": "tpu-v5-lite-podslice",
+    "ct5lp-hightpu-8t": "tpu-v5-lite-podslice",
+    "ct4p-hightpu-4t": "tpu-v4-podslice",
+    "v4": "tpu-v4-podslice",
+}
+
+
+class PathwaysJobSet:
+  """Generates JobSet configuration for Pathways (Skeleton)."""
+
+  def __init__(
+      self,
+      name: str,
+      namespace: str,
+      pathways_dir: str,
+      tpu_type: str,
+      topology: str,
+      num_slices: int,
+      user_pod_template: Optional[dict[str, Any]] = None,
+      max_restarts: int = 0,
+      jobset_api_version: str = "v1alpha2",
+      labels: Optional[dict[str, str]] = None,
+      annotations: Optional[dict[str, str]] = None,
+  ):
+    """Initializes the JobSet configuration using K8s API models."""
+    self._name = name
+    self._namespace = namespace
+    self._jobset_api_version = jobset_api_version
+    self._max_restarts = max_restarts
+    self._worker_replicas = num_slices
+    self._labels = labels or {}
+    self._annotations = annotations or {}
+
+    tpu_version = MACHINE_TYPE_TO_TPU_VERSION_MAP.get(tpu_type.lower())
+    if not tpu_version:
+      raise ValueError(f"Unsupported TPU type: {tpu_type}")
+
+    # Build minimal head template (placeholder)
+    self._head_job_template = self._build_minimal_job_template("head")
+
+    # Build minimal worker template (placeholder)
+    self._worker_job_template = self._build_minimal_job_template("worker")
+
+    self._success_policy = None
+    if user_pod_template:
+      self._success_policy = {
+          "operator": "All",
+          "targetReplicatedJobs": [PATHWAYS_HEAD_JOB_NAME],
+      }
+
+  def _build_minimal_job_template(self, role: str) -> client.V1JobTemplateSpec:
+    pod_spec = client.V1PodSpec(
+        containers=[client.V1Container(name=f"placeholder-{role}", image="ubuntu")]
+    )
+    job_spec = client.V1JobSpec(
+        template=client.V1PodTemplateSpec(
+            metadata=client.V1ObjectMeta(labels={"role": role}),
+            spec=pod_spec
+        )
+    )
+    return client.V1JobTemplateSpec(spec=job_spec)
+
+  def _compile_config(self) -> dict[str, Any]:
+    api_client = client.ApiClient()
+    serialized_head = api_client.sanitize_for_serialization(
+        self._head_job_template
+    )
+    serialized_worker = api_client.sanitize_for_serialization(
+        self._worker_job_template
+    )
+
+    replicated_jobs = [
+        {
+            "name": PATHWAYS_HEAD_JOB_NAME,
+            "replicas": 1,
+            "template": serialized_head,
+        },
+        {
+            "name": PATHWAYS_WORKER_JOB_NAME,
+            "replicas": self._worker_replicas,
+            "template": serialized_worker,
+        },
+    ]
+
+    jobset_config = {
+        "apiVersion": f"jobset.sigs.k8s.io/{self._jobset_api_version}",
+        "kind": "JobSet",
+        "metadata": {
+            "name": self._name,
+            "namespace": self._namespace,
+        },
+        "spec": {
+            "failurePolicy": {"maxRestarts": self._max_restarts},
+            "replicatedJobs": replicated_jobs,
+        },
+    }
+    if self._labels:
+      jobset_config["metadata"]["labels"] = self._labels
+    if self._annotations:
+      jobset_config["metadata"]["annotations"] = self._annotations
+    if self._success_policy:
+      jobset_config["spec"]["successPolicy"] = self._success_policy
+
+    return jobset_config
+
+  def to_dict(self) -> dict[str, Any]:
+    return self._compile_config()

pathwaysutils/test/experimental/gke/jobset_test.py

@@ -0,0 +1,61 @@
+from absl.testing import absltest
+from absl.testing import parameterized
+from pathwaysutils.experimental.gke import jobset
+
+class PathwaysJobSetTest(parameterized.TestCase):
+
+  def test_invalid_tpu_type(self):
+    with self.assertRaisesRegex(ValueError, "Unsupported TPU type"):
+      jobset.PathwaysJobSet(
+          name="test-jobset",
+          namespace="default",
+          pathways_dir="gs://test-bucket",
+          tpu_type="invalid-tpu",
+          topology="4x4",
+          num_slices=1,
+      )
+
+  def test_basic_jobset_structure(self):
+    js = jobset.PathwaysJobSet(
+        name="test-jobset",
+        namespace="default",
+        pathways_dir="gs://test-bucket",
+        tpu_type="v5e",
+        topology="4x8",
+        num_slices=2,
+        labels={"app": "pathways"},
+        annotations={"example.com/annotation": "value"},
+    )
+    config = js.to_dict()
+
+    self.assertEqual(config["apiVersion"], "jobset.sigs.k8s.io/v1alpha2")
+    self.assertEqual(config["kind"], "JobSet")
+    self.assertEqual(config["metadata"]["name"], "test-jobset")
+    self.assertEqual(config["metadata"]["namespace"], "default")
+    self.assertEqual(config["metadata"]["labels"]["app"], "pathways")
+    self.assertEqual(
+        config["metadata"]["annotations"]["example.com/annotation"], "value"
+    )
+
+    self.assertEqual(config["spec"]["failurePolicy"]["maxRestarts"], 0)
+    
+    replicated_jobs = config["spec"]["replicatedJobs"]
+    self.assertEqual(len(replicated_jobs), 2)
+    
+    head_job = replicated_jobs[0]
+    self.assertEqual(head_job["name"], "pathways-head")
+    self.assertEqual(head_job["replicas"], 1)
+    
+    # In K8s API models, V1JobTemplateSpec -> V1JobSpec -> V1PodTemplateSpec -> V1PodSpec
+    # When serialized, they match this structure.
+    head_pod_spec = head_job["template"]["spec"]["template"]["spec"]
+    self.assertEqual(head_pod_spec["containers"][0]["name"], "placeholder-head")
+
+    worker_job = replicated_jobs[1]
+    self.assertEqual(worker_job["name"], "pathways-worker")
+    self.assertEqual(worker_job["replicas"], 2)
+    worker_pod_spec = worker_job["template"]["spec"]["template"]["spec"]
+    self.assertEqual(worker_pod_spec["containers"][0]["name"], "placeholder-worker")
+
+if __name__ == "__main__":
+  absltest.main()