skypilot-nightly 1.0.0.dev20251203__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,17 @@ 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
 
 
 @dataclasses.dataclass

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
@@ -781,7 +782,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 +866,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':
@@ -1289,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
@@ -29,6 +29,7 @@ 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
@@ -151,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/common.py

@@ -6,6 +6,7 @@ import os
 from typing import Any, Dict, List, Optional, Tuple
 
 from sky import sky_logging
+from sky.utils import config_utils
 from sky.utils import env_options
 from sky.utils import resources_utils
 
@@ -36,6 +37,13 @@ class StopFailoverError(Exception):
     """
 
 
+# These fields are sensitive and should be redacted from the config for logging
+# purposes.
+SENSITIVE_FIELDS = [
+    ('docker_config', 'docker_login_config', 'password'),
+]
+
+
 @dataclasses.dataclass
 class ProvisionConfig:
     """Configuration for provisioning."""
@@ -56,6 +64,18 @@ class ProvisionConfig:
     # Optional ports to open on launch of the cluster.
     ports_to_open_on_launch: Optional[List[int]]
 
+    def get_redacted_config(self) -> Dict[str, Any]:
+        """Get the redacted config."""
+        config = dataclasses.asdict(self)
+
+        config_copy = config_utils.Config(config)
+
+        for field_list in SENSITIVE_FIELDS:
+            val = config_copy.get_nested(field_list, default_value=None)
+            if val is not None:
+                config_copy.set_nested(field_list, '<redacted>')
+        return dict(**config_copy)
+
 
 # -------------------- output data model -------------------- #
 

sky/provision/docker_utils.py

@@ -176,6 +176,17 @@ def _with_interactive(cmd):
     return ['bash', '--login', '-c', '-i', shlex.quote(force_interactive)]
 
 
+def _redact_docker_password(cmd: str) -> str:
+    parts = shlex.split(cmd)
+    for i, part in enumerate(parts):
+        if part.startswith('--password'):
+            if part.startswith('--password='):
+                parts[i] = '--password=<redacted>'
+            elif i + 1 < len(parts):
+                parts[i + 1] = '<redacted>'
+    return ' '.join(parts)
+
+
 # SkyPilot: New class to initialize docker containers on a remote node.
 # Adopted from ray.autoscaler._private.command_runner.DockerCommandRunner.
 class DockerInitializer:
@@ -219,7 +230,9 @@ class DockerInitializer:
             cmd = (f'flock {flock_args} /tmp/{flock_name} '
                    f'-c {shlex.quote(cmd)}')
 
-        logger.debug(f'+ {cmd}')
+        # Redact the password in the login command.
+        redacted_cmd = _redact_docker_password(cmd)
+        logger.debug(f'+ {redacted_cmd}')
         start = time.time()
         while True:
             rc, stdout, stderr = self.runner.run(
@@ -251,7 +264,7 @@ class DockerInitializer:
             break
         subprocess_utils.handle_returncode(
             rc,
-            cmd,
+            redacted_cmd,
             error_msg='Failed to run docker setup commands.',
             stderr=stdout + stderr,
             # Print out the error message if the command failed.