Support SkyPilot Storage configurations in file_mounts for automatic cloud sync

[andylizf]

Problem

Currently, SkypilotExecutor only supports simple string-to-string mappings in file_mounts, which limits users to local file syncing. However, SkyPilot's YAML specification supports dictionary configurations for cloud storage mounts, enabling automatic synchronization with GCS, S3, and other cloud storage backends.

Users cannot leverage SkyPilot's powerful storage mounting features when using NeMo-Run, forcing them to manually sync outputs after training completion.

Solution

This PR enhances SkypilotExecutor to handle both string paths and dictionary configurations in file_mounts, matching SkyPilot's YAML capabilities:

1. Parse mixed file_mounts: Separate string paths from dictionary configurations
2. Create Storage objects: Convert dictionary configs to sky.Storage objects using Storage.from_yaml_config()
3. Use appropriate APIs: Call task.set_file_mounts() for strings and task.set_storage_mounts() for Storage objects

Changes

• Modified SkypilotExecutor.to_task() method to handle mixed file_mount types
• Changed file_mounts type hint from dict[str, str] to dict[str, Any] to support both formats
• Added logic to separate and process different mount types

Usage Example

from nemo_run.core.execution.skypilot import SkypilotExecutor

executor = SkypilotExecutor(
    num_nodes=2,
    gpus="H100",
    gpus_per_node=8,
    file_mounts={
        # Regular file mount (string path)
        "/workspace/code": "/local/path/to/code",
        
        # Cloud storage mount (dictionary config)
        "/workspace/outputs": {
            "name": "my-training-outputs",
            "store": "gcs",  # or "s3", "azure", etc.
            "mode": "MOUNT",
            "persistent": True,
        },
        "/workspace/checkpoints": {
            "name": "model-checkpoints",
            "store": "s3",
            "mode": "MOUNT_CACHED",  # Async sync for better performance
            "persistent": True,
        }
    }
)

Benefits

• ✅ Automatic cloud sync: Training outputs are automatically saved to cloud storage
• ✅ No manual rsync needed: Eliminates post-training manual sync operations
• ✅ Support multiple backends: Works with GCS, S3, Azure, R2, etc.
• ✅ Backward compatible: Existing string-only file_mounts continue to work
• ✅ Feature parity with SkyPilot YAML: Brings Python API closer to YAML capabilities

Testing

• Tested with GCS backend on 2-node H100 cluster
• Verified backward compatibility with string-only file_mounts
• Confirmed Storage objects are correctly created and mounted

[hemildesai]

Hi @andylizf, thank you so much for your contribution. Could you fix the formatting check and add some unit tests for your change? Should be good to merge after that.

[andylizf]

@hemildesai Formatted. Can you take a look again?

[hemildesai]

Formatting looks good, can you add unit tests to https://github.com/NVIDIA-NeMo/Run/blob/main/test/core/execution/test_skypilot.py?

[andylizf]

Formatting looks good, can you add unit tests to https://github.com/NVIDIA-NeMo/Run/blob/main/test/core/execution/test_skypilot.py?

Added. PTAL!