skypilot-nightly 1.0.0.dev20250220__py3-none-any.whl → 1.0.0.dev20250222__py3-none-any.whl

sky/jobs/server/core.py

@@ -19,6 +19,7 @@ from sky import sky_logging
 from sky import task as task_lib
 from sky.backends import backend_utils
 from sky.clouds.service_catalog import common as service_catalog_common
+from sky.data import storage as storage_lib
 from sky.jobs import constants as managed_job_constants
 from sky.jobs import utils as managed_job_utils
 from sky.provision import common as provision_common
@@ -101,9 +102,35 @@ def launch(
 
     with rich_utils.safe_status(
             ux_utils.spinner_message('Initializing managed job')):
-        for task_ in dag.tasks:
-            controller_utils.maybe_translate_local_file_mounts_and_sync_up(
-                task_, task_type='jobs')
+
+        local_to_controller_file_mounts = {}
+
+        if storage_lib.get_cached_enabled_storage_clouds_or_refresh():
+            for task_ in dag.tasks:
+                controller_utils.maybe_translate_local_file_mounts_and_sync_up(
+                    task_, task_type='jobs')
+
+        else:
+            # We do not have any cloud storage available, so fall back to
+            # two-hop file_mount uploading.
+            # Note: we can't easily hack sync_storage_mounts() to upload
+            # directly to the controller, because the controller may not
+            # even be up yet.
+            for task_ in dag.tasks:
+                if task_.storage_mounts:
+                    # Technically, we could convert COPY storage_mounts that
+                    # have a local source and do not specify `store`, but we
+                    # will not do that for now. Only plain file_mounts are
+                    # supported.
+                    raise exceptions.NotSupportedError(
+                        'Cloud-based file_mounts are specified, but no cloud '
+                        'storage is available. Please specify local '
+                        'file_mounts only.')
+
+                # Merge file mounts from all tasks.
+                local_to_controller_file_mounts.update(
+                    controller_utils.translate_local_file_mounts_to_two_hop(
+                        task_))
 
     with tempfile.NamedTemporaryFile(prefix=f'managed-dag-{dag.name}-',
                                      mode='w') as f:
@@ -120,6 +147,7 @@ def launch(
         vars_to_fill = {
             'remote_user_yaml_path': remote_user_yaml_path,
             'user_yaml_path': f.name,
+            'local_to_controller_file_mounts': local_to_controller_file_mounts,
             'jobs_controller': controller_name,
             # Note: actual cluster name will be <task.name>-<managed job ID>
             'dag_name': dag.name,

sky/provision/__init__.py

@@ -20,6 +20,7 @@ from sky.provision import fluidstack
 from sky.provision import gcp
 from sky.provision import kubernetes
 from sky.provision import lambda_cloud
+from sky.provision import nebius
 from sky.provision import oci
 from sky.provision import runpod
 from sky.provision import vast

sky/provision/kubernetes/instance.py

@@ -797,7 +797,11 @@ def _create_pods(region: str, cluster_name_on_cloud: str,
                 'value': 'present',
                 'effect': 'NoSchedule'
             }
-            pod_spec_copy['spec']['tolerations'] = [tpu_toleration]
+            # Preserve existing tolerations if any
+            existing_tolerations = pod_spec_copy['spec'].get('tolerations', [])
+            pod_spec_copy['spec']['tolerations'] = existing_tolerations + [
+                tpu_toleration
+            ]
 
         return _create_namespaced_pod_with_retries(namespace, pod_spec_copy,
                                                    context)

sky/provision/kubernetes/utils.py

@@ -1302,13 +1302,13 @@ class KubernetesInstanceType:
         - Accelerators
     The name format is "{n}CPU--{k}GB" where n is the number of vCPUs and
     k is the amount of memory in GB. Accelerators can be specified by
-    appending "--{a}{type}" where a is the number of accelerators and
-    type is the accelerator type.
+    appending "--{type}:{a}" where type is the accelerator type and a
+    is the number of accelerators.
     CPU and memory can be specified as floats. Accelerator count must be int.
     Examples:
         - 4CPU--16GB
         - 0.5CPU--1.5GB
-        - 4CPU--16GB--1V100
+        - 4CPU--16GB--V100:1
     """
 
     def __init__(self,
@@ -1333,13 +1333,14 @@ class KubernetesInstanceType:
             # valid logical instance type name.
             assert self.accelerator_type is not None, self.accelerator_count
             acc_name = self.accelerator_type.replace(' ', '_')
-            name += f'--{self.accelerator_count}{acc_name}'
+            name += f'--{acc_name}:{self.accelerator_count}'
         return name
 
     @staticmethod
     def is_valid_instance_type(name: str) -> bool:
         """Returns whether the given name is a valid instance type."""
-        pattern = re.compile(r'^(\d+(\.\d+)?CPU--\d+(\.\d+)?GB)(--\d+\S+)?$')
+        pattern = re.compile(
+            r'^(\d+(\.\d+)?CPU--\d+(\.\d+)?GB)(--[\w\d-]+:\d+)?$')
         return bool(pattern.match(name))
 
     @classmethod
@@ -1354,7 +1355,7 @@ class KubernetesInstanceType:
             accelerator_type | str: Type of accelerator
         """
         pattern = re.compile(
-            r'^(?P<cpus>\d+(\.\d+)?)CPU--(?P<memory>\d+(\.\d+)?)GB(?:--(?P<accelerator_count>\d+)(?P<accelerator_type>\S+))?$'  # pylint: disable=line-too-long
+            r'^(?P<cpus>\d+(\.\d+)?)CPU--(?P<memory>\d+(\.\d+)?)GB(?:--(?P<accelerator_type>[\w\d-]+):(?P<accelerator_count>\d+))?$'  # pylint: disable=line-too-long
         )
         match = pattern.match(name)
         if match:
@@ -1400,7 +1401,7 @@ class KubernetesInstanceType:
         # Round up accelerator_count if it is not an int.
         accelerator_count = math.ceil(accelerator_count)
         if accelerator_count > 0:
-            name += f'--{accelerator_count}{accelerator_type}'
+            name += f'--{accelerator_type}:{accelerator_count}'
         return cls(cpus=cpus,
                    memory=memory,
                    accelerator_count=accelerator_count,

sky/provision/nebius/__init__.py

@@ -0,0 +1,11 @@
+"""Nebius provisioner for SkyPilot."""
+
+from sky.provision.nebius.config import bootstrap_instances
+from sky.provision.nebius.instance import cleanup_ports
+from sky.provision.nebius.instance import get_cluster_info
+from sky.provision.nebius.instance import open_ports
+from sky.provision.nebius.instance import query_instances
+from sky.provision.nebius.instance import run_instances
+from sky.provision.nebius.instance import stop_instances
+from sky.provision.nebius.instance import terminate_instances
+from sky.provision.nebius.instance import wait_instances

sky/provision/nebius/config.py

@@ -0,0 +1,11 @@
+"""Nebius configuration bootstrapping."""
+
+from sky.provision import common
+
+
+def bootstrap_instances(
+        region: str, cluster_name: str,
+        config: common.ProvisionConfig) -> common.ProvisionConfig:
+    """Bootstraps instances for the given cluster."""
+    del region, cluster_name  # unused
+    return config

sky/provision/nebius/instance.py

@@ -0,0 +1,285 @@
+"""Nebius instance provisioning."""
+import time
+from typing import Any, Dict, List, Optional
+
+from sky import sky_logging
+from sky.provision import common
+from sky.provision.nebius import utils
+from sky.utils import common_utils
+from sky.utils import status_lib
+from sky.utils import ux_utils
+
+PENDING_STATUS = ['STARTING', 'DELETING', 'STOPPING']
+
+MAX_RETRIES_TO_LAUNCH = 120  # Maximum number of retries
+
+logger = sky_logging.init_logger(__name__)
+
+
+def _filter_instances(region: str,
+                      cluster_name_on_cloud: str,
+                      status_filters: Optional[List[str]],
+                      head_only: bool = False) -> Dict[str, Any]:
+    project_id = utils.get_project_by_region(region)
+    instances = utils.list_instances(project_id)
+    filtered_instances = {}
+    for instance_id, instance in instances.items():
+        if (status_filters is not None and
+                instance['status'] not in status_filters):
+            continue
+
+        if instance['name'] and instance['name'].startswith(
+                f'{cluster_name_on_cloud}-'):
+            if head_only and instance['name'].endswith('-worker'):
+                continue
+            else:
+                filtered_instances[instance_id] = instance
+    return filtered_instances
+
+
+def _get_head_instance_id(instances: Dict[str, Any]) -> Optional[str]:
+    head_instance_id = None
+    for inst_id, inst in instances.items():
+        if inst['name'].endswith('-head'):
+            head_instance_id = inst_id
+            break
+    return head_instance_id
+
+
+def _wait_until_no_pending(region: str, cluster_name_on_cloud: str) -> None:
+    retry_count = 0
+    while retry_count < MAX_RETRIES_TO_LAUNCH:
+        instances = _filter_instances(region, cluster_name_on_cloud,
+                                      PENDING_STATUS)
+        if not instances:
+            break
+        logger.info(f'Waiting for {len(instances)} instances to be ready '
+                    f'(Attempt {retry_count + 1}/{MAX_RETRIES_TO_LAUNCH}).')
+        time.sleep(utils.POLL_INTERVAL)
+        retry_count += 1
+
+    if retry_count == MAX_RETRIES_TO_LAUNCH:
+        raise TimeoutError(f'Exceeded maximum retries '
+                           f'({MAX_RETRIES_TO_LAUNCH * utils.POLL_INTERVAL}'
+                           f' seconds) while waiting for instances'
+                           f' to be ready.')
+
+
+def run_instances(region: str, cluster_name_on_cloud: str,
+                  config: common.ProvisionConfig) -> common.ProvisionRecord:
+    """Runs instances for the given cluster."""
+    _wait_until_no_pending(region, cluster_name_on_cloud)
+    running_instances = _filter_instances(region, cluster_name_on_cloud,
+                                          ['RUNNING'])
+    head_instance_id = _get_head_instance_id(running_instances)
+    to_start_count = config.count - len(running_instances)
+    if to_start_count < 0:
+        raise RuntimeError(
+            f'Cluster {cluster_name_on_cloud} already has '
+            f'{len(running_instances)} nodes, but {config.count} are required.')
+    if to_start_count == 0:
+        if head_instance_id is None:
+            raise RuntimeError(
+                f'Cluster {cluster_name_on_cloud} has no head node.')
+        logger.info(f'Cluster {cluster_name_on_cloud} already has '
+                    f'{len(running_instances)} nodes, no need to start more.')
+        return common.ProvisionRecord(provider_name='nebius',
+                                      cluster_name=cluster_name_on_cloud,
+                                      region=region,
+                                      zone=None,
+                                      head_instance_id=head_instance_id,
+                                      resumed_instance_ids=[],
+                                      created_instance_ids=[])
+
+    created_instance_ids = []
+    resumed_instance_ids = []
+    stopped_instances = _filter_instances(region, cluster_name_on_cloud,
+                                          ['STOPPED'])
+    if config.resume_stopped_nodes and len(stopped_instances) > to_start_count:
+
+        raise RuntimeError(
+            'The number of running/stopped/stopping instances combined '
+            f'({len(stopped_instances) + len(running_instances)}) in '
+            f'cluster "{cluster_name_on_cloud}" is greater than the '
+            f'number requested by the user ({config.count}). '
+            'This is likely a resource leak. '
+            'Use "sky down" to terminate the cluster.')
+
+    for stopped_instance_id, _ in stopped_instances.items():
+        if to_start_count > 0:
+            try:
+                utils.start(stopped_instance_id)
+                resumed_instance_ids.append(stopped_instance_id)
+                to_start_count -= 1
+                if stopped_instances[stopped_instance_id]['name'].endswith(
+                        '-head'):
+                    head_instance_id = stopped_instance_id
+            except Exception as e:  # pylint: disable=broad-except
+                logger.warning(f'Start instance error: {e}')
+                raise
+            time.sleep(utils.POLL_INTERVAL)  # to avoid fake STOPPED status
+            logger.info(f'Started instance {stopped_instance_id}.')
+
+    for _ in range(to_start_count):
+        node_type = 'head' if head_instance_id is None else 'worker'
+        try:
+            platform, preset = config.node_config['InstanceType'].split('_')
+            instance_id = utils.launch(
+                cluster_name_on_cloud=cluster_name_on_cloud,
+                node_type=node_type,
+                platform=platform,
+                preset=preset,
+                region=region,
+                image_family=config.node_config['ImageId'],
+                disk_size=config.node_config['DiskSize'],
+                user_data=config.node_config['UserData'])
+        except Exception as e:  # pylint: disable=broad-except
+            logger.warning(f'run_instances error: {e}')
+            raise
+        logger.info(f'Launched instance {instance_id}.')
+        created_instance_ids.append(instance_id)
+        if head_instance_id is None:
+            head_instance_id = instance_id
+    assert head_instance_id is not None, 'head_instance_id should not be None'
+    return common.ProvisionRecord(provider_name='nebius',
+                                  cluster_name=cluster_name_on_cloud,
+                                  region=region,
+                                  zone=None,
+                                  head_instance_id=head_instance_id,
+                                  resumed_instance_ids=resumed_instance_ids,
+                                  created_instance_ids=created_instance_ids)
+
+
+def wait_instances(region: str, cluster_name_on_cloud: str,
+                   state: Optional[status_lib.ClusterStatus]) -> None:
+    _wait_until_no_pending(region, cluster_name_on_cloud)
+    if state is not None:
+        if state == status_lib.ClusterStatus.UP:
+            stopped_instances = _filter_instances(region, cluster_name_on_cloud,
+                                                  ['STOPPED'])
+            if stopped_instances:
+                raise RuntimeError(
+                    f'Cluster {cluster_name_on_cloud} is in UP state, but '
+                    f'{len(stopped_instances)} instances are stopped.')
+        if state == status_lib.ClusterStatus.STOPPED:
+            running_instances = _filter_instances(region, cluster_name_on_cloud,
+                                                  ['RUNNIG'])
+
+            if running_instances:
+                raise RuntimeError(
+                    f'Cluster {cluster_name_on_cloud} is in STOPPED state, but '
+                    f'{len(running_instances)} instances are running.')
+
+
+def stop_instances(
+    cluster_name_on_cloud: str,
+    provider_config: Optional[Dict[str, Any]] = None,
+    worker_only: bool = False,
+) -> None:
+    assert provider_config is not None
+    exist_instances = _filter_instances(provider_config['region'],
+                                        cluster_name_on_cloud, ['RUNNING'])
+    for instance in exist_instances:
+        if worker_only and instance.endswith('-head'):
+            continue
+        utils.stop(instance)
+
+
+def terminate_instances(
+    cluster_name_on_cloud: str,
+    provider_config: Optional[Dict[str, Any]] = None,
+    worker_only: bool = False,
+) -> None:
+    """See sky/provision/__init__.py"""
+
+    assert provider_config is not None
+    instances = _filter_instances(provider_config['region'],
+                                  cluster_name_on_cloud,
+                                  status_filters=None)
+    for inst_id, inst in instances.items():
+        logger.debug(f'Terminating instance {inst_id}: {inst}')
+        if worker_only and inst['name'].endswith('-head'):
+            continue
+        try:
+            utils.remove(inst_id)
+        except Exception as e:  # pylint: disable=broad-except
+            with ux_utils.print_exception_no_traceback():
+                raise RuntimeError(
+                    f'Failed to terminate instance {inst_id}: '
+                    f'{common_utils.format_exception(e, use_bracket=False)}'
+                ) from e
+    utils.delete_cluster(cluster_name_on_cloud, provider_config['region'])
+
+
+def get_cluster_info(
+        region: str,
+        cluster_name_on_cloud: str,
+        provider_config: Optional[Dict[str, Any]] = None) -> common.ClusterInfo:
+    _wait_until_no_pending(region, cluster_name_on_cloud)
+    running_instances = _filter_instances(region, cluster_name_on_cloud,
+                                          ['RUNNING'])
+    instances: Dict[str, List[common.InstanceInfo]] = {}
+    head_instance_id = None
+    for instance_id, instance_info in running_instances.items():
+        instances[instance_id] = [
+            common.InstanceInfo(
+                instance_id=instance_id,
+                internal_ip=instance_info['internal_ip'],
+                external_ip=instance_info['external_ip'],
+                tags={},
+            )
+        ]
+        if instance_info['name'].endswith('-head'):
+            head_instance_id = instance_id
+    assert head_instance_id is not None
+    return common.ClusterInfo(
+        instances=instances,
+        head_instance_id=head_instance_id,
+        provider_name='nebius',
+        provider_config=provider_config,
+    )
+
+
+def query_instances(
+    cluster_name_on_cloud: str,
+    provider_config: Optional[Dict[str, Any]] = None,
+    non_terminated_only: bool = True,
+) -> Dict[str, Optional[status_lib.ClusterStatus]]:
+    """See sky/provision/__init__.py"""
+    assert provider_config is not None, (cluster_name_on_cloud, provider_config)
+    instances = _filter_instances(provider_config['region'],
+                                  cluster_name_on_cloud, None)
+
+    status_map = {
+        'STARTING': status_lib.ClusterStatus.INIT,
+        'RUNNING': status_lib.ClusterStatus.UP,
+        'STOPPED': status_lib.ClusterStatus.STOPPED,
+        'STOPPING': status_lib.ClusterStatus.STOPPED,
+        'DELETING': status_lib.ClusterStatus.STOPPED,
+    }
+    statuses: Dict[str, Optional[status_lib.ClusterStatus]] = {}
+    for inst_id, inst in instances.items():
+        status = status_map[inst['status']]
+        if non_terminated_only and status is None:
+            continue
+        statuses[inst_id] = status
+    return statuses
+
+
+def open_ports(
+    cluster_name_on_cloud: str,
+    ports: List[str],
+    provider_config: Optional[Dict[str, Any]] = None,
+) -> None:
+    """See sky/provision/__init__.py"""
+    logger.debug(f'Skip opening ports {ports} for Nebius instances, as all '
+                 'ports are open by default.')
+    del cluster_name_on_cloud, provider_config, ports
+
+
+def cleanup_ports(
+    cluster_name_on_cloud: str,
+    ports: List[str],
+    provider_config: Optional[Dict[str, Any]] = None,
+) -> None:
+    del cluster_name_on_cloud, ports, provider_config  # Unused.