From 38a822ac6b553df0e784e559715ee4269c21f780 Mon Sep 17 00:00:00 2001
From: Zhanghao Wu <zhanghao.wu@outlook.com>
Date: Sun, 5 Jan 2025 22:51:04 -0800
Subject: [PATCH] [GCP] Activate service account for storage and controller
 (#4529)

* Activate service account for storage

* disable logging if not using service account

* Activate for controller as well.

* revert controller activate

* Add comments

* format

* fix smoke
---
 sky/cloud_stores.py                   | 12 ++++++++++--
 sky/data/data_utils.py                | 12 ++++++++----
 tests/smoke_tests/test_managed_job.py |  2 +-
 3 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/sky/cloud_stores.py b/sky/cloud_stores.py
index 108f33f2c1f..e24c4f3ad03 100644
--- a/sky/cloud_stores.py
+++ b/sky/cloud_stores.py
@@ -113,8 +113,16 @@ class GcsCloudStorage(CloudStorage):
     @property
     def _gsutil_command(self):
         gsutil_alias, alias_gen = data_utils.get_gsutil_command()
-        return (f'{alias_gen}; GOOGLE_APPLICATION_CREDENTIALS='
-                f'{gcp.DEFAULT_GCP_APPLICATION_CREDENTIAL_PATH} {gsutil_alias}')
+        return (
+            f'{alias_gen}; GOOGLE_APPLICATION_CREDENTIALS='
+            f'{gcp.DEFAULT_GCP_APPLICATION_CREDENTIAL_PATH}; '
+            # Explicitly activate service account. Unlike the gcp packages
+            # and other GCP commands, gsutil does not automatically pick up
+            # the default credential keys when it is a service account.
+            'gcloud auth activate-service-account '
+            '--key-file=$GOOGLE_APPLICATION_CREDENTIALS '
+            '2> /dev/null || true; '
+            f'{gsutil_alias}')
 
     def is_directory(self, url: str) -> bool:
         """Returns whether 'url' is a directory.
diff --git a/sky/data/data_utils.py b/sky/data/data_utils.py
index 05c2b42c844..e8dcaa83017 100644
--- a/sky/data/data_utils.py
+++ b/sky/data/data_utils.py
@@ -523,10 +523,14 @@ def get_gsutil_command() -> Tuple[str, str]:
 
 def run_upload_cli(command: str, access_denied_message: str, bucket_name: str,
                    log_path: str):
-    returncode, stdout, stderr = log_lib.run_with_log(command,
-                                                      log_path,
-                                                      shell=True,
-                                                      require_outputs=True)
+    returncode, stdout, stderr = log_lib.run_with_log(
+        command,
+        log_path,
+        shell=True,
+        require_outputs=True,
+        # We need to use bash as some of the cloud commands uses bash syntax,
+        # such as [[ ... ]]
+        executable='/bin/bash')
     if access_denied_message in stderr:
         with ux_utils.print_exception_no_traceback():
             raise PermissionError('Failed to upload files to '
diff --git a/tests/smoke_tests/test_managed_job.py b/tests/smoke_tests/test_managed_job.py
index 22381fc45e3..5c930724523 100644
--- a/tests/smoke_tests/test_managed_job.py
+++ b/tests/smoke_tests/test_managed_job.py
@@ -365,7 +365,7 @@ def test_managed_jobs_pipeline_recovery_gcp():
             # separated by `-`.
             (f'MANAGED_JOB_ID=`cat /tmp/{name}-run-id | rev | '
              f'cut -d\'_\' -f1 | rev | cut -d\'-\' -f1`; {terminate_cmd}'),
-            smoke_tests_utils.zJOB_WAIT_NOT_RUNNING.format(job_name=name),
+            smoke_tests_utils.JOB_WAIT_NOT_RUNNING.format(job_name=name),
             f'{smoke_tests_utils.GET_JOB_QUEUE} | grep {name} | head -n1 | grep "RECOVERING"',
             smoke_tests_utils.
             get_cmd_wait_until_managed_job_status_contains_matching_job_name(