skypilot-nightly 1.0.0.dev20250509__py3-none-any.whl → 1.0.0.dev20251107__py3-none-any.whl

sky/provision/nebius/utils.py

@@ -1,17 +1,21 @@
 """Nebius library wrapper for SkyPilot."""
 import time
-from typing import Any, Dict
+from typing import Any, Dict, List, Optional
 import uuid
 
 from sky import sky_logging
 from sky import skypilot_config
 from sky.adaptors import nebius
+from sky.provision.nebius import constants as nebius_constants
 from sky.utils import common_utils
+from sky.utils import resources_utils
 
 logger = sky_logging.init_logger(__name__)
 
 POLL_INTERVAL = 5
 
+_MAX_OPERATIONS_TO_FETCH = 1000
+
 
 def retry(func):
     """Decorator to retry a function."""
@@ -34,12 +38,14 @@ def retry(func):
 
 def get_project_by_region(region: str) -> str:
     service = nebius.iam().ProjectServiceClient(nebius.sdk())
-    projects = service.list(nebius.iam().ListProjectsRequest(
-        parent_id=nebius.get_tenant_id())).wait()
+    projects = nebius.sync_call(
+        service.list(
+            nebius.iam().ListProjectsRequest(parent_id=nebius.get_tenant_id()),
+            timeout=nebius.READ_TIMEOUT))
 
     #  Check is there project if in config
-    project_id = skypilot_config.get_nested(('nebius', region, 'project_id'),
-                                            None)
+    project_id = skypilot_config.get_effective_region_config(
+        cloud='nebius', region=region, keys=('project_id',), default_value=None)
     if project_id is not None:
         return project_id
     for project in projects.items:
@@ -54,19 +60,21 @@ def get_or_create_gpu_cluster(name: str, project_id: str, fabric: str) -> str:
     """
     service = nebius.compute().GpuClusterServiceClient(nebius.sdk())
     try:
-        cluster = service.get_by_name(nebius.nebius_common().GetByNameRequest(
-            parent_id=project_id,
-            name=name,
-        )).wait()
-        cluster_id = cluster.metadata.id
-    except nebius.request_error():
-        cluster = service.create(nebius.compute().CreateGpuClusterRequest(
-            metadata=nebius.nebius_common().ResourceMetadata(
+        cluster = nebius.sync_call(
+            service.get_by_name(nebius.nebius_common().GetByNameRequest(
                 parent_id=project_id,
                 name=name,
-            ),
-            spec=nebius.compute().GpuClusterSpec(
-                infiniband_fabric=fabric))).wait()
+            )))
+        cluster_id = cluster.metadata.id
+    except nebius.request_error():
+        cluster = nebius.sync_call(
+            service.create(nebius.compute().CreateGpuClusterRequest(
+                metadata=nebius.nebius_common().ResourceMetadata(
+                    parent_id=project_id,
+                    name=name,
+                ),
+                spec=nebius.compute().GpuClusterSpec(
+                    infiniband_fabric=fabric))))
         cluster_id = cluster.resource_id
     return cluster_id
 
@@ -76,14 +84,16 @@ def delete_cluster(name: str, region: str) -> None:
     project_id = get_project_by_region(region)
     service = nebius.compute().GpuClusterServiceClient(nebius.sdk())
     try:
-        cluster = service.get_by_name(nebius.nebius_common().GetByNameRequest(
-            parent_id=project_id,
-            name=name,
-        )).wait()
+        cluster = nebius.sync_call(
+            service.get_by_name(nebius.nebius_common().GetByNameRequest(
+                parent_id=project_id,
+                name=name,
+            )))
         cluster_id = cluster.metadata.id
         logger.debug(f'Found GPU Cluster : {cluster_id}.')
-        service.delete(
-            nebius.compute().DeleteGpuClusterRequest(id=cluster_id)).wait()
+        nebius.sync_call(
+            service.delete(
+                nebius.compute().DeleteGpuClusterRequest(id=cluster_id)))
         logger.debug(f'Deleted GPU Cluster : {cluster_id}.')
     except nebius.request_error():
         logger.debug('GPU Cluster does not exist.')
@@ -92,13 +102,23 @@ def delete_cluster(name: str, region: str) -> None:
 def list_instances(project_id: str) -> Dict[str, Dict[str, Any]]:
     """Lists instances associated with API key."""
     service = nebius.compute().InstanceServiceClient(nebius.sdk())
-    result = service.list(
-        nebius.compute().ListInstancesRequest(parent_id=project_id)).wait()
-
-    instances = result
+    page_token = ''
+    instances = []
+    while True:
+        result = nebius.sync_call(
+            service.list(nebius.compute().ListInstancesRequest(
+                parent_id=project_id,
+                page_size=100,
+                page_token=page_token,
+            ),
+                         timeout=nebius.READ_TIMEOUT))
+        instances.extend(result.items)
+        if not result.next_page_token:  # "" means no more pages
+            break
+        page_token = result.next_page_token
 
     instance_dict: Dict[str, Dict[str, Any]] = {}
-    for instance in instances.items:
+    for instance in instances:
         info = {}
         info['status'] = instance.status.state.name
         info['name'] = instance.metadata.name
@@ -114,12 +134,13 @@ def list_instances(project_id: str) -> Dict[str, Dict[str, Any]]:
 
 def stop(instance_id: str) -> None:
     service = nebius.compute().InstanceServiceClient(nebius.sdk())
-    service.stop(nebius.compute().StopInstanceRequest(id=instance_id)).wait()
+    nebius.sync_call(
+        service.stop(nebius.compute().StopInstanceRequest(id=instance_id)))
     retry_count = 0
     while retry_count < nebius.MAX_RETRIES_TO_INSTANCE_STOP:
         service = nebius.compute().InstanceServiceClient(nebius.sdk())
-        instance = service.get(nebius.compute().GetInstanceRequest(
-            id=instance_id,)).wait()
+        instance = nebius.sync_call(
+            service.get(nebius.compute().GetInstanceRequest(id=instance_id,)))
         if instance.status.state.name == 'STOPPED':
             break
         time.sleep(POLL_INTERVAL)
@@ -136,12 +157,13 @@ def stop(instance_id: str) -> None:
 
 def start(instance_id: str) -> None:
     service = nebius.compute().InstanceServiceClient(nebius.sdk())
-    service.start(nebius.compute().StartInstanceRequest(id=instance_id)).wait()
+    nebius.sync_call(
+        service.start(nebius.compute().StartInstanceRequest(id=instance_id)))
     retry_count = 0
     while retry_count < nebius.MAX_RETRIES_TO_INSTANCE_START:
         service = nebius.compute().InstanceServiceClient(nebius.sdk())
-        instance = service.get(nebius.compute().GetInstanceRequest(
-            id=instance_id,)).wait()
+        instance = nebius.sync_call(
+            service.get(nebius.compute().GetInstanceRequest(id=instance_id,)))
         if instance.status.state.name == 'RUNNING':
             break
         time.sleep(POLL_INTERVAL)
@@ -156,9 +178,19 @@ def start(instance_id: str) -> None:
             f' to be ready.')
 
 
-def launch(cluster_name_on_cloud: str, node_type: str, platform: str,
-           preset: str, region: str, image_family: str, disk_size: int,
-           user_data: str) -> str:
+def launch(cluster_name_on_cloud: str,
+           node_type: str,
+           platform: str,
+           preset: str,
+           region: str,
+           image_family: str,
+           disk_size: int,
+           user_data: str,
+           associate_public_ip_address: bool,
+           filesystems: List[Dict[str, Any]],
+           use_static_ip_address: bool = False,
+           use_spot: bool = False,
+           network_tier: Optional[resources_utils.NetworkTier] = None) -> str:
     # Each node must have a unique name to avoid conflicts between
     # multiple worker VMs. To ensure uniqueness,a UUID is appended
     # to the node name.
@@ -172,11 +204,26 @@ def launch(cluster_name_on_cloud: str, node_type: str, platform: str,
     # 8 GPU virtual machines can be grouped into a GPU cluster.
     # The GPU clusters are built with InfiniBand secure high-speed networking.
     # https://docs.nebius.com/compute/clusters/gpu
-    if platform in ('gpu-h100-sxm', 'gpu-h200-sxm'):
+    if platform in nebius_constants.INFINIBAND_INSTANCE_PLATFORMS:
         if preset == '8gpu-128vcpu-1600gb':
-            #  Check is there fabric in config
-            fabric = skypilot_config.get_nested(('nebius', region, 'fabric'),
-                                                None)
+            fabric = skypilot_config.get_effective_region_config(
+                cloud='nebius',
+                region=region,
+                keys=('fabric',),
+                default_value=None)
+
+            # Auto-select fabric if network_tier=best and no fabric configured
+            if (fabric is None and
+                    str(network_tier) == str(resources_utils.NetworkTier.BEST)):
+                try:
+                    fabric = nebius_constants.get_default_fabric(
+                        platform, region)
+                    logger.info(f'Auto-selected InfiniBand fabric {fabric} '
+                                f'for {platform} in {region}')
+                except ValueError as e:
+                    logger.warning(
+                        f'InfiniBand fabric auto-selection failed: {e}')
+
             if fabric is None:
                 logger.warning(
                     f'Set up fabric for region {region} in ~/.sky/config.yaml '
@@ -186,24 +233,26 @@ def launch(cluster_name_on_cloud: str, node_type: str, platform: str,
                                                        project_id, fabric)
 
     service = nebius.compute().DiskServiceClient(nebius.sdk())
-    disk = service.create(nebius.compute().CreateDiskRequest(
-        metadata=nebius.nebius_common().ResourceMetadata(
-            parent_id=project_id,
-            name=disk_name,
-        ),
-        spec=nebius.compute().DiskSpec(
-            source_image_family=nebius.compute().SourceImageFamily(
-                image_family=image_family),
-            size_gibibytes=disk_size,
-            type=nebius.compute().DiskSpec.DiskType.NETWORK_SSD,
-        ))).wait()
+    disk = nebius.sync_call(
+        service.create(nebius.compute().CreateDiskRequest(
+            metadata=nebius.nebius_common().ResourceMetadata(
+                parent_id=project_id,
+                name=disk_name,
+            ),
+            spec=nebius.compute().DiskSpec(
+                source_image_family=nebius.compute().SourceImageFamily(
+                    image_family=image_family),
+                size_gibibytes=disk_size,
+                type=nebius.compute().DiskSpec.DiskType.NETWORK_SSD,
+            ))))
     disk_id = disk.resource_id
     retry_count = 0
     while retry_count < nebius.MAX_RETRIES_TO_DISK_CREATE:
-        disk = service.get_by_name(nebius.nebius_common().GetByNameRequest(
-            parent_id=project_id,
-            name=disk_name,
-        )).wait()
+        disk = nebius.sync_call(
+            service.get_by_name(nebius.nebius_common().GetByNameRequest(
+                parent_id=project_id,
+                name=disk_name,
+            )))
         if disk.status.state.name == 'READY':
             break
         logger.debug(f'Waiting for disk {disk_name} to be ready.')
@@ -217,46 +266,102 @@ def launch(cluster_name_on_cloud: str, node_type: str, platform: str,
             f' seconds) while waiting for disk {disk_name}'
             f' to be ready.')
 
+    filesystems_spec = []
+    if filesystems:
+        for fs in filesystems:
+            filesystems_spec.append(nebius.compute().AttachedFilesystemSpec(
+                mount_tag=fs['filesystem_mount_tag'],
+                attach_mode=nebius.compute().AttachedFilesystemSpec.AttachMode[
+                    fs['filesystem_attach_mode']],
+                existing_filesystem=nebius.compute().ExistingFilesystem(
+                    id=fs['filesystem_id'])))
+
     service = nebius.vpc().SubnetServiceClient(nebius.sdk())
-    sub_net = service.list(nebius.vpc().ListSubnetsRequest(
-        parent_id=project_id,)).wait()
+    sub_net = nebius.sync_call(
+        service.list(nebius.vpc().ListSubnetsRequest(parent_id=project_id,)))
 
     service = nebius.compute().InstanceServiceClient(nebius.sdk())
-    service.create(nebius.compute().CreateInstanceRequest(
-        metadata=nebius.nebius_common().ResourceMetadata(
-            parent_id=project_id,
-            name=instance_name,
-        ),
-        spec=nebius.compute().InstanceSpec(
-            gpu_cluster=nebius.compute().InstanceGpuClusterSpec(id=cluster_id,)
-            if cluster_id is not None else None,
-            boot_disk=nebius.compute().AttachedDiskSpec(
-                attach_mode=nebius.compute(
-                ).AttachedDiskSpec.AttachMode.READ_WRITE,
-                existing_disk=nebius.compute().ExistingDisk(id=disk_id)),
-            cloud_init_user_data=user_data,
-            resources=nebius.compute().ResourcesSpec(platform=platform,
-                                                     preset=preset),
-            network_interfaces=[
-                nebius.compute().NetworkInterfaceSpec(
-                    subnet_id=sub_net.items[0].metadata.id,
-                    ip_address=nebius.compute().IPAddress(),
-                    name='network-interface-0',
-                    public_ip_address=nebius.compute().PublicIPAddress())
-            ]))).wait()
+    logger.debug(f'Creating instance {instance_name} in project {project_id}.')
+    nebius.sync_call(
+        service.create(nebius.compute().CreateInstanceRequest(
+            metadata=nebius.nebius_common().ResourceMetadata(
+                parent_id=project_id,
+                name=instance_name,
+            ),
+            spec=nebius.compute().InstanceSpec(
+                gpu_cluster=nebius.compute().InstanceGpuClusterSpec(
+                    id=cluster_id,) if cluster_id is not None else None,
+                boot_disk=nebius.compute().AttachedDiskSpec(
+                    attach_mode=nebius.compute(
+                    ).AttachedDiskSpec.AttachMode.READ_WRITE,
+                    existing_disk=nebius.compute().ExistingDisk(id=disk_id)),
+                cloud_init_user_data=user_data,
+                resources=nebius.compute().ResourcesSpec(platform=platform,
+                                                         preset=preset),
+                filesystems=filesystems_spec if filesystems_spec else None,
+                network_interfaces=[
+                    nebius.compute().NetworkInterfaceSpec(
+                        subnet_id=sub_net.items[0].metadata.id,
+                        ip_address=nebius.compute().IPAddress(),
+                        name='network-interface-0',
+                        public_ip_address=nebius.compute().PublicIPAddress(
+                            static=use_static_ip_address)
+                        if associate_public_ip_address else None,
+                    )
+                ],
+                recovery_policy=nebius.compute().InstanceRecoveryPolicy.FAIL
+                if use_spot else None,
+                preemptible=nebius.compute().PreemptibleSpec(
+                    priority=1,
+                    on_preemption=nebius.compute().PreemptibleSpec.
+                    PreemptionPolicy.STOP) if use_spot else None,
+            ))))
     instance_id = ''
     retry_count = 0
     while retry_count < nebius.MAX_RETRIES_TO_INSTANCE_READY:
         service = nebius.compute().InstanceServiceClient(nebius.sdk())
-        instance = service.get_by_name(nebius.nebius_common().GetByNameRequest(
-            parent_id=project_id,
-            name=instance_name,
-        )).wait()
+        instance = nebius.sync_call(
+            service.get_by_name(nebius.nebius_common().GetByNameRequest(
+                parent_id=project_id,
+                name=instance_name,
+            )))
+        instance_id = instance.metadata.id
         if instance.status.state.name == 'STARTING':
-            instance_id = instance.metadata.id
             break
+
+        # All Instances initially have state=STOPPED and reconciling=True,
+        # so we need to wait until reconciling is False.
+        if instance.status.state.name == 'STOPPED' and \
+                not instance.status.reconciling:
+            next_token = ''
+            total_operations = 0
+            while True:
+                operations_response = nebius.sync_call(
+                    service.list_operations_by_parent(
+                        nebius.compute().ListOperationsByParentRequest(
+                            parent_id=project_id,
+                            page_size=100,
+                            page_token=next_token,
+                        )))
+                total_operations += len(operations_response.operations)
+                for operation in operations_response.operations:
+                    # Find the most recent operation for the instance.
+                    if operation.resource_id == instance_id:
+                        error_msg = operation.description
+                        if operation.status:
+                            error_msg += f' {operation.status.message}'
+                        raise RuntimeError(error_msg)
+                # If we've fetched too many operations, or there are no more
+                # operations to fetch, just raise a generic error.
+                if total_operations > _MAX_OPERATIONS_TO_FETCH or \
+                        not operations_response.next_page_token:
+                    raise RuntimeError(
+                        f'Instance {instance_name} failed to start.')
+                next_token = operations_response.next_page_token
         time.sleep(POLL_INTERVAL)
-        logger.debug(f'Waiting for instance {instance_name} start running.')
+        logger.debug(f'Waiting for instance {instance_name} to start running. '
+                     f'State: {instance.status.state.name}, '
+                     f'Reconciling: {instance.status.reconciling}')
         retry_count += 1
 
     if retry_count == nebius.MAX_RETRIES_TO_INSTANCE_READY:
@@ -271,19 +376,19 @@ def launch(cluster_name_on_cloud: str, node_type: str, platform: str,
 def remove(instance_id: str) -> None:
     """Terminates the given instance."""
     service = nebius.compute().InstanceServiceClient(nebius.sdk())
-    result = service.get(
-        nebius.compute().GetInstanceRequest(id=instance_id)).wait()
+    result = nebius.sync_call(
+        service.get(nebius.compute().GetInstanceRequest(id=instance_id)))
     disk_id = result.spec.boot_disk.existing_disk.id
-    service.delete(
-        nebius.compute().DeleteInstanceRequest(id=instance_id)).wait()
+    nebius.sync_call(
+        service.delete(nebius.compute().DeleteInstanceRequest(id=instance_id)))
     retry_count = 0
     # The instance begins deleting and attempts to delete the disk.
     # Must wait until the disk is unlocked and becomes deletable.
     while retry_count < nebius.MAX_RETRIES_TO_DISK_DELETE:
         try:
             service = nebius.compute().DiskServiceClient(nebius.sdk())
-            service.delete(
-                nebius.compute().DeleteDiskRequest(id=disk_id)).wait()
+            nebius.sync_call(
+                service.delete(nebius.compute().DeleteDiskRequest(id=disk_id)))
             break
         except nebius.request_error():
             logger.debug('Waiting for disk deletion.')

sky/provision/oci/instance.py

@@ -10,7 +10,7 @@ import copy
 from datetime import datetime
 import time
 import typing
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
 from sky import exceptions
 from sky import sky_logging
@@ -32,10 +32,12 @@ logger = sky_logging.init_logger(__name__)
 @query_utils.debug_enabled(logger)
 @common_utils.retry
 def query_instances(
+    cluster_name: str,
     cluster_name_on_cloud: str,
     provider_config: Optional[Dict[str, Any]] = None,
     non_terminated_only: bool = True,
-) -> Dict[str, Optional['status_lib.ClusterStatus']]:
+    retry_if_missing: bool = False,
+) -> Dict[str, Tuple[Optional['status_lib.ClusterStatus'], Optional[str]]]:
     """Query instances.
 
     Returns a dictionary of instance IDs and status.
@@ -43,11 +45,13 @@ def query_instances(
     A None status means the instance is marked as "terminated"
     or "terminating".
     """
+    del cluster_name, retry_if_missing  # unused
     assert provider_config is not None, cluster_name_on_cloud
     region = provider_config['region']
 
     status_map = oci_utils.oci_config.STATE_MAPPING_OCI_TO_SKY
-    statuses: Dict[str, Optional['status_lib.ClusterStatus']] = {}
+    statuses: Dict[str, Tuple[Optional['status_lib.ClusterStatus'],
+                              Optional[str]]] = {}
     filters = {constants.TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud}
 
     instances = _get_filtered_nodes(region, filters)
@@ -56,15 +60,16 @@ def query_instances(
         sky_status = status_map[vm_status]
         if non_terminated_only and sky_status is None:
             continue
-        statuses[node['inst_id']] = sky_status
+        statuses[node['inst_id']] = (sky_status, None)
 
     return statuses
 
 
 @query_utils.debug_enabled(logger)
-def run_instances(region: str, cluster_name_on_cloud: str,
+def run_instances(region: str, cluster_name: str, cluster_name_on_cloud: str,
                   config: common.ProvisionConfig) -> common.ProvisionRecord:
     """Start instances with bootstrapped configuration."""
+    del cluster_name  # unused
     tags = dict(sorted(copy.deepcopy(config.tags).items()))
 
     start_time = round(time.time() * 1000)

sky/provision/paperspace/instance.py

@@ -1,7 +1,7 @@
 """Paperspace instance provisioning."""
 
 import time
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
 from sky import sky_logging
 from sky.provision import common
@@ -48,10 +48,10 @@ def _get_head_instance_id(instances: Dict[str, Any]) -> Optional[str]:
     return head_instance_id
 
 
-def run_instances(region: str, cluster_name_on_cloud: str,
+def run_instances(region: str, cluster_name: str, cluster_name_on_cloud: str,
                   config: common.ProvisionConfig) -> common.ProvisionRecord:
     """Runs instances for the given cluster."""
-
+    del cluster_name  # unused
     pending_status = [
         'starting', 'restarting', 'upgrading', 'provisioning', 'stopping'
     ]
@@ -277,12 +277,14 @@ def get_cluster_info(
 
 
 def query_instances(
+    cluster_name: str,
     cluster_name_on_cloud: str,
     provider_config: Optional[Dict[str, Any]] = None,
     non_terminated_only: bool = True,
-) -> Dict[str, Optional[status_lib.ClusterStatus]]:
+    retry_if_missing: bool = False,
+) -> Dict[str, Tuple[Optional['status_lib.ClusterStatus'], Optional[str]]]:
     """See sky/provision/__init__.py"""
-    del non_terminated_only
+    del cluster_name, non_terminated_only, retry_if_missing  #unused
     assert provider_config is not None, (cluster_name_on_cloud, provider_config)
     instances = _filter_instances(cluster_name_on_cloud, None)
 
@@ -297,10 +299,11 @@ def query_instances(
         'ready': status_lib.ClusterStatus.UP,
         'off': status_lib.ClusterStatus.STOPPED,
     }
-    statuses: Dict[str, Optional[status_lib.ClusterStatus]] = {}
+    statuses: Dict[str, Tuple[Optional['status_lib.ClusterStatus'],
+                              Optional[str]]] = {}
     for inst_id, inst in instances.items():
         status = status_map[inst['state']]
-        statuses[inst_id] = status
+        statuses[inst_id] = (status, None)
     return statuses
 
 

sky/provision/paperspace/utils.py

@@ -8,7 +8,7 @@ from typing import Any, Dict, List, Optional, Union
 
 from sky import sky_logging
 from sky.adaptors import common as adaptors_common
-import sky.provision.paperspace.constants as constants
+from sky.provision.paperspace import constants
 from sky.utils import common_utils
 
 if typing.TYPE_CHECKING:

sky/provision/primeintellect/__init__.py

@@ -0,0 +1,10 @@
+"""Prime Intellect provisioner for SkyPilot."""
+
+from sky.provision.primeintellect.config import bootstrap_instances
+from sky.provision.primeintellect.instance import cleanup_ports
+from sky.provision.primeintellect.instance import get_cluster_info
+from sky.provision.primeintellect.instance import query_instances
+from sky.provision.primeintellect.instance import run_instances
+from sky.provision.primeintellect.instance import stop_instances
+from sky.provision.primeintellect.instance import terminate_instances
+from sky.provision.primeintellect.instance import wait_instances

sky/provision/primeintellect/config.py

@@ -0,0 +1,11 @@
+"""Prime Intellect 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