update skypilot to use volume mounts and launcher

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

Author: ansjindal

docs/source/guides/execution.md

@@ -205,12 +205,23 @@ 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",
+        infra="k8s/mycontext",
         # 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 import backends
     from sky.utils import status_lib
+    from sky.volumes import volume as volume_lib
+    from sky import models
 
     _SKYPILOT_AVAILABLE = True
 except ImportError:
@@ -95,6 +97,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
@@ -343,6 +347,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
@@ -354,6 +366,53 @@ def _setup_launcher(self):
             launcher.rdzv_backend = "static"
             launcher.rdzv_port = 49500
 
+    def supports_launcher_transform(self) -> bool:
+        return True
+
+    def _parse_infra_for_volume_config(self) -> dict:
+        """Parse infra string and return volume config parameters."""
+        config = {}
+
+        if self.infra is not None:
+            # Parse infra string to extract cloud, region, zone components
+            # Format: cloud, cloud/region, cloud/region/zone, k8s/context
+            infra_parts = self.infra.split("/")
+            cloud = infra_parts[0] if infra_parts else None
+
+            if cloud:
+                # Special handling for Kubernetes
+                if cloud == "k8s":
+                    # VolumeConfig region and zone required even though they are marked as optional
+                    # validation fails otherwise
+                    config["cloud"] = "kubernetes"
+                    config["region"] = "kubernetes"
+                    config["zone"] = "kubernetes"
+                else:
+                    # Handle regular cloud providers
+                    config["cloud"] = cloud
+
+                    # Handle region for non-k8s clouds
+                    if len(infra_parts) >= 2:
+                        region = infra_parts[1]
+                        if region and region != "*":  # Skip wildcards
+                            config["region"] = region
+
+                    # Handle zone for non-k8s clouds
+                    if len(infra_parts) >= 3:
+                        zone = infra_parts[2]
+                        if zone and zone != "*":  # Skip wildcards
+                            config["zone"] = zone
+        else:
+            # Fall back to individual cloud, region, zone parameters
+            if self.cloud:
+                config["cloud"] = self.cloud
+            if self.region:
+                config["region"] = self.region
+            if self.zone:
+                config["zone"] = self.zone
+
+        return config
+
     def to_task(
         self,
         name: str,
@@ -377,16 +436,43 @@ def to_task(
 
 {" ".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:
+                # Configure volume based on infra if specified, otherwise use cloud/region/zone
+                volume_config_kwargs = {
+                    "name": volume_mount["volume_name"],
+                    "type": volume_mount["type"],
+                    "name_on_cloud": volume_mount["volume_name"],
+                    "size": volume_mount["size"],
+                }
+
+                # Add parsed infra configuration
+                volume_config_kwargs.update(self._parse_infra_for_volume_config())
+
+                volume_mounts.append(
+                    volume_lib.VolumeMount(
+                        path=volume_mount["path"],
+                        volume_name=volume_mount["volume_name"],
+                        volume_config=models.VolumeConfig(**volume_config_kwargs),
+                    )
+                )
+        task.volume_mounts = volume_mounts
         task.set_resources(self.to_resources())
 
         if env_vars:

test/core/execution/test_skypilot.py

@@ -561,3 +561,29 @@ def test_to_task(self, mock_task, mock_skypilot_imports, executor):
 
                 # Verify the returned task is our mock
                 assert result == mock_task_instance
+
+    def test_parse_infra_for_volume_config(self, mock_skypilot_imports):
+        """Test the _parse_infra_for_volume_config helper method."""
+
+        # Test k8s infra
+        executor_k8s = SkypilotExecutor(infra="k8s/my-context")
+        config = executor_k8s._parse_infra_for_volume_config()
+        assert config["cloud"] == "kubernetes"
+        assert config["region"] == "kubernetes"
+        assert config["zone"] == "kubernetes"
+
+        # Test AWS infra with region and zone
+        executor_aws = SkypilotExecutor(infra="aws/us-east-1/us-east-1a")
+        config = executor_aws._parse_infra_for_volume_config()
+        assert config["cloud"] == "aws"
+        assert config["region"] == "us-east-1"
+        assert config["zone"] == "us-east-1a"
+
+        # Test fallback to individual parameters
+        executor_fallback = SkypilotExecutor(
+            cloud="gcp", region="us-central1", zone="us-central1-a"
+        )
+        config = executor_fallback._parse_infra_for_volume_config()
+        assert config["cloud"] == "gcp"
+        assert config["region"] == "us-central1"
+        assert config["zone"] == "us-central1-a"