skypilot-nightly 1.0.0.dev20250521__py3-none-any.whl → 1.0.0.dev20250523__py3-none-any.whl

sky/templates/gcp-ray.yml.j2

@@ -109,12 +109,27 @@ available_node_types:
 {%- if tpu_vm %}
       acceleratorType: {{tpu_type}}
       runtimeVersion: {{runtime_version}}
+    {%- if volumes %}
+      dataDisks:
+      {%- for volume in volumes %}
+        {%- if volume.source %}
+        - sourceDisk: {{volume.source}}
+        {%- endif %}
+        {%- if volume.attach_mode %}
+          mode: {{volume.attach_mode}}
+        {%- endif %}
+      {%- endfor %}
+    {%- endif %}
       metadata:
         # TPU VM's metadata has different format than normal VMs.
         # After replacing the variables, this will become username:ssh_public_key_content.
         # This is a specific syntax required by GCP https://cloud.google.com/compute/docs/connect/add-ssh-keys
         ssh-keys: |-
               skypilot:ssh_user:skypilot:ssh_public_key_content
+      {%- if user_data is not none %}
+        startup-script: |-
+          {{ user_data | indent(10) }}
+      {%- endif %}
   {%- if use_spot %}
       schedulingConfig:
         preemptible: true
@@ -138,6 +153,34 @@ available_node_types:
           {%- if disk_iops %}
             provisionedIops: {{disk_iops}}
           {%- endif %}
+      {%- for volume in volumes %}
+        - boot: false
+          autoDelete: {{volume.auto_delete}}
+          type: {{volume.storage_type}}
+          deviceName: {{volume.device_name}}
+        {%- if volume.source %}
+          source: {{volume.source}}
+        {%- endif %}
+        {%- if volume.attach_mode %}
+          mode: {{volume.attach_mode}}
+        {%- endif %}
+        {%- if volume.interface_type %}
+          interface: {{volume.interface_type}}
+        {%- endif %}
+        {%- if volume.disk_tier %}
+          initializeParams:
+            diskType: zones/{{zones}}/diskTypes/{{volume.disk_tier}}
+        {%- endif %}
+          {%- if volume.disk_name %}
+            diskName: {{volume.disk_name}}
+          {%- endif %}
+          {%- if volume.disk_size %}
+            diskSizeGb: {{volume.disk_size}}
+          {%- endif %}
+          {%- if volume.iops %}
+            provisionedIops: {{volume.iops}}
+          {%- endif %}
+      {%- endfor %}
   {%- if gpu is not none %}
       guestAccelerators:
         - acceleratorType: projects/{{gcp_project_id}}/zones/{{zones}}/acceleratorTypes/{{gpu}}
@@ -157,7 +200,7 @@ available_node_types:
   {%- if user_data is not none %}
           - key: user-data
             value: |-
-              {{ user_data | indent(10) }}
+              {{ user_data | indent(14) }}
   {%- endif %}
   {%- if use_spot or gpu is not none %}
       scheduling:

sky/templates/nebius-ray.yml.j2

@@ -46,12 +46,17 @@ available_node_types:
       InstanceType: {{instance_type}}
       ImageId: {{image_id}}
       DiskSize: {{disk_size}}
+      filesystems:
+        {%- for fs in filesystems %}
+        - filesystem_id: {{ fs.filesystem_id }}
+          filesystem_mount_tag: {{ fs.filesystem_mount_tag }}
+          filesystem_attach_mode: {{ fs.filesystem_attach_mode }}
+          filesystem_mount_path: {{ fs.filesystem_mount_path }}
+        {%- endfor %}
       UserData: |
-        {%- if docker_image is not none %}
         runcmd:
           - sudo sed -i 's/^#\?AllowTcpForwarding.*/AllowTcpForwarding yes/' /etc/ssh/sshd_config
           - systemctl restart sshd
-        {%- endif %}
 
         {# Two available OS images:
                1. ubuntu22.04-driverless - requires Docker installation
@@ -132,6 +137,11 @@ setup_commands:
   - {%- for initial_setup_command in initial_setup_commands %}
     {{ initial_setup_command }}
     {%- endfor %}
+    {%- for fs in filesystems %}
+    sudo mkdir {{ fs.filesystem_mount_path }};
+    sudo mount -t virtiofs {{ fs.filesystem_mount_tag }} {{ fs.filesystem_mount_path }};
+    sudo chmod a+w {{ fs.filesystem_mount_path }};
+    {%- endfor %}
     sudo systemctl stop unattended-upgrades || true;
     sudo systemctl disable unattended-upgrades || true;
     sudo sed -i 's/Unattended-Upgrade "1"/Unattended-Upgrade "0"/g' /etc/apt/apt.conf.d/20auto-upgrades || true;

sky/utils/admin_policy_utils.py

@@ -1,9 +1,8 @@
 """Admin policy utils."""
+import contextlib
 import copy
 import importlib
-import os
-import tempfile
-from typing import Optional, Tuple, Union
+from typing import Iterator, Optional, Tuple, Union
 
 import colorama
 
@@ -52,9 +51,31 @@ def _get_policy_cls(
     return policy_cls
 
 
+@contextlib.contextmanager
+def apply_and_use_config_in_current_request(
+    entrypoint: Union['dag_lib.Dag', 'task_lib.Task'],
+    request_options: Optional[admin_policy.RequestOptions] = None,
+) -> Iterator['dag_lib.Dag']:
+    """Applies an admin policy and override SkyPilot config for current request
+
+    This is a helper function of `apply()` that applies an admin policy and
+    overrides the SkyPilot config for the current request as a context manager.
+    The original SkyPilot config will be restored when the context manager is
+    exited.
+
+    Refer to `apply()` for more details.
+    """
+    original_config = skypilot_config.to_dict()
+    dag, mutated_config = apply(entrypoint, request_options)
+    if mutated_config != original_config:
+        with skypilot_config.replace_skypilot_config(mutated_config):
+            yield dag
+    else:
+        yield dag
+
+
 def apply(
     entrypoint: Union['dag_lib.Dag', 'task_lib.Task'],
-    use_mutated_config_in_current_request: bool = True,
     request_options: Optional[admin_policy.RequestOptions] = None,
 ) -> Tuple['dag_lib.Dag', config_utils.Config]:
     """Applies an admin policy (if registered) to a DAG or a task.
@@ -85,8 +106,7 @@ def apply(
         return dag, skypilot_config.to_dict()
 
     logger.info(f'Applying policy: {policy}')
-    original_config = skypilot_config.to_dict()
-    config = copy.deepcopy(original_config)
+    config = copy.deepcopy(skypilot_config.to_dict())
     mutated_dag = dag_lib.Dag()
     mutated_dag.name = dag.name
 
@@ -126,22 +146,6 @@ def apply(
         mutated_dag.graph.add_edge(mutated_dag.tasks[u_idx],
                                    mutated_dag.tasks[v_idx])
 
-    if (use_mutated_config_in_current_request and
-            original_config != mutated_config):
-        with tempfile.NamedTemporaryFile(
-                delete=False,
-                mode='w',
-                prefix='policy-mutated-skypilot-config-',
-                suffix='.yaml') as temp_file:
-
-            common_utils.dump_yaml(temp_file.name, dict(**mutated_config))
-            os.environ[skypilot_config.ENV_VAR_SKYPILOT_CONFIG] = temp_file.name
-            logger.debug(f'Updated SkyPilot config: {temp_file.name}')
-            # TODO(zhwu): This is not a clean way to update the SkyPilot config,
-            # because we are resetting the global context for a single DAG,
-            # which is conceptually weird.
-            importlib.reload(skypilot_config)
-
     logger.debug(f'Mutated user request: {mutated_user_request}')
     mutated_dag.policy_applied = True
     return mutated_dag, mutated_config

sky/utils/context.py

@@ -57,6 +57,7 @@ class Context(object):
         self._log_file = None
         self._log_file_handle = None
         self.env_overrides = {}
+        self.config_context = None
 
     def cancel(self):
         """Cancel the context."""
@@ -159,17 +160,25 @@ class ContextualEnviron(MutableMapping):
         ctx = get()
         if ctx is not None:
             if key in ctx.env_overrides:
-                return ctx.env_overrides[key]
+                value = ctx.env_overrides[key]
+                # None is used to indicate that the key is deleted in the
+                # context.
+                if value is None:
+                    raise KeyError(key)
+                return value
         return self._environ[key]
 
     def __iter__(self):
         ctx = get()
+        deleted_keys = set()
         if ctx is not None:
-            for key in ctx.env_overrides:
+            for key, value in ctx.env_overrides.items():
+                if value is None:
+                    deleted_keys.add(key)
                 yield key
             for key in self._environ:
                 # Deduplicate the keys
-                if key not in ctx.env_overrides:
+                if key not in ctx.env_overrides and key not in deleted_keys:
                     yield key
         else:
             return self._environ.__iter__()
@@ -178,10 +187,27 @@ class ContextualEnviron(MutableMapping):
         return len(dict(self))
 
     def __setitem__(self, key, value):
-        return self._environ.__setitem__(key, value)
+        ctx = get()
+        if ctx is not None:
+            ctx.env_overrides[key] = value
+        else:
+            self._environ.__setitem__(key, value)
 
     def __delitem__(self, key):
-        return self._environ.__delitem__(key)
+        ctx = get()
+        if ctx is not None:
+            if key in ctx.env_overrides:
+                del ctx.env_overrides[key]
+            elif key in self._environ:
+                # If the key is not set in the context but set in the environ
+                # of the process, we mark it as deleted in the context by
+                # setting the value to None.
+                ctx.env_overrides[key] = None
+            else:
+                # The key is not set in the context nor the process.
+                raise KeyError(key)
+        else:
+            self._environ.__delitem__(key)
 
     def __repr__(self):
         return self._environ.__repr__()
@@ -190,7 +216,11 @@ class ContextualEnviron(MutableMapping):
         copied = self._environ.copy()
         ctx = get()
         if ctx is not None:
-            copied.update(ctx.env_overrides)
+            for key in ctx.env_overrides:
+                if ctx.env_overrides[key] is None:
+                    copied.pop(key)
+                else:
+                    copied[key] = ctx.env_overrides[key]
         return copied
 
     def setdefault(self, key, default=None):

sky/utils/context_utils.py

@@ -1,5 +1,6 @@
 """Utilities for SkyPilot context."""
 import asyncio
+import contextvars
 import functools
 import io
 import multiprocessing
@@ -170,3 +171,17 @@ def cancellation_guard(func: F) -> F:
         return func(*args, **kwargs)
 
     return typing.cast(F, wrapper)
+
+
+# TODO(aylei): replace this with asyncio.to_thread once we drop support for
+# python 3.8
+def to_thread(func, /, *args, **kwargs):
+    """Asynchronously run function *func* in a separate thread.
+
+    This is same as asyncio.to_thread added in python 3.9
+    """
+    loop = asyncio.get_running_loop()
+    # This is critical to pass the current coroutine context to the new thread
+    pyctx = contextvars.copy_context()
+    func_call = functools.partial(pyctx.run, func, *args, **kwargs)
+    return loop.run_in_executor(None, func_call)

sky/utils/infra_utils.py

@@ -86,6 +86,16 @@ class InfraInfo:
             cloud_name = 'kubernetes'  # Normalize k8s to kubernetes
             region = '/'.join(parts[1:]) if len(parts) >= 2 else None
             zone = None
+        elif cloud_name == 'ssh':
+            # For SSH, the entire string after "ssh/" is the
+            # node pool name. We prepend 'ssh-' for the internal implementation
+            # which reuses the context name.
+            # TODO(romilb): This is a workaround while we use the global
+            # kubeconfig to store the ssh contexts.
+            region = '/'.join(parts[1:]) if len(parts) >= 2 else None
+            if region:
+                region = f'ssh-{region}'
+            zone = None
         else:
             # For non-Kubernetes clouds, continue with regular parsing
             # but be careful to only split into max 3 parts
@@ -133,6 +143,12 @@ class InfraInfo:
         if zone is None:
             zone = '*'
 
+        # If the cloud is ssh, we remove the ssh- prefix from the region
+        # TODO(romilb): This is a workaround while we use the global
+        # kubeconfig to store the ssh contexts.
+        if region and region.startswith('ssh-'):
+            region = region[4:]
+
         # Build the parts list and filter out trailing wildcards
         parts = [cloud.lower(), region, zone]
         while parts and parts[-1] == '*':
@@ -160,7 +176,11 @@ class InfraInfo:
         if self.zone is not None and self.zone != '*':
             region_or_zone = self.zone
         elif self.region is not None and self.region != '*':
-            region_or_zone = self.region
+            # If using region, we remove the ssh- prefix if it exists for SSH
+            # Node Pools.
+            # TODO(romilb): This is a workaround while we use the global
+            # kubeconfig to store the ssh contexts.
+            region_or_zone = self.region.lstrip('ssh-')
 
         if region_or_zone is not None and truncate:
             region_or_zone = common_utils.truncate_long_string(

sky/utils/kubernetes/cleanup-tunnel.sh

@@ -0,0 +1,62 @@
+#!/bin/bash
+# cleanup-tunnel.sh - Script to clean up SSH tunnels for a Kubernetes context
+
+# Usage: cleanup-tunnel.sh CONTEXT_NAME
+
+CONTEXT="${1:-default}"
+TUNNEL_DIR="$HOME/.sky/ssh_node_pools_info"
+PID_FILE="$TUNNEL_DIR/$CONTEXT-tunnel.pid"
+LOG_FILE="$TUNNEL_DIR/$CONTEXT-tunnel.log"
+LOCK_FILE="$TUNNEL_DIR/$CONTEXT-tunnel.lock"
+
+# Get the port from kubeconfig if available
+KUBE_PORT=$(kubectl config view --minify --context="$CONTEXT" -o jsonpath='{.clusters[0].cluster.server}' 2>/dev/null | grep -o ":[0-9]\+" | tr -d ":" || echo "")
+
+if [[ -z "$KUBE_PORT" ]]; then
+    # Default to 6443 if we can't determine the port
+    KUBE_PORT=6443
+    echo "$(date): Could not determine port from kubeconfig, using default port $KUBE_PORT" >> "$LOG_FILE"
+else
+    echo "$(date): Found port $KUBE_PORT in kubeconfig for context $CONTEXT" >> "$LOG_FILE"
+fi
+
+# Check if PID file exists
+if [[ -f "$PID_FILE" ]]; then
+    OLD_PID=$(cat "$PID_FILE")
+    
+    # Log the cleanup attempt
+    echo "$(date): Attempting to clean up tunnel for context $CONTEXT (PID: $OLD_PID, Port: $KUBE_PORT)" >> "$LOG_FILE"
+    
+    # Try to kill the process
+    if kill -0 "$OLD_PID" 2>/dev/null; then
+        # Process exists, kill it
+        kill "$OLD_PID" 2>/dev/null
+        
+        # Wait a moment and check if it's really gone
+        sleep 1
+        if kill -0 "$OLD_PID" 2>/dev/null; then
+            # Still running, force kill
+            kill -9 "$OLD_PID" 2>/dev/null
+            echo "$(date): Forcefully terminated tunnel process $OLD_PID" >> "$LOG_FILE"
+        else
+            echo "$(date): Successfully terminated tunnel process $OLD_PID" >> "$LOG_FILE"
+        fi
+    else
+        echo "$(date): No running process found with PID $OLD_PID" >> "$LOG_FILE"
+    fi
+    
+    # Remove PID file
+    rm -f "$PID_FILE"
+else
+    echo "$(date): No PID file found for context $CONTEXT. Nothing to clean up." >> "$LOG_FILE"
+fi
+
+# Clean up lock file if it exists
+rm -f "$LOCK_FILE"
+
+# Check if port is still in use
+if nc -z localhost "$KUBE_PORT" 2>/dev/null; then
+    echo "$(date): Warning: Port $KUBE_PORT is still in use after cleanup. Another process might be using it." >> "$LOG_FILE"
+fi
+
+echo "$(date): Cleanup complete for context $CONTEXT" >> "$LOG_FILE" 

sky/utils/kubernetes/create_cluster.sh

@@ -85,6 +85,7 @@ fi
 if kind get clusters | grep -q skypilot; then
     echo "Local cluster already exists. Exiting."
     # Switch context to the local cluster
+    kind export kubeconfig --name skypilot
     kubectl config use-context kind-skypilot
     exit 100
 fi