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

sky/provision/hyperbolic/instance.py

@@ -0,0 +1,437 @@
+"""Hyperbolic instance provisioning."""
+import time
+from typing import Any, Dict, List, Optional, Tuple
+
+from sky import sky_logging
+from sky.provision import common
+from sky.provision.hyperbolic import utils
+from sky.utils import status_lib
+
+PROVIDER_NAME = 'hyperbolic'
+POLL_INTERVAL = 5
+QUERY_PORTS_TIMEOUT_SECONDS = 30
+#TODO come up with a reasonable value for this timeout
+TIMEOUT = 300
+
+logger = sky_logging.init_logger(__name__)
+
+
+def _filter_instances(cluster_name_on_cloud: str,
+                      status_filters: Optional[List[str]],
+                      head_only: bool = False) -> Dict[str, Dict[str, Any]]:
+    logger.debug(f'Filtering instances: cluster={cluster_name_on_cloud}, '
+                 f'status={status_filters}')
+    _ = head_only  # Mark as intentionally unused
+
+    # Filter by cluster name using metadata
+    instances = utils.list_instances(
+        metadata={'skypilot': {
+            'cluster_name': cluster_name_on_cloud
+        }})
+
+    # Normalize status filters to lowercase
+    if status_filters is not None:
+        status_filters = [s.lower() for s in status_filters]
+
+    filtered_instances: Dict[str, Dict[str, Any]] = {}
+    for instance_id, instance in instances.items():
+        try:
+            # Check status filter
+            instance_status = instance.get('status', '').lower()
+            if (status_filters is not None and
+                    instance_status not in status_filters):
+                logger.debug(
+                    f'Skipping instance {instance_id} '
+                    f'- status {instance_status} not in {status_filters}')
+                continue
+
+            filtered_instances[instance_id] = instance
+            logger.debug(f'Including instance {instance_id} '
+                         f'with status {instance_status}')
+
+        except Exception as e:  # pylint: disable=broad-except
+            logger.warning(f'Error processing instance {instance_id}: {str(e)}')
+            continue
+
+    logger.info(f'Found {len(filtered_instances)} instances matching filters')
+    return filtered_instances
+
+
+def _get_head_instance_id(instances: Dict[str, Any]) -> Optional[str]:
+    """Get the instance ID from the instances dict."""
+    if not instances:
+        return None
+    return next(iter(instances.keys()))
+
+
+def run_instances(region: str, cluster_name: str, cluster_name_on_cloud: str,
+                  config: common.ProvisionConfig) -> common.ProvisionRecord:
+    del cluster_name  # unused
+    logger.info(f'Starting run_instances with region={region}, '
+                f'cluster={cluster_name_on_cloud}')
+    logger.debug(f'Config: {config}')
+    start_time = time.time()
+
+    # Define pending statuses for Hyperbolic
+    pending_status = [
+        utils.HyperbolicInstanceStatus.CREATING.value,
+        utils.HyperbolicInstanceStatus.STARTING.value
+    ]
+    logger.debug(
+        f'Looking for instances with pending statuses: {pending_status}')
+
+    # Wait for any pending instances to be ready
+    while True:
+        if time.time() - start_time > TIMEOUT:
+            logger.error(
+                f'Timed out after {TIMEOUT}s waiting for instances to be ready')
+            raise TimeoutError(
+                f'Timed out after {TIMEOUT}s waiting for instances to be ready')
+
+        instances = _filter_instances(cluster_name_on_cloud, pending_status)
+        logger.debug(f'Found {len(instances)} instances with pending status')
+        if not instances:
+            break
+        logger.info(
+            f'Waiting for instance to be ready. Current instances: {instances}')
+        time.sleep(POLL_INTERVAL)
+
+    # Check existing running instance
+    logger.info('Checking for existing running instances')
+    exist_instances = _filter_instances(
+        cluster_name_on_cloud, [utils.HyperbolicInstanceStatus.ONLINE.value])
+    logger.debug(
+        f'Found {len(exist_instances)} running instances: {exist_instances}')
+    instance_id = _get_head_instance_id(exist_instances)
+    logger.debug(f'Head instance ID: {instance_id}')
+
+    # Calculate if we need to start a new instance
+    to_start_count = 1 - len(exist_instances)  # Always 1 for single node
+    logger.info(f'Need to start {to_start_count} new instances')
+    if to_start_count < 0:
+        logger.error(
+            f'Cluster {cluster_name_on_cloud} already has an instance running')
+        raise RuntimeError(
+            f'Cluster {cluster_name_on_cloud} already has an instance running.')
+    if to_start_count == 0:
+        if instance_id is None:
+            logger.error(
+                f'Cluster {cluster_name_on_cloud} has no running instance')
+            raise RuntimeError(
+                f'Cluster {cluster_name_on_cloud} has no running instance.')
+        logger.info(
+            f'Cluster {cluster_name_on_cloud} already has a running instance')
+        return common.ProvisionRecord(provider_name=PROVIDER_NAME,
+                                      cluster_name=cluster_name_on_cloud,
+                                      region='default',
+                                      zone=None,
+                                      head_instance_id=instance_id,
+                                      resumed_instance_ids=[],
+                                      created_instance_ids=[])
+
+    try:
+        # Get instance type from node_config
+        instance_type = config.node_config.get('InstanceType')
+        logger.debug(f'Instance type from config: {instance_type}')
+        if not instance_type:
+            logger.error('InstanceType is not set in node_config')
+            raise RuntimeError(
+                'InstanceType is not set in node_config. '
+                'Please specify an instance type for Hyperbolic.')
+
+        # Parse gpu_model configuration from instance type
+        # Format: {gpu_count}x-{gpu_model}-{cpu}-{memory}
+        # Example: 1x-A100-24-271
+        try:
+            parts = instance_type.split('-')
+            if len(parts) != 4:
+                raise ValueError(
+                    f'Invalid instance type format: {instance_type}. '
+                    'Expected format: {gpu_count}x-{gpu_model}-{cpu}-{memory}')
+
+            gpu_count = int(parts[0].rstrip('x'))
+            gpu_model = parts[1]
+            logger.info(f'Parsed GPU config from instance type: '
+                        f'model={gpu_model}, count={gpu_count}')
+
+            # Launch instance
+            instance_id, ssh_command = utils.launch_instance(
+                gpu_model, gpu_count, cluster_name_on_cloud)
+            logger.info(f'Launched instance {instance_id} with SSH command: '
+                        f'{ssh_command}')
+            created_instance_ids = [instance_id]
+
+            # Wait for instance to be ready
+            if not utils.wait_for_instance(
+                    instance_id, utils.HyperbolicInstanceStatus.ONLINE.value):
+                raise RuntimeError(
+                    f'Instance {instance_id} failed to reach ONLINE state')
+
+        except ValueError as e:
+            logger.error(f'Failed to parse instance type: {e}')
+            raise RuntimeError(str(e)) from e
+        except Exception as e:
+            logger.error(f'Failed to launch instance: {e}')
+            raise RuntimeError(str(e)) from e
+
+    except Exception as e:
+        logger.error(f'Unexpected error: {e}')
+        raise
+
+    # Wait for instance to be ready
+    logger.info(f'Waiting for instance {instance_id} to be ready')
+    while True:
+        instances = _filter_instances(
+            cluster_name_on_cloud,
+            [utils.HyperbolicInstanceStatus.ONLINE.value])
+        logger.debug(f'Current instances: {instances}')
+        if len(instances) == 1:
+            logger.info(f'Instance {instance_id} is ready')
+            break
+        if time.time() - start_time > TIMEOUT:
+            logger.error(
+                f'Timed out after {TIMEOUT}s waiting for instance to be ready')
+            raise TimeoutError(
+                f'Timed out after {TIMEOUT}s waiting for instance to be ready')
+        logger.info('Waiting for instance to be ready...')
+        time.sleep(POLL_INTERVAL)
+
+    logger.info(f'Returning ProvisionRecord for instance {instance_id}')
+    return common.ProvisionRecord(provider_name=PROVIDER_NAME,
+                                  cluster_name=cluster_name_on_cloud,
+                                  region='default',
+                                  zone=None,
+                                  head_instance_id=instance_id,
+                                  resumed_instance_ids=[],
+                                  created_instance_ids=created_instance_ids)
+
+
+def terminate_instances(
+    cluster_name_on_cloud: str,
+    provider_config: Optional[dict] = None,
+    worker_only: bool = False,
+) -> None:
+    """Terminate all instances in the cluster."""
+    del provider_config, worker_only  # unused
+    logger.info(
+        f'Terminating all instances for cluster {cluster_name_on_cloud}')
+
+    # First check if instances exist
+    instances = _filter_instances(cluster_name_on_cloud, None)
+    if not instances:
+        logger.info(f'No instances found for cluster {cluster_name_on_cloud}')
+        return
+
+    # Terminate each instance
+    for instance_id in instances:
+        try:
+            utils.terminate_instance(instance_id)
+            logger.info(f'Terminated instance {instance_id}')
+        except Exception as e:  # pylint: disable=broad-except
+            logger.warning(f'Failed to terminate instance {instance_id}: {e}')
+            continue
+
+    # Wait for instances to be terminated
+    start_time = time.time()
+    while True:
+        if time.time() - start_time > TIMEOUT:
+            logger.error(
+                f'Timed out after {TIMEOUT}s waiting for instances to terminate'
+            )
+            break
+
+        instances = _filter_instances(
+            cluster_name_on_cloud,
+            [utils.HyperbolicInstanceStatus.TERMINATED.value])
+        if not instances:
+            logger.info('All instances terminated successfully')
+            break
+
+        logger.info('Waiting for instances to terminate...')
+        time.sleep(POLL_INTERVAL)
+
+
+def get_cluster_info(
+        region: str,
+        cluster_name_on_cloud: str,
+        provider_config: Optional[Dict[str, Any]] = None) -> common.ClusterInfo:
+    """Returns information about the cluster."""
+    del region  # unused
+    running_instances = _filter_instances(
+        cluster_name_on_cloud, [utils.HyperbolicInstanceStatus.ONLINE.value])
+    instances: Dict[str, List[common.InstanceInfo]] = {}
+    head_instance_id = None
+
+    for instance_id, instance_info in running_instances.items():
+        # Extract hostname and port from sshCommand
+        ssh_command = instance_info.get('sshCommand', '')
+        if ssh_command:
+            # Format: ssh user@hostname -p port
+            parts = ssh_command.split()
+            if len(parts) >= 4:
+                user_host = parts[1]  # user@hostname
+                if '@' in user_host:
+                    ssh_user = user_host.split('@')[0]
+                    hostname = user_host.split('@')[1]
+                else:
+                    hostname = user_host
+                port = int(parts[3])
+            else:
+                hostname = instance_id
+                port = 22
+        else:
+            hostname = instance_id
+            port = 22
+
+        instances[instance_id] = [
+            common.InstanceInfo(
+                instance_id=instance_id,
+                internal_ip=hostname,
+                external_ip=hostname,
+                ssh_port=port,
+                tags={},
+            )
+        ]
+        if head_instance_id is None:
+            head_instance_id = instance_id
+
+    return common.ClusterInfo(
+        instances=instances,
+        head_instance_id=head_instance_id,
+        provider_name=PROVIDER_NAME,
+        provider_config=provider_config,
+        ssh_user=ssh_user,
+    )
+
+
+def query_instances(
+    cluster_name: str,
+    cluster_name_on_cloud: str,
+    provider_config: Optional[dict] = None,
+    non_terminated_only: bool = True,
+    retry_if_missing: bool = False,
+) -> Dict[str, Tuple[Optional['status_lib.ClusterStatus'], Optional[str]]]:
+    """Returns the status of the specified instances for Hyperbolic."""
+    del cluster_name, provider_config, retry_if_missing  # unused
+    # Fetch all instances for this cluster
+    instances = utils.list_instances(
+        metadata={'skypilot': {
+            'cluster_name': cluster_name_on_cloud
+        }})
+    if not instances:
+        # No instances found: return empty dict to indicate fully deleted
+        return {}
+
+    statuses: Dict[str, Tuple[Optional['status_lib.ClusterStatus'],
+                              Optional[str]]] = {}
+    for instance_id, instance in instances.items():
+        try:
+            raw_status = instance.get('status', 'unknown').lower()
+            hyperbolic_status = utils.HyperbolicInstanceStatus.from_raw_status(
+                raw_status)
+            status = hyperbolic_status.to_cluster_status()
+            if non_terminated_only and status is None:
+                continue
+            statuses[instance_id] = (status, None)
+        except utils.HyperbolicError as e:
+            logger.warning(
+                f'Failed to parse status for instance {instance_id}: {e}')
+            continue
+    return statuses
+
+
+def wait_instances(region: str, cluster_name_on_cloud: str,
+                   state: Optional[status_lib.ClusterStatus]) -> None:
+    """Wait for instances to reach the desired state."""
+    del region  # unused
+    if state == status_lib.ClusterStatus.UP:
+        # Check if any instances are in ONLINE state
+        instances = _filter_instances(
+            cluster_name_on_cloud,
+            [utils.HyperbolicInstanceStatus.ONLINE.value])
+        if not instances:
+            # Check if any instances are in a failed state
+            failed_instances = _filter_instances(cluster_name_on_cloud, [
+                utils.HyperbolicInstanceStatus.FAILED.value,
+                utils.HyperbolicInstanceStatus.ERROR.value
+            ])
+            if failed_instances:
+                raise RuntimeError(
+                    f'Cluster {cluster_name_on_cloud} has failed instances: '
+                    f'{failed_instances}')
+            raise RuntimeError(f'No running instances found for cluster '
+                               f'{cluster_name_on_cloud}')
+        # Check if any instances are in TERMINATED state
+        terminated_instances = _filter_instances(
+            cluster_name_on_cloud,
+            [utils.HyperbolicInstanceStatus.TERMINATED.value])
+        if terminated_instances:
+            error_msg = (
+                f'Cluster {cluster_name_on_cloud} is in UP state, but '
+                f'{len(terminated_instances)} instances are terminated.')
+            raise RuntimeError(error_msg)
+    elif state == status_lib.ClusterStatus.STOPPED:
+        # Check if any instances are in TERMINATED state
+        instances = _filter_instances(
+            cluster_name_on_cloud,
+            [utils.HyperbolicInstanceStatus.TERMINATED.value])
+        if not instances:
+            # Check if any instances are in a failed state
+            failed_instances = _filter_instances(cluster_name_on_cloud, [
+                utils.HyperbolicInstanceStatus.FAILED.value,
+                utils.HyperbolicInstanceStatus.ERROR.value
+            ])
+            if failed_instances:
+                raise RuntimeError(
+                    f'Cluster {cluster_name_on_cloud} has failed instances: '
+                    f'{failed_instances}')
+            raise RuntimeError(f'No terminated instances found for cluster '
+                               f'{cluster_name_on_cloud}')
+        # Check if any instances are in ONLINE state
+        running_instances = _filter_instances(
+            cluster_name_on_cloud,
+            [utils.HyperbolicInstanceStatus.ONLINE.value])
+        if running_instances:
+            error_msg = (
+                f'Cluster {cluster_name_on_cloud} is in STOPPED state, but '
+                f'{len(running_instances)} instances are running.')
+            raise RuntimeError(error_msg)
+    else:
+        raise RuntimeError(f'Unsupported state: {state}')
+
+
+def stop_instances(
+    cluster_name_on_cloud: str,
+    provider_config: Optional[Dict[str, Any]] = None,
+    worker_only: bool = False,
+) -> None:
+    """Stop running instances. Not supported for Hyperbolic."""
+    raise NotImplementedError('stop_instances is not supported for Hyperbolic')
+
+
+def cleanup_ports(
+    cluster_name_on_cloud: str,
+    provider_config: Optional[dict] = None,
+    ports: Optional[list] = None,
+) -> None:
+    """Cleanup ports. Not supported for Hyperbolic."""
+    raise NotImplementedError('cleanup_ports is not supported for Hyperbolic')
+
+
+def cleanup_custom_multi_network(
+    cluster_name_on_cloud: str,
+    provider_config: Dict[str, Any],
+    failover: bool = False,
+) -> None:
+    """Cleanup custom multi-network. Not supported for Hyperbolic."""
+    raise NotImplementedError(
+        'cleanup_custom_multi_network is not supported for Hyperbolic')
+
+
+def open_ports(
+    cluster_name_on_cloud: str,
+    ports: list,
+    provider_config: Optional[dict] = None,
+) -> None:
+    """Open ports. Not supported for Hyperbolic."""
+    raise NotImplementedError('open_ports is not supported for Hyperbolic')