skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/utils/cli_utils/status_utils.py

@@ -1,14 +1,22 @@
 """Utilities for sky status."""
+import typing
 from typing import Any, Callable, Dict, List, Optional
 
 import click
 import colorama
 
 from sky import backends
-from sky import status_lib
 from sky.skylet import constants
+from sky.utils import common_utils
 from sky.utils import log_utils
 from sky.utils import resources_utils
+from sky.utils import status_lib
+
+if typing.TYPE_CHECKING:
+    from sky.provision.kubernetes import utils as kubernetes_utils
+
+if typing.TYPE_CHECKING:
+    from sky.provision.kubernetes import utils as kubernetes_utils
 
 COMMAND_TRUNC_LENGTH = 25
 NUM_COST_REPORT_LINES = 5
@@ -19,25 +27,6 @@ _ClusterRecord = Dict[str, Any]
 _ClusterCostReportRecord = Dict[str, Any]
 
 
-def truncate_long_string(s: str, max_length: int = 35) -> str:
-    if len(s) <= max_length:
-        return s
-    splits = s.split(' ')
-    if len(splits[0]) > max_length:
-        return splits[0][:max_length] + '...'  # Use '…'?
-    # Truncate on word boundary.
-    i = 0
-    total = 0
-    for i, part in enumerate(splits):
-        total += len(part)
-        if total >= max_length:
-            break
-    prefix = ' '.join(splits[:i])
-    if len(prefix) < max_length:
-        prefix += s[len(prefix):max_length]
-    return prefix + '...'
-
-
 class StatusColumn:
     """One column of the displayed cluster table"""
 
@@ -54,12 +43,14 @@ class StatusColumn:
     def calc(self, record):
         val = self.calc_func(record)
         if self.trunc_length != 0:
-            val = truncate_long_string(str(val), self.trunc_length)
+            val = common_utils.truncate_long_string(str(val), self.trunc_length)
         return val
 
 
 def show_status_table(cluster_records: List[_ClusterRecord],
-                      show_all: bool) -> int:
+                      show_all: bool,
+                      show_user: bool,
+                      query_clusters: Optional[List[str]] = None) -> int:
     """Compute cluster table values and display.
 
     Returns:
@@ -70,6 +61,13 @@ def show_status_table(cluster_records: List[_ClusterRecord],
 
     status_columns = [
         StatusColumn('NAME', _get_name),
+    ]
+    if show_user:
+        status_columns.append(StatusColumn('USER', _get_user_name))
+        status_columns.append(
+            StatusColumn('USER_ID', _get_user_hash, show_by_default=False))
+
+    status_columns += [
         StatusColumn('LAUNCHED', _get_launched),
         StatusColumn('RESOURCES',
                      _get_resources,
@@ -101,7 +99,21 @@ def show_status_table(cluster_records: List[_ClusterRecord],
 
     if cluster_records:
         click.echo(cluster_table)
-    else:
+
+    if query_clusters:
+        cluster_names = {record['name'] for record in cluster_records}
+        not_found_clusters = [
+            repr(cluster)
+            for cluster in query_clusters
+            if cluster not in cluster_names
+        ]
+        cluster_str = 'Cluster'
+        if len(not_found_clusters) > 1:
+            cluster_str += 's'
+        cluster_str += ' '
+        cluster_str += ', '.join(not_found_clusters)
+        click.echo(f'{cluster_str} not found.')
+    elif not cluster_records:
         click.echo('No existing clusters.')
     return num_pending_autostop
 
@@ -202,6 +214,8 @@ def show_cost_report_table(cluster_records: List[_ClusterCostReportRecord],
 # _ClusterCostReportRecord, which is okay as we guarantee the queried fields
 # exist in those cases.
 _get_name = (lambda cluster_record: cluster_record['name'])
+_get_user_hash = (lambda cluster_record: cluster_record['user_hash'])
+_get_user_name = (lambda cluster_record: cluster_record.get('user_name', '-'))
 _get_launched = (lambda cluster_record: log_utils.readable_time_duration(
     cluster_record['launched_at']))
 _get_region = (
@@ -220,6 +234,8 @@ def _get_status_colored(cluster_record: _ClusterRecord) -> str:
 
 
 def _get_resources(cluster_record: _ClusterRecord) -> str:
+    if 'resources_str' in cluster_record:
+        return cluster_record['resources_str']
     handle = cluster_record['handle']
     if isinstance(handle, backends.LocalDockerResourceHandle):
         resources_str = 'docker'
@@ -316,3 +332,45 @@ def _get_estimated_cost_for_cost_report(
         return '-'
 
     return f'$ {cost:.2f}'
+
+
+def show_kubernetes_cluster_status_table(
+        clusters: List['kubernetes_utils.KubernetesSkyPilotClusterInfo'],
+        show_all: bool) -> None:
+    """Compute cluster table values and display for Kubernetes clusters."""
+    status_columns = [
+        StatusColumn('USER', lambda c: c.user),
+        StatusColumn('NAME', lambda c: c.cluster_name),
+        StatusColumn('LAUNCHED',
+                     lambda c: log_utils.readable_time_duration(c.launched_at)),
+        StatusColumn('RESOURCES',
+                     lambda c: c.resources_str,
+                     trunc_length=70 if not show_all else 0),
+        StatusColumn('STATUS', lambda c: c.status.colored_str()),
+        # TODO(romilb): We should consider adding POD_NAME field here when --all
+        #  is passed to help users fetch pod name programmatically.
+    ]
+
+    columns = [
+        col.name for col in status_columns if col.show_by_default or show_all
+    ]
+    cluster_table = log_utils.create_table(columns)
+
+    # Sort table by user, then by cluster name
+    sorted_clusters = sorted(clusters, key=lambda c: (c.user, c.cluster_name))
+
+    for cluster in sorted_clusters:
+        row = []
+        for status_column in status_columns:
+            if status_column.show_by_default or show_all:
+                row.append(status_column.calc(cluster))
+        cluster_table.add_row(row)
+
+    if clusters:
+        click.echo(f'{colorama.Fore.CYAN}{colorama.Style.BRIGHT}'
+                   f'SkyPilot clusters'
+                   f'{colorama.Style.RESET_ALL}')
+        click.echo(cluster_table)
+    else:
+        click.echo('No SkyPilot resources found in the '
+                   'active Kubernetes context.')

sky/utils/cluster_utils.py

@@ -0,0 +1,356 @@
+"""Utility functions for cluster yaml file."""
+
+import functools
+import glob
+import os
+import re
+import textwrap
+from typing import Dict, List, Optional
+import uuid
+
+from sky.skylet import constants
+from sky.utils import command_runner
+from sky.utils import common_utils
+from sky.utils import timeline
+
+# The cluster yaml used to create the current cluster where the module is
+# called.
+SKY_CLUSTER_YAML_REMOTE_PATH = '~/.sky/sky_ray.yml'
+
+
+def get_provider_name(config: dict) -> str:
+    """Return the name of the provider."""
+
+    provider_module = config['provider']['module']
+    # Examples:
+    #   'sky.skylet.providers.aws.AWSNodeProviderV2' -> 'aws'
+    #   'sky.provision.aws' -> 'aws'
+    provider_search = re.search(r'(?:providers|provision)\.(\w+)\.?',
+                                provider_module)
+    assert provider_search is not None, config
+    provider_name = provider_search.group(1).lower()
+    # Special handling for lambda_cloud as Lambda cloud is registered as lambda.
+    if provider_name == 'lambda_cloud':
+        provider_name = 'lambda'
+    return provider_name
+
+
+class SSHConfigHelper(object):
+    """Helper for handling local SSH configuration."""
+
+    ssh_conf_path = '~/.ssh/config'
+    ssh_conf_lock_path = os.path.expanduser('~/.sky/locks/.ssh_config.lock')
+    ssh_conf_per_cluster_lock_path = os.path.expanduser(
+        '~/.sky/locks/.ssh_config_{}.lock')
+    ssh_cluster_path = constants.SKY_USER_FILE_PATH + '/ssh/{}'
+    ssh_cluster_key_path = constants.SKY_USER_FILE_PATH + '/ssh-keys/{}.key'
+
+    @classmethod
+    def _get_generated_config(cls, autogen_comment: str, host_name: str,
+                              ip: str, username: str, ssh_key_path: str,
+                              proxy_command: Optional[str], port: int,
+                              docker_proxy_command: Optional[str]):
+        if proxy_command is not None:
+            # Already checked in resources
+            assert docker_proxy_command is None, (
+                'Cannot specify both proxy_command and docker_proxy_command.')
+            proxy = f'ProxyCommand {proxy_command}'
+        elif docker_proxy_command is not None:
+            proxy = f'ProxyCommand {docker_proxy_command}'
+        else:
+            proxy = ''
+        # StrictHostKeyChecking=no skips the host key check for the first
+        # time. UserKnownHostsFile=/dev/null and GlobalKnownHostsFile/dev/null
+        # prevent the host key from being added to the known_hosts file and
+        # always return an empty file for known hosts, making the ssh think
+        # this is a first-time connection, and thus skipping the host key
+        # check.
+        # Not adding SSH agent forwarding by default here to avoid implicitly
+        # using users' SSH keys in their local agent. Plus on sky launch side we
+        # are not default adding SSH agent forwarding either.
+        codegen = textwrap.dedent(f"""\
+            {autogen_comment}
+            Host {host_name}
+              HostName {ip}
+              User {username}
+              IdentityFile {ssh_key_path}
+              IdentitiesOnly yes
+              StrictHostKeyChecking no
+              UserKnownHostsFile=/dev/null
+              GlobalKnownHostsFile=/dev/null
+              Port {port}
+              {proxy}
+            """.rstrip())
+        codegen = codegen + '\n'
+        return codegen
+
+    @classmethod
+    def generate_local_key_file(cls, cluster_name: str,
+                                auth_config: Dict[str, str]) -> str:
+        key_content = auth_config.pop('ssh_private_key_content', None)
+        if key_content is not None:
+            cluster_private_key_path = cls.ssh_cluster_key_path.format(
+                cluster_name)
+            expanded_cluster_private_key_path = os.path.expanduser(
+                cluster_private_key_path)
+            expanded_cluster_private_key_dir = os.path.dirname(
+                expanded_cluster_private_key_path)
+            os.makedirs(expanded_cluster_private_key_dir,
+                        exist_ok=True,
+                        mode=0o700)
+            with open(expanded_cluster_private_key_path,
+                      'w',
+                      encoding='utf-8',
+                      opener=functools.partial(os.open, mode=0o600)) as f:
+                f.write(key_content)
+            auth_config['ssh_private_key'] = cluster_private_key_path
+        return auth_config['ssh_private_key']
+
+    @classmethod
+    @timeline.FileLockEvent(ssh_conf_lock_path)
+    def add_cluster(
+        cls,
+        cluster_name: str,
+        ips: List[str],
+        auth_config: Dict[str, str],
+        ports: List[int],
+        docker_user: Optional[str] = None,
+        ssh_user: Optional[str] = None,
+    ):
+        """Add authentication information for cluster to local SSH config file.
+
+        If a host with `cluster_name` already exists and the configuration was
+        not added by sky, then `ip` is used to identify the host instead in the
+        file.
+
+        If a host with `cluster_name` already exists and the configuration was
+        added by sky (e.g. a spot instance), then the configuration is
+        overwritten.
+
+        Args:
+            cluster_name: Cluster name (see `sky status`)
+            ips: List of public IP addresses in the cluster. First IP is head
+              node.
+            auth_config: `auth` in cluster yaml.
+            ports: List of port numbers for SSH corresponding to ips
+            docker_user: If not None, use this user to ssh into the docker
+            ssh_user: Override the ssh_user in auth_config
+        """
+        if ssh_user is None:
+            username = auth_config['ssh_user']
+        else:
+            username = ssh_user
+        if docker_user is not None:
+            username = docker_user
+
+        key_path = cls.generate_local_key_file(cluster_name, auth_config)
+        key_path = os.path.expanduser(key_path)
+        sky_autogen_comment = ('# Added by sky (use `sky stop/down '
+                               f'{cluster_name}` to remove)')
+        ip = ips[0]
+        if docker_user is not None:
+            ip = 'localhost'
+
+        config_path = os.path.expanduser(cls.ssh_conf_path)
+        os.makedirs(os.path.dirname(config_path), exist_ok=True, mode=0o700)
+
+        if not os.path.exists(config_path):
+            config = ['\n']
+            with open(config_path,
+                      'w',
+                      encoding='utf-8',
+                      opener=functools.partial(os.open, mode=0o644)) as f:
+                f.writelines(config)
+
+        with open(config_path, 'r', encoding='utf-8') as f:
+            config = f.readlines()
+
+        ssh_dir = cls.ssh_cluster_path.format('')
+        os.makedirs(os.path.expanduser(ssh_dir), exist_ok=True, mode=0o700)
+
+        # Handle Include on top of Config file
+        include_str = f'Include {cls.ssh_cluster_path.format("*")}'
+        found = False
+        for i, line in enumerate(config):
+            config_str = line.strip()
+            if config_str == include_str:
+                found = True
+                break
+            if 'Host' in config_str:
+                break
+        if not found:
+            # Did not find Include string. Insert `Include` lines.
+            with open(config_path, 'w', encoding='utf-8') as f:
+                config.insert(
+                    0, '# Added by SkyPilot for ssh config of all clusters\n'
+                    f'{include_str}\n')
+                f.write(''.join(config).strip())
+                f.write('\n' * 2)
+
+        proxy_command = auth_config.get('ssh_proxy_command', None)
+
+        docker_proxy_command_generator = None
+        if docker_user is not None:
+            docker_proxy_command_generator = lambda ip, port: ' '.join(
+                ['ssh'] + command_runner.ssh_options_list(
+                    key_path, ssh_control_name=None, port=port) +
+                ['-W', '%h:%p', f'{auth_config["ssh_user"]}@{ip}'])
+
+        codegen = ''
+        # Add the nodes to the codegen
+        for i, ip in enumerate(ips):
+            docker_proxy_command = None
+            port = ports[i]
+            if docker_proxy_command_generator is not None:
+                docker_proxy_command = docker_proxy_command_generator(ip, port)
+                ip = 'localhost'
+                port = constants.DEFAULT_DOCKER_PORT
+            node_name = cluster_name if i == 0 else cluster_name + f'-worker{i}'
+            # TODO(romilb): Update port number when k8s supports multinode
+            codegen += cls._get_generated_config(
+                sky_autogen_comment, node_name, ip, username, key_path,
+                proxy_command, port, docker_proxy_command) + '\n'
+
+        cluster_config_path = os.path.expanduser(
+            cls.ssh_cluster_path.format(cluster_name))
+
+        with open(cluster_config_path,
+                  'w',
+                  encoding='utf-8',
+                  opener=functools.partial(os.open, mode=0o644)) as f:
+            f.write(codegen)
+
+    @classmethod
+    def _remove_stale_cluster_config_for_backward_compatibility(
+        cls,
+        cluster_name: str,
+        ip: str,
+        auth_config: Dict[str, str],
+        docker_user: Optional[str] = None,
+    ):
+        """Remove authentication information for cluster from local SSH config.
+
+        If no existing host matching the provided specification is found, then
+        nothing is removed.
+
+        Args:
+            ip: Head node's IP address.
+            auth_config: `auth` in cluster yaml.
+            docker_user: If not None, use this user to ssh into the docker
+        """
+        username = auth_config['ssh_user']
+        config_path = os.path.expanduser(cls.ssh_conf_path)
+        cluster_config_path = os.path.expanduser(
+            cls.ssh_cluster_path.format(cluster_name))
+        if not os.path.exists(config_path):
+            return
+
+        with open(config_path, 'r', encoding='utf-8') as f:
+            config = f.readlines()
+
+        start_line_idx = None
+
+        # Scan the config for the cluster name.
+        for i, line in enumerate(config):
+            next_line = config[i + 1] if i + 1 < len(config) else ''
+            if docker_user is None:
+                found = (line.strip() == f'HostName {ip}' and
+                         next_line.strip() == f'User {username}')
+            else:
+                found = (line.strip() == 'HostName localhost' and
+                         next_line.strip() == f'User {docker_user}')
+                if found:
+                    # Find the line starting with ProxyCommand and contains ip
+                    found = False
+                    for idx in range(i, len(config)):
+                        # Stop if we reach an empty line, which means a new host
+                        if not config[idx].strip():
+                            break
+                        if config[idx].strip().startswith('ProxyCommand'):
+                            proxy_command_line = config[idx].strip()
+                            if proxy_command_line.endswith(f'@{ip}'):
+                                found = True
+                                break
+            if found:
+                start_line_idx = i - 1
+                break
+
+        if start_line_idx is not None:
+            # Scan for end of previous config.
+            cursor = start_line_idx
+            while cursor > 0 and len(config[cursor].strip()) > 0:
+                cursor -= 1
+            prev_end_line_idx = cursor
+
+            # Scan for end of the cluster config.
+            end_line_idx = None
+            cursor = start_line_idx + 1
+            start_line_idx -= 1  # remove auto-generated comment
+            while cursor < len(config):
+                if config[cursor].strip().startswith(
+                        '# ') or config[cursor].strip().startswith('Host '):
+                    end_line_idx = cursor
+                    break
+                cursor += 1
+
+            # Remove sky-generated config and update the file.
+            config[prev_end_line_idx:end_line_idx] = [
+                '\n'
+            ] if end_line_idx is not None else []
+            with open(config_path, 'w', encoding='utf-8') as f:
+                f.write(''.join(config).strip())
+                f.write('\n' * 2)
+
+        # Delete include statement if it exists in the config.
+        sky_autogen_comment = ('# Added by sky (use `sky stop/down '
+                               f'{cluster_name}` to remove)')
+        with open(config_path, 'r', encoding='utf-8') as f:
+            config = f.readlines()
+
+        for i, line in enumerate(config):
+            config_str = line.strip()
+            if f'Include {cluster_config_path}' in config_str:
+                with open(config_path, 'w', encoding='utf-8') as f:
+                    if i < len(config) - 1 and config[i + 1] == '\n':
+                        del config[i + 1]
+                    # Delete Include string
+                    del config[i]
+                    # Delete Sky Autogen Comment
+                    if i > 0 and sky_autogen_comment in config[i - 1].strip():
+                        del config[i - 1]
+                    f.write(''.join(config))
+                break
+            if 'Host' in config_str:
+                break
+
+    @classmethod
+    def remove_cluster(cls, cluster_name: str):
+        """Remove auth information for cluster from ~/.sky/ssh/<cluster_name>.
+
+        If no existing host matching the provided specification is found, then
+        nothing is removed.
+
+        Args:
+            cluster_name: Cluster name.
+        """
+
+        with timeline.FileLockEvent(
+                cls.ssh_conf_per_cluster_lock_path.format(cluster_name)):
+            cluster_config_path = os.path.expanduser(
+                cls.ssh_cluster_path.format(cluster_name))
+            common_utils.remove_file_if_exists(cluster_config_path)
+
+    @classmethod
+    def list_cluster_names(cls) -> List[str]:
+        """List all names of clusters with SSH config set up."""
+        cluster_config_dir = os.path.expanduser(cls.ssh_cluster_path.format(''))
+        return [
+            os.path.basename(path)
+            for path in glob.glob(os.path.join(cluster_config_dir, '*'))
+        ]
+
+
+def generate_cluster_name():
+    # TODO: change this ID formatting to something more pleasant.
+    # User name is helpful in non-isolated accounts, e.g., GCP, Azure.
+    return f'sky-{uuid.uuid4().hex[:4]}-{common_utils.get_cleaned_username()}'