Merge pull request #1551 from transformerlab/claude/skypilot-container-investigation-YdrpB

Add SkyPilot container and cloud settings (global + per-job)

Author: aliasaria

api/transformerlab/compute_providers/models.py

@@ -47,6 +47,10 @@ class ClusterConfig(BaseModel):
     zone: Optional[str] = None
     use_spot: bool = False
 
+    # Container / VM image override (SkyPilot only).
+    # Use "docker:<image>" for Docker containers, e.g. "docker:nvcr.io/nvidia/pytorch:23.10-py3".
+    image_id: Optional[str] = None
+
     # Cluster settings
     idle_minutes_to_autostop: Optional[int] = None
     run: Optional[str] = None  # Initial run command

api/transformerlab/compute_providers/skypilot.py

@@ -341,6 +341,8 @@ def launch_cluster(self, cluster_name: str, config: ClusterConfig) -> Dict[str,
             resources_kwargs["zone"] = config.zone
         if config.use_spot:
             resources_kwargs["use_spot"] = True
+        if config.image_id:
+            resources_kwargs["image_id"] = config.image_id
 
         if resources_kwargs:
             task.set_resources(sky_resources.Resources(**resources_kwargs))

api/transformerlab/routers/compute_provider.py

@@ -1447,6 +1447,26 @@ async def _launch_sweep_jobs(
 
                 # When file_mounts is True we use lab.copy_file_mounts() in setup; do not send to provider
                 file_mounts_for_provider = request.file_mounts if isinstance(request.file_mounts, dict) else {}
+
+                # Resolve SkyPilot-specific settings from provider config for sweep child jobs
+                sweep_image_id: str | None = None
+                sweep_region: str | None = None
+                sweep_zone: str | None = None
+                sweep_use_spot: bool = False
+                if provider.type == ProviderType.SKYPILOT.value:
+                    prov_cfg = provider.config or {}
+                    sweep_image_id = prov_cfg.get("docker_image") or None
+                    sweep_region = prov_cfg.get("default_region") or None
+                    sweep_zone = prov_cfg.get("default_zone") or None
+                    sweep_use_spot = prov_cfg.get("use_spot", False) is True
+                    if request.config:
+                        if request.config.get("docker_image"):
+                            sweep_image_id = str(request.config["docker_image"]).strip()
+                        if request.config.get("region"):
+                            sweep_region = str(request.config["region"]).strip()
+                        if request.config.get("use_spot"):
+                            sweep_use_spot = True
+
                 cluster_config = ClusterConfig(
                     cluster_name=formatted_cluster_name,
                     provider_name=provider_display_name,
@@ -1461,6 +1481,10 @@ async def _launch_sweep_jobs(
                     disk_size=disk_size,
                     file_mounts=file_mounts_for_provider,
                     provider_config={"requested_disk_space": request.disk_space},
+                    image_id=sweep_image_id,
+                    region=sweep_region,
+                    zone=sweep_zone,
+                    use_spot=sweep_use_spot,
                 )
 
                 # Launch cluster for child job
@@ -1958,6 +1982,28 @@ async def launch_template_on_provider(
     else:
         parameters_with_secrets = merged_parameters if merged_parameters else None
 
+    # For SkyPilot providers, resolve docker_image / region / use_spot.
+    # Per-job overrides (from request.config) take precedence over provider-level defaults.
+    skypilot_image_id: str | None = None
+    skypilot_region: str | None = None
+    skypilot_zone: str | None = None
+    skypilot_use_spot: bool = False
+    if provider.type == ProviderType.SKYPILOT.value:
+        prov_cfg = provider.config or {}
+        # Provider-level defaults
+        skypilot_image_id = prov_cfg.get("docker_image") or None
+        skypilot_region = prov_cfg.get("default_region") or None
+        skypilot_zone = prov_cfg.get("default_zone") or None
+        skypilot_use_spot = prov_cfg.get("use_spot", False) is True
+        # Per-job overrides from the frontend config dict
+        if request.config:
+            if request.config.get("docker_image"):
+                skypilot_image_id = str(request.config["docker_image"]).strip()
+            if request.config.get("region"):
+                skypilot_region = str(request.config["region"]).strip()
+            if request.config.get("use_spot"):
+                skypilot_use_spot = True
+
     # Build provider_config for cluster_config (and job_data for local provider)
     provider_config_dict = {"requested_disk_space": request.disk_space}
     # For SLURM, pass through any per-run custom SBATCH flags so the provider
@@ -2094,6 +2140,10 @@ async def launch_template_on_provider(
         disk_size=disk_size,
         file_mounts=file_mounts_for_provider,
         provider_config=provider_config_dict,
+        image_id=skypilot_image_id,
+        region=skypilot_region,
+        zone=skypilot_zone,
+        use_spot=skypilot_use_spot,
     )
 
     await job_service.job_update_launch_progress(

src/renderer/components/Experiment/Tasks/QueueTaskModal.tsx

@@ -112,6 +112,9 @@ export default function QueueTaskModal({
   const [sweepMetric, setSweepMetric] = React.useState('eval/loss');
   const [lowerIsBetter, setLowerIsBetter] = React.useState(true);
   const [jobSlurmFlags, setJobSlurmFlags] = React.useState<string[]>(['']);
+  const [jobDockerImage, setJobDockerImage] = React.useState('');
+  const [jobRegion, setJobRegion] = React.useState('');
+  const [jobUseSpot, setJobUseSpot] = React.useState(false);
   const [useTrackio, setUseTrackio] = React.useState(false);
   const [useProfiling, setUseProfiling] = React.useState(false);
   const [useProfilingTorch, setUseProfilingTorch] = React.useState(false);
@@ -215,6 +218,7 @@ export default function QueueTaskModal({
   );
   const isLocalProvider = selectedProvider?.type === 'local';
   const isSlurmProvider = selectedProvider?.type === 'slurm';
+  const isSkypilotProvider = selectedProvider?.type === 'skypilot';
 
   // Fetch user-specific provider settings (including default custom SBATCH flags)
   const slurmUserSettingsKey =
@@ -632,6 +636,15 @@ export default function QueueTaskModal({
     setJobSlurmFlags(lines.length > 0 ? lines : ['']);
   }, [open, isSlurmProvider, selectedProviderId, slurmUserSettings]);
 
+  // Initialize SkyPilot per-job defaults from provider config when a SkyPilot provider is selected.
+  React.useEffect(() => {
+    if (!open || !isSkypilotProvider || !selectedProvider) return;
+    const cfg = selectedProvider.config || {};
+    setJobDockerImage(cfg.docker_image || '');
+    setJobRegion(cfg.default_region || '');
+    setJobUseSpot(cfg.use_spot === true);
+  }, [open, isSkypilotProvider, selectedProviderId, selectedProvider]);
+
   // Helper function to validate constraints
   const validateParameter = (param: ProcessedParameter): string | null => {
     const { schema, value } = param;
@@ -744,6 +757,19 @@ export default function QueueTaskModal({
       }
     }
 
+    // For SkyPilot providers, add optional per-job overrides
+    if (provider?.type === 'skypilot') {
+      if (jobDockerImage.trim()) {
+        config.docker_image = jobDockerImage.trim();
+      }
+      if (jobRegion.trim()) {
+        config.region = jobRegion.trim();
+      }
+      if (jobUseSpot) {
+        config.use_spot = true;
+      }
+    }
+
     // Add sweep configuration if enabled
     if (runSweeps) {
       config.run_sweeps = true;
@@ -1580,6 +1606,53 @@ export default function QueueTaskModal({
                     the template default.
                   </FormHelperText>
 
+                  {/* SkyPilot per-job overrides */}
+                  {isSkypilotProvider && (
+                    <>
+                      <Divider />
+                      <Typography level="title-sm">
+                        SkyPilot Job Overrides
+                      </Typography>
+                      <FormControl>
+                        <FormLabel>Docker Image (optional)</FormLabel>
+                        <Input
+                          value={jobDockerImage}
+                          onChange={(e) => setJobDockerImage(e.target.value)}
+                          placeholder="docker:nvcr.io/nvidia/pytorch:23.10-py3"
+                          sx={{ fontFamily: 'monospace', fontSize: 'sm' }}
+                          disabled={isSubmitting}
+                        />
+                        <FormHelperText>
+                          Prefix with &quot;docker:&quot; to run inside a
+                          container. Defaults to the provider&apos;s global
+                          setting.
+                        </FormHelperText>
+                      </FormControl>
+                      <FormControl>
+                        <FormLabel>Region (optional)</FormLabel>
+                        <Input
+                          value={jobRegion}
+                          onChange={(e) => setJobRegion(e.target.value)}
+                          placeholder="e.g. us-east-1"
+                          disabled={isSubmitting}
+                        />
+                      </FormControl>
+                      <FormControl
+                        sx={{ flexDirection: 'row', alignItems: 'center' }}
+                      >
+                        <Switch
+                          checked={jobUseSpot}
+                          onChange={(e) => setJobUseSpot(e.target.checked)}
+                          disabled={isSubmitting}
+                          sx={{ mr: 1 }}
+                        />
+                        <FormLabel sx={{ m: 0 }}>
+                          Use Spot / Preemptible Instances
+                        </FormLabel>
+                      </FormControl>
+                    </>
+                  )}
+
                   {/* Incompatibility Warning */}
                   {selectedProvider &&
                     effectiveResources?.accelerators &&