update v1.1 pytorch job template to match parity with v1.0 change in staging repo (#228)

Author: mollyheamazon

hyperpod-pytorch-job-template/hyperpod_pytorch_job_template/v1_1/model.py

@@ -23,6 +23,8 @@
 from .quota_allocation_util import _is_valid, _get_resources_from_compute_quotas, _get_resources_from_instance, _get_limits
 
 class VolumeConfig(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
     name: str = Field(
         ..., 
         description="Volume name",
@@ -109,16 +111,15 @@ class PyTorchJobConfig(BaseModel):
         min_length=1
     )
     node_count: Optional[int] = Field(
-        default=None, 
+        default=1, 
         alias="node_count", 
         description="Number of nodes",
         ge=1
     )
-    tasks_per_node: Optional[int] = Field(
-        default=None, 
+    tasks_per_node: Optional[str] = Field(
+        default="auto", 
         alias="tasks_per_node", 
-        description="Number of tasks per node",
-        ge=1
+        description="Number of workers per node; supported values: [auto,cpu, gpu, int]",
     )
     label_selector: Optional[Dict[str, str]] = Field(
         default=None,

hyperpod-pytorch-job-template/hyperpod_pytorch_job_template/v1_1/schema.json

@@ -202,22 +202,22 @@
           "type": "null"
         }
       ],
-      "default": null,
+      "default": 1,
       "description": "Number of nodes",
       "title": "Node Count"
     },
     "tasks_per_node": {
       "anyOf": [
         {
           "minimum": 1,
-          "type": "integer"
+          "type": "string"
         },
         {
           "type": "null"
         }
       ],
-      "default": null,
-      "description": "Number of tasks per node",
+      "default": "auto",
+      "description": "Number of workers per node; supported values: [auto,cpu, gpu, int]",
       "title": "Tasks Per Node"
     },
     "label_selector": {

test/unit_tests/cli/test_training.py

@@ -441,9 +441,9 @@ def test_integer_field_validation_success(self):
         config = PyTorchJobConfig(
             job_name="test-job", 
             image="pytorch:latest", 
-            tasks_per_node=8
+            tasks_per_node="auto"
         )
-        self.assertEqual(config.tasks_per_node, 8)
+        self.assertEqual(config.tasks_per_node, "auto")
 
         # Test max_retry
         config = PyTorchJobConfig(
@@ -755,7 +755,7 @@ def test_comprehensive_valid_config(self):
             pull_policy="Always",
             instance_type="ml.p4d.24xlarge",
             node_count=2,
-            tasks_per_node=8,
+            tasks_per_node="auto",
             label_selector={"accelerator": "nvidia"},
             queue_name="training-queue",
             priority="high",
@@ -774,7 +774,7 @@ def test_comprehensive_valid_config(self):
         self.assertEqual(config.pull_policy, "Always")
         self.assertEqual(config.instance_type, "ml.p4d.24xlarge")
         self.assertEqual(config.node_count, 2)
-        self.assertEqual(config.tasks_per_node, 8)
+        self.assertEqual(config.tasks_per_node, "auto")
         self.assertEqual(config.label_selector, {"accelerator": "nvidia"})
         self.assertEqual(config.queue_name, "training-queue")
         self.assertEqual(config.priority, "high")