Repository: skypilot-org/skypilot

Issue:
[Provisioner] Parallelize wait for SSH in provisioner
<!-- Describe the bug report / feature request here -->

A user reported that we wait SSH access in provisioner sequentially, which causes inefficiency when the number of nodes increase. We should parallelize it.

https://github.com/skypilot-org/skypilot/blob/f2991b144d4b15eac55dd7f759f361b6146033b3/sky/provision/provisioner.py#L345-L382


Discussion:
asaiacai: these are some preliminary numbers but need to also test k8s and run smoke tests
before PR
```
(base) ubuntu@ip-172-31-16-12:~/skypilot$ time sky launch --cloud aws --num-nodes 20 --cpus 2+ -y -c scale --region us-east-1

real    1m13.987s
user    0m13.215s
sys     0m1.214s
```

after PR
```
(base) ubuntu@ip-172-31-16-12:~/skypilot$ time sky launch --cloud aws --num-nodes 20 --cpus 2+ -y -c scale --region us-east-1

real    0m58.341s
user    0m11.981s
sys     0m1.223s
```

Patch:
diff --git a/docs/source/reference/kubernetes/kubernetes-getting-started.rst b/docs/source/reference/kubernetes/kubernetes-getting-started.rst
index d7313fba..c1874b6b 100644
--- a/docs/source/reference/kubernetes/kubernetes-getting-started.rst
+++ b/docs/source/reference/kubernetes/kubernetes-getting-started.rst
@@ -174,7 +174,12 @@ You can also inspect the real-time GPU usage on the cluster with :code:`sky show
 
 Using Custom Images
 -------------------
-By default, we use and maintain a SkyPilot container image that has conda and a few other basic tools installed.
+By default, we maintain and use two SkyPilot container images for use on Kubernetes clusters:
+
+1. ``us-central1-docker.pkg.dev/skypilot-375900/skypilotk8s/skypilot``: used for CPU-only clusters (`Dockerfile <https://github.com/skypilot-org/skypilot/blob/master/Dockerfile_k8s>`__).
+2. ``us-central1-docker.pkg.dev/skypilot-375900/skypilotk8s/skypilot-gpu``: used for GPU clusters (`Dockerfile <https://github.com/skypilot-org/skypilot/blob/master/Dockerfile_k8s_gpu>`__).
+
+These images are pre-installed with SkyPilot dependencies for fast startup.
 
 To use your own image, add :code:`image_id: docker:<your image tag>` to the :code:`resources` section of your task YAML.
 
diff --git a/sky/__init__.py b/sky/__init__.py
index 37b5a1ca..b851775d 100644
--- a/sky/__init__.py
+++ b/sky/__init__.py
@@ -128,6 +128,7 @@ GCP = clouds.GCP
 Lambda = clouds.Lambda
 SCP = clouds.SCP
 Kubernetes = clouds.Kubernetes
+K8s = Kubernetes
 OCI = clouds.OCI
 Paperspace = clouds.Paperspace
 RunPod = clouds.RunPod
@@ -143,6 +144,7 @@ __all__ = [
     'GCP',
     'IBM',
     'Kubernetes',
+    'K8s',
     'Lambda',
     'OCI',
     'Paperspace',
diff --git a/sky/cli.py b/sky/cli.py
index fb5a38bb..7ac26fd0 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -339,7 +339,6 @@ def _get_shell_complete_args(complete_fn):
 
 
 _RELOAD_ZSH_CMD = 'source ~/.zshrc'
-_RELOAD_FISH_CMD = 'source ~/.config/fish/config.fish'
 _RELOAD_BASH_CMD = 'source ~/.bashrc'
 
 
@@ -378,7 +377,9 @@ def _install_shell_completion(ctx: click.Context, param: click.Parameter,
         cmd = '_SKY_COMPLETE=fish_source sky > \
                 ~/.config/fish/completions/sky.fish'
 
-        reload_cmd = _RELOAD_FISH_CMD
+        # Fish does not need to be reloaded and will automatically pick up
+        # completions.
+        reload_cmd = None
 
     elif value == 'zsh':
         install_cmd = f'_SKY_COMPLETE=zsh_source sky > \
@@ -398,9 +399,10 @@ def _install_shell_completion(ctx: click.Context, param: click.Parameter,
                        check=True,
                        executable=shutil.which('bash'))
         click.secho(f'Shell completion installed for {value}', fg='green')
-        click.echo(
-            'Completion will take effect once you restart the terminal: ' +
-            click.style(f'{reload_cmd}', bold=True))
+        if reload_cmd is not None:
+            click.echo(
+                'Completion will take effect once you restart the terminal: ' +
+                click.style(f'{reload_cmd}', bold=True))
     except subprocess.CalledProcessError as e:
         click.secho(f'> Installation failed with code {e.returncode}', fg='red')
     ctx.exit()
@@ -431,7 +433,9 @@ def _uninstall_shell_completion(ctx: click.Context, param: click.Parameter,
 
     elif value == 'fish':
         cmd = 'rm -f ~/.config/fish/completions/sky.fish'
-        reload_cmd = _RELOAD_FISH_CMD
+        # Fish does not need to be reloaded and will automatically pick up
+        # completions.
+        reload_cmd = None
 
     elif value == 'zsh':
         cmd = 'sed -i"" -e "/# For SkyPilot shell completion/d" ~/.zshrc && \
@@ -447,8 +451,10 @@ def _uninstall_shell_completion(ctx: click.Context, param: click.Parameter,
     try:
         subprocess.run(cmd, shell=True, check=True)
         click.secho(f'Shell completion uninstalled for {value}', fg='green')
-        click.echo('Changes will take effect once you restart the terminal: ' +
-                   click.style(f'{reload_cmd}', bold=True))
+        if reload_cmd is not None:
+            click.echo(
+                'Changes will take effect once you restart the terminal: ' +
+                click.style(f'{reload_cmd}', bold=True))
     except subprocess.CalledProcessError as e:
         click.secho(f'> Uninstallation failed with code {e.returncode}',
                     fg='red')
@@ -3056,7 +3062,8 @@ def show_gpus(
 
     # This will validate 'cloud' and raise if not found.
     cloud_obj = sky_clouds.CLOUD_REGISTRY.from_str(cloud)
-    service_catalog.validate_region_zone(region, None, clouds=cloud)
+    cloud_name = cloud_obj.canonical_name() if cloud_obj is not None else None
+    service_catalog.validate_region_zone(region, None, clouds=cloud_name)
     show_all = all
     if show_all and accelerator_str is not None:
         raise click.UsageError('--all is only allowed without a GPU name.')
@@ -3142,8 +3149,8 @@ def show_gpus(
         # Optimization - do not poll for Kubernetes API for fetching
         # common GPUs because that will be fetched later for the table after
         # common GPUs.
-        clouds_to_list = cloud
-        if cloud is None:
+        clouds_to_list = cloud_name
+        if cloud_name is None:
             clouds_to_list = [
                 c for c in service_catalog.ALL_CLOUDS if c != 'kubernetes'
             ]
@@ -3153,7 +3160,8 @@ def show_gpus(
             # Collect k8s related messages in k8s_messages and print them at end
             print_section_titles = False
             # If cloud is kubernetes, we want to show real-time capacity
-            if kubernetes_is_enabled and (cloud is None or cloud_is_kubernetes):
+            if kubernetes_is_enabled and (cloud_name is None or
+                                          cloud_is_kubernetes):
                 if region:
                     context = region
                 else:
@@ -3263,8 +3271,8 @@ def show_gpus(
                 name, quantity = accelerator_str, None
 
         print_section_titles = False
-        if (kubernetes_is_enabled and (cloud is None or cloud_is_kubernetes) and
-                not show_all):
+        if (kubernetes_is_enabled and
+            (cloud_name is None or cloud_is_kubernetes) and not show_all):
             # Print section title if not showing all and instead a specific
             # accelerator is requested
             print_section_titles = True
@@ -3336,7 +3344,7 @@ def show_gpus(
         if len(result) == 0:
             quantity_str = (f' with requested quantity {quantity}'
                             if quantity else '')
-            cloud_str = f' on {cloud_obj}.' if cloud else ' in cloud catalogs.'
+            cloud_str = f' on {cloud_obj}.' if cloud_name else ' in cloud catalogs.'
             yield f'Resources \'{name}\'{quantity_str} not found{cloud_str} '
             yield 'To show available accelerators, run: sky show-gpus --all'
             return
diff --git a/sky/clouds/cloud.py b/sky/clouds/cloud.py
index dae1d56d..3e21204f 100644
--- a/sky/clouds/cloud.py
+++ b/sky/clouds/cloud.py
@@ -819,6 +819,10 @@ class Cloud:
 
     # === End of image related methods ===
 
+    @classmethod
+    def canonical_name(cls) -> str:
+        return cls.__name__.lower()
+
     def __repr__(self):
         return self._REPR
 
diff --git a/sky/clouds/cloud_registry.py b/sky/clouds/cloud_registry.py
index 5c4b10b9..52a026aa 100644
--- a/sky/clouds/cloud_registry.py
+++ b/sky/clouds/cloud_registry.py
@@ -1,7 +1,7 @@
 """Clouds need to be registered in CLOUD_REGISTRY to be discovered"""
 
 import typing
-from typing import Optional, Type
+from typing import Callable, Dict, List, Optional, overload, Type, Union
 
 from sky.utils import ux_utils
 
@@ -12,20 +12,65 @@ if typing.TYPE_CHECKING:
 class _CloudRegistry(dict):
     """Registry of clouds."""
 
+    def __init__(self) -> None:
+        super().__init__()
+        self.aliases: Dict[str, str] = {}
+
     def from_str(self, name: Optional[str]) -> Optional['cloud.Cloud']:
+        """Returns the cloud instance from the canonical name or alias."""
         if name is None:
             return None
-        if name.lower() not in self:
-            with ux_utils.print_exception_no_traceback():
-                raise ValueError(f'Cloud {name!r} is not a valid cloud among '
-                                 f'{list(self.keys())}')
-        return self.get(name.lower())
 
+        search_name = name.lower()
+
+        if search_name in self:
+            return self[search_name]
+
+        if search_name in self.aliases:
+            return self[self.aliases[search_name]]
+
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(f'Cloud {name!r} is not a valid cloud among '
+                             f'{[*self.keys(), *self.aliases.keys()]}')
+
+    @overload
     def register(self, cloud_cls: Type['cloud.Cloud']) -> Type['cloud.Cloud']:
-        name = cloud_cls.__name__.lower()
-        assert name not in self, f'{name} already registered'
-        self[name] = cloud_cls()
-        return cloud_cls
+        ...
+
+    @overload
+    def register(
+        self,
+        cloud_cls: None = None,
+        aliases: Optional[List[str]] = None,
+    ) -> Callable[[Type['cloud.Cloud']], Type['cloud.Cloud']]:
+        ...
+
+    def register(
+        self,
+        cloud_cls: Optional[Type['cloud.Cloud']] = None,
+        aliases: Optional[List[str]] = None,
+    ) -> Union[Type['cloud.Cloud'], Callable[[Type['cloud.Cloud']],
+                                             Type['cloud.Cloud']]]:
+
+        def _register(cloud_cls: Type['cloud.Cloud']) -> Type['cloud.Cloud']:
+            name = cloud_cls.canonical_name()
+            assert name not in self, f'{name} already registered'
+            self[name] = cloud_cls()
+
+            for alias in aliases or []:
+                alias = alias.lower()
+                assert alias not in self.aliases, (
+                    f'alias {alias} already registered')
+                self.aliases[alias] = name
+
+            return cloud_cls
+
+        if cloud_cls is not None:
+            # invocation without parens (e.g. just `@register`)
+            return _register(cloud_cls)
+
+        # Invocation with parens (e.g. `@register(aliases=['alias'])`)
+        return _register
 
 
 CLOUD_REGISTRY: _CloudRegistry = _CloudRegistry()
diff --git a/sky/clouds/kubernetes.py b/sky/clouds/kubernetes.py
index da85246e..8ff4172a 100644
--- a/sky/clouds/kubernetes.py
+++ b/sky/clouds/kubernetes.py
@@ -33,7 +33,7 @@ CREDENTIAL_PATH = os.environ.get('KUBECONFIG', DEFAULT_KUBECONFIG_PATH)
 _SKYPILOT_SYSTEM_NAMESPACE = 'skypilot-system'
 
 
-@clouds.CLOUD_REGISTRY.register
+@clouds.CLOUD_REGISTRY.register(aliases=['k8s'])
 class Kubernetes(clouds.Cloud):
     """Kubernetes."""
 
diff --git a/sky/provision/kubernetes/instance.py b/sky/provision/kubernetes/instance.py
index 6663ed3f..6ce7b74d 100644
--- a/sky/provision/kubernetes/instance.py
+++ b/sky/provision/kubernetes/instance.py
@@ -18,6 +18,7 @@ from sky.provision.kubernetes import utils as kubernetes_utils
 from sky.utils import command_runner
 from sky.utils import common_utils
 from sky.utils import kubernetes_enums
+from sky.utils import subprocess_utils
 from sky.utils import ux_utils
 
 POLL_INTERVAL = 2
@@ -398,8 +399,7 @@ def _setup_ssh_in_pods(namespace: str, context: Optional[str],
         # See https://www.educative.io/answers/error-mesg-ttyname-failed-inappropriate-ioctl-for-device  # pylint: disable=line-too-long
         '$(prefix_cmd) sed -i "s/mesg n/tty -s \\&\\& mesg n/" ~/.profile;')
 
-    # TODO(romilb): Parallelize the setup of SSH in pods for multi-node clusters
-    for new_node in new_nodes:
+    def _setup_ssh_thread(new_node):
         pod_name = new_node.metadata.name
         runner = command_runner.KubernetesCommandRunner(
             ((namespace, context), pod_name))
@@ -411,6 +411,8 @@ def _setup_ssh_in_pods(namespace: str, context: Optional[str],
                                      stdout)
         logger.info(f'{"-"*20}End: Set up SSH in pod {pod_name!r} {"-"*20}')
 
+    subprocess_utils.run_in_parallel(_setup_ssh_thread, new_nodes)
+
 
 def _label_pod(namespace: str, context: Optional[str], pod_name: str,
                label: Dict[str, str]) -> None:
diff --git a/sky/provision/provisioner.py b/sky/provision/provisioner.py
index 7706a3d4..b3e96576 100644
--- a/sky/provision/provisioner.py
+++ b/sky/provision/provisioner.py
@@ -28,6 +28,7 @@ from sky.skylet import constants
 from sky.utils import common_utils
 from sky.utils import resources_utils
 from sky.utils import rich_utils
+from sky.utils import subprocess_utils
 from sky.utils import ux_utils
 
 # Do not use __name__ as we do not want to propagate logs to sky.provision,
@@ -365,14 +366,13 @@ def wait_for_ssh(cluster_info: provision_common.ClusterInfo,
     # use a queue for SSH querying
     ips = collections.deque(ip_list)
     ssh_ports = collections.deque(port_list)
-    while ips:
-        ip = ips.popleft()
-        ssh_port = ssh_ports.popleft()
-        success, stderr = waiter(ip, ssh_port, **ssh_credentials)
-        if not success:
-            ips.append(ip)
-            ssh_ports.append(ssh_port)
-            if time.time() - start > timeout:
+
+    def _retry_ssh_thread(ip_ssh_port: Tuple[str, int]):
+        ip, ssh_port = ip_ssh_port
+        success = False
+        while not success:
+            success, stderr = waiter(ip, ssh_port, **ssh_credentials)
+            if not success and time.time() - start > timeout:
                 with ux_utils.print_exception_no_traceback():
                     raise RuntimeError(
                         f'Failed to SSH to {ip} after timeout {timeout}s, with '
@@ -380,6 +380,14 @@ def wait_for_ssh(cluster_info: provision_common.ClusterInfo,
             logger.debug('Retrying in 1 second...')
             time.sleep(1)
 
+    # try one node and multiprocess the rest
+    if ips:
+        ip = ips.popleft()
+        ssh_port = ssh_ports.popleft()
+        _retry_ssh_thread((ip, ssh_port))
+    subprocess_utils.run_in_parallel(_retry_ssh_thread,
+                                     list(zip(ips, ssh_ports)))
+
 
 def _post_provision_setup(
         cloud_name: str, cluster_name: resources_utils.ClusterName,


Test patch:
diff --git a/tests/common.py b/tests/common.py
index c6f08588..d41ff3be 100644
--- a/tests/common.py
+++ b/tests/common.py
@@ -70,6 +70,9 @@ def enable_all_clouds_in_monkeypatch(
                         lambda *_args, **_kwargs: [True, ''])
     monkeypatch.setattr('sky.provision.kubernetes.utils.get_spot_label',
                         lambda *_args, **_kwargs: [None, None])
+    monkeypatch.setattr(
+        'sky.provision.kubernetes.utils.is_kubeconfig_exec_auth',
+        lambda *_args, **_kwargs: [False, None])
 
     # monkeypatch class Kubernetes.
     monkeypatch.setattr(
diff --git a/tests/test_api.py b/tests/test_api.py
index 4d6658fc..5a33336d 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -1,7 +1,20 @@
 import sky
+from sky.clouds.cloud import Cloud
 
 
 def test_sky_launch(enable_all_clouds):
     task = sky.Task()
     job_id, handle = sky.launch(task, dryrun=True)
     assert job_id is None and handle is None
+
+
+def test_k8s_alias(enable_all_clouds):
+
+    def dryrun_task_with_cloud(cloud: Cloud):
+        task = sky.Task()
+        task.set_resources_override({'cloud': cloud})
+        sky.launch(task, dryrun=True)
+
+    dryrun_task_with_cloud(sky.K8s())
+
+    dryrun_task_with_cloud(sky.Kubernetes())
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 3a2417a6..36f2a6ea 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -3,7 +3,6 @@ import textwrap
 
 from click import testing as cli_testing
 
-import sky
 from sky import exceptions
 import sky.cli as cli
 
@@ -103,3 +102,40 @@ def test_show_gpus():
         result = cli_runner.invoke(cli.show_gpus,
                                    ['V100:4', '--cloud', cloud, '--all'])
         assert isinstance(result.exception, SystemExit)
+
+
+def test_k8s_alias_check():
+    cli_runner = cli_testing.CliRunner()
+
+    result = cli_runner.invoke(cli.check, ['k8s'])
+    assert not result.exit_code
+
+    result = cli_runner.invoke(cli.check, ['kubernetes'])
+    assert not result.exit_code
+
+    result = cli_runner.invoke(cli.check, ['notarealcloud'])
+    assert isinstance(result.exception, ValueError)
+
+
+def test_k8s_alias(enable_all_clouds):
+    cli_runner = cli_testing.CliRunner()
+
+    result = cli_runner.invoke(cli.launch, ['--cloud', 'k8s', '--dryrun'])
+    assert not result.exit_code
+
+    result = cli_runner.invoke(cli.launch,
+                               ['--cloud', 'kubernetes', '--dryrun'])
+    assert not result.exit_code
+
+    result = cli_runner.invoke(cli.launch,
+                               ['--cloud', 'notarealcloud', '--dryrun'])
+    assert isinstance(result.exception, ValueError)
+
+    result = cli_runner.invoke(cli.show_gpus, ['--cloud', 'k8s'])
+    assert not result.exit_code
+
+    result = cli_runner.invoke(cli.show_gpus, ['--cloud', 'kubernetes'])
+    assert not result.exit_code
+
+    result = cli_runner.invoke(cli.show_gpus, ['--cloud', 'notarealcloud'])
+    assert isinstance(result.exception, ValueError)