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

sky/jobs/utils.py

@@ -80,9 +80,8 @@ JOB_STARTED_STATUS_CHECK_GAP_SECONDS = 5
 
 _LOG_STREAM_CHECK_CONTROLLER_GAP_SECONDS = 5
 
-_JOB_STATUS_FETCH_MAX_RETRIES = 3
-_JOB_K8S_TRANSIENT_NW_MSG = 'Unable to connect to the server: dial tcp'
 _JOB_STATUS_FETCH_TIMEOUT_SECONDS = 30
+JOB_STATUS_FETCH_TOTAL_TIMEOUT_SECONDS = 60
 
 _JOB_WAITING_STATUS_MESSAGE = ux_utils.spinner_message(
     'Waiting for task to start[/]'
@@ -329,13 +328,21 @@ def ha_recovery_for_consolidation_mode() -> None:
 
 
 async def get_job_status(
-        backend: 'backends.CloudVmRayBackend', cluster_name: str,
-        job_id: Optional[int]) -> Optional['job_lib.JobStatus']:
+    backend: 'backends.CloudVmRayBackend', cluster_name: str,
+    job_id: Optional[int]
+) -> Tuple[Optional['job_lib.JobStatus'], Optional[str]]:
     """Check the status of the job running on a managed job cluster.
 
     It can be None, INIT, RUNNING, SUCCEEDED, FAILED, FAILED_DRIVER,
     FAILED_SETUP or CANCELLED.
+
+    Returns:
+        job_status: The status of the job.
+        transient_error_reason: None if successful or fatal error; otherwise,
+            the detailed reason for the transient error.
     """
+    # TODO(zhwu, cooperc): Make this get job status aware of cluster status, so
+    # that it can exit retry early if the cluster is down.
     # TODO(luca) make this async
     handle = await context_utils.to_thread(
         global_user_state.get_handle_from_cluster_name, cluster_name)
@@ -343,85 +350,68 @@ async def get_job_status(
         # This can happen if the cluster was preempted and background status
         # refresh already noticed and cleaned it up.
         logger.info(f'Cluster {cluster_name} not found.')
-        return None
+        return None, None
     assert isinstance(handle, backends.CloudVmRayResourceHandle), handle
     job_ids = None if job_id is None else [job_id]
-    for i in range(_JOB_STATUS_FETCH_MAX_RETRIES):
-        try:
-            logger.info('=== Checking the job status... ===')
-            statuses = await asyncio.wait_for(
-                context_utils.to_thread(backend.get_job_status,
-                                        handle,
-                                        job_ids=job_ids,
-                                        stream_logs=False),
-                timeout=_JOB_STATUS_FETCH_TIMEOUT_SECONDS)
-            status = list(statuses.values())[0]
-            if status is None:
-                logger.info('No job found.')
+    try:
+        logger.info('=== Checking the job status... ===')
+        statuses = await asyncio.wait_for(
+            context_utils.to_thread(backend.get_job_status,
+                                    handle,
+                                    job_ids=job_ids,
+                                    stream_logs=False),
+            timeout=_JOB_STATUS_FETCH_TIMEOUT_SECONDS)
+        status = list(statuses.values())[0]
+        if status is None:
+            logger.info('No job found.')
+        else:
+            logger.info(f'Job status: {status}')
+        logger.info('=' * 34)
+        return status, None
+    except (exceptions.CommandError, grpc.RpcError, grpc.FutureTimeoutError,
+            ValueError, TypeError, asyncio.TimeoutError) as e:
+        # Note: Each of these exceptions has some additional conditions to
+        # limit how we handle it and whether or not we catch it.
+        potential_transient_error_reason = None
+        if isinstance(e, exceptions.CommandError):
+            returncode = e.returncode
+            potential_transient_error_reason = (f'Returncode: {returncode}. '
+                                                f'{e.detailed_reason}')
+        elif isinstance(e, grpc.RpcError):
+            potential_transient_error_reason = e.details()
+        elif isinstance(e, grpc.FutureTimeoutError):
+            potential_transient_error_reason = 'grpc timeout'
+        elif isinstance(e, asyncio.TimeoutError):
+            potential_transient_error_reason = (
+                'Job status check timed out after '
+                f'{_JOB_STATUS_FETCH_TIMEOUT_SECONDS}s')
+        # TODO(cooperc): Gracefully handle these exceptions in the backend.
+        elif isinstance(e, ValueError):
+            # If the cluster yaml is deleted in the middle of getting the
+            # SSH credentials, we could see this. See
+            # sky/global_user_state.py get_cluster_yaml_dict.
+            if re.search(r'Cluster yaml .* not found', str(e)):
+                potential_transient_error_reason = 'Cluster yaml was deleted'
             else:
-                logger.info(f'Job status: {status}')
-            logger.info('=' * 34)
-            return status
-        except (exceptions.CommandError, grpc.RpcError, grpc.FutureTimeoutError,
-                ValueError, TypeError, asyncio.TimeoutError) as e:
-            # Note: Each of these exceptions has some additional conditions to
-            # limit how we handle it and whether or not we catch it.
-            # Retry on k8s transient network errors. This is useful when using
-            # coreweave which may have transient network issue sometimes.
-            is_transient_error = False
-            detailed_reason = None
-            if isinstance(e, exceptions.CommandError):
-                detailed_reason = e.detailed_reason
-                if (detailed_reason is not None and
-                        _JOB_K8S_TRANSIENT_NW_MSG in detailed_reason):
-                    is_transient_error = True
-            elif isinstance(e, grpc.RpcError):
-                detailed_reason = e.details()
-                if e.code() in [
-                        grpc.StatusCode.UNAVAILABLE,
-                        grpc.StatusCode.DEADLINE_EXCEEDED
-                ]:
-                    is_transient_error = True
-            elif isinstance(e, grpc.FutureTimeoutError):
-                detailed_reason = 'Timeout'
-            elif isinstance(e, asyncio.TimeoutError):
-                detailed_reason = ('Job status check timed out after '
-                                   f'{_JOB_STATUS_FETCH_TIMEOUT_SECONDS}s')
-            # TODO(cooperc): Gracefully handle these exceptions in the backend.
-            elif isinstance(e, ValueError):
-                # If the cluster yaml is deleted in the middle of getting the
-                # SSH credentials, we could see this. See
-                # sky/global_user_state.py get_cluster_yaml_dict.
-                if re.search(r'Cluster yaml .* not found', str(e)):
-                    detailed_reason = 'Cluster yaml was deleted'
-                else:
-                    raise
-            elif isinstance(e, TypeError):
-                # We will grab the SSH credentials from the cluster yaml, but if
-                # handle.cluster_yaml is None, we will just return an empty dict
-                # for the credentials. See
-                # backend_utils.ssh_credential_from_yaml. Then, the credentials
-                # are passed as kwargs to SSHCommandRunner.__init__ - see
-                # cloud_vm_ray_backend.get_command_runners. So we can hit this
-                # TypeError if the cluster yaml is removed from the handle right
-                # when we pull it before the cluster is fully deleted.
-                error_msg_to_check = (
-                    'SSHCommandRunner.__init__() missing 2 required positional '
-                    'arguments: \'ssh_user\' and \'ssh_private_key\'')
-                if str(e) == error_msg_to_check:
-                    detailed_reason = 'SSH credentials were already cleaned up'
-                else:
-                    raise
-            if is_transient_error:
-                logger.info('Failed to connect to the cluster. Retrying '
-                            f'({i + 1}/{_JOB_STATUS_FETCH_MAX_RETRIES})...')
-                logger.info('=' * 34)
-                await asyncio.sleep(1)
+                raise
+        elif isinstance(e, TypeError):
+            # We will grab the SSH credentials from the cluster yaml, but if
+            # handle.cluster_yaml is None, we will just return an empty dict
+            # for the credentials. See
+            # backend_utils.ssh_credential_from_yaml. Then, the credentials
+            # are passed as kwargs to SSHCommandRunner.__init__ - see
+            # cloud_vm_ray_backend.get_command_runners. So we can hit this
+            # TypeError if the cluster yaml is removed from the handle right
+            # when we pull it before the cluster is fully deleted.
+            error_msg_to_check = (
+                'SSHCommandRunner.__init__() missing 2 required positional '
+                'arguments: \'ssh_user\' and \'ssh_private_key\'')
+            if str(e) == error_msg_to_check:
+                potential_transient_error_reason = ('SSH credentials were '
+                                                    'already cleaned up')
             else:
-                logger.info(f'Failed to get job status: {detailed_reason}')
-                logger.info('=' * 34)
-                return None
-    return None
+                raise
+        return None, potential_transient_error_reason
 
 
 def controller_process_alive(record: managed_job_state.ControllerPidRecord,
@@ -1570,6 +1560,7 @@ def get_managed_job_queue(
                     handle.launched_resources.region,
                     handle.launched_resources.zone).formatted_str()
                 job['accelerators'] = handle.launched_resources.accelerators
+                job['labels'] = handle.launched_resources.labels
             else:
                 # FIXME(zongheng): display the last cached values for these.
                 job['cluster_resources'] = '-'
@@ -1578,6 +1569,7 @@ def get_managed_job_queue(
                 job['region'] = '-'
                 job['zone'] = '-'
                 job['infra'] = '-'
+                job['labels'] = None
 
         if not fields or 'details' in fields:
             # Add details about schedule state / backoff.
@@ -1821,7 +1813,8 @@ def format_job_table(
             for replica in replica_info:
                 used_by = replica.get('used_by')
                 if used_by is not None:
-                    job_to_worker[used_by] = replica.get('replica_id')
+                    for job_id in used_by:
+                        job_to_worker[job_id] = replica.get('replica_id')
         return job_to_worker
 
     # Create mapping from job_id to worker replica_id

sky/models.py

@@ -68,6 +68,15 @@ class KubernetesNodeInfo:
     free: Dict[str, int]
     # IP address of the node (external IP preferred, fallback to internal IP)
     ip_address: Optional[str] = None
+    # CPU count (total CPUs available on the node)
+    cpu_count: Optional[float] = None
+    # Memory in GB (total memory available on the node)
+    memory_gb: Optional[float] = None
+    # Free CPU count (free CPUs available on the node after pod allocations)
+    cpu_free: Optional[float] = None
+    # Free memory in GB (free memory available on the node after pod
+    # allocations)
+    memory_free_gb: Optional[float] = None
     # Whether the node is ready (all conditions are satisfied)
     is_ready: bool = True
 

sky/optimizer.py

@@ -20,6 +20,7 @@ from sky.adaptors import common as adaptors_common
 from sky.clouds import cloud as sky_cloud
 from sky.usage import usage_lib
 from sky.utils import common
+from sky.utils import common_utils
 from sky.utils import env_options
 from sky.utils import log_utils
 from sky.utils import registry
@@ -1290,7 +1291,7 @@ def _check_specified_regions(task: task_lib.Task) -> None:
         msg = f'Task{task_name} requires '
         if region not in existing_contexts:
             if is_ssh:
-                infra_str = f'SSH/{region.lstrip("ssh-")}'
+                infra_str = f'SSH/{common_utils.removeprefix(region, "ssh-")}'
             else:
                 infra_str = f'Kubernetes/{region}'
             logger.warning(f'{infra_str} is not enabled.')

sky/provision/__init__.py

@@ -6,7 +6,7 @@ providers supported by SkyPilot need to follow.
 import functools
 import inspect
 import typing
-from typing import Any, Dict, List, Optional, Tuple, Type
+from typing import Any, Dict, List, Optional, Set, Tuple, Type
 
 from sky import models
 from sky import sky_logging
@@ -152,16 +152,18 @@ def get_volume_usedby(
 @_route_to_cloud_impl
 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.
+) -> Tuple[Dict[str, Any], Dict[str, Any], Set[str]]:
+    """Get the usedby of all volumes.
+
+    Args:
+        provider_name: Name of the provider.
+        configs: List of VolumeConfig objects.
 
     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.
+        usedby_pods: Dict of usedby pods.
+        usedby_clusters: Dict of usedby clusters.
+        failed_volume_names: Set of volume names whose usedby info
+          failed to fetch.
     """
     raise NotImplementedError
 

sky/provision/kubernetes/utils.py

@@ -144,6 +144,7 @@ DEFAULT_NAMESPACE = 'default'
 DEFAULT_SERVICE_ACCOUNT_NAME = 'skypilot-service-account'
 
 MEMORY_SIZE_UNITS = {
+    'm': 0.001,
     'B': 1,
     'K': 2**10,
     'M': 2**20,
@@ -1331,12 +1332,20 @@ class V1Pod:
 
 
 @_retry_on_error(resource_type='pod')
-def get_allocated_gpu_qty_by_node(
+def get_allocated_resources_by_node(
     *,
     context: Optional[str] = None,
-) -> Dict[str, int]:
-    """Gets allocated GPU quantity by each node by fetching pods in
+) -> Tuple[Dict[str, int], Dict[str, Tuple[float, float]]]:
+    """Gets allocated GPU, CPU, and memory by each node by fetching pods in
     all namespaces in kubernetes cluster indicated by context.
+
+    This function combines GPU and CPU/memory allocation tracking into a single
+    API call for better performance.
+
+    Returns:
+        Tuple of (allocated_gpu_qty_by_node, allocated_cpu_memory_by_node):
+        - allocated_gpu_qty_by_node: Dict mapping node name to allocated GPU count
+        - allocated_cpu_memory_by_node: Dict mapping node name to (allocated_cpu, allocated_memory_gb) tuple
     """
     if context is None:
         context = get_current_kube_config_context_name()
@@ -1355,29 +1364,67 @@ def get_allocated_gpu_qty_by_node(
         field_selector=field_selector)
     try:
         allocated_qty_by_node: Dict[str, int] = collections.defaultdict(int)
+        allocated_cpu_memory_by_node: Dict[str, Tuple[
+            float, float]] = collections.defaultdict(lambda: (0.0, 0.0))
         for item_dict in ijson.items(response,
                                      'items.item',
                                      buf_size=IJSON_BUFFER_SIZE):
             pod = V1Pod.from_dict(item_dict)
             if should_exclude_pod_from_gpu_allocation(pod):
                 logger.debug(
-                    f'Excluding pod {pod.metadata.name} from GPU count '
+                    f'Excluding pod {pod.metadata.name} from resource count '
                     f'calculations on node {pod.spec.node_name}')
                 continue
-            # Iterate over all the containers in the pod and sum the
-            # GPU requests
+            if not pod.spec.node_name:
+                continue
+
+            # Iterate over all the containers in the pod and sum the resources
             pod_allocated_qty = 0
+            pod_allocated_cpu = 0.0
+            pod_allocated_memory_gb = 0.0
             for container in pod.spec.containers:
                 if container.resources.requests:
+                    requests = container.resources.requests
+                    # Parse GPU
                     pod_allocated_qty += get_node_accelerator_count(
-                        context, container.resources.requests)
-            if pod_allocated_qty > 0 and pod.spec.node_name:
+                        context, requests)
+                    # Parse CPU
+                    if 'cpu' in requests:
+                        pod_allocated_cpu += parse_cpu_or_gpu_resource_to_float(
+                            requests['cpu'])
+                    # Parse memory
+                    if 'memory' in requests:
+                        pod_allocated_memory_gb += parse_memory_resource(
+                            requests['memory'], unit='G')
+
+            if pod_allocated_qty > 0:
                 allocated_qty_by_node[pod.spec.node_name] += pod_allocated_qty
-        return allocated_qty_by_node
+            if pod_allocated_cpu > 0 or pod_allocated_memory_gb > 0:
+                current_cpu, current_memory = allocated_cpu_memory_by_node[
+                    pod.spec.node_name]
+                allocated_cpu_memory_by_node[pod.spec.node_name] = (
+                    current_cpu + pod_allocated_cpu,
+                    current_memory + pod_allocated_memory_gb)
+        return allocated_qty_by_node, allocated_cpu_memory_by_node
     finally:
         response.release_conn()
 
 
+@_retry_on_error(resource_type='pod')
+def get_allocated_gpu_qty_by_node(
+    *,
+    context: Optional[str] = None,
+) -> Dict[str, int]:
+    """Gets allocated GPU quantity by each node by fetching pods in
+    all namespaces in kubernetes cluster indicated by context.
+
+    Note: For better performance when you also need CPU/memory allocation,
+    use get_allocated_resources_by_node() instead.
+    """
+    allocated_qty_by_node, _ = get_allocated_resources_by_node(context=context)
+    return allocated_qty_by_node
+
+
 def check_instance_fits(context: Optional[str],
                         instance: str) -> Tuple[bool, Optional[str]]:
     """Checks if the instance fits on the Kubernetes cluster.
@@ -2189,6 +2236,13 @@ def get_current_kube_config_context_name() -> Optional[str]:
         _, current_context = kubernetes.list_kube_config_contexts()
         return current_context['name']
     except k8s.config.config_exception.ConfigException:
+        # If kubeconfig is not available, check if running in-cluster and
+        # return the in-cluster context name. This is needed when kubeconfig
+        # is not uploaded to the pod (e.g., remote_identity: SERVICE_ACCOUNT)
+        # but we still need to know the context name for operations like
+        # port mode detection.
+        if is_incluster_config_available():
+            return kubernetes.in_cluster_context_name()
         return None
 
 
@@ -2313,7 +2367,7 @@ def parse_memory_resource(resource_qty_str: str,
     try:
         bytes_value = int(resource_str)
     except ValueError:
-        memory_size = re.sub(r'([KMGTPB]+)', r' \1', resource_str)
+        memory_size = re.sub(r'([KMGTPBm]+)', r' \1', resource_str)
         number, unit_index = [item.strip() for item in memory_size.split()]
         unit_index = unit_index[0]
         bytes_value = float(number) * MEMORY_SIZE_UNITS[unit_index]
@@ -3061,16 +3115,32 @@ def get_kubernetes_node_info(
             has_accelerator_nodes = True
             break
 
-    # Get the allocated GPU quantity by each node
+    # Get the allocated resources (GPU, CPU, memory) by each node in a single call
     allocated_qty_by_node: Dict[str, int] = collections.defaultdict(int)
-    error_on_get_allocated_gpu_qty_by_node = False
+    allocated_cpu_memory_by_node: Dict[str, Tuple[float, float]] = {}
+    error_on_get_allocated_resources = False
+    # Get resource allocation. For GPU allocation, only call if there are GPU nodes
+    # (same as master branch). For CPU/memory, we always need it for all nodes.
     if has_accelerator_nodes:
+        # When there are GPU nodes, get both GPU and CPU/memory in one call
         try:
-            allocated_qty_by_node = get_allocated_gpu_qty_by_node(
+            allocated_qty_by_node, allocated_cpu_memory_by_node = get_allocated_resources_by_node(
                 context=context)
         except kubernetes.api_exception() as e:
             if e.status == 403:
-                error_on_get_allocated_gpu_qty_by_node = True
+                error_on_get_allocated_resources = True
+                pass
+            else:
+                raise
+    else:
+        # When there are no GPU nodes, we still need CPU/memory allocation
+        # This is an extra API call compared to master branch
+        try:
+            _, allocated_cpu_memory_by_node = get_allocated_resources_by_node(
+                context=context)
+        except kubernetes.api_exception() as e:
+            if e.status == 403:
+                error_on_get_allocated_resources = True
                 pass
             else:
                 raise
@@ -3106,6 +3176,35 @@ def get_kubernetes_node_info(
 
         accelerator_count = get_node_accelerator_count(context,
                                                        node.status.allocatable)
+
+        # Parse CPU and memory from node capacity
+        cpu_count = None
+        memory_gb = None
+        try:
+            if 'cpu' in node.status.capacity:
+                cpu_count = float(
+                    parse_cpu_or_gpu_resource(node.status.capacity['cpu']))
+            if 'memory' in node.status.capacity:
+                memory_gb = parse_memory_resource(
+                    node.status.capacity['memory'], unit='G')
+        except (KeyError, ValueError) as e:
+            # If parsing fails, log but continue
+            logger.debug(f'Failed to parse CPU/memory for node '
+                         f'{node.metadata.name}: {e}')
+
+        # Calculate free CPU and memory
+        cpu_free = None
+        memory_free_gb = None
+        if cpu_count is not None or memory_gb is not None:
+            if not error_on_get_allocated_resources:
+                allocated_cpu, allocated_memory = allocated_cpu_memory_by_node.get(
+                    node.metadata.name, (0.0, 0.0))
+                if cpu_count is not None:
+                    cpu_free = max(0.0, cpu_count - allocated_cpu)
+                if memory_gb is not None:
+                    memory_free_gb = max(0.0, memory_gb - allocated_memory)
+            # If we can't get allocation info, set free to None (unknown)
+
         # Check if node is ready
         node_is_ready = node.is_ready()
 
@@ -3116,13 +3215,17 @@ def get_kubernetes_node_info(
                 total={'accelerator_count': 0},
                 free={'accelerators_available': 0},
                 ip_address=node_ip,
+                cpu_count=cpu_count,
+                memory_gb=memory_gb,
+                cpu_free=cpu_free,
+                memory_free_gb=memory_free_gb,
                 is_ready=node_is_ready)
             continue
 
         if not node_is_ready:
             # If node is not ready, report 0 available GPUs
             accelerators_available = 0
-        elif not has_accelerator_nodes or error_on_get_allocated_gpu_qty_by_node:
+        elif not has_accelerator_nodes or error_on_get_allocated_resources:
             accelerators_available = -1
         else:
             allocated_qty = allocated_qty_by_node[node.metadata.name]
@@ -3141,6 +3244,10 @@ def get_kubernetes_node_info(
             total={'accelerator_count': int(accelerator_count)},
             free={'accelerators_available': int(accelerators_available)},
             ip_address=node_ip,
+            cpu_count=cpu_count,
+            memory_gb=memory_gb,
+            cpu_free=cpu_free,
+            memory_free_gb=memory_free_gb,
             is_ready=node_is_ready)
     hint = ''
     if has_multi_host_tpu:

sky/provision/kubernetes/volume.py

@@ -45,7 +45,9 @@ def check_pvc_usage_for_pod(context: Optional[str], namespace: str,
             continue
         pvc = kubernetes.core_api(
             context).read_namespaced_persistent_volume_claim(
-                name=pvc_name, namespace=namespace)
+                name=pvc_name,
+                namespace=namespace,
+                _request_timeout=kubernetes.API_TIMEOUT)
         access_mode = pvc.spec.access_modes[0]
         if access_mode not in once_modes:
             continue
@@ -65,7 +67,8 @@ def apply_volume(config: models.VolumeConfig) -> models.VolumeConfig:
     if storage_class_name is not None:
         try:
             kubernetes.storage_api(context).read_storage_class(
-                name=storage_class_name)
+                name=storage_class_name,
+                _request_timeout=kubernetes.API_TIMEOUT)
         except kubernetes.api_exception() as e:
             raise config_lib.KubernetesError(
                 f'Check storage class {storage_class_name} error: {e}')
@@ -82,7 +85,7 @@ def delete_volume(config: models.VolumeConfig) -> models.VolumeConfig:
             context).delete_namespaced_persistent_volume_claim(
                 name=pvc_name,
                 namespace=namespace,
-                _request_timeout=config_lib.DELETION_TIMEOUT),
+                _request_timeout=kubernetes.API_TIMEOUT),
         resource_type='pvc',
         resource_name=pvc_name)
     logger.info(f'Deleted PVC {pvc_name} in namespace {namespace}')
@@ -119,7 +122,9 @@ def _get_volume_usedby(
     cloud_to_name_map = _get_cluster_name_on_cloud_to_cluster_name_map()
     # Get all pods in the namespace
     pods = kubernetes.core_api(context).list_namespaced_pod(
-        namespace=namespace, field_selector=field_selector)
+        namespace=namespace,
+        field_selector=field_selector,
+        _request_timeout=kubernetes.API_TIMEOUT)
     for pod in pods.items:
         if pod.spec.volumes is None:
             continue
@@ -164,8 +169,21 @@ def get_volume_usedby(
 
 def get_all_volumes_usedby(
     configs: List[models.VolumeConfig],
-) -> Tuple[Dict[str, Any], Dict[str, Any]]:
-    """Gets the usedby resources of all volumes."""
+) -> Tuple[Dict[str, Any], Dict[str, Any], Set[str]]:
+    """Gets the usedby resources of all volumes.
+
+    Args:
+        configs: List of VolumeConfig objects.
+
+    Returns:
+        usedby_pods: Dictionary of context to namespace to volume name to pods
+                     using the volume. These may include pods not created by
+                     SkyPilot.
+        usedby_clusters: Dictionary of context to namespace to volume name to
+                         clusters using the volume.
+        failed_volume_names: Set of volume names whose usedby info failed to
+          fetch.
+    """
     field_selector = ','.join([
         f'status.phase!={phase}'
         for phase in k8s_constants.PVC_NOT_HOLD_POD_PHASES
@@ -173,26 +191,39 @@ def get_all_volumes_usedby(
     label_selector = 'parent=skypilot'
     context_to_namespaces: Dict[str, Set[str]] = {}
     pvc_names = set()
+    original_volume_names: Dict[str, Dict[str, List[str]]] = {}
     for config in configs:
         context, namespace = _get_context_namespace(config)
-        if context not in context_to_namespaces:
-            context_to_namespaces[context] = set()
-        context_to_namespaces[context].add(namespace)
+        context_to_namespaces.setdefault(context, set()).add(namespace)
+        original_volume_names.setdefault(context,
+                                         {}).setdefault(namespace,
+                                                        []).append(config.name)
         pvc_names.add(config.name_on_cloud)
     cloud_to_name_map = _get_cluster_name_on_cloud_to_cluster_name_map()
     # Get all pods in the namespace
     used_by_pods: Dict[str, Dict[str, Dict[str, List[str]]]] = {}
     used_by_clusters: Dict[str, Dict[str, Dict[str, List[str]]]] = {}
+    failed_volume_names: Set[str] = set()
     for context, namespaces in context_to_namespaces.items():
         used_by_pods[context] = {}
         used_by_clusters[context] = {}
         for namespace in namespaces:
             used_by_pods[context][namespace] = {}
             used_by_clusters[context][namespace] = {}
-            pods = kubernetes.core_api(context).list_namespaced_pod(
-                namespace=namespace,
-                field_selector=field_selector,
-                label_selector=label_selector)
+            try:
+                pods = kubernetes.core_api(context).list_namespaced_pod(
+                    namespace=namespace,
+                    field_selector=field_selector,
+                    label_selector=label_selector,
+                    _request_timeout=kubernetes.API_TIMEOUT)
+            except Exception as e:  # pylint: disable=broad-except
+                logger.debug(f'Failed to get pods in namespace {namespace} '
+                             f'in context {context}: {e}')
+                # Mark all volumes in this namespace as failed
+                for original_volume_name in original_volume_names[context][
+                        namespace]:
+                    failed_volume_names.add(original_volume_name)
+                continue
             for pod in pods.items:
                 if pod.spec.volumes is None:
                     continue
@@ -217,7 +248,7 @@ def get_all_volumes_usedby(
                         used_by_clusters[context][namespace][cluster_name] = []
                     used_by_clusters[context][namespace][cluster_name].append(
                         cluster_name)
-    return used_by_pods, used_by_clusters
+    return used_by_pods, used_by_clusters, failed_volume_names
 
 
 def map_all_volumes_usedby(
@@ -292,7 +323,9 @@ def create_persistent_volume_claim(
     try:
         pvc = kubernetes.core_api(
             context).read_namespaced_persistent_volume_claim(
-                name=pvc_name, namespace=namespace)
+                name=pvc_name,
+                namespace=namespace,
+                _request_timeout=kubernetes.API_TIMEOUT)
         if config is not None:
             _populate_config_from_pvc(config, pvc)
         logger.debug(f'PVC {pvc_name} already exists')
@@ -305,8 +338,10 @@ def create_persistent_volume_claim(
         raise ValueError(
             f'PVC {pvc_name} does not exist while use_existing is True.')
     pvc = kubernetes.core_api(
-        context).create_namespaced_persistent_volume_claim(namespace=namespace,
-                                                           body=pvc_spec)
+        context).create_namespaced_persistent_volume_claim(
+            namespace=namespace,
+            body=pvc_spec,
+            _request_timeout=kubernetes.API_TIMEOUT)
     logger.info(f'Created PVC {pvc_name} in namespace {namespace}')
     if config is not None:
         _populate_config_from_pvc(config, pvc)

sky/provision/provisioner.py

@@ -493,7 +493,8 @@ def _post_provision_setup(
         # commands and rsync on the pods. SSH will still be ready after a while
         # for the users to SSH into the pod.
         is_k8s_cloud = cloud_name.lower() in ['kubernetes', 'ssh']
-        if not is_k8s_cloud:
+        is_slurm_cloud = cloud_name.lower() == 'slurm'
+        if not is_k8s_cloud and not is_slurm_cloud:
             logger.debug(
                 f'\nWaiting for SSH to be available for {cluster_name!r} ...')
             wait_for_ssh(cluster_info, ssh_credentials)