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

sky/__init__.py

@@ -7,7 +7,7 @@ import urllib.request
 from sky.utils import directory_utils
 
 # Replaced with the current commit when building the wheels.
-_SKYPILOT_COMMIT_SHA = 'c28d94abd967c1a7494e3c343f92eb6d02d29541'
+_SKYPILOT_COMMIT_SHA = '5f4cd3b33375c055093474b95f219d26018b7343'
 
 
 def _get_git_commit():
@@ -37,7 +37,7 @@ def _get_git_commit():
 
 
 __commit__ = _get_git_commit()
-__version__ = '1.0.0.dev20251210'
+__version__ = '1.0.0.dev20260112'
 __root_dir__ = directory_utils.get_sky_dir()
 
 
@@ -143,6 +143,7 @@ SCP = clouds.SCP
 Slurm = clouds.Slurm
 Kubernetes = clouds.Kubernetes
 K8s = Kubernetes
+SSH = clouds.SSH
 OCI = clouds.OCI
 Paperspace = clouds.Paperspace
 PrimeIntellect = clouds.PrimeIntellect
@@ -164,6 +165,7 @@ __all__ = [
     'IBM',
     'Kubernetes',
     'K8s',
+    'SSH',
     'Lambda',
     'OCI',
     'Paperspace',

sky/adaptors/slurm.py

@@ -1,11 +1,15 @@
 """Slurm adaptor for SkyPilot."""
 
+import ipaddress
 import logging
 import re
+import socket
 import time
 from typing import Dict, List, NamedTuple, Optional, Tuple
 
+from sky.adaptors import common
 from sky.utils import command_runner
+from sky.utils import common_utils
 from sky.utils import subprocess_utils
 from sky.utils import timeline
 
@@ -22,6 +26,11 @@ _PARTITION_NAME_REGEX = re.compile(r'PartitionName=(.+?)(?:\s+\w+=|$)')
 # Default timeout for waiting for job nodes to be allocated, in seconds.
 _SLURM_DEFAULT_PROVISION_TIMEOUT = 10
 
+_IMPORT_ERROR_MESSAGE = ('Failed to import dependencies for Slurm. '
+                         'Try running: pip install "skypilot[slurm]"')
+hostlist = common.LazyImport('hostlist',
+                             import_error_message=_IMPORT_ERROR_MESSAGE)
+
 
 class SlurmPartition(NamedTuple):
     """Information about the Slurm partitions."""
@@ -47,11 +56,13 @@ class SlurmClient:
 
     def __init__(
         self,
-        ssh_host: str,
-        ssh_port: int,
-        ssh_user: str,
-        ssh_key: Optional[str],
+        ssh_host: Optional[str] = None,
+        ssh_port: Optional[int] = None,
+        ssh_user: Optional[str] = None,
+        ssh_key: Optional[str] = None,
         ssh_proxy_command: Optional[str] = None,
+        ssh_proxy_jump: Optional[str] = None,
+        is_inside_slurm_cluster: bool = False,
     ):
         """Initialize SlurmClient.
 
@@ -61,21 +72,42 @@ class SlurmClient:
             ssh_user: SSH username.
             ssh_key: Path to SSH private key, or None for keyless SSH.
             ssh_proxy_command: Optional SSH proxy command.
+            ssh_proxy_jump: Optional SSH proxy jump destination.
+            is_inside_slurm_cluster: If True, uses local execution mode (for
+            when running on the Slurm cluster itself). Defaults to False.
         """
         self.ssh_host = ssh_host
         self.ssh_port = ssh_port
         self.ssh_user = ssh_user
         self.ssh_key = ssh_key
         self.ssh_proxy_command = ssh_proxy_command
-
-        # Internal runner for executing Slurm CLI commands
-        # on the controller node.
-        self._runner = command_runner.SSHCommandRunner(
-            (ssh_host, ssh_port),
-            ssh_user,
-            ssh_key,
-            ssh_proxy_command=ssh_proxy_command,
-        )
+        self.ssh_proxy_jump = ssh_proxy_jump
+
+        self._runner: command_runner.CommandRunner
+
+        if is_inside_slurm_cluster:
+            # Local execution mode - for running on the Slurm cluster itself
+            # (e.g., autodown from skylet).
+            self._runner = command_runner.LocalProcessCommandRunner()
+        else:
+            # Remote execution via SSH
+            assert ssh_host is not None
+            assert ssh_port is not None
+            assert ssh_user is not None
+            self._runner = command_runner.SSHCommandRunner(
+                (ssh_host, ssh_port),
+                ssh_user,
+                ssh_key,
+                ssh_proxy_command=ssh_proxy_command,
+                ssh_proxy_jump=ssh_proxy_jump,
+                enable_interactive_auth=True,
+            )
+
+    def _run_slurm_cmd(self, cmd: str) -> Tuple[int, str, str]:
+        return self._runner.run(cmd,
+                                require_outputs=True,
+                                separate_stderr=True,
+                                stream_logs=False)
 
     def query_jobs(
         self,
@@ -99,13 +131,11 @@ class SlurmClient:
         if job_name is not None:
             cmd += f' --name {job_name}'
 
-        rc, stdout, stderr = self._runner.run(cmd,
-                                              require_outputs=True,
-                                              stream_logs=False)
+        rc, stdout, stderr = self._run_slurm_cmd(cmd)
         subprocess_utils.handle_returncode(rc,
                                            cmd,
                                            'Failed to query Slurm jobs.',
-                                           stderr=stderr)
+                                           stderr=f'{stdout}\n{stderr}')
 
         job_ids = stdout.strip().splitlines()
         return job_ids
@@ -128,13 +158,11 @@ class SlurmClient:
             cmd += f' --signal {signal}'
         if full:
             cmd += ' --full'
-        rc, stdout, stderr = self._runner.run(cmd,
-                                              require_outputs=True,
-                                              stream_logs=False)
+        rc, stdout, stderr = self._run_slurm_cmd(cmd)
         subprocess_utils.handle_returncode(rc,
                                            cmd,
                                            f'Failed to cancel job {job_name}.',
-                                           stderr=stderr)
+                                           stderr=f'{stdout}\n{stderr}')
         logger.debug(f'Successfully cancelled job {job_name}: {stdout}')
 
     def info(self) -> str:
@@ -147,11 +175,12 @@ class SlurmClient:
             The stdout output from sinfo.
         """
         cmd = 'sinfo'
-        rc, stdout, stderr = self._runner.run(cmd,
-                                              require_outputs=True,
-                                              stream_logs=False)
+        rc, stdout, stderr = self._run_slurm_cmd(cmd)
         subprocess_utils.handle_returncode(
-            rc, cmd, 'Failed to get Slurm cluster information.', stderr=stderr)
+            rc,
+            cmd,
+            'Failed to get Slurm cluster information.',
+            stderr=f'{stdout}\n{stderr}')
         return stdout
 
     def info_nodes(self) -> List[NodeInfo]:
@@ -162,11 +191,12 @@ class SlurmClient:
         """
         cmd = (f'sinfo -h --Node -o '
                f'"%N{SEP}%t{SEP}%G{SEP}%c{SEP}%m{SEP}%P"')
-        rc, stdout, stderr = self._runner.run(cmd,
-                                              require_outputs=True,
-                                              stream_logs=False)
+        rc, stdout, stderr = self._run_slurm_cmd(cmd)
         subprocess_utils.handle_returncode(
-            rc, cmd, 'Failed to get Slurm node information.', stderr=stderr)
+            rc,
+            cmd,
+            'Failed to get Slurm node information.',
+            stderr=f'{stdout}\n{stderr}')
 
         nodes = []
         for line in stdout.splitlines():
@@ -211,31 +241,63 @@ class SlurmClient:
             return node_info
 
         cmd = f'scontrol show node {node_name}'
-        rc, node_details, _ = self._runner.run(cmd,
-                                               require_outputs=True,
-                                               stream_logs=False)
+        rc, node_details, stderr = self._run_slurm_cmd(cmd)
         subprocess_utils.handle_returncode(
             rc,
             cmd,
             f'Failed to get detailed node information for {node_name}.',
-            stderr=node_details)
+            stderr=f'{node_details}\n{stderr}')
         node_info = _parse_scontrol_node_output(node_details)
         return node_info
 
-    def get_node_jobs(self, node_name: str) -> List[str]:
-        """Get the list of jobs for a given node name.
+    def get_jobs_gres(self, node_name: str) -> List[str]:
+        """Get the list of jobs GRES for a given node name.
 
         Returns:
-            A list of job names for the current user on the node.
+            A list of GRES specs (e.g., 'gres/gpu:h100:4')
+            for jobs on the node.
         """
-        cmd = f'squeue --me -h --nodelist {node_name} -o "%b"'
-        rc, stdout, stderr = self._runner.run(cmd,
-                                              require_outputs=True,
-                                              stream_logs=False)
+        cmd = f'squeue -h --nodelist {node_name} -o "%b"'
+        rc, stdout, stderr = self._run_slurm_cmd(cmd)
         subprocess_utils.handle_returncode(
-            rc, cmd, f'Failed to get jobs for node {node_name}.', stderr=stderr)
+            rc,
+            cmd,
+            f'Failed to get jobs for node {node_name}.',
+            stderr=f'{stdout}\n{stderr}')
         return stdout.splitlines()
 
+    def get_all_jobs_gres(self) -> Dict[str, List[str]]:
+        """Get GRES allocation for all running jobs, grouped by node.
+
+        Returns:
+            Dict mapping node_name -> list of GRES strings for jobs on that
+            node.
+        """
+        cmd = f'squeue -h --states=running,completing -o "%N{SEP}%b"'
+        rc, stdout, stderr = self._run_slurm_cmd(cmd)
+        subprocess_utils.handle_returncode(rc,
+                                           cmd,
+                                           'Failed to get all jobs GRES.',
+                                           stderr=f'{stdout}\n{stderr}')
+
+        nodes_to_gres: Dict[str, List[str]] = {}
+        for line in stdout.splitlines():
+            line = line.strip()
+            if not line:
+                continue
+            parts = line.split(SEP)
+            if len(parts) != 2:
+                # We should never reach here, but just in case.
+                continue
+            nodelist_str, gres_str = parts
+            if not gres_str or gres_str == 'N/A':
+                continue
+
+            for node in hostlist.expand_hostlist(nodelist_str):
+                nodes_to_gres.setdefault(node, []).append(gres_str)
+
+        return nodes_to_gres
+
     def get_job_state(self, job_id: str) -> Optional[str]:
         """Get the state of a Slurm job.
 
@@ -249,17 +311,30 @@ class SlurmClient:
         # Use --only-job-state since we only need the job state.
         # This reduces the work required by slurmctld.
         cmd = f'squeue -h --only-job-state --jobs {job_id} -o "%T"'
-        rc, stdout, stderr = self._runner.run(cmd,
-                                              require_outputs=True,
-                                              stream_logs=False)
-        if rc != 0:
-            # Job may not exist
-            logger.debug(f'Failed to get job state for job {job_id}: {stderr}')
-            return None
+        rc, stdout, stderr = self._run_slurm_cmd(cmd)
+        subprocess_utils.handle_returncode(
+            rc,
+            cmd,
+            f'Failed to get job state for job {job_id}.',
+            stderr=f'{stdout}\n{stderr}')
 
         state = stdout.strip()
         return state if state else None
 
+    def get_jobs_state_by_name(self, job_name: str) -> List[str]:
+        """Get the states of all Slurm jobs by name.
+        """
+        cmd = f'squeue -h --name {job_name} -o "%T"'
+        rc, stdout, stderr = self._run_slurm_cmd(cmd)
+        subprocess_utils.handle_returncode(
+            rc,
+            cmd,
+            f'Failed to get job state for job {job_name}.',
+            stderr=f'{stdout}\n{stderr}')
+
+        states = stdout.splitlines()
+        return states
+
     @timeline.event
     def get_job_reason(self, job_id: str) -> Optional[str]:
         """Get the reason a job is in its current state
@@ -269,12 +344,12 @@ class SlurmClient:
         """
         # Without --states all, squeue omits terminated jobs.
         cmd = f'squeue -h --jobs {job_id} --states all -o "%r"'
-        rc, stdout, stderr = self._runner.run(cmd,
-                                              require_outputs=True,
-                                              stream_logs=False)
-        if rc != 0:
-            logger.debug(f'Failed to get job info for job {job_id}: {stderr}')
-            return None
+        rc, stdout, stderr = self._run_slurm_cmd(cmd)
+        subprocess_utils.handle_returncode(
+            rc,
+            cmd,
+            f'Failed to get job reason for job {job_id}.',
+            stderr=f'{stdout}\n{stderr}')
 
         output = stdout.strip()
         if not output:
@@ -312,9 +387,7 @@ class SlurmClient:
 
             # Check if nodes are allocated by trying to get node list
             cmd = f'squeue -h --jobs {job_id} -o "%N"'
-            rc, stdout, stderr = self._runner.run(cmd,
-                                                  require_outputs=True,
-                                                  stream_logs=False)
+            rc, stdout, stderr = self._run_slurm_cmd(cmd)
 
             if rc == 0 and stdout.strip():
                 # Nodes are allocated
@@ -322,7 +395,8 @@ class SlurmClient:
                     f'Job {job_id} has nodes allocated: {stdout.strip()}')
                 return
             elif rc != 0:
-                logger.debug(f'Failed to get nodes for job {job_id}: {stderr}')
+                logger.debug(f'Failed to get nodes for job {job_id}: '
+                             f'{stdout}\n{stderr}')
 
             # Wait before checking again
             time.sleep(2)
@@ -359,15 +433,16 @@ class SlurmClient:
             f'squeue -h --jobs {job_id} -o "%N" | tr \',\' \'\\n\' | '
             f'while read node; do '
             # TODO(kevin): Use json output for more robust parsing.
-            f'ip=$(scontrol show node=$node | grep NodeAddr= | '
+            f'node_addr=$(scontrol show node=$node | grep NodeAddr= | '
             f'awk -F= \'{{print $2}}\' | awk \'{{print $1}}\'); '
-            f'echo "$node $ip"; '
+            f'echo "$node $node_addr"; '
             f'done')
-        rc, stdout, stderr = self._runner.run(cmd,
-                                              require_outputs=True,
-                                              stream_logs=False)
+        rc, stdout, stderr = self._run_slurm_cmd(cmd)
         subprocess_utils.handle_returncode(
-            rc, cmd, f'Failed to get nodes for job {job_id}.', stderr=stderr)
+            rc,
+            cmd,
+            f'Failed to get nodes for job {job_id}.',
+            stderr=f'{stdout}\n{stderr}')
         logger.debug(f'Successfully got nodes for job {job_id}: {stdout}')
 
         node_info = {}
@@ -377,7 +452,23 @@ class SlurmClient:
                 parts = line.split()
                 if len(parts) >= 2:
                     node_name = parts[0]
-                    node_ip = parts[1]
+                    node_addr = parts[1]
+                    # Resolve hostname to IP if node_addr is not already
+                    # an IP address.
+                    try:
+                        ipaddress.ip_address(node_addr)
+                        # Already an IP address
+                        node_ip = node_addr
+                    except ValueError:
+                        # It's a hostname, resolve it to an IP
+                        try:
+                            node_ip = socket.gethostbyname(node_addr)
+                        except socket.gaierror as e:
+                            raise RuntimeError(
+                                f'Failed to resolve hostname {node_addr} to IP '
+                                f'for node {node_name}: '
+                                f'{common_utils.format_exception(e)}') from e
+
                     node_info[node_name] = node_ip
 
         nodes = list(node_info.keys())
@@ -408,9 +499,7 @@ class SlurmClient:
             The job ID of the submitted job.
         """
         cmd = f'sbatch --partition={partition} {script_path}'
-        rc, stdout, stderr = self._runner.run(cmd,
-                                              require_outputs=True,
-                                              stream_logs=False)
+        rc, stdout, stderr = self._run_slurm_cmd(cmd)
         subprocess_utils.handle_returncode(rc,
                                            cmd,
                                            'Failed to submit Slurm job.',
@@ -435,13 +524,11 @@ class SlurmClient:
             List of SlurmPartition objects.
         """
         cmd = 'scontrol show partitions -o'
-        rc, stdout, stderr = self._runner.run(cmd,
-                                              require_outputs=True,
-                                              stream_logs=False)
+        rc, stdout, stderr = self._run_slurm_cmd(cmd)
         subprocess_utils.handle_returncode(rc,
                                            cmd,
                                            'Failed to get Slurm partitions.',
-                                           stderr=stderr)
+                                           stderr=f'{stdout}\n{stderr}')
 
         partitions = []
         for line in stdout.strip().splitlines():

sky/backends/backend_utils.py

@@ -69,6 +69,7 @@ from sky.utils import timeline
 from sky.utils import ux_utils
 from sky.utils import volume as volume_utils
 from sky.utils import yaml_utils
+from sky.utils.plugin_extensions import ExternalFailureSource
 from sky.workspaces import core as workspaces_core
 
 if typing.TYPE_CHECKING:
@@ -763,7 +764,20 @@ def write_cluster_config(
                     keys=('allowed_contexts',),
                     default_value=None)
             if allowed_contexts is None:
-                excluded_clouds.add(cloud)
+                # Exclude both Kubernetes and SSH explicitly since:
+                # 1. isinstance(cloud, clouds.Kubernetes) matches both (SSH
+                #    inherits from Kubernetes)
+                # 2. Both share the same get_credential_file_mounts() which
+                #    returns the kubeconfig. So if we don't exclude both, the
+                #    unexcluded one will upload the kubeconfig.
+                # TODO(romilb): This is a workaround. The right long-term fix
+                # is to have SSH Node Pools use its own kubeconfig instead of
+                # sharing the global kubeconfig at ~/.kube/config. In the
+                # interim, SSH Node Pools' get_credential_file_mounts can filter
+                # contexts starting with ssh- and create a temp kubeconfig
+                # to upload.
+                excluded_clouds.add(clouds.Kubernetes())
+                excluded_clouds.add(clouds.SSH())
         else:
             excluded_clouds.add(cloud)
 
@@ -2262,6 +2276,12 @@ def _update_cluster_status(
                         for status in node_statuses) and
                     len(node_statuses) == handle.launched_nodes)
 
+    external_cluster_failures = ExternalFailureSource.get(
+        cluster_hash=record['cluster_hash'])
+    logger.debug(f'Cluster {cluster_name} with cluster_hash '
+                 f'{record["cluster_hash"]} has external cluster failures: '
+                 f'{external_cluster_failures}')
+
     def get_node_counts_from_ray_status(
             runner: command_runner.CommandRunner) -> Tuple[int, int, str, str]:
         rc, output, stderr = runner.run(
@@ -2401,8 +2421,9 @@ def _update_cluster_status(
 
     # For Slurm, skip Ray health check since it doesn't use Ray.
     should_check_ray = cloud is not None and cloud.uses_ray()
-    if all_nodes_up and (not should_check_ray or
-                         run_ray_status_to_check_ray_cluster_healthy()):
+    if (all_nodes_up and (not should_check_ray or
+                          run_ray_status_to_check_ray_cluster_healthy()) and
+            not external_cluster_failures):
         # NOTE: all_nodes_up calculation is fast due to calling cloud CLI;
         # run_ray_status_to_check_all_nodes_up() is slow due to calling `ray get
         # head-ip/worker-ips`.
@@ -2505,15 +2526,15 @@ def _update_cluster_status(
     #  (2) Otherwise, we will reset the autostop setting, unless the cluster is
     #      autostopping/autodowning.
     some_nodes_terminated = 0 < len(node_statuses) < handle.launched_nodes
-    # If all nodes are up and ray cluster is health, we would have returned
-    # earlier. So if all_nodes_up is True and we are here, it means the ray
-    # cluster must have been unhealthy.
-    ray_cluster_unhealthy = all_nodes_up
     some_nodes_not_stopped = any(status[0] != status_lib.ClusterStatus.STOPPED
                                  for status in node_statuses)
     is_abnormal = (some_nodes_terminated or some_nodes_not_stopped)
 
-    if is_abnormal:
+    if is_abnormal and not external_cluster_failures:
+        # If all nodes are up and ray cluster is healthy, we would have returned
+        # earlier. So if all_nodes_up is True and we are here, it means the ray
+        # cluster must have been unhealthy.
+        ray_cluster_unhealthy = all_nodes_up
         status_reason = ', '.join(
             [status[1] for status in node_statuses if status[1] is not None])
 
@@ -2641,8 +2662,25 @@ def _update_cluster_status(
             cluster_name,
             include_user_info=include_user_info,
             summary_response=summary_response)
-    # Now is_abnormal is False: either node_statuses is empty or all nodes are
-    # STOPPED.
+    # Now either:
+    # (1) is_abnormal is False: either node_statuses is empty or all nodes are
+    #                           STOPPED
+    # or
+    # (2) there are external cluster failures reported by a plugin.
+
+    # If there are external cluster failures and the cluster has not been
+    # terminated on cloud (to_terminate), we can return the cluster record as is.
+    # This is because when an external failure is detected, the cluster will be
+    # marked as INIT with a reason indicating the details of the failure. So, we
+    # do not want to modify the cluster status in this function except for in the
+    # case where the cluster has been terminated on cloud, in which case we should
+    # clean up the cluster from SkyPilot's global state.
+    if external_cluster_failures and not to_terminate:
+        return global_user_state.get_cluster_from_name(
+            cluster_name,
+            include_user_info=include_user_info,
+            summary_response=summary_response)
+
     verb = 'terminated' if to_terminate else 'stopped'
     backend = backends.CloudVmRayBackend()
     global_user_state.add_cluster_event(
@@ -3368,6 +3406,8 @@ def get_clusters(
             handle = record['handle']
             record['nodes'] = handle.launched_nodes
             if handle.launched_resources is None:
+                # Set default values when launched_resources is None
+                record['labels'] = {}
                 continue
             record['cloud'] = (f'{handle.launched_resources.cloud}'
                                if handle.launched_resources.cloud else None)
@@ -3380,6 +3420,8 @@ def get_clusters(
             record['accelerators'] = (
                 f'{handle.launched_resources.accelerators}'
                 if handle.launched_resources.accelerators else None)
+            record['labels'] = (handle.launched_resources.labels
+                                if handle.launched_resources.labels else {})
             if not include_handle:
                 record.pop('handle', None)