Repository: skypilot-org/skypilot
Pull request #1180 (2022-09-30 00:56:28)

Align gcloud installation (Resubmit #1155)
This is a minor change to align the installation code for our cloud store and spot controller.

The previous PR has a problem with the generated spot controller YAML due to the incorrect indent. It is now fixed by not using multiple lines for the installation commands.

Tested:
- [x] `sky spot launch -n test 'sleep 10'`
- [x] `sky launch --cloud gcp examples/using_file_mounts.yaml`
- [x] `sky launch examples/using_file_mounts.yaml`

Michaelvll: Tested again with the 3 tests in the description. I modified the backup logic, so that the `~/.sky/.sky_gcp_config_default_backup` will not be empty for the spot cluster when launched with `sky spot launch`. PTAL. : )
Michaelvll: Tested:
- [ ] `tests/run_smoke_tests.sh`

Patch:
diff --git a/sky/authentication.py b/sky/authentication.py
index 10a0877bf..8378801f5 100644
--- a/sky/authentication.py
+++ b/sky/authentication.py
@@ -12,6 +12,7 @@ from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import rsa
 from cryptography.hazmat.backends import default_backend
 
+from sky import clouds
 from sky import sky_logging
 from sky.adaptors import aws, gcp
 from sky.utils import common_utils
@@ -30,11 +31,6 @@ logger = sky_logging.init_logger(__name__)
 MAX_TRIALS = 64
 PRIVATE_SSH_KEY_PATH = '~/.ssh/sky-key'
 
-GCP_CONFIGURE_PATH = '~/.config/gcloud/configurations/config_default'
-# Do not place the backup under the gcloud config directory, as ray
-# autoscaler can overwrite that directory on the remote nodes.
-GCP_CONFIGURE_SKY_BACKUP_PATH = '~/.sky/.sky_gcp_config_default'
-
 
 def generate_rsa_key_pair():
     key = rsa.generate_private_key(backend=default_backend(),
@@ -212,27 +208,13 @@ def setup_gcp_authentication(config):
             'additional authentication steps.')
         # Read the account information from the credential file, since the user
         # should be set according the account, when the oslogin is enabled.
-        config_path = os.path.expanduser(GCP_CONFIGURE_PATH)
+        config_path = os.path.expanduser(clouds.gcp.GCP_CONFIG_PATH)
         sky_backup_config_path = os.path.expanduser(
-            GCP_CONFIGURE_SKY_BACKUP_PATH)
-        if not os.path.exists(sky_backup_config_path):
-            if not os.path.exists(config_path):
-                with ux_utils.print_exception_no_traceback():
-                    raise RuntimeError(
-                        'GCP authentication failed, as the oslogin is enabled '
-                        f'but the file {config_path} is not found.')
-
-            # Create a backup of the config_default file in the same folder (the
-            # folder will be uploaded by `sky launch`), as the original file can
-            # be modified on the remote cluster by ray causing failure of
-            # launching GCP cluster on a remote cluster.
-            subprocess.run(f'cp {config_path} {sky_backup_config_path}',
-                           shell=True,
-                           check=True)
-        new_file_mounts = config.get('file_mounts', {})
-        new_file_mounts[
-            GCP_CONFIGURE_SKY_BACKUP_PATH] = GCP_CONFIGURE_SKY_BACKUP_PATH
-        config['file_mounts'] = new_file_mounts
+            clouds.gcp.GCP_CONFIG_SKY_BACKUP_PATH)
+        assert os.path.exists(sky_backup_config_path), (
+            'GCP credential backup file '
+            f'{sky_backup_config_path!r} does not exist.')
+
         with open(sky_backup_config_path, 'r') as infile:
             for line in infile:
                 if line.startswith('account'):
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index 488c2b15b..676abf296 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -1910,7 +1910,8 @@ def interrupt_handler(signum, frame):
     logger.warning(f'{colorama.Fore.LIGHTBLACK_EX}The job will keep '
                    f'running after Ctrl-C.{colorama.Style.RESET_ALL}')
     kill_children_processes()
-    raise KeyboardInterrupt(exceptions.KEYBOARD_INTERRUPT_CODE)
+    with ux_utils.print_exception_no_traceback():
+        raise KeyboardInterrupt(exceptions.KEYBOARD_INTERRUPT_CODE)
 
 
 # Handle ctrl-z
@@ -1919,7 +1920,8 @@ def stop_handler(signum, frame):
     logger.warning(f'{colorama.Fore.LIGHTBLACK_EX}The job will keep '
                    f'running after Ctrl-Z.{colorama.Style.RESET_ALL}')
     kill_children_processes()
-    raise KeyboardInterrupt(exceptions.SIGTSTP_CODE)
+    with ux_utils.print_exception_no_traceback():
+        raise KeyboardInterrupt(exceptions.SIGTSTP_CODE)
 
 
 def validate_schema(obj, schema, err_msg_prefix=''):
diff --git a/sky/cli.py b/sky/cli.py
index cf10ba070..037c521c9 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -604,6 +604,7 @@ def _launch_with_confirm(
     no_confirm: bool = False,
     idle_minutes_to_autostop: Optional[int] = None,
     retry_until_up: bool = False,
+    no_setup: bool = False,
     node_type: Optional[str] = None,
     is_local_cloud: Optional[bool] = False,
 ):
@@ -644,14 +645,17 @@ def _launch_with_confirm(
 
     if node_type is None or maybe_status != global_user_state.ClusterStatus.UP:
         # No need to sky.launch again when interactive node is already up.
-        sky.launch(dag,
-                   dryrun=dryrun,
-                   stream_logs=True,
-                   cluster_name=cluster,
-                   detach_run=detach_run,
-                   backend=backend,
-                   idle_minutes_to_autostop=idle_minutes_to_autostop,
-                   retry_until_up=retry_until_up)
+        sky.launch(
+            dag,
+            dryrun=dryrun,
+            stream_logs=True,
+            cluster_name=cluster,
+            detach_run=detach_run,
+            backend=backend,
+            idle_minutes_to_autostop=idle_minutes_to_autostop,
+            retry_until_up=retry_until_up,
+            no_setup=no_setup,
+        )
 
 
 # TODO: skip installing ray to speed up provisioning.
@@ -1002,6 +1006,12 @@ def cli():
               default=False,
               required=False,
               help='Skip confirmation prompt.')
+@click.option('--no-setup',
+              '-n',
+              is_flag=True,
+              default=False,
+              required=False,
+              help='Skip setup phase when (re-)launching cluster.')
 @usage_lib.entrypoint
 def launch(
     entrypoint: str,
@@ -1024,6 +1034,7 @@ def launch(
     idle_minutes_to_autostop: Optional[int],
     retry_until_up: bool,
     yes: bool,
+    no_setup: bool,
 ):
     """Launch a task from a YAML or a command (rerun setup if cluster exists).
 
@@ -1072,6 +1083,7 @@ def launch(
         no_confirm=yes,
         idle_minutes_to_autostop=idle_minutes_to_autostop,
         retry_until_up=retry_until_up,
+        no_setup=no_setup,
         is_local_cloud=onprem_utils.check_if_local_cloud(cluster))
 
 
diff --git a/sky/cloud_stores.py b/sky/cloud_stores.py
index dc8e69574..b4f8cdd59 100644
--- a/sky/cloud_stores.py
+++ b/sky/cloud_stores.py
@@ -10,6 +10,7 @@ TODO:
 import subprocess
 import urllib.parse
 
+from sky.clouds import gcp
 from sky.data import data_utils
 from sky.adaptors import aws
 
@@ -93,26 +94,18 @@ class GcsCloudStorage(CloudStorage):
     # We use gsutil as a basic implementation.  One pro is that its -m
     # multi-threaded download is nice, which frees us from implementing
     # parellel workers on our end.
-    _GET_GSUTIL = [
-        # Skip if gsutil already exists.
-        'pushd /tmp &>/dev/null',
-        '(test -f ~/google-cloud-sdk/bin/gsutil || (wget --quiet '
-        'https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/'
-        'google-cloud-sdk-367.0.0-linux-x86_64.tar.gz && '
-        'tar xzf google-cloud-sdk-367.0.0-linux-x86_64.tar.gz && '
-        'mv google-cloud-sdk ~/ && '
-        '~/google-cloud-sdk/install.sh -q ))',
-        'popd &>/dev/null',
-    ]
+    # The gsutil command is part of the Google Cloud SDK, and we reuse
+    # the installation logic here.
+    _GET_GSUTIL = gcp.GCLOUD_INSTALLATION_COMMAND
 
-    _GSUTIL = '~/google-cloud-sdk/bin/gsutil'
+    _GSUTIL = 'gsutil'
 
     def is_directory(self, url: str) -> bool:
         """Returns whether 'url' is a directory.
         In cloud object stores, a "directory" refers to a regular object whose
         name is a prefix of other objects.
         """
-        commands = list(self._GET_GSUTIL)
+        commands = [self._GET_GSUTIL]
         commands.append(f'{self._GSUTIL} ls -d {url}')
         command = ' && '.join(commands)
         p = subprocess.run(command,
@@ -142,14 +135,14 @@ class GcsCloudStorage(CloudStorage):
         """Downloads a directory using gsutil."""
         download_via_gsutil = (
             f'{self._GSUTIL} -m rsync -r {source} {destination}')
-        all_commands = list(self._GET_GSUTIL)
+        all_commands = [self._GET_GSUTIL]
         all_commands.append(download_via_gsutil)
         return ' && '.join(all_commands)
 
     def make_sync_file_command(self, source: str, destination: str) -> str:
         """Downloads a file using gsutil."""
         download_via_gsutil = f'{self._GSUTIL} -m cp {source} {destination}'
-        all_commands = list(self._GET_GSUTIL)
+        all_commands = [self._GET_GSUTIL]
         all_commands.append(download_via_gsutil)
         return ' && '.join(all_commands)
 
diff --git a/sky/clouds/gcp.py b/sky/clouds/gcp.py
index 46ab4e2eb..480650347 100644
--- a/sky/clouds/gcp.py
+++ b/sky/clouds/gcp.py
@@ -18,6 +18,11 @@ DEFAULT_GCP_APPLICATION_CREDENTIAL_PATH = os.path.expanduser(
     '~/.config/gcloud/'
     'application_default_credentials.json')
 
+GCP_CONFIG_PATH = '~/.config/gcloud/configurations/config_default'
+# Do not place the backup under the gcloud config directory, as ray
+# autoscaler can overwrite that directory on the remote nodes.
+GCP_CONFIG_SKY_BACKUP_PATH = '~/.sky/.sky_gcp_config_default'
+
 # Minimum set of files under ~/.config/gcloud that grant GCP access.
 _CREDENTIAL_FILES = [
     'credentials.db',
@@ -29,6 +34,23 @@ _CREDENTIAL_FILES = [
 
 _IMAGE_ID_PREFIX = ('projects/deeplearning-platform-release/global/images/')
 
+_GCLOUD_INSTALLATION_LOG = '~/.sky/logs/gcloud_installation.log'
+# Need to be run with /bin/bash
+# We factor out the installation logic to keep it align in both spot
+# controller and cloud stores.
+GCLOUD_INSTALLATION_COMMAND = f'pushd /tmp &>/dev/null && \
+    gcloud --help > /dev/null 2>&1 || \
+    {{ mkdir -p {os.path.dirname(_GCLOUD_INSTALLATION_LOG)} && \
+    wget --quiet https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-sdk-382.0.0-linux-x86_64.tar.gz > {_GCLOUD_INSTALLATION_LOG} && \
+    tar xzf google-cloud-sdk-382.0.0-linux-x86_64.tar.gz >> {_GCLOUD_INSTALLATION_LOG} && \
+    rm -rf ~/google-cloud-sdk >> {_GCLOUD_INSTALLATION_LOG}  && \
+    mv google-cloud-sdk ~/ && \
+    ~/google-cloud-sdk/install.sh -q >> {_GCLOUD_INSTALLATION_LOG} 2>&1 && \
+    echo "source ~/google-cloud-sdk/path.bash.inc > /dev/null 2>&1" >> ~/.bashrc && \
+    source ~/google-cloud-sdk/path.bash.inc >> {_GCLOUD_INSTALLATION_LOG} 2>&1; }} && \
+    {{ cp {GCP_CONFIG_SKY_BACKUP_PATH} {GCP_CONFIG_PATH} > /dev/null 2>&1 || true; }} && \
+    popd &>/dev/null'
+
 
 def _run_output(cmd):
     proc = subprocess.run(cmd,
@@ -328,13 +350,30 @@ class GCP(clouds.Cloud):
         return True, None
 
     def get_credential_file_mounts(self) -> Dict[str, str]:
+        # Create a backup of the config_default file, as the original file can
+        # be modified on the remote cluster by ray causing authentication
+        # problems. The backup file will be updated to the remote cluster
+        # whenever the original file is not empty and will be applied
+        # appropriately on the remote cluster when neccessary.
+        if (os.path.exists(os.path.expanduser(GCP_CONFIG_PATH)) and
+                os.path.getsize(os.path.expanduser(GCP_CONFIG_PATH)) > 0):
+            subprocess.run(f'cp {GCP_CONFIG_PATH} {GCP_CONFIG_SKY_BACKUP_PATH}',
+                           shell=True,
+                           check=True)
+        elif not os.path.exists(os.path.expanduser(GCP_CONFIG_SKY_BACKUP_PATH)):
+            raise RuntimeError(
+                'GCP credential file is empty. Please make sure you '
+                'have run: gcloud init')
+
         # Excluding the symlink to the python executable created by the gcp
         # credential, which causes problem for ray up multiple nodes, tracked
         # in #494, #496, #483.
-        return {
+        credentials = {
             f'~/.config/gcloud/{filename}': f'~/.config/gcloud/{filename}'
             for filename in _CREDENTIAL_FILES
         }
+        credentials[GCP_CONFIG_SKY_BACKUP_PATH] = GCP_CONFIG_SKY_BACKUP_PATH
+        return credentials
 
     def instance_type_exists(self, instance_type):
         return service_catalog.instance_type_exists(instance_type, 'gcp')
diff --git a/sky/execution.py b/sky/execution.py
index da9d67ca4..d400bb7d8 100644
--- a/sky/execution.py
+++ b/sky/execution.py
@@ -28,6 +28,7 @@ from sky import optimizer
 from sky import sky_logging
 from sky import spot
 from sky.backends import backend_utils
+from sky.clouds import gcp
 from sky.data import data_utils
 from sky.data import storage
 from sky.usage import usage_lib
@@ -105,6 +106,7 @@ def _execute(
     cluster_name: Optional[str] = None,
     detach_run: bool = False,
     idle_minutes_to_autostop: Optional[int] = None,
+    no_setup: bool = False,
 ) -> None:
     """Runs a DAG.
 
@@ -134,6 +136,7 @@ def _execute(
       detach_run: bool; whether to detach the process after the job submitted.
       autostop_idle_minutes: int; if provided, the cluster will be set to
         autostop after this many minutes of idleness.
+      no_setup: bool; whether to skip setup commands or not when (re-)launching.
     """
     _type_check_dag(dag)
     assert len(dag) == 1, f'We support 1 task for now. {dag}'
@@ -201,7 +204,9 @@ def _execute(
             backend.sync_file_mounts(handle, task.file_mounts,
                                      task.storage_mounts)
 
-        if stages is None or Stage.SETUP in stages:
+        if no_setup:
+            logger.info('Setup commands skipped.')
+        elif stages is None or Stage.SETUP in stages:
             backend.setup(handle, task)
 
         if stages is None or Stage.PRE_EXEC in stages:
@@ -251,6 +256,7 @@ def launch(
     backend: Optional[backends.Backend] = None,
     optimize_target: OptimizeTarget = OptimizeTarget.COST,
     detach_run: bool = False,
+    no_setup: bool = False,
 ):
     """Launch a sky.DAG (rerun setup if cluster exists).
 
@@ -262,6 +268,7 @@ def launch(
             up.
         idle_minutes_to_autostop: if provided, the cluster will be auto-stop
             after this many minutes of idleness.
+        no_setup: if true, the cluster will not re-run setup instructions
 
     Examples:
         >>> import sky
@@ -290,6 +297,7 @@ def launch(
         cluster_name=cluster_name,
         detach_run=detach_run,
         idle_minutes_to_autostop=idle_minutes_to_autostop,
+        no_setup=no_setup,
     )
 
 
@@ -439,6 +447,7 @@ def spot_launch(
                 'user_yaml_path': f.name,
                 'spot_controller': controller_name,
                 'cluster_name': name,
+                'gcloud_installation_commands': gcp.GCLOUD_INSTALLATION_COMMAND,
                 'is_dev': env_options.Options.IS_DEVELOPER.get(),
                 'disable_logging': env_options.Options.DISABLE_LOGGING.get(),
                 'logging_user_hash': common_utils.get_user_hash(),
diff --git a/sky/templates/spot-controller.yaml.j2 b/sky/templates/spot-controller.yaml.j2
index 00e6e4878..d56a5ce14 100644
--- a/sky/templates/spot-controller.yaml.j2
+++ b/sky/templates/spot-controller.yaml.j2
@@ -25,13 +25,7 @@ setup: |
   # pip list | grep azure-cli > /dev/null 2>&1 || \
   #  pip3 install azure-cli==2.31.0 azure-core
 
-  gcloud --help > /dev/null 2>&1 || \
-  (wget --quiet https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-sdk-382.0.0-linux-x86_64.tar.gz && \
-    tar xzf google-cloud-sdk-382.0.0-linux-x86_64.tar.gz && \
-    mv google-cloud-sdk ~/ && \
-    ~/google-cloud-sdk/install.sh -q > /dev/null 2>&1 && \
-    source ~/google-cloud-sdk/path.bash.inc > /dev/null 2>&1 && \
-    echo 'source ~/google-cloud-sdk/path.bash.inc > /dev/null 2>&1' >> ~/.bashrc)
+  {{gcloud_installation_commands}}
 
 run: |
   python -u -m sky.spot.controller \


Test patch:
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 1ea5ff76d..59160d278 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -57,6 +57,10 @@ def test_accelerator_mismatch(monkeypatch):
         lambda: enabled_clouds,
     )
     monkeypatch.setattr('sky.check.check', lambda *_args, **_kwargs: None)
+    config_file_backup = tempfile.NamedTemporaryFile(
+        prefix='tmp_backup_config_default', delete=False)
+    monkeypatch.setattr('sky.clouds.gcp.GCP_CONFIG_SKY_BACKUP_PATH',
+                        config_file_backup.name)
 
     spec = textwrap.dedent("""\
         resources:
diff --git a/tests/test_optimizer_dryruns.py b/tests/test_optimizer_dryruns.py
index edd7e88e1..079d28232 100644
--- a/tests/test_optimizer_dryruns.py
+++ b/tests/test_optimizer_dryruns.py
@@ -38,6 +38,12 @@ def _make_resources(
         lambda: enabled_clouds,
     )
     monkeypatch.setattr('sky.check.check', lambda *_args, **_kwargs: None)
+
+    config_file_backup = tempfile.NamedTemporaryFile(
+        prefix='tmp_backup_config_default', delete=False)
+    monkeypatch.setattr('sky.clouds.gcp.GCP_CONFIG_SKY_BACKUP_PATH',
+                        config_file_backup.name)
+
     # Should create Resources here, since it uses the enabled clouds.
     return sky.Resources(*resources_args, **resources_kwargs)