[Mypy] Add more mypy

format.sh

@@ -74,8 +74,8 @@ format_changed() {
     # exist on both branches.
     MERGEBASE="$(git merge-base origin/master HEAD)"
 
-    if ! git diff --diff-filter=ACM --quiet --exit-code "$MERGEBASE" -- '*.py' &>/dev/null; then
-        git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.py' | xargs -P 5 \
+    if ! git diff --diff-filter=ACM --quiet --exit-code "$MERGEBASE" -- '*.py' '*.pyi' &>/dev/null; then
+        git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.py' '*.pyi' | xargs -P 5 \
              yapf --in-place "${YAPF_EXCLUDES[@]}" "${YAPF_FLAGS[@]}"
     fi
 

requirements-dev.txt

@@ -11,6 +11,7 @@ toml==0.10.2
 mypy==0.991
 types-PyYAML
 types-requests
+types-setuptools
 
 # testing
 pytest

sky/authentication.py

@@ -142,7 +142,6 @@ def _wait_for_compute_global_operation(project_name: str, operation_name: str,
 # avoid duplicated codes.
 # Retry for the GCP as sometimes there will be connection reset by peer error.
 @common_utils.retry
-@gcp.import_package
 def setup_gcp_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
     private_key_path, public_key_path = get_or_generate_keys()
     config = copy.deepcopy(config)
@@ -156,7 +155,7 @@ def setup_gcp_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
 
     try:
         project = compute.projects().get(project=project_id).execute()
-    except gcp.googleapiclient.errors.HttpError as e:
+    except gcp.http_error_exception() as e:
         # Can happen for a new project where Compute Engine API is disabled.
         #
         # Example message:
@@ -191,9 +190,10 @@ def setup_gcp_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
                      'Please check your network connection.')
         raise
 
-    project_oslogin = next(
+    project_oslogin: str = next(
         (item for item in project['commonInstanceMetadata'].get('items', [])
-         if item['key'] == 'enable-oslogin'), {}).get('value', 'False')
+         if item['key'] == 'enable-oslogin'), {}).get('value',
+                                                      'False')  # type: ignore
 
     if project_oslogin.lower() == 'true':
         # project.
@@ -245,17 +245,17 @@ def setup_gcp_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
                 'utf-8'):
             subprocess_utils.handle_returncode(proc.returncode, enable_ssh_cmd,
                                                'Failed to enable ssh port.',
-                                               proc.stderr)
+                                               proc.stderr.decode('utf-8'))
         return config
 
     # OS Login is not enabled for the project. Add the ssh key directly to the
     # metadata.
     # TODO(zhwu): Use cloud init to add ssh public key, to avoid the permission
     # issue. A blocker is that the cloud init is not installed in the debian
     # image by default.
-    project_keys = next(
+    project_keys: str = next(
         (item for item in project['commonInstanceMetadata'].get('items', [])
-         if item['key'] == 'ssh-keys'), {}).get('value', '')
+         if item['key'] == 'ssh-keys'), {}).get('value', '')  # type: ignore
     ssh_keys = project_keys.split('\n') if project_keys else []
 
     # Get public key from file.
@@ -287,8 +287,8 @@ def setup_gcp_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
         if len(ssh_key_index) == 0:
             metadata.append({'key': 'ssh-keys', 'value': new_ssh_key})
         else:
-            ssh_key_index = ssh_key_index[0]
-            metadata[ssh_key_index]['value'] += '\n' + new_ssh_key
+            first_ssh_key_index = ssh_key_index[0]
+            metadata[first_ssh_key_index]['value'] += '\n' + new_ssh_key
 
         project['commonInstanceMetadata']['items'] = metadata
 

sky/backends/__init__.py

@@ -1,6 +1,10 @@
 """Sky Backends."""
-from sky.backends.backend import Backend
-from sky.backends.cloud_vm_ray_backend import CloudVmRayBackend
-from sky.backends.local_docker_backend import LocalDockerBackend
+from sky.backends.backend import Backend, ResourceHandle
+from sky.backends.cloud_vm_ray_backend import CloudVmRayBackend, CloudVmRayResourceHandle
+from sky.backends.local_docker_backend import LocalDockerBackend, LocalDockerResourceHandle
 
-__all__ = ['Backend', 'CloudVmRayBackend', 'LocalDockerBackend']
+__all__ = [
+    'Backend', 'ResourceHandle', 'CloudVmRayBackend',
+    'CloudVmRayResourceHandle', 'LocalDockerBackend',
+    'LocalDockerResourceHandle'
+]

sky/backends/backend.py

@@ -1,6 +1,6 @@
 """Sky backend interface."""
 import typing
-from typing import Dict, Optional
+from typing import Dict, Generic, Optional
 
 import sky
 from sky.utils import timeline
@@ -12,36 +12,43 @@
     from sky.data import storage as storage_lib
 
 Path = str
+_ResourceHandleType = typing.TypeVar('_ResourceHandleType',
+                                     bound='ResourceHandle')
 
 
-class Backend:
+# Backend-specific handle to the launched resources (e.g., a cluster).
+# Examples: 'cluster.yaml'; 'ray://...', 'k8s://...'.
+class ResourceHandle:
+
+    def get_cluster_name(self) -> str:
+        raise NotImplementedError
+
+
+class Backend(Generic[_ResourceHandleType]):
     """Backend interface: handles provisioning, setup, and scheduling."""
 
     # NAME is used to identify the backend class from cli/yaml.
     NAME = 'backend'
 
-    # Backend-specific handle to the launched resources (e.g., a cluster).
-    # Examples: 'cluster.yaml'; 'ray://...', 'k8s://...'.
-    class ResourceHandle:
-
-        def get_cluster_name(self) -> str:
-            raise NotImplementedError
+    # Backward compatibility, with the old name of the handle.
+    ResourceHandle = ResourceHandle  # pylint: disable=invalid-name
 
     # --- APIs ---
-    def check_resources_fit_cluster(self, handle: ResourceHandle,
+    def check_resources_fit_cluster(self, handle: _ResourceHandleType,
                                     task: 'task_lib.Task') -> None:
         """Check whether resources of the task are satisfied by cluster."""
         raise NotImplementedError
 
     @timeline.event
     @usage_lib.messages.usage.update_runtime('provision')
-    def provision(self,
-                  task: 'task_lib.Task',
-                  to_provision: Optional['resources.Resources'],
-                  dryrun: bool,
-                  stream_logs: bool,
-                  cluster_name: Optional[str] = None,
-                  retry_until_up: bool = False) -> ResourceHandle:
+    def provision(
+            self,
+            task: 'task_lib.Task',
+            to_provision: Optional['resources.Resources'],
+            dryrun: bool,
+            stream_logs: bool,
+            cluster_name: Optional[str] = None,
+            retry_until_up: bool = False) -> Optional[_ResourceHandleType]:
         if cluster_name is None:
             cluster_name = sky.backends.backend_utils.generate_cluster_name()
         usage_lib.record_cluster_name_for_current_operation(cluster_name)
@@ -51,22 +58,22 @@ def provision(self,
 
     @timeline.event
     @usage_lib.messages.usage.update_runtime('sync_workdir')
-    def sync_workdir(self, handle: ResourceHandle, workdir: Path) -> None:
+    def sync_workdir(self, handle: _ResourceHandleType, workdir: Path) -> None:
         return self._sync_workdir(handle, workdir)
 
     @timeline.event
     @usage_lib.messages.usage.update_runtime('sync_file_mounts')
     def sync_file_mounts(
         self,
-        handle: ResourceHandle,
+        handle: _ResourceHandleType,
         all_file_mounts: Dict[Path, Path],
         storage_mounts: Dict[Path, 'storage_lib.Storage'],
     ) -> None:
         return self._sync_file_mounts(handle, all_file_mounts, storage_mounts)
 
     @timeline.event
     @usage_lib.messages.usage.update_runtime('setup')
-    def setup(self, handle: ResourceHandle, task: 'task_lib.Task',
+    def setup(self, handle: _ResourceHandleType, task: 'task_lib.Task',
               detach_setup: bool) -> None:
         return self._setup(handle, task, detach_setup)
 
@@ -75,15 +82,15 @@ def add_storage_objects(self, task: 'task_lib.Task') -> None:
 
     @timeline.event
     @usage_lib.messages.usage.update_runtime('execute')
-    def execute(self, handle: ResourceHandle, task: 'task_lib.Task',
+    def execute(self, handle: _ResourceHandleType, task: 'task_lib.Task',
                 detach_run: bool) -> None:
         usage_lib.record_cluster_name_for_current_operation(
             handle.get_cluster_name())
         usage_lib.messages.usage.update_actual_task(task)
         return self._execute(handle, task, detach_run)
 
     @timeline.event
-    def post_execute(self, handle: ResourceHandle, down: bool) -> None:
+    def post_execute(self, handle: _ResourceHandleType, down: bool) -> None:
         """Post execute(): e.g., print helpful inspection messages."""
         return self._post_execute(handle, down)
 
@@ -94,7 +101,7 @@ def teardown_ephemeral_storage(self, task: 'task_lib.Task') -> None:
     @timeline.event
     @usage_lib.messages.usage.update_runtime('teardown')
     def teardown(self,
-                 handle: ResourceHandle,
+                 handle: _ResourceHandleType,
                  terminate: bool,
                  purge: bool = False) -> None:
         self._teardown(handle, terminate, purge)
@@ -104,42 +111,43 @@ def register_info(self, **kwargs) -> None:
         pass
 
     # --- Implementations of the APIs ---
-    def _provision(self,
-                   task: 'task_lib.Task',
-                   to_provision: Optional['resources.Resources'],
-                   dryrun: bool,
-                   stream_logs: bool,
-                   cluster_name: str,
-                   retry_until_up: bool = False) -> ResourceHandle:
+    def _provision(
+            self,
+            task: 'task_lib.Task',
+            to_provision: Optional['resources.Resources'],
+            dryrun: bool,
+            stream_logs: bool,
+            cluster_name: str,
+            retry_until_up: bool = False) -> Optional[_ResourceHandleType]:
         raise NotImplementedError
 
-    def _sync_workdir(self, handle: ResourceHandle, workdir: Path) -> None:
+    def _sync_workdir(self, handle: _ResourceHandleType, workdir: Path) -> None:
         raise NotImplementedError
 
     def _sync_file_mounts(
         self,
-        handle: ResourceHandle,
+        handle: _ResourceHandleType,
         all_file_mounts: Dict[Path, Path],
         storage_mounts: Dict[Path, 'storage_lib.Storage'],
     ) -> None:
         raise NotImplementedError
 
-    def _setup(self, handle: ResourceHandle, task: 'task_lib.Task',
+    def _setup(self, handle: _ResourceHandleType, task: 'task_lib.Task',
                detach_setup: bool) -> None:
         raise NotImplementedError
 
-    def _execute(self, handle: ResourceHandle, task: 'task_lib.Task',
+    def _execute(self, handle: _ResourceHandleType, task: 'task_lib.Task',
                  detach_run: bool) -> None:
         raise NotImplementedError
 
-    def _post_execute(self, handle: ResourceHandle, down: bool) -> None:
+    def _post_execute(self, handle: _ResourceHandleType, down: bool) -> None:
         raise NotImplementedError
 
     def _teardown_ephemeral_storage(self, task: 'task_lib.Task') -> None:
         raise NotImplementedError
 
     def _teardown(self,
-                  handle: ResourceHandle,
+                  handle: _ResourceHandleType,
                   terminate: bool,
                   purge: bool = False):
         raise NotImplementedError