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
@@ -726,6 +726,7 @@ def test_runopts(self) -> None:
                 "image_repo",
                 "service_account",
                 "priority_class",
+                "validate_spec",
             },
         )
 
@@ -929,6 +930,102 @@ def test_min_replicas(self) -> None:
         ]
         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):
     """