Training CLI for Launch - Changes per SDK HyperPodPytorchJob constructor (#64)

* Training CLI for Launch

* Training CLI for Launch

---------

Co-authored-by: Roja Reddy Sareddy <[email protected]>

Author: rsareddy0329

hyperpod-pytorchjob-config-schemas/hyperpod_pytorchjob_config_schemas/v1_0/model.py

@@ -1,6 +1,6 @@
 from pydantic import BaseModel, ConfigDict, Field
 from typing import Optional, List, Dict, Union
-from sagemaker.hyperpod.training.config.hyperpod_pytorch_job_config import _HyperPodPytorchJob, ReplicaSpec, RunPolicy, Template, Metadata, Spec
+from sagemaker.hyperpod.training.config.hyperpod_pytorch_job_config import ReplicaSpec, RunPolicy, Template, Metadata, Spec
 
 
 class PyTorchJobConfig(BaseModel):
@@ -28,7 +28,7 @@ class PyTorchJobConfig(BaseModel):
 
 
 
-    def to_domain(self) -> _HyperPodPytorchJob:
+    def to_domain(self) -> Dict:
         """
         Convert flat config to domain model (HyperPodPytorchJobSpec)
         """
@@ -38,13 +38,6 @@ def to_domain(self) -> _HyperPodPytorchJob:
         "image": self.image,
         }
 
-        # Add resources if needed (could be moved to SDK default)
-        container["resources"] = {
-        "limits": {
-            "nvidia.com/gpu": 8
-        }
-        }
-
         # Add optional container fields only if they're not None
         optional_container_fields = [
         ("command", "command", self.command),
@@ -157,5 +150,5 @@ def to_domain(self) -> _HyperPodPytorchJob:
         # Create and return the domain model
         return { "name":self.job_name ,
                  "namespace":self.namespace,
-                 "spec":_HyperPodPytorchJob(**job_kwargs)
+                 "spec":job_kwargs
                  }

sagemaker-hyperpod/src/sagemaker/hyperpod/cli/commands/training.py

@@ -36,14 +36,28 @@ def pytorch_create(version, config):
         job_name = config.get("name")
         namespace = config.get("namespace")
         spec = config.get("spec")
-        # Create job with or without namespace
-        if namespace is None:
-            job = HyperPodPytorchJob(metadata=Metadata(name=job_name), spec=spec)
-        else:
-            job = HyperPodPytorchJob(
-                metadata=Metadata(name=job_name, namespace=namespace), spec=spec
-            )
 
+        # Prepare metadata
+        metadata_kwargs = {"name": job_name}
+        if namespace:
+            metadata_kwargs["namespace"] = namespace
+
+        # Prepare job kwargs
+        job_kwargs = {
+            "metadata": Metadata(**metadata_kwargs),
+            "replica_specs": spec.get("replica_specs", [])
+        }
+
+        # Add nproc_per_node if present
+        if "nproc_per_node" in spec:
+            job_kwargs["nproc_per_node"] = spec["nproc_per_node"]
+
+        # Add run_policy if present
+        if "run_policy" in spec:
+            job_kwargs["run_policy"] = spec["run_policy"]
+
+        # Create job
+        job = HyperPodPytorchJob(**job_kwargs)
         job.create()
 
     except Exception as e:
@@ -138,16 +152,14 @@ def pytorch_describe(job_name: str, namespace: str):
         click.echo("=" * 80)
         click.echo(f"Name:           {job.metadata.name}")
         click.echo(f"Namespace:      {job.metadata.namespace}")
-        click.echo(f"API Version:    {job.apiVersion}")
-        click.echo(f"Kind:           {job.kind}")
 
         # Print Spec details
         click.echo("\nSpec:")
         click.echo("-" * 80)
-        click.echo(f"Processes per Node: {job.spec.nprocPerNode}")
+        click.echo(f"Processes per Node: {job.nprocPerNode}")
 
         # Print Replica Specs
-        for replica in job.spec.replicaSpecs:
+        for replica in job.replicaSpecs:
             click.echo(f"\nReplica Spec:")
             click.echo(f"  Name:     {replica.name}")
             click.echo(f"  Replicas: {replica.replicas}")
@@ -169,9 +181,9 @@ def pytorch_describe(job_name: str, namespace: str):
         # Print Run Policy
         click.echo("\nRun Policy:")
         click.echo("-" * 80)
-        click.echo(f"Clean Pod Policy:          {job.spec.runPolicy.cleanPodPolicy}")
+        click.echo(f"Clean Pod Policy:          {job.runPolicy.cleanPodPolicy}")
         click.echo(
-            f"TTL Seconds After Finished: {job.spec.runPolicy.ttlSecondsAfterFinished}"
+            f"TTL Seconds After Finished: {job.runPolicy.ttlSecondsAfterFinished}"
         )
 
         # Print Status