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

Author: azzhipa

torchx/schedulers/kubernetes_scheduler.py

@@ -27,6 +27,90 @@
 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 providing
+a callable that receives the generated pod dict and returns the modified pod dict.
+
+.. code:: python
+
+    from torchx.specs import AppDef, Role
+
+    # Simple merge - replaces lists
+    role = Role(
+        name="trainer",
+        image="my-image:latest",
+        entrypoint="train.py",
+        metadata={
+            "kubernetes": lambda pod: {
+                **pod,
+                "spec": {
+                    **pod["spec"],
+                    "nodeSelector": {**pod["spec"].get("nodeSelector", {}), "gpu": "true"},
+                    "tolerations": [{"key": "nvidia.com/gpu", "operator": "Exists"}]
+                }
+            }
+        }
+    )
+
+    # Append to lists
+    role = Role(
+        name="trainer",
+        image="my-image:latest",
+        entrypoint="train.py",
+        metadata={
+            "kubernetes": lambda pod: {
+                **pod,
+                "spec": {
+                    **pod["spec"],
+                    "tolerations": pod["spec"].get("tolerations", []) + [
+                        {"key": "nvidia.com/gpu", "operator": "Exists"}
+                    ]
+                }
+            }
+        }
+    )
+
+    # Load from YAML file
+    import yaml
+    import fsspec
+
+    with fsspec.open("file:///path/to/overlay.yaml", "r") as f:
+        overlay_dict = yaml.safe_load(f)
+
+    role = Role(
+        name="trainer",
+        image="my-image:latest",
+        entrypoint="train.py",
+        metadata={
+            "kubernetes": lambda pod: {
+                **pod,
+                "spec": {
+                    **pod["spec"],
+                    **overlay_dict.get("spec", {}),
+                    "tolerations": pod["spec"].get("tolerations", []) + 
+                                   overlay_dict.get("spec", {}).get("tolerations", [])
+                }
+            }
+        }
+    )
+
+Example ``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
@@ -36,6 +120,7 @@
 from datetime import datetime
 from typing import (
     Any,
+    Callable,
     cast,
     Dict,
     Iterable,
@@ -45,6 +130,7 @@
     Tuple,
     TYPE_CHECKING,
     TypedDict,
+    Union,
 )
 
 import torchx
@@ -97,6 +183,30 @@
 RESERVED_MILLICPU = 100
 RESERVED_MEMMB = 1024
 
+
+def apply_pod_overlay(
+    pod: "V1Pod",
+    overlay: Callable[[Dict[str, Any]], Dict[str, Any]],
+) -> "V1Pod":
+    """Apply overlay function to V1Pod object.
+
+    Args:
+        pod: Kubernetes V1Pod object to modify
+        overlay: Callable that receives pod dict and returns modified pod dict
+
+    Returns:
+        Modified V1Pod object
+    """
+    from kubernetes import client
+
+    assert callable(overlay), f"overlay must be callable, got {type(overlay)}"
+
+    api = client.ApiClient()
+    return api._ApiClient__deserialize(
+        overlay(api.sanitize_for_serialization(pod)), "V1Pod"
+    )
+
+
 RETRY_POLICIES: Mapping[str, Iterable[Mapping[str, str]]] = {
     RetryPolicy.REPLICA: [],
     RetryPolicy.APPLICATION: [
@@ -402,6 +512,8 @@ def app_to_resource(
             replica_role.env["TORCHX_IMAGE"] = replica_role.image
 
             pod = role_to_pod(name, replica_role, service_account)
+            if k8s_overlay := role.metadata.get("kubernetes"):
+                pod = apply_pod_overlay(pod, k8s_overlay)
             pod.metadata.labels.update(
                 pod_labels(
                     app=app,
@@ -636,7 +748,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,200 @@ def test_min_replicas(self) -> None:
         ]
         self.assertEqual(min_available, [1, 1, 0])
 
+    def test_apply_pod_overlay_merge(self) -> None:
+        from kubernetes.client.models import V1Container, V1ObjectMeta, V1Pod, V1PodSpec
+        from torchx.schedulers.kubernetes_scheduler import (  # pyre-ignore[21]
+            apply_pod_overlay,
+        )
+
+        pod = V1Pod(
+            spec=V1PodSpec(
+                containers=[V1Container(name="test", image="test:latest")],
+                node_selector={"existing": "label"},
+            ),
+            metadata=V1ObjectMeta(name="test-pod"),
+        )
+
+        overlay = lambda pod_dict: {
+            **pod_dict,
+            "spec": {
+                **pod_dict["spec"],
+                "nodeSelector": {
+                    **pod_dict["spec"].get("nodeSelector", {}),
+                    "gpu": "true",
+                },
+                "tolerations": [{"key": "nvidia.com/gpu", "operator": "Exists"}],
+            },
+        }
+
+        pod = apply_pod_overlay(pod, overlay)  # pyre-ignore[16]
+
+        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_append_lists(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")],
+                tolerations=[{"key": "existing", "operator": "Exists"}],
+            ),
+            metadata=V1ObjectMeta(name="test-pod"),
+        )
+
+        overlay = lambda pod_dict: {
+            **pod_dict,
+            "spec": {
+                **pod_dict["spec"],
+                "tolerations": pod_dict["spec"].get("tolerations", [])
+                + [{"key": "nvidia.com/gpu", "operator": "Exists"}],
+            },
+        }
+
+        pod = apply_pod_overlay(pod, overlay)  # pyre-ignore[16]
+
+        self.assertEqual(len(pod.spec.tolerations), 2)
+        self.assertEqual(pod.spec.tolerations[0].key, "existing")
+        self.assertEqual(pod.spec.tolerations[1].key, "nvidia.com/gpu")
+
+    def test_submit_dryrun_with_pod_overlay(self) -> None:
+        scheduler = create_scheduler("test")
+
+        trainer_role = specs.Role(
+            name="trainer",
+            image="pytorch/torchx:latest",
+            entrypoint="main",
+            resource=specs.Resource(cpu=1, memMB=1000, gpu=0),
+            metadata={
+                "kubernetes": lambda pod: {
+                    **pod,
+                    "spec": {
+                        **pod["spec"],
+                        "nodeSelector": {
+                            **pod["spec"].get("nodeSelector", {}),
+                            "gpu": "true",
+                        },
+                    },
+                }
+            },
+        )
+        app = specs.AppDef("test", roles=[trainer_role])
+        cfg = KubernetesOpts({"queue": "testqueue"})
+
+        info = scheduler.submit_dryrun(app, cfg)
+        resource = info.request.resource
+
+        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_from_yaml(self) -> None:
+        import tempfile
+
+        import yaml
+
+        scheduler = create_scheduler("test")
+
+        overlay_dict = {"spec": {"nodeSelector": {"instance-type": "p4d.24xlarge"}}}
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
+            yaml.dump(overlay_dict, f)
+            overlay_path = f.name
+
+        try:
+            import fsspec
+
+            with fsspec.open(f"file://{overlay_path}", "r") as f:
+                loaded_overlay = yaml.safe_load(f)
+
+            trainer_role = specs.Role(
+                name="trainer",
+                image="pytorch/torchx:latest",
+                entrypoint="main",
+                resource=specs.Resource(cpu=1, memMB=1000, gpu=0),
+                metadata={
+                    "kubernetes": lambda pod: {
+                        **pod,
+                        "spec": {**pod["spec"], **loaded_overlay.get("spec", {})},
+                    }
+                },
+            )
+            app = specs.AppDef("test", roles=[trainer_role])
+            cfg = KubernetesOpts({"queue": "testqueue"})
+
+            info = scheduler.submit_dryrun(app, cfg)
+            resource = info.request.resource
+
+            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_append_from_yaml(self) -> None:
+        import tempfile
+
+        import yaml
+
+        scheduler = create_scheduler("test")
+
+        overlay_dict = {
+            "spec": {"tolerations": [{"key": "nvidia.com/gpu", "operator": "Exists"}]}
+        }
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
+            yaml.dump(overlay_dict, f)
+            overlay_path = f.name
+
+        try:
+            import fsspec
+
+            with fsspec.open(f"file://{overlay_path}", "r") as f:
+                loaded_overlay = yaml.safe_load(f)
+
+            trainer_role = specs.Role(
+                name="trainer",
+                image="pytorch/torchx:latest",
+                entrypoint="main",
+                resource=specs.Resource(cpu=1, memMB=1000, gpu=0),
+                metadata={
+                    "kubernetes": lambda pod: {
+                        **pod,
+                        "spec": {
+                            **pod["spec"],
+                            "tolerations": pod["spec"].get("tolerations", [])
+                            + loaded_overlay.get("spec", {}).get("tolerations", []),
+                        },
+                    }
+                },
+            )
+            app = specs.AppDef("test", roles=[trainer_role])
+            cfg = KubernetesOpts({"queue": "testqueue"})
+
+            info = scheduler.submit_dryrun(app, cfg)
+            resource = info.request.resource
+
+            tasks = resource["spec"]["tasks"]  # pyre-ignore[16]
+            for task in tasks:
+                pod = task["template"]
+                self.assertIsNotNone(pod.spec.tolerations)
+                self.assertTrue(
+                    any(t.key == "nvidia.com/gpu" for t in pod.spec.tolerations)
+                )
+        finally:
+            import os
+
+            os.unlink(overlay_path)
+
 
 class KubernetesSchedulerNoImportTest(unittest.TestCase):
     """