update skypilot to use volume mounts and launcher

Signed-off-by: ansjindal <[email protected]>

Author: ansjindal

docs/source/guides/execution.md

@@ -205,12 +205,25 @@ def your_skypilot_executor(nodes: int, devices: int, container_image: str):
     return SkypilotExecutor(
         gpus="RTX5880-ADA-GENERATION",
         gpus_per_node=devices,
-        nodes = nodes
-        env_vars=common_envs()
+        num_nodes = nodes,
+        env_vars=common_envs(),
         container_image=container_image,
         cloud="kubernetes",
+        region="us-central1",
+        zone="us-central1-a",
         # Optional to reuse Skypilot cluster
         cluster_name="tester",
+        volumes={
+            "nemo-workspace": "nemo-workspace"
+        },
+        volume_mounts=[
+            {
+                "path": "/data",
+                "volume_name": "nemo-workspace",
+                "size": "50Gi",
+                "type": "k8s-pvc"
+            }
+        ],
         setup="""
     conda deactivate
     nvidia-smi

nemo_run/core/execution/skypilot.py

@@ -18,7 +18,7 @@
 import subprocess
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Optional, Type, Union
+from typing import Any, Dict, List, Optional, Type, Union
 
 from invoke.context import Context
 
@@ -37,6 +37,8 @@
     import sky.task as skyt
     from sky.utils import status_lib
     from sky import backends
+    from sky.volumes import volume as volume_lib
+    from sky import models
 
     _SKYPILOT_AVAILABLE = True
 except ImportError:
@@ -94,6 +96,8 @@ class SkypilotExecutor(Executor):
     memory: Optional[Union[int | float, list[int | float]]] = None
     instance_type: Optional[Union[str, list[str]]] = None
     num_nodes: int = 1
+    volumes: Optional[Dict[str, str]] = None
+    volume_mounts: Optional[List[Any]] = None
     use_spot: Optional[Union[bool, list[bool]]] = None
     disk_size: Optional[Union[int, list[int]]] = None
     disk_tier: Optional[Union[str, list[str]]] = None
@@ -108,12 +112,12 @@ class SkypilotExecutor(Executor):
     packager: Packager = field(default_factory=lambda: GitArchivePackager())  # type: ignore  # noqa: F821
 
     def __post_init__(self):
-        assert _SKYPILOT_AVAILABLE, (
-            'Skypilot is not installed. Please install it using `pip install "nemo_run[skypilot]"`.'
-        )
-        assert isinstance(self.packager, GitArchivePackager), (
-            "Only GitArchivePackager is currently supported for SkypilotExecutor."
-        )
+        assert (
+            _SKYPILOT_AVAILABLE
+        ), 'Skypilot is not installed. Please install it using `pip install "nemo_run[skypilot]"`.'
+        assert isinstance(
+            self.packager, GitArchivePackager
+        ), "Only GitArchivePackager is currently supported for SkypilotExecutor."
 
     @classmethod
     def parse_app(cls: Type["SkypilotExecutor"], app_id: str) -> tuple[str, str, int]:
@@ -331,6 +335,14 @@ def macro_values(self) -> Optional[ExecutorMacros]:
         )
 
     def _setup_launcher(self):
+        # Auto-enable torchrun for distributed training scenarios:
+        # 1. Multi-node training (num_nodes > 1)
+        # 2. Single-node multi-GPU training (gpus_per_node > 1)
+        if self.launcher is None and (
+            self.num_nodes > 1 or (self.gpus_per_node and self.gpus_per_node > 1)
+        ):
+            self.launcher = "torchrun"
+
         super()._setup_launcher()
         launcher = self.launcher
         # Dynamic rendezvous has an error in Skypilot Kubernetes currently
@@ -342,6 +354,9 @@ def _setup_launcher(self):
             launcher.rdzv_backend = "static"
             launcher.rdzv_port = 49500
 
+    def supports_launcher_transform(self) -> bool:
+        return True
+
     def to_task(
         self,
         name: str,
@@ -361,20 +376,48 @@ def to_task(
 head_node_ip=`echo "$SKYPILOT_NODE_IPS" | head -n1`
 echo "head_node_ip=$head_node_ip"
 
+export MASTER_ADDR=$head_node_ip
+export WORLD_SIZE=$num_nodes
+export RANK=`echo "$SKYPILOT_NODE_RANK"`
+echo "MASTER_ADDR=$MASTER_ADDR"
 cd /nemo_run/code
 
 {" ".join(cmd)}
 """
+
         task = Task(
             name=name,
             setup=self.setup if self.setup else "",
             run=run_cmd,
             envs=self.env_vars,
             num_nodes=self.num_nodes,
+            volumes=self.volumes,
         )
+
         file_mounts = self.file_mounts or {}
         file_mounts["/nemo_run"] = self.job_dir
         task.set_file_mounts(file_mounts)
+        task.set_volumes(self.volumes)
+
+        volume_mounts = []
+        if self.volume_mounts:
+            for volume_mount in self.volume_mounts:
+                volume_mounts.append(
+                    volume_lib.VolumeMount(
+                        path=volume_mount["path"],
+                        volume_name=volume_mount["volume_name"],
+                        volume_config=models.VolumeConfig(
+                            name=volume_mount["volume_name"],
+                            type=volume_mount["type"],
+                            cloud=self.cloud,
+                            region=self.region,
+                            zone=self.zone,
+                            name_on_cloud=volume_mount["volume_name"],
+                            size=volume_mount["size"],
+                        ),
+                    )
+                )
+        task.volume_mounts = volume_mounts
         task.set_resources(self.to_resources())
 
         if env_vars: