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

sky/core.py

@@ -1,6 +1,4 @@
 """SDK functions for cluster/job management."""
-import os
-import shlex
 import typing
 from typing import Any, Dict, List, Optional, Tuple, Union
 
@@ -9,7 +7,6 @@ import colorama
 from sky import admin_policy
 from sky import backends
 from sky import catalog
-from sky import check as sky_check
 from sky import clouds
 from sky import dag as dag_lib
 from sky import data
@@ -20,16 +17,18 @@ from sky import optimizer
 from sky import sky_logging
 from sky import skypilot_config
 from sky import task as task_lib
+from sky.adaptors import common as adaptors_common
 from sky.backends import backend_utils
+from sky.backends import cloud_vm_ray_backend
 from sky.clouds import cloud as sky_cloud
 from sky.jobs.server import core as managed_jobs_core
 from sky.provision.kubernetes import constants as kubernetes_constants
 from sky.provision.kubernetes import utils as kubernetes_utils
 from sky.schemas.api import responses
+from sky.server.requests import request_names
 from sky.skylet import autostop_lib
 from sky.skylet import constants
 from sky.skylet import job_lib
-from sky.skylet import log_lib
 from sky.usage import usage_lib
 from sky.utils import admin_policy_utils
 from sky.utils import common
@@ -44,6 +43,9 @@ from sky.utils.kubernetes import kubernetes_deploy_utils
 
 if typing.TYPE_CHECKING:
     from sky import resources as resources_lib
+    from sky.schemas.generated import jobsv1_pb2
+else:
+    jobsv1_pb2 = adaptors_common.LazyImport('sky.schemas.generated.jobsv1_pb2')
 
 logger = sky_logging.init_logger(__name__)
 
@@ -83,7 +85,9 @@ def optimize(
     # but we do not apply the admin policy there. We should apply the admin
     # policy in the optimizer, but that will require some refactoring.
     with admin_policy_utils.apply_and_use_config_in_current_request(
-            dag, request_options=request_options) as dag:
+            dag,
+            request_name=request_names.AdminPolicyRequestName.OPTIMIZE,
+            request_options=request_options) as dag:
         dag.resolve_and_validate_volumes()
         return optimizer.Optimizer.optimize(dag=dag,
                                             minimize=minimize,
@@ -97,6 +101,8 @@ def status(
     refresh: common.StatusRefreshMode = common.StatusRefreshMode.NONE,
     all_users: bool = False,
     include_credentials: bool = False,
+    summary_response: bool = False,
+    include_handle: bool = True,
 ) -> List[responses.StatusResponse]:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
     """Gets cluster statuses.
@@ -176,16 +182,25 @@ def status(
         refresh=refresh,
         cluster_names=cluster_names,
         all_users=all_users,
-        include_credentials=include_credentials)
-    return [
-        responses.StatusResponse.model_validate(cluster) for cluster in clusters
-    ]
+        include_credentials=include_credentials,
+        summary_response=summary_response,
+        include_handle=include_handle)
+
+    status_responses = []
+    for cluster in clusters:
+        try:
+            status_responses.append(
+                responses.StatusResponse.model_validate(cluster))
+        except Exception as e:  # pylint: disable=broad-except
+            logger.warning('Failed to validate status responses for cluster '
+                           f'{cluster.get("name")}: {e}')
+    return status_responses
 
 
 def status_kubernetes(
 ) -> Tuple[List['kubernetes_utils.KubernetesSkyPilotClusterInfoPayload'],
            List['kubernetes_utils.KubernetesSkyPilotClusterInfoPayload'],
-           List[Dict[str, Any]], Optional[str]]:
+           List[responses.ManagedJobRecord], Optional[str]]:
     """Gets all SkyPilot clusters and jobs in the Kubernetes cluster.
 
     Managed jobs and services are also included in the clusters returned.
@@ -260,6 +275,7 @@ all_clusters, unmanaged_clusters, all_jobs, context
         kubernetes_utils.KubernetesSkyPilotClusterInfoPayload.from_cluster(c)
         for c in unmanaged_clusters
     ]
+    all_jobs = [responses.ManagedJobRecord(**job) for job in all_jobs]
     return all_clusters, unmanaged_clusters, all_jobs, context
 
 
@@ -288,7 +304,10 @@ def endpoints(cluster: str,
 
 
 @usage_lib.entrypoint
-def cost_report(days: Optional[int] = None) -> List[Dict[str, Any]]:
+def cost_report(
+        days: Optional[int] = None,
+        dashboard_summary_response: bool = False,
+        cluster_hashes: Optional[List[str]] = None) -> List[Dict[str, Any]]:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
     """Get all cluster cost reports, including those that have been downed.
 
@@ -334,7 +353,12 @@ def cost_report(days: Optional[int] = None) -> List[Dict[str, Any]]:
     if days is None:
         days = constants.COST_REPORT_DEFAULT_DAYS
 
-    cluster_reports = global_user_state.get_clusters_from_history(days=days)
+    abbreviate_response = dashboard_summary_response and cluster_hashes is None
+
+    cluster_reports = global_user_state.get_clusters_from_history(
+        days=days,
+        abbreviate_response=abbreviate_response,
+        cluster_hashes=cluster_hashes)
     logger.debug(
         f'{len(cluster_reports)} clusters found from history with {days} days.')
 
@@ -352,43 +376,6 @@ def cost_report(days: Optional[int] = None) -> List[Dict[str, Any]]:
             cost = (launched_resources.get_cost(duration) * launched_nodes)
             return cost
 
-        def _update_record_with_resources(record: Dict[str, Any]) -> None:
-            """Add resource fields for dashboard compatibility."""
-            if record is None:
-                return
-            resources = record.get('resources')
-            if resources is None:
-                return
-            fields = ['cloud', 'region', 'cpus', 'memory', 'accelerators']
-            for field in fields:
-                try:
-                    record[field] = str(getattr(resources, field))
-                except Exception as e:  # pylint: disable=broad-except
-                    # Ok to skip the fields as this is just for display
-                    # purposes.
-                    logger.debug(f'Failed to get resources.{field} for cluster '
-                                 f'{record["name"]}: {str(e)}')
-                    record[field] = None
-
-            # Add resources_str and resources_str_full for dashboard
-            # compatibility
-            num_nodes = record.get('num_nodes', 1)
-            try:
-                resource_str_simple = resources_utils.format_resource(
-                    resources, simplify=True)
-                resource_str_full = resources_utils.format_resource(
-                    resources, simplify=False)
-                record['resources_str'] = f'{num_nodes}x{resource_str_simple}'
-                record[
-                    'resources_str_full'] = f'{num_nodes}x{resource_str_full}'
-            except Exception as e:  # pylint: disable=broad-except
-                logger.debug(f'Failed to get resources_str for cluster '
-                             f'{record["name"]}: {str(e)}')
-                for field in fields:
-                    record[field] = None
-                record['resources_str'] = '-'
-                record['resources_str_full'] = '-'
-
         try:
             report['total_cost'] = get_total_cost(report)
         except Exception as e:  # pylint: disable=broad-except
@@ -397,17 +384,62 @@ def cost_report(days: Optional[int] = None) -> List[Dict[str, Any]]:
                            f'{report["name"]}: {str(e)}')
             report['total_cost'] = 0.0
 
-        _update_record_with_resources(report)
         return report
 
     # Process clusters in parallel
     if not cluster_reports:
         return []
 
-    processed_reports = subprocess_utils.run_in_parallel(
-        _process_cluster_report, cluster_reports)
+    if not abbreviate_response:
+        cluster_reports = subprocess_utils.run_in_parallel(
+            _process_cluster_report, cluster_reports)
+
+    def _update_record_with_resources(record: Dict[str, Any]) -> None:
+        """Add resource fields for dashboard compatibility."""
+        if record is None:
+            return
+        resources = record.get('resources')
+        if resources is None:
+            return
+        if not dashboard_summary_response:
+            fields = ['cloud', 'region', 'cpus', 'memory', 'accelerators']
+        else:
+            fields = ['cloud']
+        for field in fields:
+            try:
+                record[field] = str(getattr(resources, field))
+            except Exception as e:  # pylint: disable=broad-except
+                # Ok to skip the fields as this is just for display
+                # purposes.
+                logger.debug(f'Failed to get resources.{field} for cluster '
+                             f'{record["name"]}: {str(e)}')
+                record[field] = None
 
-    return processed_reports
+        # Add resources_str and resources_str_full for dashboard
+        # compatibility
+        num_nodes = record.get('num_nodes', 1)
+        try:
+            resource_str_simple, resource_str_full = (
+                resources_utils.format_resource(resources,
+                                                simplified_only=False))
+            record['resources_str'] = f'{num_nodes}x{resource_str_simple}'
+            record['resources_str_full'] = f'{num_nodes}x{resource_str_full}'
+        except Exception as e:  # pylint: disable=broad-except
+            logger.debug(f'Failed to get resources_str for cluster '
+                         f'{record["name"]}: {str(e)}')
+            for field in fields:
+                record[field] = None
+            record['resources_str'] = '-'
+            record['resources_str_full'] = '-'
+
+    for report in cluster_reports:
+        _update_record_with_resources(report)
+        if dashboard_summary_response:
+            report.pop('usage_intervals')
+            report.pop('user_hash')
+            report.pop('resources')
+
+    return cluster_reports
 
 
 def _start(
@@ -466,6 +498,32 @@ def _start(
                 controller_autostop_config.enabled):
             idle_minutes_to_autostop = controller_autostop_config.idle_minutes
             down = controller_autostop_config.down
+    else:
+        # For non-controller clusters, restore autostop configuration from
+        # database if not explicitly provided.
+        if idle_minutes_to_autostop is None:
+            cluster_record = global_user_state.get_cluster_from_name(
+                cluster_name, include_user_info=False, summary_response=True)
+            if cluster_record is not None:
+                stored_autostop = cluster_record.get('autostop', -1)
+                stored_to_down = cluster_record.get('to_down', False)
+                # Restore autostop if it was previously set (autostop > 0)
+                if stored_autostop > 0:
+                    logger.warning(f'Restoring cluster {cluster_name!r} with '
+                                   f'autostop set to {stored_autostop} minutes'
+                                   f'. To turn off autostop, run: '
+                                   f'`sky autostop {cluster_name} --cancel`')
+                    idle_minutes_to_autostop = stored_autostop
+                    # Only restore 'down' if it was explicitly set and we're
+                    # restoring autostop
+                    if stored_to_down:
+                        down = stored_to_down
+                elif stored_autostop == 0:
+                    logger.warning(
+                        f'Autostop was previously set to 0 minutes '
+                        f'for cluster {cluster_name!r} so it will '
+                        'not be restored. To turn on autostop, run: '
+                        f'`sky autostop {cluster_name} -i <minutes>`')
 
     usage_lib.record_cluster_name_for_current_operation(cluster_name)
 
@@ -773,7 +831,7 @@ def autostop(
 @usage_lib.entrypoint
 def queue(cluster_name: str,
           skip_finished: bool = False,
-          all_users: bool = False) -> List[dict]:
+          all_users: bool = False) -> List[responses.ClusterJobRecord]:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
     """Gets the job queue of a cluster.
 
@@ -811,7 +869,6 @@ def queue(cluster_name: str,
         user_hash = None
     else:
         user_hash = common_utils.get_current_user().id
-    code = job_lib.JobLibCodeGen.get_job_queue(user_hash, all_jobs)
 
     handle = backend_utils.check_cluster_available(
         cluster_name,
@@ -819,18 +876,49 @@ def queue(cluster_name: str,
     )
     backend = backend_utils.get_backend_from_handle(handle)
 
-    returncode, jobs_payload, stderr = backend.run_on_head(handle,
-                                                           code,
-                                                           require_outputs=True,
-                                                           separate_stderr=True)
-    subprocess_utils.handle_returncode(
-        returncode,
-        command=code,
-        error_msg=f'Failed to get job queue on cluster {cluster_name}.',
-        stderr=f'{jobs_payload + stderr}',
-        stream_logs=True)
-    jobs = job_lib.load_job_queue(jobs_payload)
-    return jobs
+    use_legacy = not handle.is_grpc_enabled_with_flag
+
+    if not use_legacy:
+        try:
+            request = jobsv1_pb2.GetJobQueueRequest(user_hash=user_hash,
+                                                    all_jobs=all_jobs)
+            response = backend_utils.invoke_skylet_with_retries(
+                lambda: cloud_vm_ray_backend.SkyletClient(
+                    handle.get_grpc_channel()).get_job_queue(request))
+            jobs = []
+            for job_info in response.jobs:
+                job_dict = {
+                    'job_id': job_info.job_id,
+                    'job_name': job_info.job_name,
+                    'submitted_at': job_info.submitted_at,
+                    'status': job_lib.JobStatus.from_protobuf(job_info.status),
+                    'run_timestamp': job_info.run_timestamp,
+                    'start_at': job_info.start_at
+                                if job_info.HasField('start_at') else None,
+                    'end_at': job_info.end_at
+                              if job_info.HasField('end_at') else None,
+                    'resources': job_info.resources,
+                    'log_path': job_info.log_path,
+                    'user_hash': job_info.username,
+                }
+                # Copied from job_lib.load_job_queue.
+                user = global_user_state.get_user(job_dict['user_hash'])
+                job_dict['username'] = user.name if user is not None else None
+                jobs.append(job_dict)
+        except exceptions.SkyletMethodNotImplementedError:
+            use_legacy = True
+    if use_legacy:
+        code = job_lib.JobLibCodeGen.get_job_queue(user_hash, all_jobs)
+        returncode, jobs_payload, stderr = backend.run_on_head(
+            handle, code, require_outputs=True, separate_stderr=True)
+        subprocess_utils.handle_returncode(
+            returncode,
+            command=code,
+            error_msg=f'Failed to get job queue on cluster {cluster_name}.',
+            stderr=f'{jobs_payload + stderr}',
+            stream_logs=True)
+        jobs = job_lib.load_job_queue(jobs_payload)
+    return [responses.ClusterJobRecord.model_validate(job) for job in jobs]
 
 
 @usage_lib.entrypoint
@@ -1070,25 +1158,25 @@ def job_status(cluster_name: str,
 # = Storage Management =
 # ======================
 @usage_lib.entrypoint
-def storage_ls() -> List[Dict[str, Any]]:
+def storage_ls() -> List[responses.StorageRecord]:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
     """Gets the storages.
 
     Returns:
-        [
-            {
-                'name': str,
-                'launched_at': int timestamp of creation,
-                'store': List[sky.StoreType],
-                'last_use': int timestamp of last use,
-                'status': sky.StorageStatus,
-            }
-        ]
+        List[responses.StorageRecord]: A list of storage records.
     """
     storages = global_user_state.get_storage()
+    storage_records = []
     for storage in storages:
-        storage['store'] = list(storage.pop('handle').sky_stores.keys())
-    return storages
+        storage_records.append(
+            responses.StorageRecord(
+                name=storage['name'],
+                launched_at=storage['launched_at'],
+                store=list(storage.pop('handle').sky_stores.keys()),
+                last_use=storage['last_use'],
+                status=storage['status'],
+            ))
+    return storage_records
 
 
 @usage_lib.entrypoint
@@ -1104,9 +1192,7 @@ def storage_delete(name: str) -> None:
     if handle is None:
         raise ValueError(f'Storage name {name!r} not found.')
     else:
-        storage_object = data.Storage(name=handle.storage_name,
-                                      source=handle.source,
-                                      sync_on_reconstruction=False)
+        storage_object = data.Storage.from_handle(handle)
         storage_object.delete()
 
 
@@ -1125,6 +1211,7 @@ def enabled_clouds(workspace: Optional[str] = None,
             return [cloud.canonical_name() for cloud in cached_clouds]
         enabled_ssh_infras = []
         enabled_k8s_infras = []
+        enabled_slurm_infras = []
         enabled_cloud_infras = []
         for cloud in cached_clouds:
             cloud_infra = cloud.expand_infras()
@@ -1132,10 +1219,16 @@ def enabled_clouds(workspace: Optional[str] = None,
                 enabled_ssh_infras.extend(cloud_infra)
             elif isinstance(cloud, clouds.Kubernetes):
                 enabled_k8s_infras.extend(cloud_infra)
+            elif isinstance(cloud, clouds.Slurm):
+                enabled_slurm_infras.extend(cloud_infra)
             else:
                 enabled_cloud_infras.extend(cloud_infra)
+        # We do not sort slurm infras alphabetically because the
+        # default partition should appear first.
+        # Ordering of slurm infras is enforced in Slurm implementation.
         all_infras = sorted(enabled_ssh_infras) + sorted(
-            enabled_k8s_infras) + sorted(enabled_cloud_infras)
+            enabled_k8s_infras) + enabled_slurm_infras + sorted(
+                enabled_cloud_infras)
         return all_infras
 
 
@@ -1146,7 +1239,14 @@ def realtime_kubernetes_gpu_availability(
     quantity_filter: Optional[int] = None,
     is_ssh: Optional[bool] = None
 ) -> List[Tuple[str, List[models.RealtimeGpuAvailability]]]:
+    """Gets the real-time Kubernetes GPU availability.
 
+    Returns:
+        A list of tuples, where each tuple contains:
+        - context (str): The Kubernetes context.
+        - availability_list (List[models.RealtimeGpuAvailability]): A list
+            of RealtimeGpuAvailability objects for that context.
+    """
     if context is None:
         # Include contexts from both Kubernetes and SSH clouds
         kubernetes_contexts = clouds.Kubernetes.existing_allowed_contexts()
@@ -1228,132 +1328,133 @@ def realtime_kubernetes_gpu_availability(
     return availability_lists
 
 
-# =================
-# = Local Cluster =
-# =================
-@usage_lib.entrypoint
-def local_up(gpus: bool,
-             ips: Optional[List[str]],
-             ssh_user: Optional[str],
-             ssh_key: Optional[str],
-             cleanup: bool,
-             context_name: Optional[str] = None,
-             password: Optional[str] = None) -> None:
-    """Creates a local or remote cluster."""
-
-    def _validate_args(ips, ssh_user, ssh_key, cleanup):
-        # If any of --ips, --ssh-user, or --ssh-key-path is specified,
-        # all must be specified
-        if bool(ips) or bool(ssh_user) or bool(ssh_key):
-            if not (ips and ssh_user and ssh_key):
-                with ux_utils.print_exception_no_traceback():
-                    raise ValueError(
-                        'All ips, ssh_user, and ssh_key must be specified '
-                        'together.')
-
-        # --cleanup can only be used if --ips, --ssh-user and --ssh-key-path
-        # are all provided
-        if cleanup and not (ips and ssh_user and ssh_key):
-            with ux_utils.print_exception_no_traceback():
-                raise ValueError(
-                    'cleanup can only be used with ips, ssh_user and ssh_key.')
-
-    _validate_args(ips, ssh_user, ssh_key, cleanup)
-
-    # If remote deployment arguments are specified, run remote up script
-    if ips:
-        assert ssh_user is not None and ssh_key is not None
-        kubernetes_deploy_utils.deploy_remote_cluster(ips, ssh_user, ssh_key,
-                                                      cleanup, context_name,
-                                                      password)
-    else:
-        # Run local deployment (kind) if no remote args are specified
-        kubernetes_deploy_utils.deploy_local_cluster(gpus)
+def realtime_slurm_gpu_availability(
+        slurm_cluster_name: Optional[str] = None,
+        name_filter: Optional[str] = None,
+        quantity_filter: Optional[int] = None,
+        env_vars: Optional[Dict[str, str]] = None,
+        **kwargs) -> List[Tuple[str, List[models.RealtimeGpuAvailability]]]:
+    """Gets Slurm real-time GPU availability grouped by partition.
 
+    This function calls the Slurm backend to fetch GPU info.
 
-def local_down() -> None:
-    """Tears down the Kubernetes cluster started by local_up."""
-    cluster_removed = False
+    Args:
+        name_filter: Optional name filter for GPUs.
+        quantity_filter: Optional quantity filter for GPUs.
+        env_vars: Environment variables (may be needed for backend).
+        kwargs: Additional keyword arguments.
 
-    path_to_package = os.path.dirname(__file__)
-    down_script_path = os.path.join(path_to_package, 'utils/kubernetes',
-                                    'delete_cluster.sh')
+    Returns:
+        A list of tuples, where each tuple contains:
+        - partition_name (str): The name of the Slurm partition.
+        - availability_list (List[models.RealtimeGpuAvailability]): A list
+            of RealtimeGpuAvailability objects for that partition.
+        Example structure:
+        [
+            ('gpu_partition_1', [
+                RealtimeGpuAvailability(gpu='V100', counts=[4, 8],
+                                        capacity=16, available=10),
+                RealtimeGpuAvailability(gpu='A100', counts=[8],
+                                        capacity=8, available=0),
+            ]),
+            ('gpu_partition_2', [
+                RealtimeGpuAvailability(gpu='V100', counts=[4],
+                                        capacity=4, available=4),
+            ])
+        ]
 
-    cwd = os.path.dirname(os.path.abspath(down_script_path))
-    run_command = shlex.split(down_script_path)
+    Raises:
+        ValueError: If Slurm is not configured or no matching GPUs are found.
+        exceptions.NotSupportedError: If Slurm is not enabled or configured.
+    """
+    del env_vars, kwargs  # Currently unused
 
-    # Setup logging paths
-    run_timestamp = sky_logging.get_run_timestamp()
-    log_path = os.path.join(constants.SKY_LOGS_DIRECTORY, run_timestamp,
-                            'local_down.log')
+    if slurm_cluster_name is None:
+        # Include contexts from both Kubernetes and SSH clouds
+        slurm_cluster_names = clouds.Slurm.existing_allowed_clusters()
+    else:
+        slurm_cluster_names = [slurm_cluster_name]
 
-    with rich_utils.safe_status(
-            ux_utils.spinner_message('Removing local cluster',
-                                     log_path=log_path,
-                                     is_local=True)):
-
-        returncode, stdout, stderr = log_lib.run_with_log(cmd=run_command,
-                                                          log_path=log_path,
-                                                          require_outputs=True,
-                                                          stream_logs=False,
-                                                          cwd=cwd)
-        stderr = stderr.replace('No kind clusters found.\n', '')
-
-        if returncode == 0:
-            cluster_removed = True
-        elif returncode == 100:
-            logger.info(ux_utils.error_message('Local cluster does not exist.'))
-        else:
-            with ux_utils.print_exception_no_traceback():
-                raise RuntimeError('Failed to create local cluster. '
-                                   f'Stdout: {stdout}'
-                                   f'\nError: {stderr}')
-    if cluster_removed:
-        # Run sky check
-        with rich_utils.safe_status(
-                ux_utils.spinner_message('Running sky check...')):
-            sky_check.check_capability(sky_cloud.CloudCapability.COMPUTE,
-                                       clouds=['kubernetes'],
-                                       quiet=True)
-        logger.info(
-            ux_utils.finishing_message('Local cluster removed.',
-                                       log_path=log_path,
-                                       is_local=True))
+    # Optional: Check if Slurm is enabled first
+    # enabled = global_user_state.get_enabled_clouds(
+    #     capability=sky_cloud.CloudCapability.COMPUTE)
+    # if not clouds.Slurm() in enabled:
+    #    raise exceptions.NotSupportedError(
+    #       "Slurm is not enabled. Run 'sky check' to enable it.")
 
+    def realtime_slurm_gpu_availability_single(
+            slurm_cluster_name: str) -> List[models.RealtimeGpuAvailability]:
+        try:
+            # This function now returns aggregated data per GPU type:
+            # Tuple[Dict[str, List[InstanceTypeInfo]], Dict[str, int],
+            # Dict[str, int]]
+            # (qtys_map, total_capacity, total_available)
+            accelerator_counts, total_capacity, total_available = (
+                catalog.list_accelerator_realtime(
+                    gpus_only=True,  # Ensure we only query for GPUs
+                    name_filter=name_filter,
+                    # Pass None for region_filter here; filtering happens
+                    # inside if needed, but we want all partitions returned
+                    # for grouping.
+                    region_filter=slurm_cluster_name,
+                    quantity_filter=quantity_filter,
+                    clouds='slurm',
+                    case_sensitive=False,
+                ))
+        except exceptions.NotSupportedError as e:
+            logger.error(f'Failed to query Slurm GPU availability: {e}')
+            raise
+        except ValueError as e:
+            # Re-raise ValueError if no GPUs are found matching the filters
+            logger.error(f'Error querying Slurm GPU availability: {e}')
+            raise
+        except Exception as e:
+            logger.error(
+                'Error querying Slurm GPU availability: '
+                f'{common_utils.format_exception(e, use_bracket=True)}')
+            raise ValueError(
+                f'Error querying Slurm GPU availability: {e}') from e
 
-@usage_lib.entrypoint
-def ssh_up(infra: Optional[str] = None, cleanup: bool = False) -> None:
-    """Deploys or tears down a Kubernetes cluster on SSH targets.
+        # --- Format the output ---
+        realtime_gpu_availability_list: List[
+            models.RealtimeGpuAvailability] = []
+        for gpu_type, _ in sorted(accelerator_counts.items()):
+            realtime_gpu_availability_list.append(
+                models.RealtimeGpuAvailability(
+                    gpu_type,
+                    accelerator_counts.pop(gpu_type),
+                    total_capacity[gpu_type],
+                    total_available[gpu_type],
+                ))
+        return realtime_gpu_availability_list
 
-    Args:
-        infra: Name of the cluster configuration in ssh_node_pools.yaml.
-            If None, the first cluster in the file is used.
-        cleanup: If True, clean up the cluster instead of deploying.
-    """
-    kubernetes_deploy_utils.deploy_ssh_cluster(
-        cleanup=cleanup,
-        infra=infra,
-    )
+    parallel_queried = subprocess_utils.run_in_parallel(
+        realtime_slurm_gpu_availability_single, slurm_cluster_names)
+    availability_lists: List[Tuple[str,
+                                   List[models.RealtimeGpuAvailability]]] = []
+    for slurm_cluster_name, queried in zip(slurm_cluster_names,
+                                           parallel_queried):
+        if len(queried) == 0:
+            logger.debug(f'No gpus found in Slurm cluster {slurm_cluster_name}')
+            continue
+        availability_lists.append((slurm_cluster_name, queried))
+    return availability_lists
 
 
+# =================
+# = Local Cluster =
+# =================
 @usage_lib.entrypoint
-def ssh_status(context_name: str) -> Tuple[bool, str]:
-    """Check the status of an SSH Node Pool context.
+def local_up(gpus: bool,
+             name: Optional[str] = None,
+             port_start: Optional[int] = None) -> None:
+    """Creates a local cluster."""
+    kubernetes_deploy_utils.deploy_local_cluster(name, port_start, gpus)
 
-    Args:
-        context_name: The SSH context name (e.g., 'ssh-my-cluster')
 
-    Returns:
-        Tuple[bool, str]: (is_ready, reason)
-            - is_ready: True if the SSH Node Pool is ready, False otherwise
-            - reason: Explanation of the status
-    """
-    try:
-        is_ready, reason = clouds.SSH.check_single_context(context_name)
-        return is_ready, reason
-    except Exception as e:  # pylint: disable=broad-except
-        return False, ('Failed to check SSH context: '
-                       f'{common_utils.format_exception(e)}')
+def local_down(name: Optional[str] = None) -> None:
+    """Tears down the Kubernetes cluster started by local_up."""
+    kubernetes_deploy_utils.teardown_local_cluster(name)
 
 
 def get_all_contexts() -> List[str]: