skypilot-nightly 1.0.0.dev20250905__py3-none-any.whl → 1.0.0.dev20251210__py3-none-any.whl

sky/provision/primeintellect/utils.py

@@ -0,0 +1,398 @@
+"""Prime Intellect library wrapper for SkyPilot."""
+
+import json
+import os
+import shlex
+import time
+from typing import Any, Dict, List, Optional, Tuple, Union
+import uuid
+
+import requests
+
+from sky.catalog import common as catalog_common
+from sky.utils import common_utils
+
+_df = None
+_lookup_dict = None
+
+# Base URL for Prime Intellect API (used as default if not configured).
+DEFAULT_BASE_URL = 'https://api.primeintellect.ai'
+CREDENTIALS_PATH = '~/.prime/config.json'
+INITIAL_BACKOFF_SECONDS = 10
+MAX_BACKOFF_FACTOR = 10
+MAX_ATTEMPTS = 6
+
+
+class PrimeintellectAPIError(Exception):
+    """Base exception for Prime Intellect API errors."""
+
+    def __init__(self,
+                 message: str,
+                 status_code: Optional[int] = None,
+                 response_data: Optional[Dict[str, Any]] = None):
+        super().__init__(message)
+        self.status_code = status_code
+        self.response_data = response_data
+
+
+class PrimeintellectResourcesUnavailableError(PrimeintellectAPIError):
+    """Exception for when resources are unavailable on Prime Intellect."""
+    pass
+
+
+def _parse_api_error(response: Any) -> Tuple[str, bool]:
+    """Parse API error response to extract meaningful error messages.
+
+    Returns:
+        Tuple[str, bool]:
+        - str: A human-readable error message parsed from the API response.
+        - bool: True if the error indicates resource unavailability (e.g.,
+          capacity issues or quota/limit exceeded), otherwise False.
+    """
+    try:
+        if hasattr(response, 'json'):
+            error_data = response.json()
+        else:
+            error_data = response
+
+        if isinstance(error_data, dict):
+            # Try to extract error message from common error response fields
+            error_message = error_data.get('message', '')
+            if not error_message:
+                error_message = error_data.get('error', '')
+            if not error_message:
+                error_message = error_data.get('detail', '')
+
+            # Check if it's a resource unavailability error
+            if any(keyword in error_message.lower() for keyword in [
+                    'no capacity', 'capacity', 'unavailable', 'out of stock',
+                    'insufficient', 'not available', 'quota exceeded',
+                    'limit exceeded'
+            ]):
+                return error_message, True
+
+            return error_message, False
+
+        return str(error_data), False
+    except Exception:  # pylint: disable=broad-except
+        return f'HTTP {response.status_code} {response.reason}', False
+
+
+def _try_request_with_backoff(
+        method: str,
+        url: str,
+        headers: Dict[str, str],
+        data: Optional[Union[str, Dict[str, Any]]] = None) -> Dict[str, Any]:
+    backoff = common_utils.Backoff(initial_backoff=INITIAL_BACKOFF_SECONDS,
+                                   max_backoff_factor=MAX_BACKOFF_FACTOR)
+    for i in range(MAX_ATTEMPTS):
+        timeout = 30
+        if method == 'get':
+            response = requests.get(url,
+                                    headers=headers,
+                                    params=data,
+                                    timeout=timeout)
+        elif method == 'post':
+            response = requests.post(url,
+                                     headers=headers,
+                                     json=data,
+                                     timeout=timeout)
+        elif method == 'put':
+            response = requests.put(url,
+                                    headers=headers,
+                                    json=data,
+                                    timeout=timeout)
+        elif method == 'patch':
+            response = requests.patch(url,
+                                      headers=headers,
+                                      json=data,
+                                      timeout=timeout)
+        elif method == 'delete':
+            response = requests.delete(url, headers=headers, timeout=timeout)
+        else:
+            raise ValueError(f'Unsupported requests method: {method}')
+        # If rate limited, wait and try again
+        if response.status_code == 429 and i != MAX_ATTEMPTS - 1:
+            time.sleep(backoff.current_backoff())
+            continue
+        if response.ok:
+            return response.json()
+        else:
+            # Parse the error response for meaningful messages
+            err, is_resource_unavailable = _parse_api_error(response)
+
+            # Create a more informative error message
+            if not err:
+                err = (f'API request failed: {method} {url}: '
+                       f'{response.status_code} {response.reason}')
+            else:
+                err = f'API request failed: {err}'
+
+            # Raise appropriate exception based on error type
+            if is_resource_unavailable:
+                raise PrimeintellectResourcesUnavailableError(
+                    err,
+                    status_code=response.status_code,
+                    response_data=response.json()
+                    if hasattr(response, 'json') else None)
+            else:
+                raise PrimeintellectAPIError(
+                    err,
+                    status_code=response.status_code,
+                    response_data=response.json()
+                    if hasattr(response, 'json') else None)
+    return {}
+
+
+def get_upstream_cloud_id(instance_type: str) -> Optional[str]:
+    global _df, _lookup_dict
+    if _df is None:
+        _df = catalog_common.read_catalog('primeintellect/vms.csv')
+        _lookup_dict = (
+            _df.set_index('InstanceType')['UpstreamCloudId'].to_dict())
+    return _lookup_dict.get(instance_type)
+
+
+class PrimeIntellectAPIClient:
+    """Client for interacting with Prime Intellect API."""
+
+    def __init__(self) -> None:
+        self.credentials = os.path.expanduser(CREDENTIALS_PATH)
+        assert os.path.exists(self.credentials), 'Credentials not found'
+        with open(self.credentials, 'r', encoding='utf-8') as f:
+            self._credentials = json.load(f)
+        self.api_key = self._credentials.get('api_key')
+        self.team_id = self._credentials.get('team_id')
+        self.base_url = self._credentials.get('base_url', DEFAULT_BASE_URL)
+        self.headers = {
+            'Authorization': f'Bearer {self.api_key}',
+            'Content-Type': 'application/json'
+        }
+
+    def list_instances(self, **search_kwargs) -> List[Dict[str, Any]]:
+        response = _try_request_with_backoff('get',
+                                             f'{self.base_url}/api/v1/pods',
+                                             headers=self.headers,
+                                             data=search_kwargs)
+        return response['data']
+
+    def get_instance_details(self, instance_id: str) -> Dict[str, Any]:
+        return _try_request_with_backoff(
+            'get',
+            f'{self.base_url}/api/v1/pods/{instance_id}',
+            headers=self.headers)
+
+    def launch(self,
+               name: str,
+               instance_type: str,
+               region: str,
+               availability_zone: str,
+               disk_size: int,
+               vcpus: int = 0,
+               memory: int = 0) -> Dict[str, Any]:
+        """Create a pod/instance via Prime Intellect API.
+
+        Args:
+            name: User-visible name of the pod.
+            instance_type: A catalog instance type string. The expected format
+                is:
+                "<provider>__<accelerator>__<vcpus>__<memory>[_SPOT]".
+
+                - <provider>: Upstream provider tag (e.g., "primecompute").
+                - <accelerator>:
+                  * GPU nodes: "<N>x<GPU_MODEL>", e.g., "8xH100_80GB".
+                  * CPU-only nodes: the literal "CPU_NODE".
+                - <vcpus>: Integer string for vCPU count (e.g., "104").
+                - <memory>: Integer string for memory in GB (e.g., "752").
+                - Optional suffix "_SPOT" may be present in the full string
+                  (ignored here; pricing/spot behavior is not controlled by
+                  this method).
+
+                Notes:
+                - Parsing: only the first two components (provider,
+                  accelerator) are needed to build the payload. The vCPU
+                  and memory values are provided via the ``vcpus`` and
+                  ``memory`` arguments.
+                - Catalog lookup: the full instance_type string is used to
+                  map to the catalog's UpstreamCloudId.
+                - CPU-only: accelerator "CPU_NODE" is a sentinel for
+                  "no GPU". We set gpuType='CPU_NODE' and gpuCount=1 to
+                  represent CPU-only pods.
+                - Spot: the optional "__SPOT" suffix (if present) is ignored
+                  here; pricing/spot behavior is handled elsewhere.
+
+            region: Country/region code used by Prime Intellect.
+            availability_zone: Data center ID (zone) within the region.
+            disk_size: Boot disk size in GB.
+            vcpus: Optional explicit vCPU override; if >0 it will be sent.
+            memory: Optional explicit memory override in GB; if >0 it will be
+                sent.
+
+        Returns:
+            The API response JSON as a dict.
+        """
+        cloud_id = get_upstream_cloud_id(instance_type)
+        assert cloud_id, 'cloudId cannot be None'
+        assert availability_zone, 'availability_zone cannot be None'
+
+        # Parse the instance_type. We only need the first two components:
+        # provider and accelerator info (see docstring above).
+        provider, gpu_parts, _, _ = instance_type.split('__', 3)
+        if 'CPU_NODE' in gpu_parts:
+            # Prime Intellect API uses the same schema for CPU-only and GPU
+            # pods. For CPU-only instances, we set gpuType='CPU_NODE' and
+            # gpuCount=1 as a sentinel to indicate "no GPU". This is how CPU
+            # instances are represented internally on our platform; the
+            # backend does not interpret this as having a physical GPU.
+            gpu_type = 'CPU_NODE'
+            gpu_count = 1
+        else:
+            parts = gpu_parts.split('x', 1)
+            gpu_count = int(parts[0])
+            gpu_type = parts[1]
+
+        payload: Dict[str, Any] = {
+            'pod': {
+                'name': name,
+                'cloudId': cloud_id,
+                'socket': 'PCIe',
+                'gpuType': gpu_type,
+                'gpuCount': int(gpu_count),
+                'diskSize': disk_size,
+                # Prime Intellect API historically required maxPrice.
+                # Set to 0 to indicate on-demand/non-spot pricing.
+                'maxPrice': 0,
+            },
+            'provider': {
+                'type': provider,
+            }
+        }
+
+        if vcpus > 0:
+            payload['pod']['vcpus'] = vcpus
+        if memory > 0:
+            payload['pod']['memory'] = memory
+
+        if region != 'UNSPECIFIED':
+            payload['pod']['country'] = region
+        if availability_zone != 'UNSPECIFIED':
+            payload['pod']['dataCenterId'] = availability_zone
+
+        if self.team_id is not None and self.team_id != '':
+            payload['team'] = {'teamId': self.team_id}
+
+        response = _try_request_with_backoff(
+            'post',
+            f'{self.base_url}/api/v1/pods',
+            headers=self.headers,
+            data=payload,
+        )
+        return response
+
+    def remove(self, instance_id: str) -> Dict[str, Any]:
+        return _try_request_with_backoff(
+            'delete',
+            f'{self.base_url}/api/v1/pods/{instance_id}',
+            headers=self.headers,
+        )
+
+    def list_ssh_keys(self) -> List[Dict[str, Any]]:
+        response = _try_request_with_backoff('get',
+                                             f'{self.base_url}/api/v1/ssh_keys',
+                                             headers=self.headers)
+        return response['data']
+
+    def get_or_add_ssh_key(self, ssh_pub_key: str = '') -> Dict[str, str]:
+        """Add ssh key if not already added."""
+        # Check if the public key is already added
+        ssh_keys = self.list_ssh_keys()
+        for key in ssh_keys:
+            if key['publicKey'].strip().split()[:2] == ssh_pub_key.strip(
+            ).split()[:2]:
+                return {'name': key['name'], 'ssh_key': ssh_pub_key}
+
+        # Add the public key to Prime Intellect account if not already added
+        ssh_key_name = 'skypilot-' + str(uuid.uuid4()).replace('-', '')[:8]
+        _try_request_with_backoff(
+            'post',
+            f'{self.base_url}/api/v1/ssh_keys',
+            headers=self.headers,
+            data={
+                'name': ssh_key_name,
+                'publicKey': ssh_pub_key
+            },
+        )
+        return {'name': ssh_key_name, 'ssh_key': ssh_pub_key}
+
+
+def parse_ssh_connection(ssh_connection: Any) -> Tuple[Optional[str], int]:
+    """Parse and extract SSH username and port from a connection field.
+
+    The provider may return the SSH connection in multiple shapes. This helper
+    robustly extracts the SSH username and port while tolerating extra flags or
+    various tokenizations.
+
+    Accepted formats (examples):
+    - String with port flag:
+      "ubuntu@1.2.3.4 -p 2222 [-o <flag> ...]"
+    - String without explicit port (defaults to 22):
+      "ubuntu@1.2.3.4"
+    - String with host:port:
+      "ubuntu@1.2.3.4:2222"
+    - List with a single target:
+      ["ubuntu@1.2.3.4"]
+    - List of tokens (e.g., split form):
+      ["ubuntu@1.2.3.4", "-p", "2222"]
+
+    Args:
+        ssh_connection: The raw field from the API; can be a string or a list
+            of strings.
+
+    Returns:
+        (ssh_user, ssh_port): username if found, else None; port if found,
+        else 22.
+    """
+    ssh_user: Optional[str] = None
+    ssh_port: int = 22
+
+    # Normalize into a list of tokens for easier processing.
+    tokens: List[str] = []
+    if isinstance(ssh_connection, str):
+        try:
+            tokens = shlex.split(ssh_connection)
+        except Exception:  # pylint: disable=broad-except
+            tokens = [ssh_connection]
+    elif isinstance(ssh_connection, list):
+        for elem in ssh_connection:
+            if isinstance(elem, str):
+                try:
+                    tokens.extend(shlex.split(elem))
+                except Exception:  # pylint: disable=broad-except
+                    tokens.append(elem)
+    else:
+        # Unknown type; return defaults.
+        return ssh_user, ssh_port
+
+    # Find the first token containing '@' as the user@host candidate.
+    user_host: Optional[str] = next((t for t in tokens if '@' in t), None)
+    if user_host:
+        ssh_user = user_host.split('@', 1)[0].strip()
+        # Try host:port format (after '@').
+        host_part = user_host.split('@', 1)[1]
+        if ':' in host_part:
+            _, maybe_port = host_part.rsplit(':', 1)
+            try:
+                ssh_port = int(maybe_port)
+            except ValueError:
+                pass
+
+    # Check for '-p <port>' pair anywhere in the tokens. This takes priority.
+    if '-p' in tokens:
+        idx = tokens.index('-p')
+        if idx + 1 < len(tokens):
+            try:
+                ssh_port = int(tokens[idx + 1])
+            except ValueError:
+                pass
+
+    return ssh_user, ssh_port

sky/provision/provisioner.py

@@ -18,6 +18,7 @@ from sky import exceptions
 from sky import global_user_state
 from sky import logs
 from sky import provision
+from sky import resources as resources_lib
 from sky import sky_logging
 from sky import skypilot_config
 from sky.adaptors import aws
@@ -27,6 +28,7 @@ from sky.provision import common as provision_common
 from sky.provision import instance_setup
 from sky.provision import logging as provision_logging
 from sky.provision import metadata_utils
+from sky.provision import volume as provision_volume
 from sky.skylet import constants
 from sky.utils import common
 from sky.utils import common_utils
@@ -58,6 +60,11 @@ def _bulk_provision(
     region_name = region.name
 
     start = time.time()
+
+    provision_volume.provision_ephemeral_volumes(cloud, region_name,
+                                                 cluster_name.name_on_cloud,
+                                                 bootstrap_config)
+
     # TODO(suquark): Should we cache the bootstrapped result?
     #  Currently it is not necessary as bootstrapping takes
     #  only ~3s, caching it seems over-engineering and could
@@ -69,6 +76,7 @@ def _bulk_provision(
 
     provision_record = provision.run_instances(provider_name,
                                                region_name,
+                                               str(cluster_name),
                                                cluster_name.name_on_cloud,
                                                config=config)
 
@@ -149,9 +157,9 @@ def bulk_provision(
             logger.debug(f'SkyPilot version: {sky.__version__}; '
                          f'commit: {sky.__commit__}')
             logger.debug(_TITLE.format('Provisioning'))
-            logger.debug(
-                'Provision config:\n'
-                f'{json.dumps(dataclasses.asdict(bootstrap_config), indent=2)}')
+            redacted_config = bootstrap_config.get_redacted_config()
+            logger.debug('Provision config:\n'
+                         f'{json.dumps(redacted_config, indent=2)}')
             return _bulk_provision(cloud, region, cluster_name,
                                    bootstrap_config)
         except exceptions.NoClusterLaunchedError:
@@ -235,6 +243,7 @@ def teardown_cluster(cloud_name: str, cluster_name: resources_utils.ClusterName,
         provision.terminate_instances(cloud_name, cluster_name.name_on_cloud,
                                       provider_config)
         metadata_utils.remove_cluster_metadata(cluster_name.name_on_cloud)
+        provision_volume.delete_ephemeral_volumes(provider_config)
     else:
         provision.stop_instances(cloud_name, cluster_name.name_on_cloud,
                                  provider_config)
@@ -427,18 +436,27 @@ def wait_for_ssh(cluster_info: provision_common.ClusterInfo,
 
 
 def _post_provision_setup(
-        cloud_name: str, cluster_name: resources_utils.ClusterName,
-        handle_cluster_yaml: str,
+        launched_resources: resources_lib.Resources,
+        cluster_name: resources_utils.ClusterName, handle_cluster_yaml: str,
         provision_record: provision_common.ProvisionRecord,
         custom_resource: Optional[str]) -> provision_common.ClusterInfo:
     config_from_yaml = global_user_state.get_cluster_yaml_dict(
         handle_cluster_yaml)
     provider_config = config_from_yaml.get('provider')
+    cloud_name = repr(launched_resources.cloud)
     cluster_info = provision.get_cluster_info(cloud_name,
                                               provision_record.region,
                                               cluster_name.name_on_cloud,
                                               provider_config=provider_config)
 
+    # Update cluster info in handle so cluster instance ids are set. This
+    # allows us to expose provision logs to debug nodes that failed during post
+    # provision setup.
+    handle = global_user_state.get_handle_from_cluster_name(
+        cluster_name.display_name)
+    handle.cached_cluster_info = cluster_info
+    global_user_state.update_cluster_handle(cluster_name.display_name, handle)
+
     if cluster_info.num_instances > 1:
         # Only worker nodes have logs in the per-instance log directory. Head
         # node's log will be redirected to the main log file.
@@ -474,12 +492,13 @@ def _post_provision_setup(
         # ready by the provisioner, and we use kubectl instead of SSH to run the
         # commands and rsync on the pods. SSH will still be ready after a while
         # for the users to SSH into the pod.
-        if cloud_name.lower() != 'kubernetes':
+        is_k8s_cloud = cloud_name.lower() in ['kubernetes', 'ssh']
+        if not is_k8s_cloud:
             logger.debug(
                 f'\nWaiting for SSH to be available for {cluster_name!r} ...')
             wait_for_ssh(cluster_info, ssh_credentials)
             logger.debug(f'SSH Connection ready for {cluster_name!r}')
-        vm_str = 'Instance' if cloud_name.lower() != 'kubernetes' else 'Pod'
+        vm_str = 'Instance' if not is_k8s_cloud else 'Pod'
         plural = '' if len(cluster_info.instances) == 1 else 's'
         verb = 'is' if len(cluster_info.instances) == 1 else 'are'
         indent_str = (ux_utils.INDENT_SYMBOL
@@ -526,6 +545,7 @@ def _post_provision_setup(
             status.update(
                 ux_utils.spinner_message(
                     'Checking controller version compatibility'))
+
             try:
                 server_jobs_utils.check_version_mismatch_and_non_terminal_jobs()
             except exceptions.ClusterNotUpError:
@@ -615,10 +635,15 @@ def _post_provision_setup(
         status.update(
             runtime_preparation_str.format(step=3, step_name='runtime'))
 
+        skip_ray_setup = False
         ray_port = constants.SKY_REMOTE_RAY_PORT
         head_ray_needs_restart = True
         ray_cluster_healthy = False
-        if (not provision_record.is_instance_just_booted(
+        if (launched_resources.cloud is not None and
+                not launched_resources.cloud.uses_ray()):
+            skip_ray_setup = True
+            logger.debug('Skip Ray cluster setup as cloud does not use Ray.')
+        elif (not provision_record.is_instance_just_booted(
                 head_instance.instance_id)):
             # Check if head node Ray is alive
             (ray_port, ray_cluster_healthy,
@@ -643,7 +668,9 @@ def _post_provision_setup(
                              'async setup to complete...')
                 time.sleep(1)
 
-        if head_ray_needs_restart:
+        if skip_ray_setup:
+            logger.debug('Skip Ray cluster setup on the head node.')
+        elif head_ray_needs_restart:
             logger.debug('Starting Ray on the entire cluster.')
             instance_setup.start_ray_on_head_node(
                 cluster_name.name_on_cloud,
@@ -666,7 +693,9 @@ def _post_provision_setup(
         # We don't need to restart ray on worker nodes if the ray cluster is
         # already healthy, i.e. the head node has expected number of nodes
         # connected to the ray cluster.
-        if cluster_info.num_instances > 1 and not ray_cluster_healthy:
+        if skip_ray_setup:
+            logger.debug('Skip Ray cluster setup on the worker nodes.')
+        elif cluster_info.num_instances > 1 and not ray_cluster_healthy:
             instance_setup.start_ray_on_worker_nodes(
                 cluster_name.name_on_cloud,
                 no_restart=not head_ray_needs_restart,
@@ -692,8 +721,9 @@ def _post_provision_setup(
                                                     cluster_info,
                                                     ssh_credentials)
 
-        instance_setup.start_skylet_on_head_node(cluster_name.name_on_cloud,
-                                                 cluster_info, ssh_credentials)
+        instance_setup.start_skylet_on_head_node(cluster_name, cluster_info,
+                                                 ssh_credentials,
+                                                 launched_resources)
 
     logger.info(
         ux_utils.finishing_message(f'Cluster launched: {cluster_name}.',
@@ -704,8 +734,8 @@ def _post_provision_setup(
 
 @timeline.event
 def post_provision_runtime_setup(
-        cloud_name: str, cluster_name: resources_utils.ClusterName,
-        handle_cluster_yaml: str,
+        launched_resources: resources_lib.Resources,
+        cluster_name: resources_utils.ClusterName, handle_cluster_yaml: str,
         provision_record: provision_common.ProvisionRecord,
         custom_resource: Optional[str],
         log_dir: str) -> provision_common.ClusterInfo:
@@ -726,7 +756,7 @@ def post_provision_runtime_setup(
         try:
             logger.debug(_TITLE.format('System Setup After Provision'))
             return _post_provision_setup(
-                cloud_name,
+                launched_resources,
                 cluster_name,
                 handle_cluster_yaml=handle_cluster_yaml,
                 provision_record=provision_record,

sky/provision/runpod/__init__.py

@@ -11,4 +11,6 @@ from sky.provision.runpod.instance import terminate_instances
 from sky.provision.runpod.instance import wait_instances
 from sky.provision.runpod.volume import apply_volume
 from sky.provision.runpod.volume import delete_volume
+from sky.provision.runpod.volume import get_all_volumes_usedby
 from sky.provision.runpod.volume import get_volume_usedby
+from sky.provision.runpod.volume import map_all_volumes_usedby

sky/provision/runpod/instance.py

@@ -44,10 +44,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 = ['CREATED', 'RESTARTING']
 
     while True:
@@ -222,9 +222,10 @@ def query_instances(
     cluster_name_on_cloud: str,
     provider_config: Optional[Dict[str, Any]] = None,
     non_terminated_only: bool = True,
+    retry_if_missing: bool = False,
 ) -> Dict[str, Tuple[Optional['status_lib.ClusterStatus'], Optional[str]]]:
     """See sky/provision/__init__.py"""
-    del cluster_name  # unused
+    del cluster_name, retry_if_missing  # unused
     assert provider_config is not None, (cluster_name_on_cloud, provider_config)
     instances = _filter_instances(cluster_name_on_cloud, None)