add typechecking for boto clients/resources (#5319)

* add typechecking for boto clients/resources

* import module instead of class

Author: cg505

requirements-dev.txt

@@ -40,3 +40,6 @@ aiohttp==3.9.3
 
 # For mocking AWS
 moto==5.1.2
+
+# boto3 type checking
+boto3-stubs[s3,ec2,sts,iam,service-quotas]

sky/adaptors/aws.py

@@ -31,19 +31,31 @@
 import logging
 import threading
 import time
-from typing import Any, Callable
+import typing
+from typing import Callable, Literal, Optional, TypeVar
 
 from sky.adaptors import common
 from sky.utils import annotations
 from sky.utils import common_utils
 
+if typing.TYPE_CHECKING:
+    import boto3
+    _ = boto3  # Supress pylint use before assignment error
+    import mypy_boto3_ec2
+    import mypy_boto3_iam
+    import mypy_boto3_s3
+    import mypy_boto3_service_quotas
+    import mypy_boto3_sts
+
 _IMPORT_ERROR_MESSAGE = ('Failed to import dependencies for AWS. '
                          'Try pip install "skypilot[aws]"')
 boto3 = common.LazyImport('boto3', import_error_message=_IMPORT_ERROR_MESSAGE)
 botocore = common.LazyImport('botocore',
                              import_error_message=_IMPORT_ERROR_MESSAGE)
 _LAZY_MODULES = (boto3, botocore)
 
+T = TypeVar('T')
+
 logger = logging.getLogger(__name__)
 _session_creation_lock = threading.RLock()
 
@@ -73,8 +85,8 @@ def _assert_kwargs_builtin_type(kwargs):
         f'kwargs should not contain none built-in types: {kwargs}')
 
 
-def _create_aws_object(creation_fn_or_cls: Callable[[], Any],
-                       object_name: str) -> Any:
+def _create_aws_object(creation_fn_or_cls: Callable[[], T],
+                       object_name: str) -> T:
     """Create an AWS object.
 
     Args:
@@ -123,6 +135,25 @@ def session(check_credentials: bool = True):
     return s
 
 
+# New typing overloads can be added as needed.
+@typing.overload
+def resource(service_name: Literal['ec2'],
+             **kwargs) -> 'mypy_boto3_ec2.ServiceResource':
+    ...
+
+
+@typing.overload
+def resource(service_name: Literal['s3'],
+             **kwargs) -> 'mypy_boto3_s3.ServiceResource':
+    ...
+
+
+@typing.overload
+def resource(service_name: Literal['iam'],
+             **kwargs) -> 'mypy_boto3_iam.ServiceResource':
+    ...
+
+
 # Avoid caching the resource/client objects. If we are using the assumed role,
 # the credentials will be automatically rotated, but the cached resource/client
 # object will only refresh the credentials with a fixed 15 minutes interval,
@@ -142,7 +173,7 @@ def resource(service_name: str, **kwargs):
     """
     _assert_kwargs_builtin_type(kwargs)
 
-    max_attempts = kwargs.pop('max_attempts', None)
+    max_attempts: Optional[int] = kwargs.pop('max_attempts', None)
     if max_attempts is not None:
         config = botocore_config().Config(
             retries={'max_attempts': max_attempts})
@@ -158,6 +189,28 @@ def resource(service_name: str, **kwargs):
             service_name, **kwargs), 'resource')
 
 
+# New typing overloads can be added as needed.
+@typing.overload
+def client(service_name: Literal['s3'], **kwargs) -> 'mypy_boto3_s3.Client':
+    pass
+
+
+@typing.overload
+def client(service_name: Literal['ec2'], **kwargs) -> 'mypy_boto3_ec2.Client':
+    pass
+
+
+@typing.overload
+def client(service_name: Literal['sts'], **kwargs) -> 'mypy_boto3_sts.Client':
+    pass
+
+
+@typing.overload
+def client(service_name: Literal['service-quotas'],
+           **kwargs) -> 'mypy_boto3_service_quotas.Client':
+    pass
+
+
 def client(service_name: str, **kwargs):
     """Create an AWS client of a certain service.
 

sky/adaptors/common.py

@@ -2,10 +2,11 @@
 import functools
 import importlib
 import threading
+import types
 from typing import Any, Callable, Optional, Tuple
 
 
-class LazyImport:
+class LazyImport(types.ModuleType):
     """Lazy importer for modules.
 
     This is mainly used in two cases:

sky/clouds/aws.py

@@ -312,13 +312,12 @@ def get_image_size(cls, image_id: str, region: Optional[str]) -> float:
             'Example: ami-0729d913a335efca7')
         try:
             client = aws.client('ec2', region_name=region)
-            image_info = client.describe_images(ImageIds=[image_id])
-            image_info = image_info.get('Images', [])
+            image_info = client.describe_images(ImageIds=[image_id]).get(
+                'Images', [])
             if not image_info:
                 with ux_utils.print_exception_no_traceback():
                     raise ValueError(image_not_found_message)
-            image_info = image_info[0]
-            image_size = image_info['BlockDeviceMappings'][0]['Ebs'][
+            image_size = image_info[0]['BlockDeviceMappings'][0]['Ebs'][
                 'VolumeSize']
         except (aws.botocore_exceptions().NoCredentialsError,
                 aws.botocore_exceptions().ProfileNotFound):

sky/clouds/service_catalog/data_fetchers/fetch_aws.py

@@ -13,7 +13,7 @@
 import textwrap
 import traceback
 import typing
-from typing import Dict, List, Optional, Set, Tuple, Union
+from typing import List, Optional, Set, Tuple, Union
 
 import numpy as np
 
@@ -24,6 +24,7 @@
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
+    from mypy_boto3_ec2 import type_defs as ec2_type_defs
     import pandas as pd
 else:
     pd = adaptors_common.LazyImport('pandas')
@@ -192,7 +193,7 @@ def _get_spot_pricing_table(region: str) -> 'pd.DataFrame':
     paginator = client.get_paginator('describe_spot_price_history')
     response_iterator = paginator.paginate(ProductDescriptions=['Linux/UNIX'],
                                            StartTime=datetime.datetime.utcnow())
-    ret: List[Dict[str, str]] = []
+    ret: List['ec2_type_defs.SpotPriceTypeDef'] = []
     for response in response_iterator:
         # response['SpotPriceHistory'] is a list of dicts, each dict is like:
         # {

sky/data/storage.py

@@ -38,8 +38,8 @@
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
-    import boto3  # type: ignore
     from google.cloud import storage  # type: ignore
+    import mypy_boto3_s3
 
 logger = sky_logging.init_logger(__name__)
 
@@ -1363,7 +1363,7 @@ def __init__(self,
                  is_sky_managed: Optional[bool] = None,
                  sync_on_reconstruction: bool = True,
                  _bucket_sub_path: Optional[str] = None):
-        self.client: 'boto3.client.Client'
+        self.client: 'mypy_boto3_s3.Client'
         self.bucket: 'StorageHandle'
         # TODO(romilb): This is purely a stopgap fix for
         #  https://github.com/skypilot-org/skypilot/issues/3405
@@ -3295,7 +3295,7 @@ def __init__(self,
                  is_sky_managed: Optional[bool] = None,
                  sync_on_reconstruction: Optional[bool] = True,
                  _bucket_sub_path: Optional[str] = None):
-        self.client: 'boto3.client.Client'
+        self.client: 'mypy_boto3_s3.Client'
         self.bucket: 'StorageHandle'
         super().__init__(name, source, region, is_sky_managed,
                          sync_on_reconstruction, _bucket_sub_path)
@@ -4700,7 +4700,7 @@ def __init__(self,
                  is_sky_managed: Optional[bool] = None,
                  sync_on_reconstruction: bool = True,
                  _bucket_sub_path: Optional[str] = None):
-        self.client: 'boto3.client.Client'
+        self.client: 'mypy_boto3_s3.Client'
         self.bucket: 'StorageHandle'
         super().__init__(name, source, region, is_sky_managed,
                          sync_on_reconstruction, _bucket_sub_path)

sky/provision/aws/config.py

@@ -11,6 +11,7 @@
 import json
 import logging
 import time
+import typing
 from typing import Any, Dict, List, Optional, Set, Tuple
 
 import colorama
@@ -23,6 +24,10 @@
 from sky.utils import annotations
 from sky.utils import common_utils
 
+if typing.TYPE_CHECKING:
+    import mypy_boto3_ec2
+    from mypy_boto3_ec2 import type_defs as ec2_type_defs
+
 logger = sky_logging.init_logger(__name__)
 
 RAY = 'ray-autoscaler'
@@ -223,7 +228,8 @@ def _get_role(role_name: str):
 
 
 @annotations.lru_cache(scope='request', maxsize=128)  # Keep bounded.
-def _get_route_tables(ec2, vpc_id: Optional[str], region: str,
+def _get_route_tables(ec2: 'mypy_boto3_ec2.ServiceResource',
+                      vpc_id: Optional[str], region: str,
                       main: bool) -> List[Any]:
     """Get route tables associated with a VPC and region
 
@@ -248,7 +254,8 @@ def _get_route_tables(ec2, vpc_id: Optional[str], region: str,
         'RouteTables', [])
 
 
-def _is_subnet_public(ec2, subnet_id, vpc_id: Optional[str]) -> bool:
+def _is_subnet_public(ec2: 'mypy_boto3_ec2.ServiceResource', subnet_id,
+                      vpc_id: Optional[str]) -> bool:
     """Checks if a subnet is public by existence of a route to an IGW.
 
     Conventionally, public subnets connect to a IGW, and private subnets to a
@@ -441,10 +448,14 @@ def _get_pruned_subnets(current_subnets: List[Any]) -> Set[str]:
     return subnets, first_subnet_vpc_id
 
 
-def _vpc_id_from_security_group_ids(ec2, sg_ids: List[str]) -> Any:
+def _vpc_id_from_security_group_ids(ec2: 'mypy_boto3_ec2.ServiceResource',
+                                    sg_ids: List[str]) -> Any:
     # sort security group IDs to support deterministic unit test stubbing
     sg_ids = sorted(set(sg_ids))
-    filters = [{'Name': 'group-id', 'Values': sg_ids}]
+    filters: List['ec2_type_defs.FilterTypeDef'] = [{
+        'Name': 'group-id',
+        'Values': sg_ids
+    }]
     security_groups = ec2.security_groups.filter(Filters=filters)
     vpc_ids = [sg.vpc_id for sg in security_groups]
     vpc_ids = list(set(vpc_ids))
@@ -462,15 +473,19 @@ def _vpc_id_from_security_group_ids(ec2, sg_ids: List[str]) -> Any:
     return vpc_ids[0]
 
 
-def _get_vpc_id_by_name(ec2, vpc_name: str, region: str) -> str:
+def _get_vpc_id_by_name(ec2: 'mypy_boto3_ec2.ServiceResource', vpc_name: str,
+                        region: str) -> str:
     """Returns the VPC ID of the unique VPC with a given name.
 
     Exits with code 1 if:
       - No VPC with the given name is found in the current region.
       - More than 1 VPC with the given name are found in the current region.
     """
     # Look in the 'Name' tag (shown as Name column in console).
-    filters = [{'Name': 'tag:Name', 'Values': [vpc_name]}]
+    filters: List['ec2_type_defs.FilterTypeDef'] = [{
+        'Name': 'tag:Name',
+        'Values': [vpc_name]
+    }]
     vpcs = list(ec2.vpcs.filter(Filters=filters))
     if not vpcs:
         _skypilot_log_error_and_exit_for_failover(
@@ -486,8 +501,9 @@ def _get_vpc_id_by_name(ec2, vpc_name: str, region: str) -> str:
     return vpcs[0].id
 
 
-def _get_subnet_and_vpc_id(ec2, security_group_ids: Optional[List[str]],
-                           region: str, availability_zone: Optional[str],
+def _get_subnet_and_vpc_id(ec2: 'mypy_boto3_ec2.ServiceResource',
+                           security_group_ids: Optional[List[str]], region: str,
+                           availability_zone: Optional[str],
                            use_internal_ips: bool,
                            vpc_name: Optional[str]) -> Tuple[Any, str]:
     if vpc_name is not None:
@@ -514,7 +530,8 @@ def _get_subnet_and_vpc_id(ec2, security_group_ids: Optional[List[str]],
     return subnets, vpc_id
 
 
-def _configure_security_group(ec2, vpc_id: str, expected_sg_name: str,
+def _configure_security_group(ec2: 'mypy_boto3_ec2.ServiceResource',
+                              vpc_id: str, expected_sg_name: str,
                               extended_ip_rules: List) -> List[str]:
     security_group = _get_or_create_vpc_security_group(ec2, vpc_id,
                                                        expected_sg_name)
@@ -551,7 +568,8 @@ def _configure_security_group(ec2, vpc_id: str, expected_sg_name: str,
     return sg_ids
 
 
-def _get_or_create_vpc_security_group(ec2, vpc_id: str,
+def _get_or_create_vpc_security_group(ec2: 'mypy_boto3_ec2.ServiceResource',
+                                      vpc_id: str,
                                       expected_sg_name: str) -> Any:
     """Find or create a security group in the specified VPC.
 
@@ -612,7 +630,8 @@ def _get_or_create_vpc_security_group(ec2, vpc_id: str,
     return security_group
 
 
-def _get_security_group_from_vpc_id(ec2, vpc_id: str,
+def _get_security_group_from_vpc_id(ec2: 'mypy_boto3_ec2.ServiceResource',
+                                    vpc_id: str,
                                     group_name: str) -> Optional[Any]:
     """Get security group by VPC ID and group name."""
     existing_groups = list(