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

sky/logs/agent.py

@@ -34,23 +34,50 @@ class FluentbitAgent(LoggingAgent):
     def get_setup_command(self,
                           cluster_name: resources_utils.ClusterName) -> str:
         install_cmd = (
-            'if ! command -v fluent-bit >/dev/null 2>&1; then '
-            'sudo apt-get install -y gnupg; '
             # pylint: disable=line-too-long
-            'curl https://raw.githubusercontent.com/fluent/fluent-bit/master/install.sh | sh; '
+            'if ! command -v fluent-bit >/dev/null 2>&1 && [ ! -f /opt/fluent-bit/bin/fluent-bit ]; then '
+            'sudo apt-get update; sudo apt-get install -y gnupg; '
+            # pylint: disable=line-too-long
+            'sudo sh -c \'curl -L https://packages.fluentbit.io/fluentbit.key | gpg --dearmor > /usr/share/keyrings/fluentbit-keyring.gpg\'; '
+            # pylint: disable=line-too-long
+            'os_id=$(grep -oP \'(?<=^ID=).*\' /etc/os-release 2>/dev/null || lsb_release -is 2>/dev/null | tr \'[:upper:]\' \'[:lower:]\'); '
+            # pylint: disable=line-too-long
+            'codename=$(grep -oP \'(?<=VERSION_CODENAME=).*\' /etc/os-release 2>/dev/null || lsb_release -cs 2>/dev/null); '
+            # pylint: disable=line-too-long
+            'echo "deb [signed-by=/usr/share/keyrings/fluentbit-keyring.gpg] https://packages.fluentbit.io/$os_id/$codename $codename main" | sudo tee /etc/apt/sources.list.d/fluent-bit.list; '
+            'sudo apt-get update; '
+            'sudo apt-get install -y fluent-bit; '
             'fi')
         cfg = self.fluentbit_config(cluster_name)
         cfg_path = os.path.join(constants.LOGGING_CONFIG_DIR, 'fluentbit.yaml')
         config_cmd = (f'mkdir -p {constants.LOGGING_CONFIG_DIR} && '
                       f'echo {shlex.quote(cfg)} > {cfg_path}')
+        kill_prior_cmd = (
+            'if [ -f "/tmp/fluentbit.pid" ]; then '
+            # pylint: disable=line-too-long
+            'echo "Killing prior fluent-bit process $(cat /tmp/fluentbit.pid)"; '
+            'kill "$(cat /tmp/fluentbit.pid)" || true; '
+            'fi')
         start_cmd = ('nohup $(command -v fluent-bit || '
                      'echo "/opt/fluent-bit/bin/fluent-bit") '
-                     f'-c {cfg_path} > /tmp/fluentbit.log 2>&1 &')
-        return f'set -e; {install_cmd}; {config_cmd}; {start_cmd}'
+                     f'-c {cfg_path} > /tmp/fluentbit.log 2>&1 & '
+                     'echo $! > /tmp/fluentbit.pid')
+        return ('set -e; '
+                f'{install_cmd}; '
+                f'{config_cmd}; '
+                f'{kill_prior_cmd}; '
+                f'{start_cmd}')
 
     def fluentbit_config(self,
                          cluster_name: resources_utils.ClusterName) -> str:
         cfg_dict = {
+            'parsers': [{
+                'name': 'sky-ray-parser',
+                'format': 'regex',
+                # pylint: disable=line-too-long
+                'regex': r'(?:\x1b\[[\d;]+m)?\((?<worker_name>[^,]+)(?:,\s*rank=(?<rank>\d+))?(?:,\s*pid=(?<pid>\d+))(?:,\s*ip=(?<ip>[\d.]+))?\)(?:\x1b\[[\d;]+m)?\s*(?<log_line>.*)',
+                'types': 'rank:integer pid:integer',
+            }],
             'pipeline': {
                 'inputs': [{
                     'name': 'tail',
@@ -62,6 +89,14 @@ class FluentbitAgent(LoggingAgent):
                     # right after the job completion.
                     'refresh_interval': 1,
                 }],
+                'filters': [{
+                    'name': 'parser',
+                    'match': '*',
+                    'key_name': 'log',
+                    'parser': 'sky-ray-parser',
+                    'preserve_key': 'on',  # preserve field for backwards compat
+                    'reserve_data': 'on',
+                }],
                 'outputs': [self.fluentbit_output_config(cluster_name)],
             }
         }

sky/logs/aws.py

@@ -5,7 +5,6 @@ from typing import Any, Dict, Optional
 import pydantic
 
 from sky.logs.agent import FluentbitAgent
-from sky.skylet import constants
 from sky.utils import resources_utils
 from sky.utils import yaml_utils
 
@@ -176,6 +175,8 @@ class CloudwatchLoggingAgent(FluentbitAgent):
         Returns:
             The Fluent Bit configuration as a YAML string.
         """
+        cfg_dict = yaml_utils.read_yaml_str(
+            super().fluentbit_config(cluster_name))
         display_name = cluster_name.display_name
         unique_name = cluster_name.name_on_cloud
         # Build tags for the log stream
@@ -197,24 +198,13 @@ class CloudwatchLoggingAgent(FluentbitAgent):
                 'value': value
             })
 
-        cfg_dict = {
-            'pipeline': {
-                'inputs': [{
-                    'name': 'tail',
-                    'path': f'{constants.SKY_LOGS_DIRECTORY}/*/*.log',
-                    'path_key': 'log_path',
-                    # Shorten the refresh interval from 60s to 1s since every
-                    # job creates a new log file and we must be responsive
-                    # for this: the VM might be autodown within a minute
-                    # right after the job completion.
-                    'refresh_interval': 1,
-                    'processors': {
-                        'logs': log_processors,
-                    }
-                }],
-                'outputs': [self.fluentbit_output_config(cluster_name)],
-            }
-        }
+        # Add log processors to config
+        processors_config = cfg_dict['pipeline']['inputs'][0].get(
+            'processors', {})
+        processors_logs_config = processors_config.get('logs', [])
+        processors_logs_config.extend(log_processors)
+        processors_config['logs'] = processors_logs_config
+        cfg_dict['pipeline']['inputs'][0]['processors'] = processors_config
 
         return yaml_utils.dump_yaml_str(cfg_dict)
 

sky/metrics/utils.py

@@ -1,11 +1,218 @@
 """Utilities for processing GPU metrics from Kubernetes clusters."""
+import contextlib
+import functools
 import os
 import re
+import select
 import subprocess
 import time
 from typing import List, Optional, Tuple
 
 import httpx
+import prometheus_client as prom
+
+from sky import sky_logging
+from sky.skylet import constants
+from sky.utils import common_utils
+from sky.utils import context_utils
+
+_SELECT_TIMEOUT = 1
+_SELECT_BUFFER_SIZE = 4096
+
+_KB = 2**10
+_MB = 2**20
+_MEM_BUCKETS = [
+    _KB,
+    256 * _KB,
+    512 * _KB,
+    _MB,
+    2 * _MB,
+    4 * _MB,
+    8 * _MB,
+    16 * _MB,
+    32 * _MB,
+    64 * _MB,
+    128 * _MB,
+    256 * _MB,
+    float('inf'),
+]
+
+logger = sky_logging.init_logger(__name__)
+
+# Whether the metrics are enabled, cannot be changed at runtime.
+METRICS_ENABLED = os.environ.get(constants.ENV_VAR_SERVER_METRICS_ENABLED,
+                                 'false').lower() == 'true'
+
+# Time spent processing a piece of code, refer to time_it().
+SKY_APISERVER_CODE_DURATION_SECONDS = prom.Histogram(
+    'sky_apiserver_code_duration_seconds',
+    'Time spent processing code',
+    ['name', 'group'],
+    buckets=(0.001, 0.005, 0.01, 0.025, 0.05, 0.075, 0.1, 0.125, 0.15, 0.25,
+             0.35, 0.5, 0.75, 1, 1.25, 1.5, 1.75, 2, 2.5, 2.75, 3, 3.5, 4, 4.5,
+             5, 7.5, 10.0, 12.5, 15.0, 17.5, 20.0, 25.0, 30.0, 35.0, 40.0, 45.0,
+             50.0, 55.0, 60.0, 80.0, 120.0, 140.0, 160.0, 180.0, 200.0, 220.0,
+             240.0, 260.0, 280.0, 300.0, 320.0, 340.0, 360.0, 380.0, 400.0,
+             420.0, 440.0, 460.0, 480.0, 500.0, 520.0, 540.0, 560.0, 580.0,
+             600.0, 620.0, 640.0, 660.0, 680.0, 700.0, 720.0, 740.0, 760.0,
+             780.0, 800.0, 820.0, 840.0, 860.0, 880.0, 900.0, 920.0, 940.0,
+             960.0, 980.0, 1000.0, float('inf')),
+)
+
+# Total number of API server requests, grouped by path, method, and status.
+SKY_APISERVER_REQUESTS_TOTAL = prom.Counter(
+    'sky_apiserver_requests_total',
+    'Total number of API server requests',
+    ['path', 'method', 'status'],
+)
+
+# Time spent processing API server requests, grouped by path, method, and
+# status.
+SKY_APISERVER_REQUEST_DURATION_SECONDS = prom.Histogram(
+    'sky_apiserver_request_duration_seconds',
+    'Time spent processing API server requests',
+    ['path', 'method', 'status'],
+    buckets=(0.001, 0.005, 0.01, 0.025, 0.05, 0.075, 0.1, 0.125, 0.15, 0.25,
+             0.35, 0.5, 0.75, 1, 1.25, 1.5, 1.75, 2, 2.5, 2.75, 3, 3.5, 4, 4.5,
+             5, 7.5, 10.0, 12.5, 15.0, 17.5, 20.0, 25.0, 30.0, 35.0, 40.0, 45.0,
+             50.0, 55.0, 60.0, 80.0, 120.0, 140.0, 160.0, 180.0, 200.0, 220.0,
+             240.0, 260.0, 280.0, 300.0, 320.0, 340.0, 360.0, 380.0, 400.0,
+             420.0, 440.0, 460.0, 480.0, 500.0, 520.0, 540.0, 560.0, 580.0,
+             600.0, 620.0, 640.0, 660.0, 680.0, 700.0, 720.0, 740.0, 760.0,
+             780.0, 800.0, 820.0, 840.0, 860.0, 880.0, 900.0, 920.0, 940.0,
+             960.0, 980.0, 1000.0, float('inf')),
+)
+
+SKY_APISERVER_EVENT_LOOP_LAG_SECONDS = prom.Histogram(
+    'sky_apiserver_event_loop_lag_seconds',
+    'Scheduling delay of the server event loop',
+    ['pid'],
+    buckets=(0.001, 0.005, 0.01, 0.025, 0.05, 0.075, 0.1, 0.125, 0.15, 0.25,
+             0.35, 0.5, 0.75, 1, 1.25, 1.5, 1.75, 2, 2.5, 2.75, 3, 3.5, 4, 4.5,
+             5, 7.5, 10.0, 12.5, 15.0, 17.5, 20.0, 25.0, 30.0, 35.0, 40.0, 45.0,
+             50.0, 55.0, 60.0, 80.0, 120.0, 140.0, 160.0, 180.0, 200.0, 220.0,
+             240.0, 260.0, 280.0, 300.0, 320.0, 340.0, 360.0, 380.0, 400.0,
+             420.0, 440.0, 460.0, 480.0, 500.0, 520.0, 540.0, 560.0, 580.0,
+             600.0, 620.0, 640.0, 660.0, 680.0, 700.0, 720.0, 740.0, 760.0,
+             780.0, 800.0, 820.0, 840.0, 860.0, 880.0, 900.0, 920.0, 940.0,
+             960.0, 980.0, 1000.0, float('inf')),
+)
+
+SKY_APISERVER_WEBSOCKET_CONNECTIONS = prom.Gauge(
+    'sky_apiserver_websocket_connections',
+    'Number of websocket connections',
+    ['pid'],
+    multiprocess_mode='livesum',
+)
+
+SKY_APISERVER_WEBSOCKET_CLOSED_TOTAL = prom.Counter(
+    'sky_apiserver_websocket_closed_total',
+    'Number of websocket closed',
+    ['pid', 'reason'],
+)
+
+# The number of execution starts in each worker process, we do not record
+# histogram here as the duration has been measured in
+# SKY_APISERVER_CODE_DURATION_SECONDS without the worker label (process id).
+# Recording histogram WITH worker label will cause high cardinality.
+SKY_APISERVER_PROCESS_EXECUTION_START_TOTAL = prom.Counter(
+    'sky_apiserver_process_execution_start_total',
+    'Total number of execution starts in each worker process',
+    ['request', 'pid'],
+)
+
+SKY_APISERVER_PROCESS_PEAK_RSS = prom.Gauge(
+    'sky_apiserver_process_peak_rss',
+    'Peak RSS we saw in each process in last 30 seconds',
+    ['pid', 'type'],
+)
+
+SKY_APISERVER_PROCESS_CPU_TOTAL = prom.Gauge(
+    'sky_apiserver_process_cpu_total',
+    'Total CPU times a worker process has been running',
+    ['pid', 'type', 'mode'],
+)
+
+SKY_APISERVER_REQUEST_MEMORY_USAGE_BYTES = prom.Histogram(
+    'sky_apiserver_request_memory_usage_bytes',
+    'Peak memory usage of requests', ['name'],
+    buckets=_MEM_BUCKETS)
+
+SKY_APISERVER_REQUEST_RSS_INCR_BYTES = prom.Histogram(
+    'sky_apiserver_request_rss_incr_bytes',
+    'RSS increment after requests', ['name'],
+    buckets=_MEM_BUCKETS)
+
+SKY_APISERVER_WEBSOCKET_SSH_LATENCY_SECONDS = prom.Histogram(
+    'sky_apiserver_websocket_ssh_latency_seconds',
+    ('Time taken for ssh message to go from client to API server and back'
+     'to the client. This does not include: latency to reach the pod, '
+     'overhead from sending through the k8s port-forward tunnel, or '
+     'ssh server lag on the destination pod.'),
+    ['pid'],
+    buckets=(0.001, 0.005, 0.01, 0.025, 0.05, 0.075, 0.1, 0.125, 0.15, 0.25,
+             0.35, 0.5, 0.75, 1, 1.25, 1.5, 1.75, 2, 2.5, 2.75, 3, 3.5, 4, 4.5,
+             5, 7.5, 10.0, 12.5, 15.0, 17.5, 20.0, 25.0, 30.0, 35.0, 40.0, 45.0,
+             50.0, 55.0, 60.0, 80.0, 120.0, 140.0, 160.0, 180.0, 200.0, 220.0,
+             240.0, 260.0, 280.0, 300.0, 320.0, 340.0, 360.0, 380.0, 400.0,
+             420.0, 440.0, 460.0, 480.0, 500.0, 520.0, 540.0, 560.0, 580.0,
+             600.0, 620.0, 640.0, 660.0, 680.0, 700.0, 720.0, 740.0, 760.0,
+             780.0, 800.0, 820.0, 840.0, 860.0, 880.0, 900.0, 920.0, 940.0,
+             960.0, 980.0, 1000.0, float('inf')),
+)
+
+SKY_APISERVER_LONG_EXECUTORS = prom.Gauge(
+    'sky_apiserver_long_executors',
+    'Total number of long-running request executors in the API server',
+)
+
+SKY_APISERVER_SHORT_EXECUTORS = prom.Gauge(
+    'sky_apiserver_short_executors',
+    'Total number of short-running request executors in the API server',
+)
+
+
+@contextlib.contextmanager
+def time_it(name: str, group: str = 'default'):
+    """Context manager to measure and record code execution duration."""
+    if not METRICS_ENABLED:
+        yield
+    else:
+        start_time = time.time()
+        try:
+            yield
+        finally:
+            duration = time.time() - start_time
+            SKY_APISERVER_CODE_DURATION_SECONDS.labels(
+                name=name, group=group).observe(duration)
+
+
+def time_me(func):
+    """Measure the duration of decorated function."""
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        if not METRICS_ENABLED:
+            return func(*args, **kwargs)
+        name = f'{func.__module__}/{func.__name__}'
+        with time_it(name, group='function'):
+            return func(*args, **kwargs)
+
+    return wrapper
+
+
+def time_me_async(func):
+    """Measure the duration of decorated async function."""
+
+    @functools.wraps(func)
+    async def async_wrapper(*args, **kwargs):
+        if not METRICS_ENABLED:
+            return await func(*args, **kwargs)
+        name = f'{func.__module__}/{func.__name__}'
+        with time_it(name, group='function'):
+            return await func(*args, **kwargs)
+
+    return async_wrapper
 
 
 def start_svc_port_forward(context: str, namespace: str, service: str,
@@ -34,46 +241,72 @@ def start_svc_port_forward(context: str, namespace: str, service: str,
     if 'KUBECONFIG' not in env:
         env['KUBECONFIG'] = os.path.expanduser('~/.kube/config')
 
-    # start the port forward process
-    port_forward_process = subprocess.Popen(cmd,
-                                            stdout=subprocess.PIPE,
-                                            stderr=subprocess.STDOUT,
-                                            text=True,
-                                            env=env)
-
+    port_forward_process = None
+    port_forward_exit = False
     local_port = None
-    start_time = time.time()
-
-    # wait for the port forward to start and extract the local port
-    while time.time() - start_time < start_port_forward_timeout:
-        if port_forward_process.poll() is not None:
-            # port forward process has terminated
-            if port_forward_process.returncode != 0:
-                raise RuntimeError(
-                    f'Port forward failed for service {service} in namespace '
-                    f'{namespace} on context {context}')
-            break
-
-        # read output line by line to find the local port
-        if port_forward_process.stdout:
-            line = port_forward_process.stdout.readline()
-            if line:
-                # look for 'Forwarding from 127.0.0.1:XXXXX -> service_port'
-                match = re.search(r'Forwarding from 127\.0\.0\.1:(\d+)', line)
+    poller = None
+    fd = None
+
+    try:
+        # start the port forward process
+        port_forward_process = subprocess.Popen(cmd,
+                                                stdout=subprocess.PIPE,
+                                                stderr=subprocess.STDOUT,
+                                                text=True,
+                                                env=env)
+
+        # Use poll() instead of select() to avoid FD_SETSIZE limit
+        poller = select.poll()
+        assert port_forward_process.stdout is not None
+        fd = port_forward_process.stdout.fileno()
+        poller.register(fd, select.POLLIN)
+
+        start_time = time.time()
+        buffer = ''
+        # wait for the port forward to start and extract the local port
+        while time.time() - start_time < start_port_forward_timeout:
+            if port_forward_process.poll() is not None:
+                # port forward process has terminated
+                if port_forward_process.returncode != 0:
+                    port_forward_exit = True
+                break
+
+            # Wait up to 1000ms for data to be available without blocking
+            # poll() takes timeout in milliseconds
+            events = poller.poll(_SELECT_TIMEOUT * 1000)
+
+            if events:
+                # Read available bytes from the FD without blocking
+                raw = os.read(fd, _SELECT_BUFFER_SIZE)
+                chunk = raw.decode(errors='ignore')
+                buffer += chunk
+                match = re.search(r'Forwarding from 127\.0\.0\.1:(\d+)', buffer)
                 if match:
                     local_port = int(match.group(1))
                     break
 
-        # sleep for 100ms to avoid busy-waiting
-        time.sleep(0.1)
-
+            # sleep for 100ms to avoid busy-waiting
+            time.sleep(0.1)
+    except BaseException:  # pylint: disable=broad-exception-caught
+        if port_forward_process:
+            stop_svc_port_forward(port_forward_process,
+                                  timeout=terminate_port_forward_timeout)
+        raise
+    finally:
+        if poller is not None and fd is not None:
+            try:
+                poller.unregister(fd)
+            except (OSError, ValueError):
+                # FD may already be unregistered or invalid
+                pass
+    if port_forward_exit:
+        raise RuntimeError(f'Port forward failed for service {service} in '
+                           f'namespace {namespace} on context {context}')
     if local_port is None:
         try:
-            port_forward_process.terminate()
-            port_forward_process.wait(timeout=terminate_port_forward_timeout)
-        except subprocess.TimeoutExpired:
-            port_forward_process.kill()
-            port_forward_process.wait()
+            if port_forward_process:
+                stop_svc_port_forward(port_forward_process,
+                                      timeout=terminate_port_forward_timeout)
         finally:
             raise RuntimeError(
                 f'Failed to extract local port for service {service} in '
@@ -82,14 +315,15 @@ def start_svc_port_forward(context: str, namespace: str, service: str,
     return port_forward_process, local_port
 
 
-def stop_svc_port_forward(port_forward_process: subprocess.Popen) -> None:
+def stop_svc_port_forward(port_forward_process: subprocess.Popen,
+                          timeout: int = 5) -> None:
     """Stops a port forward to a service in a Kubernetes cluster.
     Args:
         port_forward_process: The subprocess.Popen process to terminate
     """
     try:
         port_forward_process.terminate()
-        port_forward_process.wait(timeout=5)
+        port_forward_process.wait(timeout=timeout)
     except subprocess.TimeoutExpired:
         port_forward_process.kill()
         port_forward_process.wait()
@@ -122,8 +356,8 @@ async def send_metrics_request_with_port_forward(
     port_forward_process = None
     try:
         # Start port forward
-        port_forward_process, local_port = start_svc_port_forward(
-            context, namespace, service, service_port)
+        port_forward_process, local_port = await context_utils.to_thread(
+            start_svc_port_forward, context, namespace, service, service_port)
 
         # Build endpoint URL
         endpoint = f'http://localhost:{local_port}{endpoint_path}'
@@ -140,10 +374,15 @@ async def send_metrics_request_with_port_forward(
             response.raise_for_status()
             return response.text
 
+    except Exception as e:  # pylint: disable=broad-exception-caught
+        logger.error(f'Failed to send metrics request with port forward: '
+                     f'{common_utils.format_exception(e)}')
+        raise
     finally:
         # Always clean up port forward
         if port_forward_process:
-            stop_svc_port_forward(port_forward_process)
+            await context_utils.to_thread(stop_svc_port_forward,
+                                          port_forward_process)
 
 
 async def add_cluster_name_label(metrics_text: str, context: str) -> str:
@@ -193,7 +432,11 @@ async def get_metrics_for_context(context: str) -> str:
     """
     # Query both DCGM metrics and kube_pod_labels metrics
     # This ensures the dashboard can perform joins to filter by skypilot cluster
-    match_patterns = ['{__name__=~"DCGM_.*"}', 'kube_pod_labels']
+    match_patterns = [
+        '{__name__=~"node_memory_MemAvailable_bytes|node_memory_MemTotal_bytes|DCGM_.*"}',  # pylint: disable=line-too-long
+        'kube_pod_labels',
+        'node_cpu_seconds_total{mode="idle"}'
+    ]
 
     # TODO(rohan): don't hardcode the namespace and service name
     metrics_text = await send_metrics_request_with_port_forward(

sky/models.py

@@ -68,6 +68,8 @@ class KubernetesNodeInfo:
     free: Dict[str, int]
     # IP address of the node (external IP preferred, fallback to internal IP)
     ip_address: Optional[str] = None
+    # Whether the node is ready (all conditions are satisfied)
+    is_ready: bool = True
 
 
 @dataclasses.dataclass

sky/optimizer.py

@@ -781,7 +781,7 @@ class Optimizer:
         def _instance_type_str(resources: 'resources_lib.Resources') -> str:
             instance_type = resources.instance_type
             assert instance_type is not None, 'Instance type must be specified'
-            if isinstance(resources.cloud, clouds.Kubernetes):
+            if isinstance(resources.cloud, (clouds.Kubernetes, clouds.Slurm)):
                 instance_type = '-'
                 if resources.use_spot:
                     instance_type = ''
@@ -865,11 +865,12 @@ class Optimizer:
                 'use_spot': resources.use_spot
             }
 
-            # Handle special case for Kubernetes and SSH clouds
-            if isinstance(resources.cloud, clouds.Kubernetes):
+            # Handle special case for Kubernetes, SSH, and SLURM clouds
+            if isinstance(resources.cloud, (clouds.Kubernetes, clouds.Slurm)):
                 # Region for Kubernetes-like clouds (SSH, Kubernetes) is the
-                # context name, i.e. different Kubernetes clusters. We add
-                # region to the key to show all the Kubernetes clusters in the
+                # context name, i.e. different Kubernetes clusters.
+                # Region for SLURM is the cluster name.
+                # We add region to the key to show all the clusters in the
                 # optimizer table for better UX.
 
                 if resources.cloud.__class__.__name__ == 'SSH':
@@ -1019,7 +1020,7 @@ class Optimizer:
                             if res.instance_type is not None
                         ])
                         candidate_str = resources_utils.format_resource(
-                            best_resources, simplify=True)
+                            best_resources, simplified_only=True)[0]
 
                         logger.info(
                             f'{colorama.Style.DIM}🔍 Multiple {cloud} instances '

sky/provision/__init__.py

@@ -24,8 +24,12 @@ 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 primeintellect
 from sky.provision import runpod
 from sky.provision import scp
+from sky.provision import seeweb
+from sky.provision import shadeform
+from sky.provision import slurm
 from sky.provision import ssh
 from sky.provision import vast
 from sky.provision import vsphere
@@ -77,6 +81,7 @@ 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]]]:
     """Query instances.
 
@@ -85,6 +90,11 @@ def query_instances(
 
     A None status means the instance is marked as "terminated"
     or "terminating".
+
+    Args:
+        retry_if_missing: Whether to retry the call to the cloud api if the
+          cluster is not found when querying the live status on the cloud.
+          NOTE: This is currently only used on kubernetes.
     """
     raise NotImplementedError
 
@@ -140,7 +150,34 @@ def get_volume_usedby(
 
 
 @_route_to_cloud_impl
-def run_instances(provider_name: str, region: str, cluster_name_on_cloud: str,
+def get_all_volumes_usedby(
+    provider_name: str, configs: List[models.VolumeConfig]
+) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    """Get the usedby of a volume.
+
+    Returns:
+        usedby_pods: List of dictionaries, each containing the config keys for
+                     a volume and a key containing pods using the volume.
+                     These may include pods not created by SkyPilot.
+        usedby_clusters: List of dictionaries, each containing the config keys
+                         for a volume and a key containing clusters using
+                         the volume.
+    """
+    raise NotImplementedError
+
+
+@_route_to_cloud_impl
+def map_all_volumes_usedby(
+        provider_name: str, used_by_pods: Dict[str, Any],
+        used_by_clusters: Dict[str, Any],
+        config: models.VolumeConfig) -> Tuple[List[str], List[str]]:
+    """Map the usedby resources of a volume."""
+    raise NotImplementedError
+
+
+@_route_to_cloud_impl
+def run_instances(provider_name: str, region: str, cluster_name: str,
+                  cluster_name_on_cloud: str,
                   config: common.ProvisionConfig) -> common.ProvisionRecord:
     """Start instances with bootstrapped configuration."""
     raise NotImplementedError