skypilot-nightly 1.0.0.dev20250502__py3-none-any.whl → 1.0.0.dev20251203__py3-none-any.whl

sky/clouds/nebius.py

@@ -1,25 +1,22 @@
 """ Nebius Cloud. """
+import json
 import os
 import typing
-from typing import Dict, Iterator, List, Optional, Tuple, Union
+from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
+from sky import catalog
 from sky import clouds
+from sky import exceptions
+from sky import skypilot_config
 from sky.adaptors import nebius
-from sky.clouds import service_catalog
+from sky.provision.nebius import constants as nebius_constants
 from sky.utils import annotations
 from sky.utils import registry
 from sky.utils import resources_utils
 
 if typing.TYPE_CHECKING:
     from sky import resources as resources_lib
-
-_CREDENTIAL_FILES = [
-    # credential files for Nebius
-    nebius.NEBIUS_TENANT_ID_FILENAME,
-    nebius.NEBIUS_IAM_TOKEN_FILENAME,
-    nebius.NEBIUS_PROJECT_ID_FILENAME,
-    nebius.NEBIUS_CREDENTIALS_FILENAME
-]
+    from sky.utils import volume as volume_lib
 
 _INDENT_PREFIX = '    '
 
@@ -55,14 +52,18 @@ class Nebius(clouds.Cloud):
     _CLOUD_UNSUPPORTED_FEATURES = {
         clouds.CloudImplementationFeatures.AUTODOWN:
             ('Autodown not supported. Can\'t delete OS disk.'),
-        clouds.CloudImplementationFeatures.SPOT_INSTANCE:
-            ('Spot is not supported, as Nebius API does not implement spot.'),
         clouds.CloudImplementationFeatures.CLONE_DISK_FROM_CLUSTER:
             (f'Migrating disk is currently not supported on {_REPR}.'),
         clouds.CloudImplementationFeatures.CUSTOM_DISK_TIER:
             (f'Custom disk tier is currently not supported on {_REPR}.'),
+        clouds.CloudImplementationFeatures.CUSTOM_NETWORK_TIER:
+            ('Custom network tier is currently only supported for '
+             'H100:8 and H200:8 on Nebius.'),
         clouds.CloudImplementationFeatures.HIGH_AVAILABILITY_CONTROLLERS:
             ('High availability controllers are not supported on Nebius.'),
+        clouds.CloudImplementationFeatures.CUSTOM_MULTI_NETWORK:
+            ('Customized multiple network interfaces are not supported on '
+             f'{_REPR}.'),
     }
     # Nebius maximum instance name length defined as <= 63 as a hostname length
     # 63 - 8 - 5 = 50 characters since
@@ -77,25 +78,43 @@ class Nebius(clouds.Cloud):
 
     @classmethod
     def _unsupported_features_for_resources(
-        cls, resources: 'resources_lib.Resources'
+        cls,
+        resources: 'resources_lib.Resources',
+        region: Optional[str] = None,
     ) -> Dict[clouds.CloudImplementationFeatures, str]:
-        del resources  # unused
-        return cls._CLOUD_UNSUPPORTED_FEATURES
+        unsupported = cls._CLOUD_UNSUPPORTED_FEATURES.copy()
+
+        # Check if the accelerators support InfiniBand (H100 or H200) and 8 GPUs
+        if resources.accelerators is not None:
+            for acc_name, acc_count in resources.accelerators.items():
+                if acc_name.lower() in ('h100', 'h200') and acc_count == 8:
+                    # Remove CUSTOM_NETWORK_TIER from unsupported features for
+                    # InfiniBand-capable accelerators. Refer to:
+                    # https://docs.nebius.com/compute/clusters/gpu#fabrics
+                    unsupported.pop(
+                        clouds.CloudImplementationFeatures.CUSTOM_NETWORK_TIER,
+                        None)
+                    break
+
+        return unsupported
 
     @classmethod
     def _max_cluster_name_length(cls) -> Optional[int]:
         return cls._MAX_CLUSTER_NAME_LEN_LIMIT
 
     @classmethod
-    def regions_with_offering(cls, instance_type: str,
-                              accelerators: Optional[Dict[str, int]],
-                              use_spot: bool, region: Optional[str],
-                              zone: Optional[str]) -> List[clouds.Region]:
+    def regions_with_offering(
+        cls,
+        instance_type: str,
+        accelerators: Optional[Dict[str, int]],
+        use_spot: bool,
+        region: Optional[str],
+        zone: Optional[str],
+        resources: Optional['resources_lib.Resources'] = None,
+    ) -> List[clouds.Region]:
         assert zone is None, 'Nebius does not support zones.'
         del accelerators, zone  # unused
-        if use_spot:
-            return []
-        regions = service_catalog.get_region_zones_for_instance_type(
+        regions = catalog.get_region_zones_for_instance_type(
             instance_type, use_spot, 'nebius')
 
         if region is not None:
@@ -107,8 +126,8 @@ class Nebius(clouds.Cloud):
         cls,
         instance_type: str,
     ) -> Tuple[Optional[float], Optional[float]]:
-        return service_catalog.get_vcpus_mem_from_instance_type(instance_type,
-                                                                clouds='nebius')
+        return catalog.get_vcpus_mem_from_instance_type(instance_type,
+                                                        clouds='nebius')
 
     @classmethod
     def zones_provision_loop(
@@ -135,11 +154,11 @@ class Nebius(clouds.Cloud):
                                      use_spot: bool,
                                      region: Optional[str] = None,
                                      zone: Optional[str] = None) -> float:
-        return service_catalog.get_hourly_cost(instance_type,
-                                               use_spot=use_spot,
-                                               region=region,
-                                               zone=zone,
-                                               clouds='nebius')
+        return catalog.get_hourly_cost(instance_type,
+                                       use_spot=use_spot,
+                                       region=region,
+                                       zone=zone,
+                                       clouds='nebius')
 
     def accelerators_to_hourly_cost(self,
                                     accelerators: Dict[str, int],
@@ -161,69 +180,124 @@ class Nebius(clouds.Cloud):
         return isinstance(other, Nebius)
 
     @classmethod
-    def get_default_instance_type(
-            cls,
-            cpus: Optional[str] = None,
-            memory: Optional[str] = None,
-            disk_tier: Optional[resources_utils.DiskTier] = None
-    ) -> Optional[str]:
+    def get_default_instance_type(cls,
+                                  cpus: Optional[str] = None,
+                                  memory: Optional[str] = None,
+                                  disk_tier: Optional[
+                                      resources_utils.DiskTier] = None,
+                                  region: Optional[str] = None,
+                                  zone: Optional[str] = None) -> Optional[str]:
         """Returns the default instance type for Nebius."""
-        return service_catalog.get_default_instance_type(cpus=cpus,
-                                                         memory=memory,
-                                                         disk_tier=disk_tier,
-                                                         clouds='nebius')
+        return catalog.get_default_instance_type(cpus=cpus,
+                                                 memory=memory,
+                                                 disk_tier=disk_tier,
+                                                 region=region,
+                                                 zone=zone,
+                                                 clouds='nebius')
 
     @classmethod
     def get_accelerators_from_instance_type(
         cls,
         instance_type: str,
     ) -> Optional[Dict[str, Union[int, float]]]:
-        return service_catalog.get_accelerators_from_instance_type(
-            instance_type, clouds='nebius')
+        return catalog.get_accelerators_from_instance_type(instance_type,
+                                                           clouds='nebius')
 
     @classmethod
     def get_zone_shell_cmd(cls) -> Optional[str]:
         return None
 
     def make_deploy_resources_variables(
-            self,
-            resources: 'resources_lib.Resources',
-            cluster_name: resources_utils.ClusterName,
-            region: 'clouds.Region',
-            zones: Optional[List['clouds.Zone']],
-            num_nodes: int,
-            dryrun: bool = False) -> Dict[str, Optional[str]]:
+        self,
+        resources: 'resources_lib.Resources',
+        cluster_name: resources_utils.ClusterName,
+        region: 'clouds.Region',
+        zones: Optional[List['clouds.Zone']],
+        num_nodes: int,
+        dryrun: bool = False,
+        volume_mounts: Optional[List['volume_lib.VolumeMount']] = None,
+    ) -> Dict[str, Any]:
         del dryrun, cluster_name
         assert zones is None, ('Nebius does not support zones', zones)
 
-        r = resources
-        acc_dict = self.get_accelerators_from_instance_type(r.instance_type)
+        resources = resources.assert_launchable()
+        acc_dict = self.get_accelerators_from_instance_type(
+            resources.instance_type)
         custom_resources = resources_utils.make_ray_custom_resources_str(
             acc_dict)
         platform, _ = resources.instance_type.split('_')
 
-        if platform in ('cpu-d3', 'cpu-e2'):
-            image_family = 'ubuntu22.04-driverless'
-        elif platform in ('gpu-h100-sxm', 'gpu-h200-sxm', 'gpu-l40s-a'):
-            image_family = 'ubuntu22.04-cuda12'
+        # Selecting image_family by platform
+        # https://docs.nebius.com/compute/storage/boot-disk-images
+        if platform.startswith('cpu'):
+            image_family = 'ubuntu24.04-driverless'
+        elif platform.startswith('gpu'):
+            image_family = 'ubuntu24.04-cuda12'
         else:
             raise RuntimeError('Unsupported instance type for Nebius cloud:'
                                f' {resources.instance_type}')
 
-        resources_vars = {
+        config_fs = skypilot_config.get_effective_region_config(
+            cloud='nebius',
+            region=region.name,
+            keys=('filesystems',),
+            default_value=[])
+        resources_vars_fs = []
+        for i, fs in enumerate(config_fs):
+            resources_vars_fs.append({
+                'filesystem_id': fs['filesystem_id'],
+                'filesystem_attach_mode': fs.get('attach_mode', 'READ_WRITE'),
+                'filesystem_mount_path': fs.get(
+                    'mount_path', f'/mnt/filesystem-skypilot-{i+1}'),
+                'filesystem_mount_tag': f'filesystem-skypilot-{i+1}'
+            })
+
+        use_static_ip_address = skypilot_config.get_nested(
+            ('nebius', 'use_static_ip_address'), default_value=False)
+        resources_vars: Dict[str, Any] = {
             'instance_type': resources.instance_type,
             'custom_resources': custom_resources,
+            'use_static_ip_address': use_static_ip_address,
             'region': region.name,
             'image_id': image_family,
             # Nebius does not support specific zones.
             'zones': None,
+            'use_spot': resources.use_spot,
+            'filesystems': resources_vars_fs,
+            'network_tier': resources.network_tier
         }
 
+        docker_run_options = []
+
         if acc_dict is not None:
             # Nebius cloud's docker runtime information does not contain
             # 'nvidia-container-runtime', causing no GPU option to be added to
             # the docker run command. We patch this by adding it here.
-            resources_vars['docker_run_options'] = ['--gpus all']
+            docker_run_options.append('--gpus all')
+
+            # Check for InfiniBand support with network_tier: best
+            is_infiniband_capable = (
+                platform in nebius_constants.INFINIBAND_INSTANCE_PLATFORMS)
+            if (is_infiniband_capable and
+                    resources.network_tier == resources_utils.NetworkTier.BEST):
+                # For Docker containers, add InfiniBand device access and
+                # IPC_LOCK capability
+                if resources.extract_docker_image() is not None:
+                    docker_run_options.extend(
+                        nebius_constants.INFINIBAND_DOCKER_OPTIONS)
+
+                    # Add InfiniBand environment variables to docker run options
+                    for env_var, env_value in (
+                            nebius_constants.INFINIBAND_ENV_VARS.items()):
+                        docker_run_options.extend(
+                            ['-e', f'{env_var}={env_value}'])
+
+                # For all InfiniBand-capable instances, add env variables
+                resources_vars[
+                    'env_vars'] = nebius_constants.INFINIBAND_ENV_VARS
+
+        if docker_run_options:
+            resources_vars['docker_run_options'] = docker_run_options
 
         return resources_vars
 
@@ -255,7 +329,9 @@ class Nebius(clouds.Cloud):
             default_instance_type = Nebius.get_default_instance_type(
                 cpus=resources.cpus,
                 memory=resources.memory,
-                disk_tier=resources.disk_tier)
+                disk_tier=resources.disk_tier,
+                region=resources.region,
+                zone=resources.zone)
             if default_instance_type is None:
                 # TODO: Add hints to all return values in this method to help
                 #  users understand why the resources are not launchable.
@@ -266,15 +342,16 @@ class Nebius(clouds.Cloud):
 
         assert len(accelerators) == 1, resources
         acc, acc_count = list(accelerators.items())[0]
-        (instance_list, fuzzy_candidate_list
-        ) = service_catalog.get_instance_type_for_accelerator(
-            acc,
-            acc_count,
-            use_spot=resources.use_spot,
-            cpus=resources.cpus,
-            region=resources.region,
-            zone=resources.zone,
-            clouds='nebius')
+        (instance_list,
+         fuzzy_candidate_list) = catalog.get_instance_type_for_accelerator(
+             acc,
+             acc_count,
+             use_spot=resources.use_spot,
+             cpus=resources.cpus,
+             memory=resources.memory,
+             region=resources.region,
+             zone=resources.zone,
+             clouds='nebius')
         if instance_list is None:
             return resources_utils.FeasibleResources([], fuzzy_candidate_list,
                                                      None)
@@ -282,25 +359,25 @@ class Nebius(clouds.Cloud):
                                                  fuzzy_candidate_list, None)
 
     @classmethod
-    @annotations.lru_cache(scope='request')
-    def _check_compute_credentials(cls) -> Tuple[bool, Optional[str]]:
+    def _check_compute_credentials(
+            cls) -> Tuple[bool, Optional[Union[str, Dict[str, str]]]]:
         """Checks if the user has access credentials to
         Nebius's compute service."""
         token_cred_msg = (
             f'{_INDENT_PREFIX}Credentials can be set up by running: \n'
-            f'{_INDENT_PREFIX}  $ nebius iam get-access-token > {nebius.NEBIUS_IAM_TOKEN_PATH} \n'  # pylint: disable=line-too-long
-            f'{_INDENT_PREFIX} or generate  ~/.nebius/credentials.json \n')
+            f'{_INDENT_PREFIX}  $ nebius iam get-access-token > {nebius.iam_token_path()} \n'  # pylint: disable=line-too-long
+            f'{_INDENT_PREFIX} or generate  {nebius.credentials_path()} \n')
 
-        tenant_msg = (f'{_INDENT_PREFIX} Copy your tenat ID from the web console and save it to file \n'  # pylint: disable=line-too-long
-                      f'{_INDENT_PREFIX}  $ echo $NEBIUS_TENANT_ID_PATH > {nebius.NEBIUS_TENANT_ID_PATH} \n'  # pylint: disable=line-too-long
+        tenant_msg = (f'{_INDENT_PREFIX} Copy your tenant ID from the web console and save it to file \n'  # pylint: disable=line-too-long
+                      f'{_INDENT_PREFIX}  $ echo $NEBIUS_TENANT_ID_PATH > {nebius.tenant_id_path()} \n'  # pylint: disable=line-too-long
                       f'{_INDENT_PREFIX} Or if you have 1 tenant you can run:\n'  # pylint: disable=line-too-long
-                      f'{_INDENT_PREFIX}  $ nebius --format json iam whoami|jq -r \'.user_profile.tenants[0].tenant_id\' > {nebius.NEBIUS_TENANT_ID_PATH} \n')  # pylint: disable=line-too-long
+                      f'{_INDENT_PREFIX}  $ nebius --format json iam whoami|jq -r \'.user_profile.tenants[0].tenant_id\' > {nebius.tenant_id_path()} \n')  # pylint: disable=line-too-long
         if not nebius.is_token_or_cred_file_exist():
             return False, f'{token_cred_msg}'
-        sdk = nebius.sdk()
         tenant_id = nebius.get_tenant_id()
         if tenant_id is None:
             return False, f'{tenant_msg}'
+        sdk = nebius.sdk()
         try:
             service = nebius.iam().ProjectServiceClient(sdk)
             service.list(
@@ -314,7 +391,8 @@ class Nebius(clouds.Cloud):
 
     @classmethod
     @annotations.lru_cache(scope='request')
-    def _check_storage_credentials(cls) -> Tuple[bool, Optional[str]]:
+    def _check_storage_credentials(
+            cls) -> Tuple[bool, Optional[Union[str, Dict[str, str]]]]:
         """Checks if the user has access credentials to Nebius Object Storage.
 
         Returns:
@@ -341,8 +419,8 @@ class Nebius(clouds.Cloud):
 
     def get_credential_file_mounts(self) -> Dict[str, str]:
         credential_file_mounts = {
-            f'~/.nebius/{filename}': f'~/.nebius/{filename}'
-            for filename in _CREDENTIAL_FILES
+            filepath: filepath
+            for filepath in nebius.get_credential_file_paths()
         }
         if nebius_profile_in_aws_cred_and_config():
             credential_file_mounts['~/.aws/credentials'] = '~/.aws/credentials'
@@ -356,9 +434,60 @@ class Nebius(clouds.Cloud):
         return None
 
     def instance_type_exists(self, instance_type: str) -> bool:
-        return service_catalog.instance_type_exists(instance_type, 'nebius')
+        return catalog.instance_type_exists(instance_type, 'nebius')
 
     def validate_region_zone(self, region: Optional[str], zone: Optional[str]):
-        return service_catalog.validate_region_zone(region,
-                                                    zone,
-                                                    clouds='nebius')
+        return catalog.validate_region_zone(region, zone, clouds='nebius')
+
+    @classmethod
+    def get_user_identities(cls) -> Optional[List[List[str]]]:
+        """Returns the email address + project id of the active user."""
+        nebius_workspace_config = json.dumps(
+            skypilot_config.get_workspace_cloud('nebius'), sort_keys=True)
+        return cls._get_user_identities(nebius_workspace_config)
+
+    @classmethod
+    @annotations.lru_cache(scope='request', maxsize=5)
+    def _get_user_identities(
+            cls, workspace_config: Optional[str]) -> Optional[List[List[str]]]:
+        # We add workspace_config in args to avoid caching the identity for when
+        # different workspace configs are used.
+        del workspace_config  # Unused
+        sdk = nebius.sdk()
+        profile_client = nebius.iam().ProfileServiceClient(sdk)
+        try:
+            profile = nebius.sync_call(
+                profile_client.get(nebius.iam().GetProfileRequest(),
+                                   timeout=nebius.READ_TIMEOUT))
+        except Exception as e:
+            raise exceptions.CloudUserIdentityError(
+                f'Error getting Nebius profile: {e}')
+        if profile.user_profile is not None:
+            if profile.user_profile.attributes is None:
+                raise exceptions.CloudUserIdentityError(
+                    'Nebius profile is a UserProfile, but has no attributes: '
+                    f'{profile.user_profile}')
+            if profile.user_profile.attributes.email is None:
+                raise exceptions.CloudUserIdentityError(
+                    'Nebius profile is a UserProfile, but has no email: '
+                    f'{profile.user_profile}')
+            return [[profile.user_profile.attributes.email]]
+        if profile.service_account_profile is not None:
+            if profile.service_account_profile.info is None:
+                raise exceptions.CloudUserIdentityError(
+                    'Nebius profile is a ServiceAccountProfile, but has no '
+                    f'info: {profile.service_account_profile}')
+            if profile.service_account_profile.info.metadata is None:
+                raise exceptions.CloudUserIdentityError(
+                    'Nebius profile is a ServiceAccountProfile, but has no '
+                    f'metadata: {profile.service_account_profile}')
+            if profile.service_account_profile.info.metadata.name is None:
+                raise exceptions.CloudUserIdentityError(
+                    'Nebius profile is a ServiceAccountProfile, but has no '
+                    f'name: {profile.service_account_profile}')
+            return [[profile.service_account_profile.info.metadata.name]]
+        if profile.anonymous_profile is not None:
+            return None
+        unknown_profile_type = profile.which_field_in_oneof('profile')
+        raise exceptions.CloudUserIdentityError(
+            f'Nebius profile is of an unknown type - {unknown_profile_type}')