feat: pod overlay for kubernetes scheduler (#1067,#1068)

Author: azzhipa

torchx/schedulers/kubernetes_scheduler.py

@@ -27,6 +27,65 @@
 See the
 `Volcano Quickstart <https://github.com/volcano-sh/volcano>`_
 for more information.
+
+Pod Overlay
+===========
+
+You can overlay arbitrary Kubernetes Pod fields on generated pods by setting
+the ``kubernetes`` metadata on your role. The structure follows the Kubernetes
+Pod spec with ``metadata`` and ``spec`` fields.
+
+The metadata value can be:
+- A dict with the overlay structure
+- A resource URI pointing to a YAML file (e.g. ``file://``, ``s3://``, ``gs://``)
+
+.. code:: python
+
+    from torchx.specs import AppDef, Role
+
+    # Option 1: Dict
+    role = Role(
+        name="trainer",
+        image="my-image:latest",
+        entrypoint="train.py",
+        metadata={
+            "kubernetes": {
+                "spec": {
+                    "nodeSelector": {"gpu": "true"},
+                    "tolerations": [{
+                        "key": "nvidia.com/gpu",
+                        "operator": "Exists",
+                        "effect": "NoSchedule"
+                    }]
+                }
+            }
+        }
+    )
+
+    # Option 2: Resource URI
+    role = Role(
+        name="trainer",
+        image="my-image:latest",
+        entrypoint="train.py",
+        metadata={
+            "kubernetes": "file:///path/to/pod_overlay.yaml"
+        }
+    )
+
+Example ``pod_overlay.yaml``:
+
+.. code:: yaml
+
+    spec:
+      nodeSelector:
+        node.kubernetes.io/instance-type: p4d.24xlarge
+      tolerations:
+        - key: nvidia.com/gpu
+          operator: Exists
+          effect: NoSchedule
+
+The overlay is deep-merged with the generated pod, preserving existing fields
+and adding or overriding specified ones.
 """
 
 import json
@@ -45,6 +104,7 @@
     Tuple,
     TYPE_CHECKING,
     TypedDict,
+    Union,
 )
 
 import torchx
@@ -97,6 +157,29 @@
 RESERVED_MILLICPU = 100
 RESERVED_MEMMB = 1024
 
+
+def _apply_pod_overlay(pod: "V1Pod", overlay: Dict[str, Any]) -> None:
+    """Apply overlay dict to V1Pod object, merging nested fields."""
+    from kubernetes import client
+
+    api = client.ApiClient()
+    pod_dict = api.sanitize_for_serialization(pod)
+
+    def deep_merge(base: Dict[str, Any], overlay: Dict[str, Any]) -> None:
+        for key, value in overlay.items():
+            if isinstance(value, dict) and key in base and isinstance(base[key], dict):
+                deep_merge(base[key], value)
+            else:
+                base[key] = value
+
+    deep_merge(pod_dict, overlay)
+
+    merged_pod = api._ApiClient__deserialize(pod_dict, "V1Pod")
+    for attr in ["api_version", "kind", "metadata", "spec", "status"]:
+        if hasattr(merged_pod, attr):
+            setattr(pod, attr, getattr(merged_pod, attr))
+
+
 RETRY_POLICIES: Mapping[str, Iterable[Mapping[str, str]]] = {
     RetryPolicy.REPLICA: [],
     RetryPolicy.APPLICATION: [
@@ -402,6 +485,17 @@ def app_to_resource(
             replica_role.env["TORCHX_IMAGE"] = replica_role.image
 
             pod = role_to_pod(name, replica_role, service_account)
+            if k8s_metadata := role.metadata.get("kubernetes"):
+                if isinstance(k8s_metadata, str):
+                    import fsspec  # pyre-ignore[21]
+
+                    with fsspec.open(k8s_metadata, "r") as f:
+                        k8s_metadata = yaml.safe_load(f)
+                elif not isinstance(k8s_metadata, dict):
+                    raise ValueError(
+                        f"metadata['kubernetes'] must be a dict or resource URI, got {type(k8s_metadata)}"
+                    )
+                _apply_pod_overlay(pod, k8s_metadata)
             pod.metadata.labels.update(
                 pod_labels(
                     app=app,
@@ -636,7 +730,7 @@ def schedule(self, dryrun_info: AppDryRunInfo[KubernetesJob]) -> str:
             else:
                 raise
 
-        return f'{namespace}:{resp["metadata"]["name"]}'
+        return f"{namespace}:{resp['metadata']['name']}"
 
     def _submit_dryrun(
         self, app: AppDef, cfg: KubernetesOpts

torchx/schedulers/test/kubernetes_scheduler_test.py

@@ -929,6 +929,157 @@ def test_min_replicas(self) -> None:
         ]
         self.assertEqual(min_available, [1, 1, 0])
 
+    def test_apply_pod_overlay(self) -> None:
+        from kubernetes.client.models import V1Container, V1ObjectMeta, V1Pod, V1PodSpec
+        from torchx.schedulers.kubernetes_scheduler import _apply_pod_overlay
+
+        pod = V1Pod(
+            spec=V1PodSpec(
+                containers=[V1Container(name="test", image="test:latest")],
+                node_selector={"existing": "label"},
+            ),
+            metadata=V1ObjectMeta(name="test-pod"),
+        )
+
+        overlay = {
+            "spec": {
+                "nodeSelector": {"gpu": "true"},
+                "tolerations": [{"key": "nvidia.com/gpu", "operator": "Exists"}],
+            }
+        }
+
+        _apply_pod_overlay(pod, overlay)
+
+        self.assertEqual(pod.spec.node_selector, {"existing": "label", "gpu": "true"})
+        self.assertEqual(len(pod.spec.tolerations), 1)
+        self.assertEqual(pod.spec.tolerations[0].key, "nvidia.com/gpu")
+
+    def test_apply_pod_overlay_new_fields(self) -> None:
+        from kubernetes.client.models import V1Container, V1ObjectMeta, V1Pod, V1PodSpec
+        from torchx.schedulers.kubernetes_scheduler import _apply_pod_overlay
+
+        # Pod without nodeSelector or tolerations
+        pod = V1Pod(
+            spec=V1PodSpec(containers=[V1Container(name="test", image="test:latest")]),
+            metadata=V1ObjectMeta(name="test-pod"),
+        )
+
+        # Overlay adds fields not present in original
+        overlay = {
+            "spec": {
+                "nodeSelector": {"gpu": "true"},
+                "tolerations": [{"key": "nvidia.com/gpu", "operator": "Exists"}],
+                "affinity": {
+                    "nodeAffinity": {
+                        "requiredDuringSchedulingIgnoredDuringExecution": {
+                            "nodeSelectorTerms": [
+                                {
+                                    "matchExpressions": [
+                                        {
+                                            "key": "gpu",
+                                            "operator": "In",
+                                            "values": ["true"],
+                                        }
+                                    ]
+                                }
+                            ]
+                        }
+                    }
+                },
+            }
+        }
+
+        _apply_pod_overlay(pod, overlay)
+
+        self.assertEqual(pod.spec.node_selector, {"gpu": "true"})
+        self.assertEqual(len(pod.spec.tolerations), 1)
+        self.assertIsNotNone(pod.spec.affinity)
+        self.assertIsNotNone(pod.spec.affinity.node_affinity)
+
+    def test_submit_dryrun_with_pod_overlay(self) -> None:
+        scheduler = create_scheduler("test")
+
+        # Create app with metadata
+        trainer_role = specs.Role(
+            name="trainer",
+            image="pytorch/torchx:latest",
+            entrypoint="main",
+            resource=specs.Resource(cpu=1, memMB=1000, gpu=0),
+            metadata={"kubernetes": {"spec": {"nodeSelector": {"gpu": "true"}}}},
+        )
+        app = specs.AppDef("test", roles=[trainer_role])
+        cfg = KubernetesOpts({"queue": "testqueue"})
+
+        info = scheduler.submit_dryrun(app, cfg)
+        resource = info.request.resource
+
+        # Check that overlay was applied to all pods
+        tasks = resource["spec"]["tasks"]  # pyre-ignore[16]
+        for task in tasks:
+            pod = task["template"]
+            self.assertIn("gpu", pod.spec.node_selector)
+            self.assertEqual(pod.spec.node_selector["gpu"], "true")
+
+    def test_submit_dryrun_with_pod_overlay_file_uri(self) -> None:
+        import tempfile
+
+        import yaml
+
+        scheduler = create_scheduler("test")
+
+        # Create overlay file
+        overlay = {"spec": {"nodeSelector": {"instance-type": "p4d.24xlarge"}}}
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
+            yaml.dump(overlay, f)
+            overlay_path = f.name
+
+        try:
+            # Create app with file URI
+            trainer_role = specs.Role(
+                name="trainer",
+                image="pytorch/torchx:latest",
+                entrypoint="main",
+                resource=specs.Resource(cpu=1, memMB=1000, gpu=0),
+                metadata={"kubernetes": f"file://{overlay_path}"},
+            )
+            app = specs.AppDef("test", roles=[trainer_role])
+            cfg = KubernetesOpts({"queue": "testqueue"})
+
+            info = scheduler.submit_dryrun(app, cfg)
+            resource = info.request.resource
+
+            # Check that overlay was applied
+            tasks = resource["spec"]["tasks"]  # pyre-ignore[16]
+            for task in tasks:
+                pod = task["template"]
+                self.assertIn("instance-type", pod.spec.node_selector)
+                self.assertEqual(
+                    pod.spec.node_selector["instance-type"], "p4d.24xlarge"
+                )
+        finally:
+            import os
+
+            os.unlink(overlay_path)
+
+    def test_submit_dryrun_with_pod_overlay_invalid_type(self) -> None:
+        scheduler = create_scheduler("test")
+
+        # Create app with invalid metadata type
+        trainer_role = specs.Role(
+            name="trainer",
+            image="pytorch/torchx:latest",
+            entrypoint="main",
+            resource=specs.Resource(cpu=1, memMB=1000, gpu=0),
+            metadata={"kubernetes": 123},  # Invalid type
+        )
+        app = specs.AppDef("test", roles=[trainer_role])
+        cfg = KubernetesOpts({"queue": "testqueue"})
+
+        with self.assertRaises(ValueError) as ctx:
+            scheduler.submit_dryrun(app, cfg)
+
+        self.assertIn("must be a dict or resource URI", str(ctx.exception))
+
 
 class KubernetesSchedulerNoImportTest(unittest.TestCase):
     """