feat: add option to validate k8s spec (#1152)

Author: azzhipa

torchx/schedulers/kubernetes_scheduler.py

@@ -369,7 +369,7 @@ def app_to_resource(
     queue: str,
     service_account: Optional[str],
     priority_class: Optional[str] = None,
-) -> Dict[str, object]:
+) -> Dict[str, Any]:
     """
     app_to_resource creates a volcano job kubernetes resource definition from
     the provided AppDef. The resource definition can be used to launch the
@@ -444,7 +444,7 @@ def app_to_resource(
     if priority_class is not None:
         job_spec["priorityClassName"] = priority_class
 
-    resource: Dict[str, object] = {
+    resource: Dict[str, Any] = {
         "apiVersion": "batch.volcano.sh/v1alpha1",
         "kind": "Job",
         "metadata": {"name": f"{unique_app_id}"},
@@ -456,7 +456,7 @@ def app_to_resource(
 @dataclass
 class KubernetesJob:
     images_to_push: Dict[str, Tuple[str, str]]
-    resource: Dict[str, object]
+    resource: Dict[str, Any]
 
     def __str__(self) -> str:
         return yaml.dump(sanitize_for_serialization(self.resource))
@@ -471,6 +471,7 @@ class KubernetesOpts(TypedDict, total=False):
     image_repo: Optional[str]
     service_account: Optional[str]
     priority_class: Optional[str]
+    validate_spec: Optional[bool]
 
 
 class KubernetesScheduler(
@@ -659,6 +660,36 @@ def _submit_dryrun(
         ), "priority_class must be a str"
 
         resource = app_to_resource(app, queue, service_account, priority_class)
+
+        if cfg.get("validate_spec"):
+            try:
+                self._custom_objects_api().create_namespaced_custom_object(
+                    group="batch.volcano.sh",
+                    version="v1alpha1",
+                    namespace=cfg.get("namespace") or "default",
+                    plural="jobs",
+                    body=resource,
+                    dry_run="All",
+                )
+            except Exception as e:
+                from kubernetes.client.rest import ApiException
+
+                if isinstance(e, ApiException):
+                    raise ValueError(f"Invalid job spec: {e.reason}") from e
+                raise
+
+            job_name = resource["metadata"]["name"]
+            for task in resource["spec"]["tasks"]:
+                task_name = task["name"]
+                replicas = task.get("replicas", 1)
+                max_index = replicas - 1
+                pod_name = f"{job_name}-{task_name}-{max_index}"
+                if len(pod_name) > 63:
+                    raise ValueError(
+                        f"Pod name '{pod_name}' ({len(pod_name)} chars) exceeds 63 character limit. "
+                        f"Shorten app.name or role names"
+                    )
+
         req = KubernetesJob(
             resource=resource,
             images_to_push=images_to_push,
@@ -703,6 +734,12 @@ def _run_opts(self) -> runopts:
             type_=str,
             help="The name of the PriorityClass to set on the job specs",
         )
+        opts.add(
+            "validate_spec",
+            type_=bool,
+            help="Validate job spec using Kubernetes API dry-run before submission",
+            default=True,
+        )
         return opts
 
     def describe(self, app_id: str) -> Optional[DescribeAppResponse]:

torchx/schedulers/test/kubernetes_scheduler_test.py

@@ -11,7 +11,7 @@
 import sys
 import unittest
 from datetime import datetime
-from typing import Any, Dict
+from typing import Any, cast, Dict
 from unittest.mock import MagicMock, patch
 
 import torchx
@@ -111,7 +111,6 @@ def test_app_to_resource_resolved_macros(self) -> None:
             make_unique_ctx.return_value = unique_app_name
             resource = app_to_resource(app, "test_queue", service_account=None)
             actual_cmd = (
-                # pyre-ignore [16]
                 resource["spec"]["tasks"][0]["template"]
                 .spec.containers[0]
                 .command
@@ -135,7 +134,6 @@ def test_retry_policy_not_set(self) -> None:
                 {"event": "PodEvicted", "action": "RestartJob"},
                 {"event": "PodFailed", "action": "RestartJob"},
             ],
-            # pyre-ignore [16]
             resource["spec"]["tasks"][0]["policies"],
         )
         for role in app.roles:
@@ -517,7 +515,7 @@ def test_rank0_env(self) -> None:
             make_unique_ctx.return_value = "app-name-42"
             info = scheduler.submit_dryrun(app, cfg)
 
-        tasks = info.request.resource["spec"]["tasks"]  # pyre-ignore[16]
+        tasks = info.request.resource["spec"]["tasks"]
         container0 = tasks[0]["template"].spec.containers[0]
         self.assertIn("TORCHX_RANK0_HOST", container0.command)
         self.assertIn(
@@ -726,6 +724,7 @@ def test_runopts(self) -> None:
                 "image_repo",
                 "service_account",
                 "priority_class",
+                "validate_spec",
             },
         )
 
@@ -925,10 +924,106 @@ def test_min_replicas(self) -> None:
         resource = app_to_resource(app, "test_queue", service_account=None)
         min_available = [
             task["minAvailable"]
-            for task in resource["spec"]["tasks"]  # pyre-ignore[16]
+            for task in resource["spec"]["tasks"]
         ]
         self.assertEqual(min_available, [1, 1, 0])
 
+    @patch(
+        "torchx.schedulers.kubernetes_scheduler.KubernetesScheduler._custom_objects_api"
+    )
+    def test_validate_spec_invalid_name(self, mock_api: MagicMock) -> None:
+        from kubernetes.client.rest import ApiException
+
+        scheduler = create_scheduler("test")
+        app = _test_app()
+        app.name = "Invalid_Name"
+
+        mock_api_instance = MagicMock()
+        mock_api_instance.create_namespaced_custom_object.side_effect = ApiException(
+            status=422,
+            reason="Invalid",
+        )
+        mock_api.return_value = mock_api_instance
+
+        cfg = cast(KubernetesOpts, {"queue": "testqueue", "validate_spec": True})
+
+        with self.assertRaises(ValueError) as ctx:
+            scheduler.submit_dryrun(app, cfg)
+
+        self.assertIn("Invalid job spec", str(ctx.exception))
+        mock_api_instance.create_namespaced_custom_object.assert_called_once()
+        call_kwargs = mock_api_instance.create_namespaced_custom_object.call_args[1]
+        self.assertEqual(call_kwargs["dry_run"], "All")
+
+    def test_validate_spec_enabled_by_default(self) -> None:
+        scheduler = create_scheduler("test")
+        app = _test_app()
+
+        cfg = KubernetesOpts({"queue": "testqueue"})
+
+        with patch(
+            "torchx.schedulers.kubernetes_scheduler.KubernetesScheduler._custom_objects_api"
+        ) as mock_api:
+            mock_api_instance = MagicMock()
+            mock_api_instance.create_namespaced_custom_object.return_value = {}
+            mock_api.return_value = mock_api_instance
+
+            info = scheduler.submit_dryrun(app, cfg)
+
+        self.assertIsNotNone(info)
+        mock_api_instance.create_namespaced_custom_object.assert_called_once()
+        call_kwargs = mock_api_instance.create_namespaced_custom_object.call_args[1]
+        self.assertEqual(call_kwargs["dry_run"], "All")
+
+    @patch(
+        "torchx.schedulers.kubernetes_scheduler.KubernetesScheduler._custom_objects_api"
+    )
+    def test_validate_spec_invalid_task_name(self, mock_api: MagicMock) -> None:
+        from kubernetes.client.rest import ApiException
+
+        scheduler = create_scheduler("test")
+        app = _test_app()
+        app.roles[0].name = "Invalid-Task-Name"
+
+        mock_api_instance = MagicMock()
+        mock_api_instance.create_namespaced_custom_object.side_effect = ApiException(
+            status=422,
+            reason="Invalid",
+        )
+        mock_api.return_value = mock_api_instance
+
+        cfg = cast(KubernetesOpts, {"queue": "testqueue", "validate_spec": True})
+
+        with self.assertRaises(ValueError) as ctx:
+            scheduler.submit_dryrun(app, cfg)
+
+        self.assertIn("Invalid job spec", str(ctx.exception))
+
+    @patch(
+        "torchx.schedulers.kubernetes_scheduler.KubernetesScheduler._custom_objects_api"
+    )
+    def test_validate_spec_long_pod_name(self, mock_api: MagicMock) -> None:
+        scheduler = create_scheduler("test")
+        app = _test_app()
+        app.name = "x" * 50
+        app.roles[0].name = "y" * 20
+
+        mock_api_instance = MagicMock()
+        mock_api_instance.create_namespaced_custom_object.return_value = {}
+        mock_api.return_value = mock_api_instance
+
+        cfg = cast(KubernetesOpts, {"queue": "testqueue", "validate_spec": True})
+
+        with patch(
+            "torchx.schedulers.kubernetes_scheduler.make_unique"
+        ) as make_unique_ctx:
+            make_unique_ctx.return_value = "x" * 50
+            with self.assertRaises(ValueError) as ctx:
+                scheduler.submit_dryrun(app, cfg)
+
+        self.assertIn("Pod name", str(ctx.exception))
+        self.assertIn("exceeds 63 character limit", str(ctx.exception))
+
 
 class KubernetesSchedulerNoImportTest(unittest.TestCase):
     """