GCP: make sky check enable the 3 required APIs & optional TPU API.

sky/authentication.py

@@ -4,13 +4,17 @@
 import hashlib
 import os
 import pathlib
+import re
 import subprocess
+import sys
 import time
 import uuid
 
+import colorama
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import rsa
 from cryptography.hazmat.backends import default_backend
+import googleapiclient
 
 from sky import clouds
 from sky import sky_logging
@@ -19,8 +23,7 @@
 from sky.utils import subprocess_utils
 from sky.utils import ux_utils
 
-logger = sky_logging.init_logger(__name__)
-
+colorama.init()
 logger = sky_logging.init_logger(__name__)
 
 # TODO: Should tolerate if gcloud is not installed. Also,
@@ -195,7 +198,37 @@ def setup_gcp_authentication(config):
                         credentials=None,
                         cache_discovery=False)
     user = config['auth']['ssh_user']
-    project = compute.projects().get(project=project_id).execute()
+
+    try:
+        project = compute.projects().get(project=project_id).execute()
+    except googleapiclient.errors.HttpError as e:
+        # Can happen for a new project where Compute Engine API is disabled.
+        # Ex:
+        # 'Compute Engine API has not been used in project 123456 before
+        # or it is disabled. Enable it by visiting
+        # https://console.developers.google.com/apis/api/compute.googleapis.com/overview?project=123456
+        # then retry. If you enabled this API recently, wait a few minutes for
+        # the action to propagate to our systems and retry.'
+        if ' API has not been used in project' in e.reason:
+            match = re.fullmatch(r'(.+)(https://.*project=\d+) (.+)', e.reason)
+            if match is None:
+                raise  # This should not happen.
+            yellow = colorama.Fore.YELLOW
+            reset = colorama.Style.RESET_ALL
+            bright = colorama.Style.BRIGHT
+            dim = colorama.Style.DIM
+            logger.error(
+                f'{yellow}Certain GCP APIs are disabled for the GCP project '
+                f'{project_id}.{reset}')
+            logger.error(f'{yellow}Enable them by running:{reset} '
+                         f'{bright}sky check{reset}')
+            logger.error('Details:')
+            logger.error(f'{dim}{match.group(1)}{reset}\n'
+                         f'{dim}    {match.group(2)}{reset}\n'
+                         f'{dim}{match.group(3)}{reset}')
+            sys.exit(1)
+        else:
+            raise
 
     project_oslogin = next(
         (item for item in project['commonInstanceMetadata'].get('items', [])

sky/clouds/gcp.py

@@ -2,6 +2,7 @@
 import json
 import os
 import subprocess
+import time
 import typing
 from typing import Dict, Iterator, List, Optional, Tuple
 
@@ -61,6 +62,16 @@ def _run_output(cmd):
     return proc.stdout.decode('ascii')
 
 
+def is_api_disabled(endpoint: str, project_id: str) -> bool:
+    proc = subprocess.run((f'gcloud services list --project {project_id} '
+                           f' | grep {endpoint}.googleapis.com'),
+                          check=False,
+                          shell=True,
+                          stderr=subprocess.PIPE,
+                          stdout=subprocess.PIPE)
+    return proc.returncode != 0
+
+
 @clouds.CLOUD_REGISTRY.register
 class GCP(clouds.Cloud):
     """Google Cloud Platform."""
@@ -347,6 +358,46 @@ def check_credentials(self) -> Tuple[bool, Optional[str]]:
                 'For more info: '
                 'https://skypilot.readthedocs.io/en/latest/getting-started/installation.html'  # pylint: disable=line-too-long
             )
+
+        # Check APIs.
+        project_id = self.get_project_id()
+        apis = (
+            ('compute', 'Compute Engine'),
+            ('cloudresourcemanager', 'Cloud Resource Manager'),
+            ('iam', 'Identity and Access Management (IAM)'),
+        )
+        enabled_api = False
+        for endpoint, display_name in apis:
+            if is_api_disabled(endpoint, project_id):
+                # For 'compute': ~55-60 seconds for the first run. If already
+                # enabled, ~1s. Other API endpoints take ~1-5s to enable.
+                if endpoint == 'compute':
+                    suffix = ' (may take a minute)'
+                else:
+                    suffix = ''
+                print(f'\nEnabling {display_name} API{suffix}...')
+                t1 = time.time()
+                proc = subprocess.run(
+                    f'gcloud services enable {endpoint}.googleapis.com '
+                    f'--project {project_id}',
+                    check=False,
+                    shell=True,
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.STDOUT)
+                if proc.returncode == 0:
+                    enabled_api = True
+                    print(f'Done. Took {time.time() - t1:.1f} secs.')
+                else:
+                    print('Failed; please manually enable the API. Log:')
+                    print(proc.stdout.decode())
+                    return False, (
+                        f'{display_name} API is disabled. Please retry '
+                        '`sky check` in a few minutes, or manually enable it.')
+        if enabled_api:
+            print('\nHint: Enabled GCP API(s) may take a few minutes to take '
+                  'effect. If any SkyPilot commands/calls failed, retry after '
+                  'some time.')
+
         return True, None
 
     def get_credential_file_mounts(self) -> Dict[str, str]: