skypilot-nightly 1.0.0.dev20250624__py3-none-any.whl → 1.0.0.dev20250626__py3-none-any.whl

sky/task.py

@@ -24,6 +24,7 @@ from sky.skylet import constants
 from sky.utils import common_utils
 from sky.utils import schemas
 from sky.utils import ux_utils
+from sky.volumes import volume as volume_lib
 
 if typing.TYPE_CHECKING:
     import yaml
@@ -246,12 +247,14 @@ class Task:
         secrets: Optional[Dict[str, str]] = None,
         workdir: Optional[str] = None,
         num_nodes: Optional[int] = None,
+        volumes: Optional[Dict[str, str]] = None,
         # Advanced:
         docker_image: Optional[str] = None,
         event_callback: Optional[str] = None,
         blocked_resources: Optional[Iterable['resources_lib.Resources']] = None,
         # Internal use only.
         file_mounts_mapping: Optional[Dict[str, str]] = None,
+        volume_mounts: Optional[List[volume_lib.VolumeMount]] = None,
     ):
         """Initializes a Task.
 
@@ -319,6 +322,7 @@ class Task:
         self.setup = setup
         self._envs = envs or {}
         self._secrets = secrets or {}
+        self._volumes = volumes or {}
 
         # Validate Docker login configuration early if both envs and secrets
         # contain Docker variables
@@ -361,7 +365,9 @@ class Task:
         self.best_resources: Optional[sky.Resources] = None
 
         # For internal use only.
-        self.file_mounts_mapping = file_mounts_mapping
+        self.file_mounts_mapping: Optional[Dict[str, str]] = file_mounts_mapping
+        self.volume_mounts: Optional[List[volume_lib.VolumeMount]] = (
+            volume_mounts)
 
         dag = sky.dag.get_current_dag()
         if dag is not None:
@@ -442,12 +448,9 @@ class Task:
         if self.file_mounts is None:
             return
         for target, source in self.file_mounts.items():
-            if target.endswith('/') or source.endswith('/'):
-                with ux_utils.print_exception_no_traceback():
-                    raise ValueError(
-                        'File mount paths cannot end with a slash '
-                        '(try "/mydir: /mydir" or "/myfile: /myfile"). '
-                        f'Found: target={target} source={source}')
+            location = f'file_mounts.{target}: {source}'
+            self._validate_mount_path(target, location)
+            self._validate_path(source, location)
             if data_utils.is_cloud_store_url(target):
                 with ux_utils.print_exception_no_traceback():
                     raise ValueError(
@@ -462,17 +465,25 @@ class Task:
                             f'File mount source {source!r} does not exist '
                             'locally. To fix: check if it exists, and correct '
                             'the path.')
-            # TODO(zhwu): /home/username/sky_workdir as the target path need
-            # to be filtered out as well.
-            if (target == constants.SKY_REMOTE_WORKDIR and
-                    self.workdir is not None):
-                with ux_utils.print_exception_no_traceback():
-                    raise ValueError(
-                        f'Cannot use {constants.SKY_REMOTE_WORKDIR!r} as a '
-                        'destination path of a file mount, as it will be used '
-                        'by the workdir. If uploading a file/folder to the '
-                        'workdir is needed, please specify the full path to '
-                        'the file/folder.')
+
+    def _validate_mount_path(self, path: str, location: str):
+        self._validate_path(path, location)
+        # TODO(zhwu): /home/username/sky_workdir as the target path need
+        # to be filtered out as well.
+        if (path == constants.SKY_REMOTE_WORKDIR and self.workdir is not None):
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    f'Cannot use {constants.SKY_REMOTE_WORKDIR!r} as a '
+                    'destination path of a file mount, as it will be used '
+                    'by the workdir. If uploading a file/folder to the '
+                    'workdir is needed, please specify the full path to '
+                    'the file/folder.')
+
+    def _validate_path(self, path: str, location: str):
+        if path.endswith('/'):
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError('Mount paths cannot end with a slash '
+                                 f'Found: {path} in {location}')
 
     def expand_and_validate_workdir(self):
         """Expand workdir to absolute path and validate it.
@@ -587,6 +598,7 @@ class Task:
             secrets=config.pop('secrets', None),
             event_callback=config.pop('event_callback', None),
             file_mounts_mapping=config.pop('file_mounts_mapping', None),
+            volumes=config.pop('volumes', None),
         )
 
         # Create lists to store storage objects inlined in file_mounts.
@@ -711,6 +723,16 @@ class Task:
             service = service_spec.SkyServiceSpec.from_yaml_config(service)
         task.set_service(service)
 
+        volume_mounts = config.pop('volume_mounts', None)
+        if volume_mounts is not None:
+            task.volume_mounts = []
+            for vol in volume_mounts:
+                common_utils.validate_schema(vol,
+                                             schemas.get_volume_mount_schema(),
+                                             'Invalid volume mount config: ')
+                volume_mount = volume_lib.VolumeMount.from_yaml_config(vol)
+                task.volume_mounts.append(volume_mount)
+
         assert not config, f'Invalid task args: {config.keys()}'
         return task
 
@@ -745,6 +767,97 @@ class Task:
             config = {}
         return Task.from_yaml_config(config)
 
+    def resolve_and_validate_volumes(self) -> None:
+        """Resolve volumes config to volume mounts and validate them.
+
+        Raises:
+            exceptions.VolumeNotFoundError: if any volume is not found.
+            exceptions.VolumeTopologyConflictError: if there is conflict in the
+              volumes and compute topology.
+        """
+        # Volumes has been resolved, a typical case is that the API server
+        # has resolved the volumes and the dag was then submitted to
+        # controllers.
+        if self.volume_mounts is not None:
+            return None
+        if not self._volumes:
+            return None
+        volume_mounts: List[volume_lib.VolumeMount] = []
+        for dst_path, vol in self._volumes.items():
+            self._validate_mount_path(dst_path, location='volumes')
+            # Shortcut for `dst_path: volume_name`
+            if isinstance(vol, str):
+                volume_mount = volume_lib.VolumeMount.resolve(dst_path, vol)
+            elif isinstance(vol, dict):
+                assert 'name' in vol, 'Volume name must be set.'
+                volume_mount = volume_lib.VolumeMount.resolve(
+                    dst_path, vol['name'])
+            else:
+                raise ValueError(f'Invalid volume config: {dst_path}: {vol}')
+            volume_mounts.append(volume_mount)
+        # Disable certain access modes
+        disabled_modes = {}
+        if self.num_nodes > 1:
+            disabled_modes[
+                volume_lib.VolumeAccessMode.READ_WRITE_ONCE.value] = (
+                    'access mode ReadWriteOnce is not supported for '
+                    'multi-node tasks.')
+            disabled_modes[
+                volume_lib.VolumeAccessMode.READ_WRITE_ONCE_POD.value] = (
+                    'access mode ReadWriteOncePod is not supported for '
+                    'multi-node tasks.')
+        # TODO(aylei): generalize access mode to all volume types
+        # Record the required topology and the volume that requires it, e.g.
+        # {'cloud': ('volume_name', 'aws')}
+        topology: Dict[str, Tuple[str, Optional[str]]] = {
+            'cloud': ('', None),
+            'region': ('', None),
+            'zone': ('', None),
+        }
+        for vol in volume_mounts:
+            # Check access mode
+            access_mode = vol.volume_config.config.get('access_mode', '')
+            if access_mode in disabled_modes:
+                raise ValueError(f'Volume {vol.volume_name} with '
+                                 f'{disabled_modes[access_mode]}')
+            # Check topology
+            for key, (vol_name, previous_req) in topology.items():
+                req = getattr(vol.volume_config, key)
+                if req is not None:
+                    if previous_req is not None and req != previous_req:
+                        raise exceptions.VolumeTopologyConflictError(
+                            f'Volume {vol.volume_name} can only be attached on '
+                            f'{key}:{req}, which conflicts with another volume '
+                            f'{vol_name} that requires {key}:{previous_req}.'
+                            f'Please use different volumes and retry.')
+                    topology[key] = (vol_name, req)
+        # Now we have the topology requirements from the intersection of all
+        # volumes. Check if there is topology conflict with the resources.
+        # Volume must have no conflict with ALL resources even if user
+        # specifies 'any_of' resources to ensure no resources will conflict
+        # with the volumes during failover.
+
+        for res in self.resources:
+            for key, (vol_name, vol_req) in topology.items():
+                req = getattr(res, key)
+                if (req is not None and vol_req is not None and
+                        str(req) != vol_req):
+                    raise exceptions.VolumeTopologyConflictError(
+                        f'The task requires {key}:{req}, which conflicts with '
+                        f'the volume constraint {key}:{vol_req}. Please '
+                        f'use different volumes and retry.')
+        # No topology conflict, we safely override the topology of resources to
+        # satisfy the volume constraints.
+        override_params = {}
+        for key, (vol_name, vol_req) in topology.items():
+            if vol_req is not None:
+                if key == 'cloud':
+                    override_params[key] = sky.CLOUD_REGISTRY.from_str(vol_req)
+                else:
+                    override_params[key] = vol_req
+        self.set_resources_override(override_params)
+        self.volume_mounts = volume_mounts
+
     @property
     def num_nodes(self) -> int:
         return self._num_nodes
@@ -767,6 +880,10 @@ class Task:
     def secrets(self) -> Dict[str, str]:
         return self._secrets
 
+    @property
+    def volumes(self) -> Dict[str, str]:
+        return self._volumes
+
     def update_envs(
             self, envs: Union[None, List[Tuple[str, str]],
                               Dict[str, str]]) -> 'Task':
@@ -1453,6 +1570,12 @@ class Task:
             })
 
         add_if_not_none('file_mounts_mapping', self.file_mounts_mapping)
+        add_if_not_none('volumes', self.volumes)
+        if self.volume_mounts is not None:
+            config['volume_mounts'] = [
+                volume_mount.to_yaml_config()
+                for volume_mount in self.volume_mounts
+            ]
         return config
 
     def get_required_cloud_features(

sky/templates/kubernetes-ray.yml.j2

@@ -243,6 +243,22 @@ provider:
             # This selector must match the head node pod's selector below.
             selector:
                 component: {{cluster_name_on_cloud}}-head
+      # Headless service mapping hostnames to rest of the worker nodes
+      {% for worker_id in range(1, num_nodes) %}
+      - apiVersion: v1
+        kind: Service
+        metadata:
+          labels:
+            parent: skypilot
+            skypilot-cluster: {{cluster_name_on_cloud}}
+            skypilot-user: {{ user }}
+          name: {{cluster_name_on_cloud}}-worker{{ worker_id }}
+        spec:
+          selector:
+              component: {{cluster_name_on_cloud}}-worker{{ worker_id }}
+          clusterIP: None
+      {% endfor %}
+
 
 # Specify the pod type for the ray head node (as configured below).
 head_node_type: ray_head_default
@@ -255,7 +271,7 @@ available_node_types:
       metadata:
         # name will be filled in the provisioner
         # head node name will be {{cluster_name_on_cloud}}-head, which will match the head node service selector above if a head node
-        # service is required.
+        # service is required. Worker nodes are named {{cluster_name_on_cloud}}-worker{{ node_id }}
         labels:
             parent: skypilot
             # component will be set for the head node pod to be the same as the head node service selector above if a
@@ -287,6 +303,10 @@ available_node_types:
         serviceAccountName: {{k8s_service_account_name}}
         automountServiceAccountToken: {{k8s_automount_sa_token}}
         restartPolicy: {{ "Always" if high_availability else "Never" }}
+        {% if volume_mounts %}
+        securityContext:
+          fsGroup: 1000
+        {% endif %}
 
         # Add node selector if GPU/TPUs are requested:
         {% if (k8s_topology_label_key is not none and k8s_topology_label_value is not none) or (k8s_spot_label_key is not none) %}
@@ -365,6 +385,11 @@ available_node_types:
           persistentVolumeClaim:
             claimName: {{cluster_name_on_cloud}}-{{k8s_high_availability_deployment_volume_mount_name}}
         {% endif %}
+        {% for volume_mount in volume_mounts %}
+        - name: {{volume_mount.name}}
+          persistentVolumeClaim:
+            claimName: {{volume_mount.volume_name_on_cloud}}
+        {% endfor %}
         containers:
         - name: ray-node
           imagePullPolicy: IfNotPresent
@@ -734,6 +759,10 @@ available_node_types:
           - name: fusermount-shared-dir
             mountPath: {{k8s_fusermount_shared_dir}}
           {% endif %}
+          {% for volume_mount in volume_mounts %}
+          - name: {{volume_mount.name}}
+            mountPath: {{volume_mount.path}}
+          {% endfor %}
           resources:
             requests:
               cpu: {{cpus}}

sky/users/permission.py

@@ -18,6 +18,8 @@ from sky.utils import common_utils
 
 logging.getLogger('casbin.policy').setLevel(sky_logging.ERROR)
 logging.getLogger('casbin.role').setLevel(sky_logging.ERROR)
+logging.getLogger('casbin.model').setLevel(sky_logging.ERROR)
+logging.getLogger('casbin.rbac').setLevel(sky_logging.ERROR)
 logger = sky_logging.init_logger(__name__)
 
 # Filelocks for the policy update.

sky/utils/context.py

@@ -254,7 +254,9 @@ class Popen(subprocess.Popen):
     def __init__(self, *args, **kwargs):
         env = kwargs.pop('env', None)
         if env is None:
-            env = os.environ
+            # Pass a copy of current context.environ to avoid race condition
+            # when the context is updated after the Popen is created.
+            env = os.environ.copy()
         super().__init__(*args, env=env, **kwargs)
 
 

sky/utils/kubernetes/deploy_remote_cluster.py

@@ -11,11 +11,12 @@ import shutil
 import subprocess
 import sys
 import tempfile
-from typing import Any, Dict, List, Optional, Set
+from typing import List, Set
 
 import yaml
 
 from sky.utils import ux_utils
+from sky.utils.kubernetes import ssh_utils
 
 # Colors for nicer UX
 RED = '\033[0;31m'
@@ -24,7 +25,6 @@ YELLOW = '\033[1;33m'
 WARNING_YELLOW = '\x1b[33m'
 NC = '\033[0m'  # No color
 
-DEFAULT_SSH_NODE_POOLS_PATH = os.path.expanduser('~/.sky/ssh_node_pools.yaml')
 DEFAULT_KUBECONFIG_PATH = os.path.expanduser('~/.kube/config')
 SSH_CONFIG_PATH = os.path.expanduser('~/.ssh/config')
 NODE_POOLS_INFO_DIR = os.path.expanduser('~/.sky/ssh_node_pools_info')
@@ -33,29 +33,6 @@ NODE_POOLS_INFO_DIR = os.path.expanduser('~/.sky/ssh_node_pools_info')
 SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
 
 
-class UniqueKeySafeLoader(yaml.SafeLoader):
-    """Custom YAML loader that raises an error if there are duplicate keys."""
-
-    def construct_mapping(self, node, deep=False):
-        mapping = {}
-        for key_node, value_node in node.value:
-            key = self.construct_object(key_node, deep=deep)
-            if key in mapping:
-                raise yaml.constructor.ConstructorError(
-                    note=(f'Duplicate cluster config for cluster {key!r}.\n'
-                          'Please remove one of them from: '
-                          f'{DEFAULT_SSH_NODE_POOLS_PATH}'))
-            value = self.construct_object(value_node, deep=deep)
-            mapping[key] = value
-        return mapping
-
-
-# Register the custom constructor inside the class
-UniqueKeySafeLoader.add_constructor(
-    yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG,
-    UniqueKeySafeLoader.construct_mapping)
-
-
 def parse_args():
     parser = argparse.ArgumentParser(
         description='Deploy a Kubernetes cluster on remote machines.')
@@ -64,9 +41,9 @@ def parse_args():
     parser.add_argument(
         '--ssh-node-pools-file',
         dest='ssh_node_pools_file',
-        default=DEFAULT_SSH_NODE_POOLS_PATH,
+        default=ssh_utils.DEFAULT_SSH_NODE_POOLS_PATH,
         help=
-        f'Path to SSH node pools YAML file (default: {DEFAULT_SSH_NODE_POOLS_PATH})'
+        f'Path to SSH node pools YAML file (default: {ssh_utils.DEFAULT_SSH_NODE_POOLS_PATH})'
     )
     parser.add_argument(
         '--kubeconfig-path',
@@ -117,156 +94,6 @@ def parse_args():
     return parser.parse_args()
 
 
-def load_ssh_targets(file_path: str) -> Dict[str, Any]:
-    """Load SSH targets from YAML file."""
-    if not os.path.exists(file_path):
-        with ux_utils.print_exception_no_traceback():
-            raise ValueError(f'SSH Node Pools file not found: {file_path}')
-
-    try:
-        with open(file_path, 'r', encoding='utf-8') as f:
-            targets = yaml.load(f, Loader=UniqueKeySafeLoader)
-        return targets
-    except yaml.constructor.ConstructorError as e:
-        with ux_utils.print_exception_no_traceback():
-            raise ValueError(e.note) from e
-    except (yaml.YAMLError, IOError, OSError) as e:
-        with ux_utils.print_exception_no_traceback():
-            raise ValueError(f'Error loading SSH Node Pools file: {e}') from e
-
-
-def check_host_in_ssh_config(hostname: str) -> bool:
-    """Return True iff *hostname* matches at least one `Host`/`Match` stanza
-    in the user's OpenSSH client configuration (including anything pulled in
-    via Include).
-
-    It calls:  ssh -vvG <hostname> -o ConnectTimeout=0
-    which:
-      • -G  expands the effective config without connecting
-      • -vv prints debug lines that show which stanzas are applied
-      • ConnectTimeout=0 avoids a DNS lookup if <hostname> is a FQDN/IP
-
-    No config files are opened or parsed manually.
-
-    Parameters
-    ----------
-    hostname : str
-        The alias/IP/FQDN you want to test.
-
-    Returns
-    -------
-    bool
-        True  – a specific stanza matched the host
-        False – nothing but the global defaults (`Host *`) applied
-    """
-    # We direct stderr→stdout because debug output goes to stderr.
-    proc = subprocess.run(
-        ['ssh', '-vvG', hostname, '-o', 'ConnectTimeout=0'],
-        text=True,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.STDOUT,
-        check=False,  # we only want the text, not to raise
-    )
-
-    # Look for lines like:
-    #   debug1: ~/.ssh/config line 42: Applying options for <hostname>
-    # Anything other than "*"
-    pattern = re.compile(r'^debug\d+: .*Applying options for ([^*].*)$',
-                         re.MULTILINE)
-
-    return bool(pattern.search(proc.stdout))
-
-
-def get_cluster_config(targets: Dict[str, Any],
-                       cluster_name: Optional[str] = None,
-                       file_path: Optional[str] = None) -> Dict[str, Any]:
-    """Get configuration for specific clusters or all clusters."""
-    if not targets:
-        with ux_utils.print_exception_no_traceback():
-            raise ValueError(
-                f'No clusters defined in SSH Node Pools file {file_path}')
-
-    if cluster_name:
-        if cluster_name not in targets:
-            with ux_utils.print_exception_no_traceback():
-                raise ValueError(f'Cluster {cluster_name!r} not found in '
-                                 f'SSH Node Pools file {file_path}')
-        return {cluster_name: targets[cluster_name]}
-
-    # Return all clusters if no specific cluster is specified
-    return targets
-
-
-def prepare_hosts_info(cluster_name: str,
-                       cluster_config: Dict[str, Any]) -> List[Dict[str, str]]:
-    """Prepare list of hosts with resolved user, identity_file, and password."""
-    if 'hosts' not in cluster_config or not cluster_config['hosts']:
-        with ux_utils.print_exception_no_traceback():
-            raise ValueError(
-                f'No hosts defined in cluster {cluster_name} configuration')
-
-    # Get cluster-level defaults
-    cluster_user = cluster_config.get('user', '')
-    cluster_identity_file = os.path.expanduser(
-        cluster_config.get('identity_file', ''))
-    cluster_password = cluster_config.get('password', '')
-
-    # Check if cluster identity file exists
-    if cluster_identity_file and not os.path.isfile(cluster_identity_file):
-        with ux_utils.print_exception_no_traceback():
-            raise ValueError(
-                f'SSH Identity File Missing: {cluster_identity_file}')
-
-    hosts_info = []
-    for host in cluster_config['hosts']:
-        # Host can be a string (IP or SSH config hostname) or a dict
-        if isinstance(host, str):
-            # Check if this is an SSH config hostname
-            is_ssh_config_host = check_host_in_ssh_config(host)
-
-            hosts_info.append({
-                'ip': host,
-                'user': '' if is_ssh_config_host else cluster_user,
-                'identity_file': '' if is_ssh_config_host else
-                                 cluster_identity_file,
-                'password': cluster_password,
-                'use_ssh_config': is_ssh_config_host
-            })
-        else:
-            # It's a dict with potential overrides
-            if 'ip' not in host:
-                print(
-                    f'{RED}Warning: Host missing \'ip\' field, skipping: {host}{NC}'
-                )
-                continue
-
-            # Check if this is an SSH config hostname
-            is_ssh_config_host = check_host_in_ssh_config(host['ip'])
-
-            # Use host-specific values or fall back to cluster defaults
-            host_user = '' if is_ssh_config_host else host.get(
-                'user', cluster_user)
-            host_identity_file = os.path.expanduser(
-                '' if is_ssh_config_host else host.
-                get('identity_file', cluster_identity_file))
-            host_password = host.get('password', cluster_password)
-
-            if host_identity_file and not os.path.isfile(host_identity_file):
-                with ux_utils.print_exception_no_traceback():
-                    raise ValueError(
-                        f'SSH Identity File Missing: {host_identity_file}')
-
-            hosts_info.append({
-                'ip': host['ip'],
-                'user': host_user,
-                'identity_file': host_identity_file,
-                'password': host_password,
-                'use_ssh_config': is_ssh_config_host
-            })
-
-    return hosts_info
-
-
 def run_command(cmd, shell=False):
     """Run a local command and return the output."""
     process = subprocess.run(cmd,
@@ -675,10 +502,10 @@ def main():
         password = args.password
 
         # Check if hosts are in SSH config
-        head_use_ssh_config = global_use_ssh_config or check_host_in_ssh_config(
+        head_use_ssh_config = global_use_ssh_config or ssh_utils.check_host_in_ssh_config(
             head_node)
         worker_use_ssh_config = [
-            global_use_ssh_config or check_host_in_ssh_config(node)
+            global_use_ssh_config or ssh_utils.check_host_in_ssh_config(node)
             for node in worker_nodes
         ]
 
@@ -688,10 +515,9 @@ def main():
                        kubeconfig_path, args.cleanup)
     else:
         # Using YAML configuration
-        targets = load_ssh_targets(args.ssh_node_pools_file)
-        clusters_config = get_cluster_config(targets,
-                                             args.infra,
-                                             file_path=args.ssh_node_pools_file)
+        targets = ssh_utils.load_ssh_targets(args.ssh_node_pools_file)
+        clusters_config = ssh_utils.get_cluster_config(
+            targets, args.infra, file_path=args.ssh_node_pools_file)
 
         # Print information about clusters being processed
         num_clusters = len(clusters_config)
@@ -705,7 +531,8 @@ def main():
                 print(f'SKYPILOT_CURRENT_CLUSTER: {cluster_name}')
                 print(
                     f'{YELLOW}==== Deploying cluster: {cluster_name} ====${NC}')
-                hosts_info = prepare_hosts_info(cluster_name, cluster_config)
+                hosts_info = ssh_utils.prepare_hosts_info(
+                    cluster_name, cluster_config)
 
                 if not hosts_info:
                     print(
@@ -744,7 +571,7 @@ def main():
                                 f'Cluster configuration has changed for field {key!r}. '
                                 f'Previous value: {history.get(key)}, '
                                 f'Current value: {cluster_config.get(key)}')
-                    history_hosts_info = prepare_hosts_info(
+                    history_hosts_info = ssh_utils.prepare_hosts_info(
                         cluster_name, history)
                     if not args.cleanup and history_hosts_info[0] != hosts_info[
                             0]: