Repository: skypilot-org/skypilot
Pull request #1564 (2022-12-30 07:40:32)

[Mypy] Add more mypy
<!-- Describe the changes in this PR -->

This PR adds mypy checks for almost all the modules excluding the following files:
```
--exclude sky/benchmark
--exclude sky/callbacks
--exclude sky/skylet/providers # blocked by the ray upstream
--exclude sky/resources.py # blocked by #1429 
--exclude sky/global_user_state.py
--exclude sky/optimizer.py
--exclude sky/task.py
```

Except for the modules with comments above, we will add mypy checks for the rest of the files in a future PR.

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [x] Any manual or new tests for this PR (please specify below)
  - mypy checks pass
- [x] All smoke tests: `bash tests/run_smoke_tests.sh` 
- [x] Backward compatibility tests: `bash tests/backward_comaptibility_tests.sh`


Michaelvll: With the latest changes, I avoid the nested class of the `ResourceHandle` as we are using the type of the handle outside the outer class. There is no reason to keep it nested in the Backend.
Michaelvll: Thanks for the review @romilbhardwaj! Tested:
- [ ] `pytest tests/test_smoke.py`

Patch:
diff --git a/format.sh b/format.sh
index 51d8aa7d7..281a51f55 100755
--- a/format.sh
+++ b/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
 
diff --git a/requirements-dev.txt b/requirements-dev.txt
index 7ccdbf937..8f29a5b46 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -11,6 +11,7 @@ toml==0.10.2
 mypy==0.991
 types-PyYAML
 types-requests
+types-setuptools
 
 # testing
 pytest
diff --git a/sky/authentication.py b/sky/authentication.py
index ae1aaa66c..95f089ebf 100644
--- a/sky/authentication.py
+++ b/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,7 +245,7 @@ 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
@@ -253,9 +253,9 @@ def setup_gcp_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
     # 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
 
diff --git a/sky/backends/__init__.py b/sky/backends/__init__.py
index 4f524943f..81007e4e2 100644
--- a/sky/backends/__init__.py
+++ b/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'
+]
diff --git a/sky/backends/backend.py b/sky/backends/backend.py
index 245f98a8f..dc9f5f1e0 100644
--- a/sky/backends/backend.py
+++ b/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 @@ if typing.TYPE_CHECKING:
     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,14 +58,14 @@ class Backend:
 
     @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:
@@ -66,7 +73,7 @@ class Backend:
 
     @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,7 +82,7 @@ class Backend:
 
     @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())
@@ -83,7 +90,7 @@ class Backend:
         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 @@ class Backend:
     @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 @@ class Backend:
         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
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index a349a352b..6485b9833 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -14,6 +14,7 @@ import threading
 import time
 import typing
 from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, Union
+from typing_extensions import Literal
 import uuid
 
 import colorama
@@ -55,6 +56,8 @@ from sky.usage import usage_lib
 if typing.TYPE_CHECKING:
     from sky import resources
     from sky import task as task_lib
+    from sky.backends import cloud_vm_ray_backend
+    from sky.backends import local_docker_backend
 
 logger = sky_logging.init_logger(__name__)
 console = rich_console.Console()
@@ -94,7 +97,7 @@ DEFAULT_TASK_CPU_DEMAND = 0.5
 # Mapping from reserved cluster names to the corresponding group name (logging
 # purpose).
 # NOTE: each group can only have one reserved cluster name for now.
-SKY_RESERVED_CLUSTER_NAMES = {
+SKY_RESERVED_CLUSTER_NAMES: Dict[str, str] = {
     spot_lib.SPOT_CONTROLLER_NAME: 'Managed spot controller'
 }
 
@@ -153,15 +156,15 @@ def fill_template(template_name: str,
     with open(template_path) as fin:
         template = fin.read()
     if output_path is None:
-        assert ('cluster_name' in variables), ('cluster_name is required.')
         cluster_name = variables.get('cluster_name')
+        assert isinstance(cluster_name, str), cluster_name
         output_path = _get_yaml_path_from_cluster_name(cluster_name,
                                                        output_prefix)
     output_path = os.path.abspath(output_path)
 
     # Write out yaml config.
-    template = jinja2.Template(template)
-    content = template.render(**variables)
+    j2_template = jinja2.Template(template)
+    content = j2_template.render(**variables)
     with open(output_path, 'w') as fout:
         fout.write(content)
     return output_path
@@ -501,8 +504,9 @@ class SSHConfigHelper(object):
         external_worker_ips = list(sorted(external_worker_ips))
 
         overwrites = [False] * len(external_worker_ips)
-        overwrite_begin_idxs = [None] * len(external_worker_ips)
-        codegens = [None] * len(external_worker_ips)
+        overwrite_begin_idxs: List[Optional[int]] = [None
+                                                    ] * len(external_worker_ips)
+        codegens: List[Optional[str]] = [None] * len(external_worker_ips)
         worker_names = []
         extra_path_name = cls.ssh_multinode_path.format(cluster_name)
 
@@ -583,6 +587,7 @@ class SSHConfigHelper(object):
             overwrite = overwrites[idx]
             overwrite_begin_idx = overwrite_begin_idxs[idx]
             codegen = codegens[idx]
+            assert codegen is not None, (codegens, idx)
             if overwrite:
                 assert overwrite_begin_idx is not None
                 updated_lines = codegen.splitlines(keepends=True) + ['\n']
@@ -702,9 +707,9 @@ class SSHConfigHelper(object):
                 break
 
 
-def _replace_yaml_dicts(new_yaml: str, old_yaml: str,
-                        restore_key_names: Set[str],
-                        restore_key_names_exceptions: List[List[str]]) -> str:
+def _replace_yaml_dicts(
+        new_yaml: str, old_yaml: str, restore_key_names: Set[str],
+        restore_key_names_exceptions: Sequence[Sequence[str]]) -> str:
     """Replaces 'new' with 'old' for all keys in restore_key_names.
 
     The replacement will be applied recursively and only for the blocks
@@ -766,7 +771,6 @@ def write_cluster_config(
         wheel_hash: str,
         region: Optional[clouds.Region] = None,
         zones: Optional[List[clouds.Zone]] = None,
-        auth_config: Optional[Dict[str, str]] = None,
         dryrun: bool = False,
         keep_launch_fields_in_existing_config: bool = True) -> Dict[str, str]:
     """Fills in cluster configuration templates and writes them out.
@@ -1008,7 +1012,7 @@ def wait_until_ray_cluster_ready(
 ) -> bool:
     """Returns whether the entire ray cluster is ready."""
     if num_nodes <= 1:
-        return
+        return True
 
     # Manually fetching head ip instead of using `ray exec` to avoid the bug
     # that `ray exec` fails to connect to the head node after some workers
@@ -1134,7 +1138,7 @@ def ssh_credential_from_yaml(cluster_yaml: str) -> Dict[str, str]:
 
 def parallel_data_transfer_to_nodes(
     runners: List[command_runner.SSHCommandRunner],
-    source: str,
+    source: Optional[str],
     target: str,
     cmd: Optional[str],
     run_rsync: bool,
@@ -1147,9 +1151,9 @@ def parallel_data_transfer_to_nodes(
     """Runs a command on all nodes and optionally runs rsync from src->dst.
 
     Args:
-        runners: A list of SSHCommandRunner's that represent multiple nodes.
-        source_target: Tuple[str, str]; Source for rsync on local node and
-            Destination on remote node for rsync
+        runners: A list of SSHCommandRunner objects that represent multiple nodes.
+        source: Optional[str]; Source for rsync on local node
+        target: str; Destination on remote node for rsync
         cmd: str; Command to be executed on all nodes
         action_message: str; Message to be printed while the command runs
         log_path: str; Path to the log file
@@ -1173,8 +1177,9 @@ def parallel_data_transfer_to_nodes(
                 stderr=stdout + stderr)
 
         if run_rsync:
+            assert source is not None
             # TODO(zhwu): Optimize for large amount of files.
-            # zip / transfer/ unzip
+            # zip / transfer / unzip
             runner.rsync(
                 source=source,
                 target=target,
@@ -1261,8 +1266,8 @@ def _query_head_ip_with_retries(cluster_yaml: str,
                 f'ray get-head-ip {full_cluster_yaml!r}',
                 stdout=subprocess.PIPE,
                 stderr=subprocess.DEVNULL).stdout.decode().strip()
-            head_ip = re.findall(IP_ADDR_REGEX, out)
-            if len(head_ip) > 1:
+            head_ip_list = re.findall(IP_ADDR_REGEX, out)
+            if len(head_ip_list) > 1:
                 # This could be triggered if e.g., some logging is added in
                 # skypilot_config, a module that has some code executed
                 # whenever `sky` is imported.
@@ -1271,12 +1276,12 @@ def _query_head_ip_with_retries(cluster_yaml: str,
                     'the `ray get-head-ip` command. This could '
                     'happen if there is extra output from it, '
                     'which should be inspected below.\nProceeding with '
-                    f'the last detected IP ({head_ip[-1]}) as head IP.'
+                    f'the last detected IP ({head_ip_list[-1]}) as head IP.'
                     f'\n== Output ==\n{out}'
                     f'\n== Output ends ==')
-                head_ip = head_ip[-1:]
-            assert 1 == len(head_ip), (out, head_ip)
-            head_ip = head_ip[0]
+                head_ip_list = head_ip_list[-1:]
+            assert 1 == len(head_ip_list), (out, head_ip_list)
+            head_ip = head_ip_list[0]
             break
         except subprocess.CalledProcessError as e:
             if i == max_attempts - 1:
@@ -1290,7 +1295,8 @@ def _query_head_ip_with_retries(cluster_yaml: str,
 @timeline.event
 def get_node_ips(cluster_yaml: str,
                  expected_num_nodes: int,
-                 handle: Optional[backends.Backend.ResourceHandle] = None,
+                 handle: Optional[
+                     'cloud_vm_ray_backend.CloudVmRayResourceHandle'] = None,
                  head_ip_max_attempts: int = 1,
                  worker_ip_max_attempts: int = 1,
                  get_internal_ips: bool = False) -> List[str]:
@@ -1304,6 +1310,7 @@ def get_node_ips(cluster_yaml: str,
     if use_tpu_vm:
         assert expected_num_nodes == 1, (
             'TPU VM only supports single node for now.')
+        assert handle is not None, 'handle is required for TPU VM.'
         try:
             ips = _get_tpu_vm_pod_ips(ray_config, get_internal_ips)
         except exceptions.CommandError as e:
@@ -1329,7 +1336,7 @@ def get_node_ips(cluster_yaml: str,
     except RuntimeError as e:
         raise exceptions.FetchIPError(
             exceptions.FetchIPError.Reason.HEAD) from e
-    head_ip = [head_ip]
+    head_ip_list = [head_ip]
     if expected_num_nodes > 1:
         backoff = common_utils.Backoff(initial_backoff=5, max_backoff_factor=5)
 
@@ -1363,7 +1370,7 @@ def get_node_ips(cluster_yaml: str,
             worker_ips = re.findall(IP_ADDR_REGEX, out)
             # Remove head ip from worker ip list.
             for i, ip in enumerate(worker_ips):
-                if ip == head_ip[0]:
+                if ip == head_ip_list[0]:
                     del worker_ips[i]
                     break
         if len(worker_ips) != expected_num_nodes - 1:
@@ -1387,7 +1394,7 @@ def get_node_ips(cluster_yaml: str,
                     exceptions.FetchIPError.Reason.WORKER)
     else:
         worker_ips = []
-    return head_ip + worker_ips
+    return head_ip_list + worker_ips
 
 
 @timeline.event
@@ -1460,7 +1467,7 @@ def _get_tpu_vm_pod_ips(ray_config: Dict[str, Any],
 
 @timeline.event
 def get_head_ip(
-    handle: backends.Backend.ResourceHandle,
+    handle: 'cloud_vm_ray_backend.CloudVmRayResourceHandle',
     use_cached_head_ip: bool = True,
     max_attempts: int = 1,
 ) -> str:
@@ -1479,87 +1486,6 @@ def get_head_ip(
     return head_ip
 
 
-def run_command_and_handle_ssh_failure(
-        runner: command_runner.SSHCommandRunner,
-        command: str,
-        failure_message: Optional[str] = None) -> str:
-    """Runs command remotely and returns output with proper error handling."""
-    rc, stdout, stderr = runner.run(command,
-                                    require_outputs=True,
-                                    stream_logs=False)
-    if rc == 255:
-        # SSH failed
-        raise RuntimeError(
-            f'SSH with user {runner.ssh_user} and key {runner.ssh_private_key} '
-            f'to {runner.ip} failed. This is most likely due to incorrect '
-            'credentials or incorrect permissions for the key file. Check '
-            'your credentials and try again.')
-    subprocess_utils.handle_returncode(rc,
-                                       command,
-                                       failure_message,
-                                       stderr=stderr)
-    return stdout
-
-
-def do_filemounts_and_setup_on_local_workers(
-        cluster_config_file: str,
-        worker_ips: List[str] = None,
-        extra_setup_cmds: List[str] = None):
-    """Completes filemounting and setup on worker nodes.
-
-    Syncs filemounts and runs setup on worker nodes for a local cluster. This
-    is a workaround for a Ray Autoscaler bug where `ray up` does not perform
-    filemounting or setup for local cluster worker nodes.
-    """
-    config = common_utils.read_yaml(cluster_config_file)
-
-    ssh_credentials = ssh_credential_from_yaml(cluster_config_file)
-    if worker_ips is None:
-        worker_ips = config['provider']['worker_ips']
-    file_mounts = config['file_mounts']
-
-    setup_cmds = config['setup_commands']
-    if extra_setup_cmds is not None:
-        setup_cmds += extra_setup_cmds
-    setup_script = log_lib.make_task_bash_script('\n'.join(setup_cmds))
-
-    worker_runners = command_runner.SSHCommandRunner.make_runner_list(
-        worker_ips, **ssh_credentials)
-
-    # Uploads setup script to the worker node
-    with tempfile.NamedTemporaryFile('w', prefix='sky_setup_') as f:
-        f.write(setup_script)
-        f.flush()
-        setup_sh_path = f.name
-        setup_file = os.path.basename(setup_sh_path)
-        file_mounts[f'/tmp/{setup_file}'] = setup_sh_path
-
-        # Ray Autoscaler Bug: Filemounting + Ray Setup
-        # does not happen on workers.
-        def _setup_local_worker(runner: command_runner.SSHCommandRunner):
-            for dst, src in file_mounts.items():
-                mkdir_dst = f'mkdir -p {os.path.dirname(dst)}'
-                run_command_and_handle_ssh_failure(
-                    runner,
-                    mkdir_dst,
-                    failure_message=f'Failed to run {mkdir_dst} on remote.')
-                if os.path.isdir(src):
-                    src = os.path.join(src, '')
-                runner.rsync(source=src, target=dst, up=True, stream_logs=False)
-
-            setup_cmd = f'/bin/bash -i /tmp/{setup_file} 2>&1'
-            rc, stdout, _ = runner.run(setup_cmd,
-                                       stream_logs=False,
-                                       require_outputs=True)
-            subprocess_utils.handle_returncode(
-                rc,
-                setup_cmd,
-                'Failed to setup Ray autoscaler commands on remote.',
-                stderr=stdout)
-
-        subprocess_utils.run_in_parallel(_setup_local_worker, worker_runners)
-
-
 def check_network_connection():
     # Tolerate 3 retries as it is observed that connections can fail.
     adapter = adapters.HTTPAdapter(max_retries=retry_lib.Retry(total=3))
@@ -1794,7 +1720,7 @@ def check_owner_identity(cluster_name: str) -> None:
     if record is None:
         return
     handle = record['handle']
-    if not isinstance(handle, backends.CloudVmRayBackend.ResourceHandle):
+    if not isinstance(handle, backends.CloudVmRayResourceHandle):
         return
 
     cloud = handle.launched_resources.cloud
@@ -1822,7 +1748,7 @@ def check_owner_identity(cluster_name: str) -> None:
 
 
 def _get_cluster_status_via_cloud_cli(
-    handle: 'backends.CloudVmRayBackend.ResourceHandle'
+    handle: 'cloud_vm_ray_backend.CloudVmRayResourceHandle'
 ) -> List[global_user_state.ClusterStatus]:
     """Returns the status of the cluster."""
     resources: sky.Resources = handle.launched_resources
@@ -1837,7 +1763,7 @@ def _update_cluster_status_no_lock(
     if record is None:
         return None
     handle = record['handle']
-    if not isinstance(handle, backends.CloudVmRayBackend.ResourceHandle):
+    if not isinstance(handle, backends.CloudVmRayResourceHandle):
         return record
 
     cluster_name = handle.cluster_name
@@ -2056,7 +1982,7 @@ def _refresh_cluster_record(
     check_owner_identity(cluster_name)
 
     handle = record['handle']
-    if isinstance(handle, backends.CloudVmRayBackend.ResourceHandle):
+    if isinstance(handle, backends.CloudVmRayResourceHandle):
         use_spot = handle.launched_resources.use_spot
         has_autostop = (
             record['status'] != global_user_state.ClusterStatus.STOPPED and
@@ -2075,7 +2001,7 @@ def refresh_cluster_status_handle(
     force_refresh: bool = False,
     acquire_per_cluster_status_lock: bool = True,
 ) -> Tuple[Optional[global_user_state.ClusterStatus],
-           Optional[backends.Backend.ResourceHandle]]:
+           Optional[backends.ResourceHandle]]:
     """Refresh the cluster, and return the possibly updated status and handle.
 
     This is a wrapper of refresh_cluster_record, which returns the status and
@@ -2091,12 +2017,32 @@ def refresh_cluster_status_handle(
     return record['status'], record['handle']
 
 
+@typing.overload
+def check_cluster_available(
+    cluster_name: str,
+    *,
+    operation: str,
+    check_cloud_vm_ray_backend: Literal[True] = True,
+) -> 'cloud_vm_ray_backend.CloudVmRayResourceHandle':
+    ...
+
+
+@typing.overload
+def check_cluster_available(
+    cluster_name: str,
+    *,
+    operation: str,
+    check_cloud_vm_ray_backend: Literal[False],
+) -> backends.ResourceHandle:
+    ...
+
+
 def check_cluster_available(
     cluster_name: str,
     *,
     operation: str,
     check_cloud_vm_ray_backend: bool = True,
-) -> backends.Backend.ResourceHandle:
+) -> backends.ResourceHandle:
     """Check if the cluster is available.
 
     Raises:
@@ -2177,7 +2123,7 @@ class CloudFilter(enum.Enum):
 def get_clusters(
     include_reserved: bool,
     refresh: bool,
-    cloud_filter: str = CloudFilter.CLOUDS_AND_DOCKER,
+    cloud_filter: CloudFilter = CloudFilter.CLOUDS_AND_DOCKER,
     cluster_names: Optional[Union[str, Sequence[str]]] = None,
 ) -> List[Dict[str, Any]]:
     """Returns a list of cached or optionally refreshed cluster records.
@@ -2233,7 +2179,7 @@ def get_clusters(
 
     def _is_local_cluster(record):
         handle = record['handle']
-        if isinstance(handle, backends.LocalDockerBackend.ResourceHandle):
+        if isinstance(handle, backends.LocalDockerResourceHandle):
             return False
         cluster_resources = handle.launched_resources
         return isinstance(cluster_resources.cloud, clouds.Local)
@@ -2317,15 +2263,30 @@ def get_clusters(
     return kept_records
 
 
+@typing.overload
 def get_backend_from_handle(
-        handle: backends.Backend.ResourceHandle) -> backends.Backend:
+    handle: 'cloud_vm_ray_backend.CloudVmRayResourceHandle'
+) -> 'cloud_vm_ray_backend.CloudVmRayBackend':
+    ...
+
+
+@typing.overload
+def get_backend_from_handle(
+    handle: 'local_docker_backend.LocalDockerResourceHandle'
+) -> 'local_docker_backend.LocalDockerBackend':
+    ...
+
+
+def get_backend_from_handle(
+        handle: backends.ResourceHandle) -> backends.Backend:
     """Gets a Backend object corresponding to a handle.
 
     Inspects handle type to infer the backend used for the resource.
     """
-    if isinstance(handle, backends.CloudVmRayBackend.ResourceHandle):
+    backend: backends.Backend
+    if isinstance(handle, backends.CloudVmRayResourceHandle):
         backend = backends.CloudVmRayBackend()
-    elif isinstance(handle, backends.LocalDockerBackend.ResourceHandle):
+    elif isinstance(handle, backends.LocalDockerResourceHandle):
         backend = backends.LocalDockerBackend()
     else:
         raise NotImplementedError(
@@ -2350,8 +2311,7 @@ def safe_console_status(msg: str):
     return NoOpConsole()
 
 
-def get_task_demands_dict(
-        task: 'task_lib.Task') -> Optional[Tuple[Optional[str], int]]:
+def get_task_demands_dict(task: 'task_lib.Task') -> Optional[Dict[str, float]]:
     """Returns the accelerator dict of the task"""
     # TODO: CPU and other memory resources are not supported yet.
     accelerator_dict = None
@@ -2465,3 +2425,24 @@ def check_public_cloud_enabled():
             raise RuntimeError(
                 'Cloud access is not set up. Run: '
                 f'{colorama.Style.BRIGHT}sky check{colorama.Style.RESET_ALL}')
+
+
+def run_command_and_handle_ssh_failure(runner: command_runner.SSHCommandRunner,
+                                       command: str,
+                                       failure_message: str) -> str:
+    """Runs command remotely and returns output with proper error handling."""
+    rc, stdout, stderr = runner.run(command,
+                                    require_outputs=True,
+                                    stream_logs=False)
+    if rc == 255:
+        # SSH failed
+        raise RuntimeError(
+            f'SSH with user {runner.ssh_user} and key {runner.ssh_private_key} '
+            f'to {runner.ip} failed. This is most likely due to incorrect '
+            'credentials or incorrect permissions for the key file. Check '
+            'your credentials and try again.')
+    subprocess_utils.handle_returncode(rc,
+                                       command,
+                                       failure_message,
+                                       stderr=stderr)
+    return stdout
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index d3959e155..c38af365d 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -54,7 +54,6 @@ from sky.utils import ux_utils
 if typing.TYPE_CHECKING:
     from sky import dag
 
-OptimizeTarget = optimizer.OptimizeTarget
 Path = str
 
 SKY_REMOTE_APP_DIR = backend_utils.SKY_REMOTE_APP_DIR
@@ -303,7 +302,7 @@ class RayCodeGen:
     def add_gang_scheduling_placement_group(
         self,
         num_nodes: int,
-        accelerator_dict: Dict[str, int],
+        accelerator_dict: Optional[Dict[str, float]],
         stable_cluster_internal_ips: List[str],
     ) -> None:
         """Create the gang scheduling placement group for a Task.
@@ -332,7 +331,7 @@ class RayCodeGen:
             # FIXME: This is a hack to make sure that we do not reserve
             # GPU when requesting TPU.
             if 'tpu' in acc_name.lower():
-                gpu_dict = dict()
+                gpu_dict = {}
             for bundle in bundles:
                 bundle.update({
                     **accelerator_dict,
@@ -402,12 +401,12 @@ class RayCodeGen:
         ]
 
     def add_ray_task(self,
-                     bash_script: str,
+                     bash_script: Optional[str],
                      task_name: Optional[str],
                      job_run_id: Optional[str],
                      ray_resources_dict: Optional[Dict[str, float]],
                      log_dir: str,
-                     env_vars: Dict[str, str] = None,
+                     env_vars: Optional[Dict[str, str]] = None,
                      gang_scheduling_id: int = 0,
                      use_sudo: bool = False) -> None:
         """Generates code for a ray remote task that runs a bash command."""
@@ -422,7 +421,7 @@ class RayCodeGen:
         cpu_str = f', num_cpus={backend_utils.DEFAULT_TASK_CPU_DEMAND}'
 
         resources_str = ''
-        num_gpus = 0
+        num_gpus = 0.0
         num_gpus_str = ''
         if ray_resources_dict is not None:
             assert len(ray_resources_dict) == 1, \
@@ -445,7 +444,7 @@ class RayCodeGen:
 
         sky_env_vars_dict_str = [
             textwrap.dedent("""\
-            sky_env_vars_dict = dict()
+            sky_env_vars_dict = {}
             sky_env_vars_dict['SKYPILOT_NODE_IPS'] = job_ip_list_str
             # Environment starting with `SKY_` is deprecated.
             sky_env_vars_dict['SKY_NODE_IPS'] = job_ip_list_str
@@ -555,8 +554,8 @@ class RetryingVmProvisioner(object):
         """Resources to be provisioned."""
 
         def __init__(
-            self, cluster_name: str,
-            resources: Optional[resources_lib.Resources], num_nodes: int,
+            self, cluster_name: str, resources: resources_lib.Resources,
+            num_nodes: int,
             prev_cluster_status: Optional[global_user_state.ClusterStatus]
         ) -> None:
             assert cluster_name is not None, 'cluster_name must be specified.'
@@ -572,10 +571,10 @@ class RetryingVmProvisioner(object):
         HEAD_FAILED = 2
 
     def __init__(self, log_dir: str, dag: 'dag.Dag',
-                 optimize_target: OptimizeTarget,
+                 optimize_target: 'optimizer.OptimizeTarget',
                  requested_features: Set[clouds.CloudImplementationFeatures],
                  local_wheel_path: pathlib.Path, wheel_hash: str):
-        self._blocked_resources = set()
+        self._blocked_resources: Set[resources_lib.Resources] = set()
 
         self.log_dir = os.path.expanduser(log_dir)
         self._dag = dag
@@ -593,17 +592,17 @@ class RetryingVmProvisioner(object):
 
         del region  # unused
         style = colorama.Style
-        assert len(zones) == 1, zones
+        assert zones and len(zones) == 1, zones
         zone = zones[0]
         splits = stderr.split('\n')
-        exception_str = [s for s in splits if s.startswith('Exception: ')]
+        exception_list = [s for s in splits if s.startswith('Exception: ')]
         httperror_str = [
             s for s in splits
             if s.startswith('googleapiclient.errors.HttpError: ')
         ]
-        if len(exception_str) == 1:
+        if len(exception_list) == 1:
             # Parse structured response {'errors': [...]}.
-            exception_str = exception_str[0][len('Exception: '):]
+            exception_str = exception_list[0][len('Exception: '):]
             try:
                 exception_dict = ast.literal_eval(exception_str)
             except Exception as e:
@@ -672,7 +671,7 @@ class RetryingVmProvisioner(object):
                     launchable_resources.copy(zone=zone.name))
         else:
             # No such structured error response found.
-            assert not exception_str, stderr
+            assert not exception_list, stderr
             if 'was not found' in stderr:
                 # Example: The resource
                 # 'projects/<id>/zones/zone/acceleratorTypes/nvidia-tesla-v100'
@@ -696,7 +695,8 @@ class RetryingVmProvisioner(object):
             self, launchable_resources: 'resources_lib.Resources',
             region: 'clouds.Region', zones: Optional[List['clouds.Zone']],
             stdout: str, stderr: str):
-        assert zones is not None
+        assert zones is not None, 'AWS should always have zones.'
+
         style = colorama.Style
         stdout_splits = stdout.split('\n')
         stderr_splits = stderr.split('\n')
@@ -742,6 +742,7 @@ class RetryingVmProvisioner(object):
             zone=None)
         assert len(region_with_zones_list) == 1, region_with_zones_list
         region_with_zones = region_with_zones_list[0]
+        assert region_with_zones.zones is not None, region_with_zones
         if set(zones) == set(region_with_zones.zones):
             # The underlying AWS NodeProvider will try all specified zones of a
             # region. (Each boto3 request takes one zone.)
@@ -883,6 +884,7 @@ class RetryingVmProvisioner(object):
                     self._blocked_resources.add(
                         launchable_resources.copy(zone=zone.name))
             return False  # definitely_no_nodes_launched
+        assert stdout is not None and stderr is not None, (stdout, stderr)
 
         # TODO(zongheng): refactor into Cloud interface?
         handlers = {
@@ -1136,7 +1138,7 @@ class RetryingVmProvisioner(object):
         dryrun: bool,
         stream_logs: bool,
         cluster_name: str,
-        cloud_user_identity: str,
+        cloud_user_identity: Optional[str],
         prev_cluster_status: Optional[global_user_state.ClusterStatus],
     ):
         """The provision retry loop."""
@@ -1216,7 +1218,7 @@ class RetryingVmProvisioner(object):
             # Record early, so if anything goes wrong, 'sky status' will show
             # the cluster name and users can appropriately 'sky down'.  It also
             # means a second 'sky launch -c <name>' will attempt to reuse.
-            handle = CloudVmRayBackend.ResourceHandle(
+            handle = CloudVmRayResourceHandle(
                 cluster_name=cluster_name,
                 cluster_yaml=cluster_config_file,
                 launched_nodes=num_nodes,
@@ -1362,7 +1364,7 @@ class RetryingVmProvisioner(object):
             message, no_failover=is_prev_cluster_healthy)
 
     def _tpu_pod_setup(self, cluster_yaml: str,
-                       cluster_handle: 'backends.Backend.ResourceHandle'):
+                       cluster_handle: 'backends.CloudVmRayResourceHandle'):
         """Completes setup and start Ray cluster on TPU VM Pod nodes.
 
         This is a workaround for Ray Autoscaler where `ray up` does not
@@ -1382,8 +1384,8 @@ class RetryingVmProvisioner(object):
                                                       **ssh_credentials)
         cmd_str = 'python3 -c \"import ray; print(ray._private.services.get_node_ip_address())\"'  # pylint: disable=line-too-long
         rc, stdout, stderr = head_runner.run(cmd_str,
-                                             stream_logs=False,
-                                             require_outputs=True)
+                                             require_outputs=True,
+                                             stream_logs=False)
         subprocess_utils.handle_returncode(
             rc,
             cmd_str,
@@ -1396,7 +1398,7 @@ class RetryingVmProvisioner(object):
         worker_start_ray_commands += ray_config['worker_start_ray_commands']
 
         # Setup TPU VM Pod workers and launch Ray cluster.
-        backend_utils.do_filemounts_and_setup_on_local_workers(
+        onprem_utils.do_filemounts_and_setup_on_local_workers(
             cluster_yaml,
             worker_ips=all_ips[1:],
             extra_setup_cmds=worker_start_ray_commands)
@@ -1404,7 +1406,7 @@ class RetryingVmProvisioner(object):
     @timeline.event
     def _gang_schedule_ray_up(
             self, to_provision_cloud: clouds.Cloud, cluster_config_file: str,
-            cluster_handle: 'backends.Backend.ResourceHandle',
+            cluster_handle: 'backends.CloudVmRayResourceHandle',
             log_abs_path: str, stream_logs: bool, logging_info: dict,
             use_spot: bool
     ) -> Tuple[GangSchedulingStatus, str, str, Optional[str]]:
@@ -1588,7 +1590,7 @@ class RetryingVmProvisioner(object):
         # nodes. Hence, this method here replicates what the Ray autoscaler
         # would do were it for public cloud.
         if isinstance(to_provision_cloud, clouds.Local):
-            backend_utils.do_filemounts_and_setup_on_local_workers(
+            onprem_utils.do_filemounts_and_setup_on_local_workers(
                 cluster_config_file)
 
         # FIXME(zongheng): the below requires ray processes are up on head. To
@@ -1623,8 +1625,7 @@ class RetryingVmProvisioner(object):
         # gang_succeeded = False, if head OK, but workers failed.
         return cluster_status, '', '', None
 
-    def _ensure_cluster_ray_started(self,
-                                    handle: 'CloudVmRayBackend.ResourceHandle',
+    def _ensure_cluster_ray_started(self, handle: 'CloudVmRayResourceHandle',
                                     log_abs_path) -> None:
         """Ensures ray processes are up on a just-provisioned cluster."""
         if handle.launched_nodes > 1:
@@ -1801,243 +1802,242 @@ class RetryingVmProvisioner(object):
         return config_dict
 
 
-class CloudVmRayBackend(backends.Backend):
-    """Backend: runs on cloud virtual machines, managed by Ray.
+class CloudVmRayResourceHandle(backends.backend.ResourceHandle):
+    """A pickle-able tuple of:
 
-    Changing this class may also require updates to:
-      * Cloud providers' templates under config/
-      * Cloud providers' implementations under clouds/
+    - (required) Cluster name.
+    - (required) Path to a cluster.yaml file.
+    - (optional) A cached head node public IP.  Filled in after a
+        successful provision().
+    - (optional) A cached stable list of (internal IP, external IP) tuples
+        for all nodes in a cluster. Filled in after successful task execution.
+    - (optional) Launched num nodes
+    - (optional) Launched resources
+    - (optional) If TPU(s) are managed, a path to a deletion script.
     """
-
-    NAME = 'cloudvmray'
-
-    class ResourceHandle(object):
-        """A pickle-able tuple of:
-
-        - (required) Cluster name.
-        - (required) Path to a cluster.yaml file.
-        - (optional) A cached head node public IP.  Filled in after a
-            successful provision().
-        - (optional) A cached stable list of (internal IP, external IP) tuples
-         for all nodes in a cluster. Filled in after successful task execution.
-        - (optional) Launched num nodes
-        - (optional) Launched resources
-        - (optional) If TPU(s) are managed, a path to a deletion script.
+    _VERSION = 3
+
+    def __init__(self,
+                 *,
+                 cluster_name: str,
+                 cluster_yaml: str,
+                 launched_nodes: int,
+                 launched_resources: resources_lib.Resources,
+                 stable_internal_external_ips: Optional[List[Tuple[
+                     str, str]]] = None,
+                 tpu_create_script: Optional[str] = None,
+                 tpu_delete_script: Optional[str] = None) -> None:
+        self._version = self._VERSION
+        self.cluster_name = cluster_name
+        self._cluster_yaml = cluster_yaml.replace(os.path.expanduser('~'), '~',
+                                                  1)
+        # List of (internal_ip, external_ip) tuples for all the nodes
+        # in the cluster, sorted by the external ips.
+        self.stable_internal_external_ips = stable_internal_external_ips
+        self.launched_nodes = launched_nodes
+        self.launched_resources = launched_resources
+        self.tpu_create_script = tpu_create_script
+        self.tpu_delete_script = tpu_delete_script
+        self._maybe_make_local_handle()
+
+    def __repr__(self):
+        return (f'ResourceHandle('
+                f'\n\tcluster_name={self.cluster_name},'
+                f'\n\thead_ip={self.head_ip},'
+                '\n\tstable_internal_external_ips='
+                f'{self.stable_internal_external_ips},'
+                '\n\tcluster_yaml='
+                f'{self.cluster_yaml}, '
+                f'\n\tlaunched_resources={self.launched_nodes}x '
+                f'{self.launched_resources}, '
+                f'\n\ttpu_create_script={self.tpu_create_script}, '
+                f'\n\ttpu_delete_script={self.tpu_delete_script})')
+
+    def get_cluster_name(self):
+        return self.cluster_name
+
+    def _maybe_make_local_handle(self):
+        """Adds local handle for the local cloud case.
+
+        For public cloud, the first time sky launch is ran, task resources
+        = cluster resources. For the local cloud case, sky launch is ran,
+        task resources != cluster resources; hence, this method is needed
+        to correct this.
         """
-        _VERSION = 3
+        self.local_handle = None
+        local_file = os.path.expanduser(
+            onprem_utils.SKY_USER_LOCAL_CONFIG_PATH.format(self.cluster_name))
+        # Local cluster case requires several modifications:
+        #   1) Create local_handle to store local cluster IPs and
+        #      custom accelerators for each node.
+        #   2) Replace launched_resources to represent a generic local
+        #      node (without accelerator specifications).
+        #   3) Replace launched_nodes to represent the total nodes in the
+        #      local cluster.
+        if os.path.isfile(local_file):
+            config = onprem_utils.get_local_cluster_config_or_error(
+                self.cluster_name)
+            self.local_handle = {}
+            cluster_config = config['cluster']
+            auth_config = config['auth']
+            ips = cluster_config['ips']
+            local_region = clouds.Local.regions()[0].name
+            # Convert existing ResourceHandle fields to specify local
+            # cluster resources.
+            self.launched_resources = resources_lib.Resources(
+                cloud=clouds.Local(), region=local_region)
+            self.launched_nodes = len(ips)
+            self.local_handle['ips'] = ips
+            cluster_accs = onprem_utils.get_local_cluster_accelerators(
+                ips, auth_config)
+            self.local_handle['cluster_resources'] = [
+                resources_lib.Resources(
+                    cloud=clouds.Local(),
+                    accelerators=acc_dict if acc_dict else None,
+                    region=local_region) for acc_dict in cluster_accs
+            ]
 
-        def __init__(
-                self,
-                *,
-                cluster_name: str,
-                cluster_yaml: str,
-                stable_internal_external_ips: Optional[List[Tuple[str,
-                                                                  str]]] = None,
-                launched_nodes: Optional[int] = None,
-                launched_resources: Optional[resources_lib.Resources] = None,
-                tpu_create_script: Optional[str] = None,
-                tpu_delete_script: Optional[str] = None) -> None:
-            self._version = self._VERSION
-            self.cluster_name = cluster_name
-            self._cluster_yaml = cluster_yaml.replace(os.path.expanduser('~'),
-                                                      '~', 1)
-            # List of (internal_ip, external_ip) tuples for all the nodes
-            # in the cluster, sorted by the external ips.
-            self.stable_internal_external_ips = stable_internal_external_ips
-            self.launched_nodes = launched_nodes
-            self.launched_resources = launched_resources
-            self.tpu_create_script = tpu_create_script
-            self.tpu_delete_script = tpu_delete_script
-            self._maybe_make_local_handle()
-
-        def __repr__(self):
-            return (f'ResourceHandle('
-                    f'\n\tcluster_name={self.cluster_name},'
-                    f'\n\thead_ip={self.head_ip},'
-                    '\n\tstable_internal_external_ips='
-                    f'{self.stable_internal_external_ips},'
-                    '\n\tcluster_yaml='
-                    f'{self.cluster_yaml}, '
-                    f'\n\tlaunched_resources={self.launched_nodes}x '
-                    f'{self.launched_resources}, '
-                    f'\n\ttpu_create_script={self.tpu_create_script}, '
-                    f'\n\ttpu_delete_script={self.tpu_delete_script})')
-
-        def get_cluster_name(self):
-            return self.cluster_name
-
-        def _maybe_make_local_handle(self):
-            """Adds local handle for the local cloud case.
-
-            For public cloud, the first time sky launch is ran, task resources
-            = cluster resources. For the local cloud case, sky launch is ran,
-            task resources != cluster resources; hence, this method is needed
-            to correct this.
-            """
-            self.local_handle = None
-            local_file = os.path.expanduser(
-                onprem_utils.SKY_USER_LOCAL_CONFIG_PATH.format(
-                    self.cluster_name))
-            # Local cluster case requires several modifications:
-            #   1) Create local_handle to store local cluster IPs and
-            #      custom accelerators for each node.
-            #   2) Replace launched_resources to represent a generic local
-            #      node (without accelerator specifications).
-            #   3) Replace launched_nodes to represent the total nodes in the
-            #      local cluster.
-            if os.path.isfile(local_file):
-                config = onprem_utils.get_local_cluster_config_or_error(
-                    self.cluster_name)
-                self.local_handle = {}
-                cluster_config = config['cluster']
-                auth_config = config['auth']
-                ips = cluster_config['ips']
-                local_region = clouds.Local.regions()[0].name
-                # Convert existing ResourceHandle fields to specify local
-                # cluster resources.
-                self.launched_resources = resources_lib.Resources(
-                    cloud=clouds.Local(), region=local_region)
-                self.launched_nodes = len(ips)
-                self.local_handle['ips'] = ips
-                cluster_accs = onprem_utils.get_local_cluster_accelerators(
-                    ips, auth_config)
-                self.local_handle['cluster_resources'] = [
-                    resources_lib.Resources(
-                        cloud=clouds.Local(),
-                        accelerators=acc_dict if acc_dict else None,
-                        region=local_region) for acc_dict in cluster_accs
-                ]
-
-        def _update_cluster_region(self):
-            if self.launched_resources.region is not None:
-                return
+    def _update_cluster_region(self):
+        if self.launched_resources.region is not None:
+            return
+
+        config = common_utils.read_yaml(self.cluster_yaml)
+        provider = config['provider']
+        cloud = self.launched_resources.cloud
+        if cloud.is_same_cloud(clouds.Azure()):
+            region = provider['location']
+        elif cloud.is_same_cloud(clouds.GCP()) or cloud.is_same_cloud(
+                clouds.AWS()):
+            region = provider['region']
+        elif cloud.is_same_cloud(clouds.Local()):
+            # There is only 1 region for Local cluster, 'Local'.
+            local_regions = clouds.Local.regions()
+            region = local_regions[0].name
+
+        self.launched_resources = self.launched_resources.copy(region=region)
+
+    def _update_stable_cluster_ips(self, max_attempts: int = 1) -> None:
+        cluster_external_ips = backend_utils.get_node_ips(
+            self.cluster_yaml,
+            self.launched_nodes,
+            handle=self,
+            head_ip_max_attempts=max_attempts,
+            worker_ip_max_attempts=max_attempts,
+            get_internal_ips=False)
+
+        if self.external_ips() == cluster_external_ips:
+            # Optimization: If the cached external IPs are the same as the
+            # retrieved external IPs, then we can skip retrieving internal
+            # IPs since the cached IPs are up-to-date.
+            return
 
-            config = common_utils.read_yaml(self.cluster_yaml)
-            provider = config['provider']
-            cloud = self.launched_resources.cloud
-            if cloud.is_same_cloud(clouds.Azure()):
-                region = provider['location']
-            elif cloud.is_same_cloud(clouds.GCP()) or cloud.is_same_cloud(
-                    clouds.AWS()):
-                region = provider['region']
-            elif cloud.is_same_cloud(clouds.Local()):
-                # There is only 1 region for Local cluster, 'Local'.
-                local_regions = clouds.Local.regions()
-                region = local_regions[0].name
-
-            self.launched_resources = self.launched_resources.copy(
-                region=region)
-
-        def _update_stable_cluster_ips(self,
-                                       max_attempts: int = 1) -> List[str]:
-            cluster_external_ips = backend_utils.get_node_ips(
+        is_cluster_aws = (self.launched_resources is not None and
+                          isinstance(self.launched_resources.cloud, clouds.AWS))
+        if is_cluster_aws and skypilot_config.get_nested(
+                keys=('aws', 'use_internal_ips'), default_value=False):
+            # Optimization: if we know use_internal_ips is True (currently
+            # only exposed for AWS), then our AWS NodeProvider is
+            # guaranteed to pick subnets that will not assign public IPs,
+            # thus the first list of IPs returned above are already private
+            # IPs. So skip the second query.
+            cluster_internal_ips = list(cluster_external_ips)
+        else:
+            cluster_internal_ips = backend_utils.get_node_ips(
                 self.cluster_yaml,
                 self.launched_nodes,
                 handle=self,
                 head_ip_max_attempts=max_attempts,
                 worker_ip_max_attempts=max_attempts,
-                get_internal_ips=False)
+                get_internal_ips=True)
+
+        assert len(cluster_external_ips) == len(cluster_internal_ips), (
+            f'Cluster {self.cluster_name!r}:'
+            f'Expected same number of internal IPs {cluster_internal_ips}'
+            f' and external IPs {cluster_external_ips}.')
+
+        internal_external_ips = list(
+            zip(cluster_internal_ips, cluster_external_ips))
+
+        # Ensure head node is the first element, then sort based on the
+        # external IPs for stableness
+        stable_internal_external_ips = [internal_external_ips[0]] + sorted(
+            internal_external_ips[1:], key=lambda x: x[1])
+        self.stable_internal_external_ips = stable_internal_external_ips
+
+    def internal_ips(self,
+                     max_attempts: int = _FETCH_IP_MAX_ATTEMPTS,
+                     use_cached_ips: bool = True) -> Optional[List[str]]:
+        if not use_cached_ips:
+            self._update_stable_cluster_ips(max_attempts=max_attempts)
+        if self.stable_internal_external_ips is not None:
+            return [ips[0] for ips in self.stable_internal_external_ips]
+        return None
+
+    def external_ips(self,
+                     max_attempts: int = _FETCH_IP_MAX_ATTEMPTS,
+                     use_cached_ips: bool = True) -> Optional[List[str]]:
+        if not use_cached_ips:
+            self._update_stable_cluster_ips(max_attempts=max_attempts)
+        if self.stable_internal_external_ips is not None:
+            return [ips[1] for ips in self.stable_internal_external_ips]
+        return None
+
+    def get_hourly_price(self) -> float:
+        hourly_cost = (self.launched_resources.get_cost(3600) *
+                       self.launched_nodes)
+        return hourly_cost
+
+    @property
+    def cluster_yaml(self):
+        return os.path.expanduser(self._cluster_yaml)
+
+    @property
+    def head_ip(self):
+        external_ips = self.external_ips()
+        if external_ips is not None:
+            return external_ips[0]
+        return None
+
+    def __setstate__(self, state):
+        self._version = self._VERSION
+
+        version = state.pop('_version', None)
+        if version is None:
+            version = -1
+            state.pop('cluster_region', None)
+        if version < 2:
+            state['_cluster_yaml'] = state.pop('cluster_yaml')
+        if version < 3:
+            head_ip = state.pop('head_ip', None)
+            state['stable_internal_external_ips'] = None
+
+        self.__dict__.update(state)
+
+        # Because the _update_stable_cluster_ips function uses the handle,
+        # we call it on the current instance after the state is updated
+        if version < 3 and head_ip is not None:
+            try:
+                self._update_stable_cluster_ips()
+            except exceptions.FetchIPError:
+                # This occurs when an old cluster from was autostopped,
+                # so the head IP in the database is not updated.
+                pass
 
-            if self.external_ips() == cluster_external_ips:
-                # Optimization: If the cached external IPs are the same as the
-                # retrieved external IPs, then we can skip retrieving internal
-                # IPs since the cached IPs are up-to-date.
-                return
+        self._update_cluster_region()
 
-            is_cluster_aws = (self.launched_resources is not None and
-                              isinstance(self.launched_resources.cloud,
-                                         clouds.AWS))
-            if is_cluster_aws and skypilot_config.get_nested(
-                    keys=('aws', 'use_internal_ips'), default_value=False):
-                # Optimization: if we know use_internal_ips is True (currently
-                # only exposed for AWS), then our AWS NodeProvider is
-                # guaranteed to pick subnets that will not assign public IPs,
-                # thus the first list of IPs returned above are already private
-                # IPs. So skip the second query.
-                cluster_internal_ips = list(cluster_external_ips)
-            else:
-                cluster_internal_ips = backend_utils.get_node_ips(
-                    self.cluster_yaml,
-                    self.launched_nodes,
-                    handle=self,
-                    head_ip_max_attempts=max_attempts,
-                    worker_ip_max_attempts=max_attempts,
-                    get_internal_ips=True)
-
-            assert len(cluster_external_ips) == len(cluster_internal_ips), (
-                f'Cluster {self.cluster_name!r}:'
-                f'Expected same number of internal IPs {cluster_internal_ips}'
-                f' and external IPs {cluster_external_ips}.')
-
-            internal_external_ips = list(
-                zip(cluster_internal_ips, cluster_external_ips))
-
-            # Ensure head node is the first element, then sort based on the
-            # external IPs for stableness
-            stable_internal_external_ips = [internal_external_ips[0]] + sorted(
-                internal_external_ips[1:], key=lambda x: x[1])
-            self.stable_internal_external_ips = stable_internal_external_ips
-
-        def internal_ips(self,
-                         max_attempts: int = _FETCH_IP_MAX_ATTEMPTS,
-                         use_cached_ips: bool = True) -> Optional[List[str]]:
-            if not use_cached_ips:
-                self._update_stable_cluster_ips(max_attempts=max_attempts)
-            if self.stable_internal_external_ips is not None:
-                return [ips[0] for ips in self.stable_internal_external_ips]
-            return None
-
-        def external_ips(self,
-                         max_attempts: int = _FETCH_IP_MAX_ATTEMPTS,
-                         use_cached_ips: bool = True) -> Optional[List[str]]:
-            if not use_cached_ips:
-                self._update_stable_cluster_ips(max_attempts=max_attempts)
-            if self.stable_internal_external_ips is not None:
-                return [ips[1] for ips in self.stable_internal_external_ips]
-            return None
-
-        def get_hourly_price(self) -> float:
-            hourly_cost = (self.launched_resources.get_cost(3600) *
-                           self.launched_nodes)
-            return hourly_cost
-
-        @property
-        def cluster_yaml(self):
-            return os.path.expanduser(self._cluster_yaml)
-
-        @property
-        def head_ip(self):
-            external_ips = self.external_ips()
-            if external_ips is not None:
-                return external_ips[0]
-            return None
-
-        def __setstate__(self, state):
-            self._version = self._VERSION
-
-            version = state.pop('_version', None)
-            if version is None:
-                version = -1
-                state.pop('cluster_region', None)
-            if version < 2:
-                state['_cluster_yaml'] = state.pop('cluster_yaml')
-            if version < 3:
-                head_ip = state.pop('head_ip', None)
-                state['stable_internal_external_ips'] = None
-
-            self.__dict__.update(state)
-
-            # Because the _update_stable_cluster_ips function uses the handle,
-            # we call it on the current instance after the state is updated
-            if version < 3 and head_ip is not None:
-                try:
-                    self._update_stable_cluster_ips()
-                except exceptions.FetchIPError:
-                    # This occurs when an old cluster from was autostopped,
-                    # so the head IP in the database is not updated.
-                    pass
 
-            self._update_cluster_region()
+class CloudVmRayBackend(backends.Backend['CloudVmRayResourceHandle']):
+    """Backend: runs on cloud virtual machines, managed by Ray.
+
+    Changing this class may also require updates to:
+      * Cloud providers' templates under config/
+      * Cloud providers' implementations under clouds/
+    """
+
+    NAME = 'cloudvmray'
+
+    # Backward compatibility, with the old name of the handle.
+    ResourceHandle = CloudVmRayResourceHandle  # pylint: disable=invalid-name
 
     def __init__(self):
         self.run_timestamp = backend_utils.get_run_timestamp()
@@ -2061,12 +2061,13 @@ class CloudVmRayBackend(backends.Backend):
     def register_info(self, **kwargs) -> None:
         self._dag = kwargs.pop('dag', self._dag)
         self._optimize_target = kwargs.pop(
-            'optimize_target', self._optimize_target) or OptimizeTarget.COST
+            'optimize_target',
+            self._optimize_target) or optimizer.OptimizeTarget.COST
         self._requested_features = kwargs.pop('requested_features',
                                               self._requested_features)
         assert len(kwargs) == 0, f'Unexpected kwargs: {kwargs}'
 
-    def check_resources_fit_cluster(self, handle: ResourceHandle,
+    def check_resources_fit_cluster(self, handle: CloudVmRayResourceHandle,
                                     task: task_lib.Task):
         """Check if resources requested by the task fit the cluster.
 
@@ -2133,13 +2134,14 @@ class CloudVmRayBackend(backends.Backend):
                     f'{handle.launched_resources}\n'
                     f'{mismatch_str}')
 
-    def _provision(self,
-                   task: task_lib.Task,
-                   to_provision: Optional[resources_lib.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_lib.Resources],
+            dryrun: bool,
+            stream_logs: bool,
+            cluster_name: str,
+            retry_until_up: bool = False) -> Optional[CloudVmRayResourceHandle]:
         """Provisions using 'ray up'.
 
         Raises:
@@ -2250,7 +2252,7 @@ class CloudVmRayBackend(backends.Backend):
                             error_message,
                             failover_history=e.failover_history) from None
             if dryrun:
-                return
+                return None
             cluster_config_file = config_dict['ray']
 
             handle = self.ResourceHandle(
@@ -2314,7 +2316,7 @@ class CloudVmRayBackend(backends.Backend):
             return handle
 
     def _update_after_cluster_provisioned(
-            self, handle: ResourceHandle, task: task_lib.Task,
+            self, handle: CloudVmRayResourceHandle, task: task_lib.Task,
             prev_cluster_status: global_user_state.ClusterStatus,
             ip_list: List[str], lock_path: str) -> None:
         usage_lib.messages.usage.update_cluster_resources(
@@ -2369,7 +2371,8 @@ class CloudVmRayBackend(backends.Backend):
 
             common_utils.remove_file_if_exists(lock_path)
 
-    def _sync_workdir(self, handle: ResourceHandle, workdir: Path) -> None:
+    def _sync_workdir(self, handle: CloudVmRayResourceHandle,
+                      workdir: Path) -> None:
         # Even though provision() takes care of it, there may be cases where
         # this function is called in isolation, without calling provision(),
         # e.g., in CLI.  So we should rerun rsync_up.
@@ -2431,7 +2434,7 @@ class CloudVmRayBackend(backends.Backend):
 
     def _sync_file_mounts(
         self,
-        handle: ResourceHandle,
+        handle: CloudVmRayResourceHandle,
         all_file_mounts: Dict[Path, Path],
         storage_mounts: Dict[Path, storage_lib.Storage],
     ) -> None:
@@ -2439,8 +2442,8 @@ class CloudVmRayBackend(backends.Backend):
         self._execute_file_mounts(handle, all_file_mounts)
         self._execute_storage_mounts(handle, storage_mounts)
 
-    def _setup(self, handle: ResourceHandle, task: task_lib.Task,
-               detach_setup: bool) -> Optional[str]:
+    def _setup(self, handle: CloudVmRayResourceHandle, task: task_lib.Task,
+               detach_setup: bool) -> None:
         start = time.time()
         style = colorama.Style
         fore = colorama.Fore
@@ -2470,7 +2473,7 @@ class CloudVmRayBackend(backends.Backend):
             # Need this `-i` option to make sure `source ~/.bashrc` work
             setup_cmd = f'/bin/bash -i /tmp/{setup_file} 2>&1'
 
-            def _setup_node(runner: command_runner.SSHCommandRunner) -> int:
+            def _setup_node(runner: command_runner.SSHCommandRunner) -> None:
                 runner.rsync(source=setup_sh_path,
                              target=f'/tmp/{setup_file}',
                              up=True,
@@ -2535,7 +2538,7 @@ class CloudVmRayBackend(backends.Backend):
 
     def _exec_code_on_head(
         self,
-        handle: ResourceHandle,
+        handle: CloudVmRayResourceHandle,
         codegen: str,
         job_id: int,
         executable: str,
@@ -2638,7 +2641,7 @@ class CloudVmRayBackend(backends.Backend):
 
     def _setup_and_create_job_cmd_on_local_head(
         self,
-        handle: ResourceHandle,
+        handle: CloudVmRayResourceHandle,
         ray_command: str,
         ray_job_id: str,
     ):
@@ -2674,7 +2677,7 @@ class CloudVmRayBackend(backends.Backend):
             f'--no-wait -- {switch_user_cmd}')
         return job_submit_cmd
 
-    def _add_job(self, handle: ResourceHandle, job_name: str,
+    def _add_job(self, handle: CloudVmRayResourceHandle, job_name: str,
                  resources_str: str) -> int:
         username = getpass.getuser()
         code = job_lib.JobLibCodeGen.add_job(job_name, username,
@@ -2703,7 +2706,7 @@ class CloudVmRayBackend(backends.Backend):
 
     def _execute(
         self,
-        handle: ResourceHandle,
+        handle: CloudVmRayResourceHandle,
         task: task_lib.Task,
         detach_run: bool,
     ) -> None:
@@ -2725,7 +2728,8 @@ class CloudVmRayBackend(backends.Backend):
             # Case: task_lib.Task(run, num_nodes=1)
             self._execute_task_one_node(handle, task, job_id, detach_run)
 
-    def _post_execute(self, handle: ResourceHandle, down: bool) -> None:
+    def _post_execute(self, handle: CloudVmRayResourceHandle,
+                      down: bool) -> None:
         colorama.init()
         fore = colorama.Fore
         style = colorama.Style
@@ -2760,7 +2764,7 @@ class CloudVmRayBackend(backends.Backend):
                     storage.delete()
 
     def _teardown(self,
-                  handle: ResourceHandle,
+                  handle: CloudVmRayResourceHandle,
                   terminate: bool,
                   purge: bool = False):
         cluster_name = handle.cluster_name
@@ -2789,10 +2793,10 @@ class CloudVmRayBackend(backends.Backend):
 
     def get_job_status(
         self,
-        handle: ResourceHandle,
-        job_ids: Optional[List[str]] = None,
+        handle: CloudVmRayResourceHandle,
+        job_ids: Optional[List[int]] = None,
         stream_logs: bool = True
-    ) -> Dict[Optional[str], Optional[job_lib.JobStatus]]:
+    ) -> Dict[Optional[int], Optional[job_lib.JobStatus]]:
         code = job_lib.JobLibCodeGen.get_job_status(job_ids)
         returncode, stdout, stderr = self.run_on_head(handle,
                                                       code,
@@ -2804,7 +2808,8 @@ class CloudVmRayBackend(backends.Backend):
         statuses = job_lib.load_statuses_payload(stdout)
         return statuses
 
-    def cancel_jobs(self, handle: ResourceHandle, jobs: Optional[List[int]]):
+    def cancel_jobs(self, handle: CloudVmRayResourceHandle,
+                    jobs: Optional[List[int]]):
         job_owner = onprem_utils.get_job_owner(handle.cluster_yaml)
         code = job_lib.JobLibCodeGen.cancel_jobs(job_owner, jobs)
 
@@ -2819,8 +2824,8 @@ class CloudVmRayBackend(backends.Backend):
 
     def sync_down_logs(
             self,
-            handle: ResourceHandle,
-            job_ids: Optional[str],
+            handle: CloudVmRayResourceHandle,
+            job_ids: Optional[List[str]],
             local_dir: str = constants.SKY_LOGS_DIRECTORY) -> Dict[str, str]:
         """Sync down logs for the given job_ids.
 
@@ -2841,7 +2846,7 @@ class CloudVmRayBackend(backends.Backend):
             logger.info(f'{colorama.Fore.YELLOW}'
                         'No matching log directories found'
                         f'{colorama.Style.RESET_ALL}')
-            return
+            return {}
 
         job_ids = list(run_timestamps.keys())
         run_timestamps = list(run_timestamps.values())
@@ -2897,7 +2902,7 @@ class CloudVmRayBackend(backends.Backend):
         return dict(zip(job_ids, local_log_dirs))
 
     def tail_logs(self,
-                  handle: ResourceHandle,
+                  handle: CloudVmRayResourceHandle,
                   job_id: Optional[int],
                   spot_job_id: Optional[int] = None,
                   follow: bool = True) -> int:
@@ -2933,7 +2938,7 @@ class CloudVmRayBackend(backends.Backend):
         return returncode
 
     def tail_spot_logs(self,
-                       handle: ResourceHandle,
+                       handle: CloudVmRayResourceHandle,
                        job_id: Optional[int] = None,
                        job_name: Optional[str] = None,
                        follow: bool = True) -> None:
@@ -2960,7 +2965,7 @@ class CloudVmRayBackend(backends.Backend):
         )
 
     def teardown_no_lock(self,
-                         handle: ResourceHandle,
+                         handle: CloudVmRayResourceHandle,
                          terminate: bool,
                          purge: bool = False,
                          post_teardown_cleanup: bool = True,
@@ -3120,7 +3125,7 @@ class CloudVmRayBackend(backends.Backend):
             return True
 
     def post_teardown_cleanup(self,
-                              handle: ResourceHandle,
+                              handle: CloudVmRayResourceHandle,
                               terminate: bool,
                               purge: bool = False) -> bool:
         """Cleanup local configs/caches and delete TPUs after teardown.
@@ -3186,7 +3191,7 @@ class CloudVmRayBackend(backends.Backend):
         return True
 
     def set_autostop(self,
-                     handle: ResourceHandle,
+                     handle: CloudVmRayResourceHandle,
                      idle_minutes_to_autostop: Optional[int],
                      down: bool = False,
                      stream_logs: bool = True) -> None:
@@ -3206,7 +3211,7 @@ class CloudVmRayBackend(backends.Backend):
                 handle.cluster_name, idle_minutes_to_autostop, down)
 
     def is_definitely_autostopping(self,
-                                   handle: ResourceHandle,
+                                   handle: CloudVmRayResourceHandle,
                                    stream_logs: bool = True) -> bool:
         """Check if the cluster is autostopping.
 
@@ -3235,10 +3240,10 @@ class CloudVmRayBackend(backends.Backend):
     @timeline.event
     def run_on_head(
         self,
-        handle: ResourceHandle,
+        handle: CloudVmRayResourceHandle,
         cmd: str,
         *,
-        port_forward: Optional[List[str]] = None,
+        port_forward: Optional[List[int]] = None,
         log_path: str = '/dev/null',
         process_stream: bool = True,
         stream_logs: bool = False,
@@ -3330,7 +3335,8 @@ class CloudVmRayBackend(backends.Backend):
                                                        task.num_nodes,
                                                        prev_cluster_status=None)
 
-    def _set_tpu_name(self, handle: ResourceHandle, tpu_name: str) -> None:
+    def _set_tpu_name(self, handle: CloudVmRayResourceHandle,
+                      tpu_name: str) -> None:
         """Sets TPU_NAME on all nodes."""
         ip_list = handle.external_ips()
         assert ip_list is not None, 'external_ips is not cached in handle'
@@ -3352,7 +3358,7 @@ class CloudVmRayBackend(backends.Backend):
 
         subprocess_utils.run_in_parallel(_setup_tpu_name_on_node, runners)
 
-    def _execute_file_mounts(self, handle: ResourceHandle,
+    def _execute_file_mounts(self, handle: CloudVmRayResourceHandle,
                              file_mounts: Dict[Path, Path]):
         """Executes file mounts - rsyncing local files and
         copying from remote stores."""
@@ -3481,7 +3487,7 @@ class CloudVmRayBackend(backends.Backend):
         end = time.time()
         logger.debug(f'File mount sync took {end - start} seconds.')
 
-    def _execute_storage_mounts(self, handle: ResourceHandle,
+    def _execute_storage_mounts(self, handle: CloudVmRayResourceHandle,
                                 storage_mounts: Dict[Path,
                                                      storage_lib.Storage]):
         """Executes storage mounts: installing mounting tools and mounting."""
@@ -3527,6 +3533,8 @@ class CloudVmRayBackend(backends.Backend):
             mount_cmd = store.mount_command(dst)
             src_print = (storage_obj.source
                          if storage_obj.source else storage_obj.name)
+            if isinstance(src_print, list):
+                src_print = ', '.join(src_print)
             backend_utils.parallel_data_transfer_to_nodes(
                 runners,
                 source=src_print,
@@ -3539,7 +3547,7 @@ class CloudVmRayBackend(backends.Backend):
         end = time.time()
         logger.debug(f'Storage mount sync took {end - start} seconds.')
 
-    def _execute_task_one_node(self, handle: ResourceHandle,
+    def _execute_task_one_node(self, handle: CloudVmRayResourceHandle,
                                task: task_lib.Task, job_id: int,
                                detach_run: bool) -> None:
         # Launch the command as a Ray task.
@@ -3571,7 +3579,7 @@ class CloudVmRayBackend(backends.Backend):
             constants.JOB_ID_ENV_VAR,
             common_utils.get_global_job_id(self.run_timestamp,
                                            cluster_name=handle.cluster_name,
-                                           job_id=job_id))
+                                           job_id=str(job_id)))
 
         command_for_node = task.run if isinstance(task.run, str) else None
         use_sudo = isinstance(handle.launched_resources.cloud, clouds.Local)
@@ -3592,8 +3600,9 @@ class CloudVmRayBackend(backends.Backend):
                                 executable='python3',
                                 detach_run=detach_run)
 
-    def _execute_task_n_nodes(self, handle: ResourceHandle, task: task_lib.Task,
-                              job_id: int, detach_run: bool) -> None:
+    def _execute_task_n_nodes(self, handle: CloudVmRayResourceHandle,
+                              task: task_lib.Task, job_id: int,
+                              detach_run: bool) -> None:
         # Strategy:
         #   ray.init(...)
         #   for node:
@@ -3611,6 +3620,7 @@ class CloudVmRayBackend(backends.Backend):
                 handle.launched_resources)
         else:
             num_actual_nodes = task.num_nodes
+        assert isinstance(num_actual_nodes, int), num_actual_nodes
 
         codegen = RayCodeGen()
         is_local = isinstance(handle.launched_resources.cloud, clouds.Local)
@@ -3636,7 +3646,7 @@ class CloudVmRayBackend(backends.Backend):
             constants.JOB_ID_ENV_VAR,
             common_utils.get_global_job_id(self.run_timestamp,
                                            cluster_name=handle.cluster_name,
-                                           job_id=job_id))
+                                           job_id=str(job_id)))
 
         # TODO(zhwu): The resources limitation for multi-node ray.tune and
         # horovod should be considered.
diff --git a/sky/backends/docker_utils.py b/sky/backends/docker_utils.py
index 3b5629f51..20cce2244 100644
--- a/sky/backends/docker_utils.py
+++ b/sky/backends/docker_utils.py
@@ -4,7 +4,7 @@ import shutil
 import subprocess
 import tempfile
 import textwrap
-from typing import Optional
+from typing import Dict, Optional, Tuple
 
 import colorama
 
@@ -40,8 +40,8 @@ def create_dockerfile(
     setup_command: str,
     copy_path: str,
     build_dir: str,
-    run_command: str = None,
-) -> str:
+    run_command: Optional[str] = None,
+) -> Tuple[str, Dict[str, str]]:
     """Writes a valid dockerfile to the specified path.
 
     performs three operations:
@@ -78,7 +78,8 @@ def create_dockerfile(
         dockerfile_contents += '\n' + DOCKERFILE_COPYCMD.format(
             copy_command=copy_docker_cmd)
 
-    def add_script_to_dockerfile(dockerfile_contents: str, multiline_cmds: str,
+    def add_script_to_dockerfile(dockerfile_contents: str,
+                                 multiline_cmds: Optional[str],
                                  out_filename: str):
         # Converts multiline commands to a script and adds the script to the
         # dockerfile. You still need to add the docker command to run the
@@ -144,7 +145,8 @@ def _execute_build(tag, context_path):
         raise
 
 
-def build_dockerimage(task, tag):
+def build_dockerimage(task: task_mod.Task,
+                      tag: str) -> Tuple[str, Dict[str, str]]:
     """
     Builds a docker image for the given task.
 
@@ -182,7 +184,8 @@ def build_dockerimage(task, tag):
     return tag, img_metadata
 
 
-def build_dockerimage_from_task(task: task_mod.Task):
+def build_dockerimage_from_task(
+        task: task_mod.Task) -> Tuple[str, Dict[str, str]]:
     """ Builds a docker image from a Task"""
     tag, img_metadata = build_dockerimage(task, tag=task.name)
     return tag, img_metadata
diff --git a/sky/backends/local_docker_backend.py b/sky/backends/local_docker_backend.py
index ced2cb41a..27f7110e7 100644
--- a/sky/backends/local_docker_backend.py
+++ b/sky/backends/local_docker_backend.py
@@ -2,7 +2,7 @@
 import subprocess
 import tempfile
 import typing
-from typing import Dict, Optional, Union
+from typing import Any, Dict, Optional, Tuple, Union
 
 import colorama
 from rich import console as rich_console
@@ -30,7 +30,21 @@ _DOCKER_LABEL_PREFIX = 'skymeta_'
 _DOCKER_HANDLE_PREFIX = 'skydocker-'
 
 
-class LocalDockerBackend(backends.Backend):
+class LocalDockerResourceHandle(str, backends.ResourceHandle):
+    """The name of the cluster/container prefixed with the handle prefix."""
+
+    def __new__(cls, s, **kw):
+        if s.startswith(_DOCKER_HANDLE_PREFIX):
+            prefixed_str = s
+        else:
+            prefixed_str = _DOCKER_HANDLE_PREFIX + s
+        return str.__new__(cls, prefixed_str, **kw)
+
+    def get_cluster_name(self):
+        return self.lstrip(_DOCKER_HANDLE_PREFIX)
+
+
+class LocalDockerBackend(backends.Backend['LocalDockerResourceHandle']):
     """Local docker backend for debugging.
 
     Ignores resource demands when allocating. Optionally uses GPU if required.
@@ -78,18 +92,8 @@ class LocalDockerBackend(backends.Backend):
 
     NAME = 'localdocker'
 
-    class ResourceHandle(str):
-        """The name of the cluster/container prefixed with the handle prefix."""
-
-        def __new__(cls, s, **kw):
-            if s.startswith(_DOCKER_HANDLE_PREFIX):
-                prefixed_str = s
-            else:
-                prefixed_str = _DOCKER_HANDLE_PREFIX + s
-            return str.__new__(cls, prefixed_str, **kw)
-
-        def get_cluster_name(self):
-            return self.lstrip(_DOCKER_HANDLE_PREFIX)
+    # Backward compatibility, with the old name of the handle.
+    ResourceHandle = LocalDockerResourceHandle  # pylint: disable=invalid-name
 
     # Define the Docker-in-Docker mount
     _dind_mount = {
@@ -110,33 +114,41 @@ class LocalDockerBackend(backends.Backend):
         """
         self._use_gpu = backend_utils.check_local_gpus() if use_gpu == 'auto' \
             else use_gpu
-        self.volume_mounts = {}  # Stores the ResourceHandle->volume mounts map
-        self.images = {}  # Stores the ResourceHandle->[image_tag, metadata] map
-        self.containers = {}
+        self.volume_mounts: Dict[LocalDockerResourceHandle, Dict[str, Any]] = {
+        }  # Stores the LocalDockerResourceHandle->volume mounts map
+        self.images: Dict[LocalDockerResourceHandle, Tuple[str, Dict[
+            str, str]]] = {
+            }  # Stores the LocalDockerResourceHandle->[image_tag, metadata] map
+        self.containers: Dict[LocalDockerResourceHandle, Any] = {}
         self.client = docker.from_env()
         self._update_state()
 
     # --- Implementation of Backend APIs ---
 
-    def check_resources_fit_cluster(self, handle: ResourceHandle,
+    def check_resources_fit_cluster(self, handle: 'LocalDockerResourceHandle',
                                     task: 'task_lib.Task') -> None:
         pass
 
-    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[LocalDockerResourceHandle]:
         """
         Builds docker image for the task and returns the cluster name as handle.
 
         Since resource demands are ignored, There's no provisioning in local
         docker.
         """
+        del to_provision  # Unused
         assert task.name is not None, ('Task name cannot be None - have you '
                                        'specified a task name?')
+        if dryrun:
+            return None
         if retry_until_up:
             logger.warning(
                 f'Retrying until up is not supported in backend: {self.NAME}. '
@@ -145,12 +157,12 @@ class LocalDockerBackend(backends.Backend):
             logger.info(
                 'Streaming build logs is not supported in LocalDockerBackend. '
                 'Build logs will be shown on failure.')
-        handle = LocalDockerBackend.ResourceHandle(cluster_name)
+        handle = LocalDockerResourceHandle(cluster_name)
         logger.info(f'Building docker image for task {task.name}. '
                     'This might take some time.')
         with console.status('[bold cyan]Building Docker image[/]'):
             image_tag, metadata = docker_utils.build_dockerimage_from_task(task)
-        self.images[handle] = [image_tag, metadata]
+        self.images[handle] = (image_tag, metadata)
         logger.info(f'Image {image_tag} built.')
         logger.info('Provisioning complete.')
         global_user_state.add_or_update_cluster(
@@ -160,18 +172,20 @@ class LocalDockerBackend(backends.Backend):
             ready=False)
         return handle
 
-    def _sync_workdir(self, handle: ResourceHandle, workdir: Path) -> None:
+    def _sync_workdir(self, handle: LocalDockerResourceHandle,
+                      workdir: Path) -> None:
         """Workdir is sync'd by adding to the docker image.
 
         This happens in the execute step.
         """
+        del handle, workdir  # Unused
         logger.info('Since the workdir is synced at build time, sync_workdir is'
                     ' a NoOp. If you are running sky exec, your workdir has not'
                     ' been updated.')
 
     def _sync_file_mounts(
         self,
-        handle: ResourceHandle,
+        handle: LocalDockerResourceHandle,
         all_file_mounts: Dict[Path, Path],
         storage_mounts: Dict[Path, storage_lib.Storage],
     ) -> None:
@@ -192,15 +206,15 @@ class LocalDockerBackend(backends.Backend):
                 }
         self.volume_mounts[handle] = docker_mounts
 
-    def _setup(self, handle: ResourceHandle, task: 'task_lib.Task') -> None:
-        """
-        Launches a container and runs a sleep command on it.
+    def _setup(self, handle: LocalDockerResourceHandle, task: 'task_lib.Task',
+               detach_setup: bool) -> None:
+        """Launches a container and runs a sleep command on it.
 
         setup() in LocalDockerBackend runs the container with a sleep job
         so that the container is kept alive and we can issue docker exec cmds
         to it to handle sky exec commands.
         """
-        del task  # unused
+        del task, detach_setup  # unused
         colorama.init()
         style = colorama.Style
         assert handle in self.images, \
@@ -252,7 +266,7 @@ class LocalDockerBackend(backends.Backend):
             requested_resources=task.resources,
             ready=True)
 
-    def _execute(self, handle: ResourceHandle, task: 'task_lib.Task',
+    def _execute(self, handle: LocalDockerResourceHandle, task: 'task_lib.Task',
                  detach_run: bool) -> None:
         """ Launches the container."""
 
@@ -272,7 +286,8 @@ class LocalDockerBackend(backends.Backend):
 
         self._execute_task_one_node(handle, task)
 
-    def _post_execute(self, handle: ResourceHandle, down: bool) -> None:
+    def _post_execute(self, handle: LocalDockerResourceHandle,
+                      down: bool) -> None:
         del down  # unused
         colorama.init()
         style = colorama.Style
@@ -296,7 +311,7 @@ class LocalDockerBackend(backends.Backend):
             f'{container.image.tags[0]} /bin/bash{style.RESET_ALL}')
 
     def _teardown(self,
-                  handle: ResourceHandle,
+                  handle: LocalDockerResourceHandle,
                   terminate: bool,
                   purge: bool = False):
         """Teardown kills the container."""
@@ -343,7 +358,7 @@ class LocalDockerBackend(backends.Backend):
             self.images[c.name] = [c.image, metadata]
             self.containers[c.name] = c
 
-    def _execute_task_one_node(self, handle: ResourceHandle,
+    def _execute_task_one_node(self, handle: LocalDockerResourceHandle,
                                task: 'task_lib.Task') -> None:
         container = self.containers[handle]
         _, image_metadata = self.images[handle]
diff --git a/sky/backends/onprem_utils.py b/sky/backends/onprem_utils.py
index 3d21bc720..56d2891f2 100644
--- a/sky/backends/onprem_utils.py
+++ b/sky/backends/onprem_utils.py
@@ -16,6 +16,7 @@ from sky import global_user_state
 from sky import sky_logging
 from sky.backends import backend_utils
 from sky.skylet import constants
+from sky.skylet import log_lib
 from sky.utils import command_runner
 from sky.utils import common_utils
 from sky.utils import schemas
@@ -67,7 +68,7 @@ def check_and_get_local_clusters(suppress_error: bool = False) -> List[str]:
     ]
 
     local_cluster_names = []
-    name_to_path_dict = {}
+    name_to_path_dict: Dict[str, str] = {}
 
     for path in local_cluster_paths:
         with open(path, 'r') as f:
@@ -127,7 +128,7 @@ def get_local_ips(cluster_name: str) -> List[str]:
     return ips
 
 
-def get_local_auth_config(cluster_name: str) -> List[str]:
+def get_local_auth_config(cluster_name: str) -> Dict[str, str]:
     """Returns IP addresses of the local cluster."""
     config = get_local_cluster_config_or_error(cluster_name)
     return config['auth']
@@ -139,7 +140,7 @@ def get_python_executable(cluster_name: str) -> str:
     return config['python']
 
 
-def get_job_owner(cluster_yaml: dict) -> str:
+def get_job_owner(cluster_yaml: str) -> str:
     """Get the owner of the job."""
     cluster_config = common_utils.read_yaml(os.path.expanduser(cluster_yaml))
     # User name is guaranteed to exist (on all jinja files)
@@ -318,9 +319,8 @@ def get_local_cluster_accelerators(
     return custom_resources
 
 
-def launch_ray_on_local_cluster(
-        cluster_config: Dict[str, Dict[str, Any]],
-        custom_resources: List[Dict[str, int]] = None) -> None:
+def launch_ray_on_local_cluster(cluster_config: Dict[str, Dict[str, Any]],
+                                custom_resources: List[Dict[str, int]]) -> None:
     """Launches Ray on all nodes for local cluster.
 
     Launches Ray on the root user of all nodes and opens the Ray dashboard port
@@ -388,8 +388,9 @@ def launch_ray_on_local_cluster(
     # to worker node.
     remote_ssh_key = f'~/.ssh/{os.path.basename(ssh_key)}'
     dashboard_remote_path = '~/.sky/dashboard_portforward.sh'
-    worker_runners = [(runner, idx) for idx, runner in enumerate(worker_runners)
-                     ]
+    worker_runner_idxs = [
+        (runner, idx) for idx, runner in enumerate(worker_runners)
+    ]
     # Connect head node's Ray dashboard to worker nodes
     # Worker nodes need access to Ray dashboard to poll the
     # JobSubmissionClient (in subprocess_daemon.py) for completed,
@@ -443,10 +444,10 @@ def launch_ray_on_local_cluster(
                 failure_message=
                 f'Failed to connect ray dashboard to worker node {runner.ip}.')
 
-        subprocess_utils.run_in_parallel(_start_ray_workers, worker_runners)
+        subprocess_utils.run_in_parallel(_start_ray_workers, worker_runner_idxs)
 
 
-def save_distributable_yaml(cluster_config: Dict[str, Dict[str, Any]]) -> None:
+def save_distributable_yaml(cluster_config: Dict[str, Any]) -> None:
     """Generates a distributable yaml for the system admin to send to users.
 
     Args:
@@ -511,3 +512,63 @@ def check_local_cloud_args(cloud: Optional[str] = None,
                     'See `sky status` for local cluster name(s).')
 
         return False
+
+
+def do_filemounts_and_setup_on_local_workers(
+        cluster_config_file: str,
+        worker_ips: Optional[List[str]] = None,
+        extra_setup_cmds: Optional[List[str]] = None):
+    """Completes filemounting and setup on worker nodes.
+
+    Syncs filemounts and runs setup on worker nodes for a local cluster. This
+    is a workaround for a Ray Autoscaler bug where `ray up` does not perform
+    filemounting or setup for local cluster worker nodes.
+    """
+    config = common_utils.read_yaml(cluster_config_file)
+
+    ssh_credentials = backend_utils.ssh_credential_from_yaml(
+        cluster_config_file)
+    if worker_ips is None:
+        worker_ips = config['provider']['worker_ips']
+    file_mounts = config['file_mounts']
+
+    setup_cmds = config['setup_commands']
+    if extra_setup_cmds is not None:
+        setup_cmds += extra_setup_cmds
+    setup_script = log_lib.make_task_bash_script('\n'.join(setup_cmds))
+
+    worker_runners = command_runner.SSHCommandRunner.make_runner_list(
+        worker_ips, **ssh_credentials)
+
+    # Uploads setup script to the worker node
+    with tempfile.NamedTemporaryFile('w', prefix='sky_setup_') as f:
+        f.write(setup_script)
+        f.flush()
+        setup_sh_path = f.name
+        setup_file = os.path.basename(setup_sh_path)
+        file_mounts[f'/tmp/{setup_file}'] = setup_sh_path
+
+        # Ray Autoscaler Bug: Filemounting + Ray Setup
+        # does not happen on workers.
+        def _setup_local_worker(runner: command_runner.SSHCommandRunner):
+            for dst, src in file_mounts.items():
+                mkdir_dst = f'mkdir -p {os.path.dirname(dst)}'
+                backend_utils.run_command_and_handle_ssh_failure(
+                    runner,
+                    mkdir_dst,
+                    failure_message=f'Failed to run {mkdir_dst} on remote.')
+                if os.path.isdir(src):
+                    src = os.path.join(src, '')
+                runner.rsync(source=src, target=dst, up=True, stream_logs=False)
+
+            setup_cmd = f'/bin/bash -i /tmp/{setup_file} 2>&1'
+            rc, stdout, _ = runner.run(setup_cmd,
+                                       stream_logs=False,
+                                       require_outputs=True)
+            subprocess_utils.handle_returncode(
+                rc,
+                setup_cmd,
+                'Failed to setup Ray autoscaler commands on remote.',
+                stderr=stdout)
+
+        subprocess_utils.run_in_parallel(_setup_local_worker, worker_runners)
diff --git a/sky/benchmark/benchmark_state.py b/sky/benchmark/benchmark_state.py
index a59f2187e..1b6f4a0e6 100644
--- a/sky/benchmark/benchmark_state.py
+++ b/sky/benchmark/benchmark_state.py
@@ -133,9 +133,8 @@ def add_benchmark(benchmark_name: str, task_name: Optional[str],
     _BENCHMARK_DB.conn.commit()
 
 
-def add_benchmark_result(
-        benchmark_name: str,
-        cluster_handle: 'backend_lib.Backend.ResourceHandle') -> None:
+def add_benchmark_result(benchmark_name: str,
+                         cluster_handle: 'backend_lib.ResourceHandle') -> None:
     name = cluster_handle.cluster_name
     num_nodes = cluster_handle.launched_nodes
     resources = pickle.dumps(cluster_handle.launched_resources)
diff --git a/sky/cli.py b/sky/cli.py
index c931d059b..5f6faf9d0 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -36,7 +36,7 @@ import subprocess
 import sys
 import textwrap
 import typing
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Sequence, Tuple
 
 import click
 import colorama
@@ -99,7 +99,7 @@ _INTERACTIVE_NODE_DEFAULT_RESOURCES = {
 }
 
 
-def _get_glob_clusters(clusters: List[str]) -> List[str]:
+def _get_glob_clusters(clusters: Sequence[str]) -> List[str]:
     """Returns a list of clusters that match the glob pattern."""
     glob_clusters = []
     for cluster in clusters:
@@ -115,7 +115,7 @@ def _get_glob_clusters(clusters: List[str]) -> List[str]:
     return list(set(glob_clusters))
 
 
-def _get_glob_storages(storages: List[str]) -> List[str]:
+def _get_glob_storages(storages: Sequence[str]) -> List[str]:
     """Returns a list of storages that match the glob pattern."""
     glob_storages = []
     for storage_object in storages:
@@ -300,7 +300,12 @@ def _parse_env_var(env_var: str) -> Tuple[str, str]:
             raise click.UsageError(
                 f'{env_var} is not set in local environment.')
         return (env_var, value)
-    return tuple(env_var.split('=', 1))
+    ret = tuple(env_var.split('=', 1))
+    if len(ret) != 2:
+        raise click.UsageError(
+            f'Invalid env var: {env_var}. Must be in the form of KEY=VAL '
+            'or KEY.')
+    return ret[0], ret[1]
 
 
 _TASK_OPTIONS = [
@@ -570,7 +575,7 @@ def _parse_override_params(cloud: Optional[str] = None,
                            image_id: Optional[str] = None,
                            disk_size: Optional[int] = None) -> Dict[str, Any]:
     """Parses the override parameters into a dictionary."""
-    override_params = {}
+    override_params: Dict[str, Any] = {}
     if cloud is not None:
         if cloud.lower() == 'none':
             override_params['cloud'] = None
@@ -674,7 +679,7 @@ def _check_resources_match(backend: backends.Backend,
 def _launch_with_confirm(
     task: sky.Task,
     backend: backends.Backend,
-    cluster: str,
+    cluster: Optional[str],
     *,
     dryrun: bool,
     detach_run: bool,
@@ -685,7 +690,6 @@ def _launch_with_confirm(
     retry_until_up: bool = False,
     no_setup: bool = False,
     node_type: Optional[str] = None,
-    is_local_cloud: Optional[bool] = False,
 ):
     """Launch a cluster with a Task."""
     with sky.Dag() as dag:
@@ -714,7 +718,7 @@ def _launch_with_confirm(
         prompt = None
         if maybe_status is None:
             cluster_str = '' if cluster is None else f' {cluster!r}'
-            if is_local_cloud:
+            if onprem_utils.check_if_local_cloud(cluster):
                 prompt = f'Initializing local cluster{cluster_str}. Proceed?'
             else:
                 prompt = f'Launching a new cluster{cluster_str}. Proceed?'
@@ -811,6 +815,10 @@ def _create_and_ssh_into_node(
     task.set_resources(resources)
 
     backend = backend if backend is not None else backends.CloudVmRayBackend()
+    if not isinstance(backend, backends.CloudVmRayBackend):
+        raise click.UsageError('Interactive nodes are only supported for '
+                               f'{backends.CloudVmRayBackend.__name__} '
+                               f'backend. Got {type(backend).__name__}.')
     maybe_status, _ = backend_utils.refresh_cluster_status_handle(cluster_name)
     if maybe_status is not None and user_requested_resources:
         name_arg = ''
@@ -866,7 +874,7 @@ def _create_and_ssh_into_node(
     click.secho(f'rsync -rP {cluster_name}:/remote/path /local/path', bold=True)
 
 
-def _check_yaml(entrypoint: str) -> Tuple[bool, dict]:
+def _check_yaml(entrypoint: str) -> Tuple[bool, Optional[Dict[str, Any]]]:
     """Checks if entrypoint is a readable YAML file.
 
     Args:
@@ -928,12 +936,13 @@ def _make_task_from_entrypoint_with_overrides(
     use_spot: Optional[bool] = None,
     image_id: Optional[str] = None,
     disk_size: Optional[int] = None,
-    env: List[Dict[str, str]] = None,
+    env: Optional[List[Tuple[str, str]]] = None,
     # spot launch specific
     spot_recovery: Optional[str] = None,
 ) -> sky.Task:
     entrypoint = ' '.join(entrypoint)
     is_yaml, yaml_config = _check_yaml(entrypoint)
+    entrypoint: Optional[str]
     if is_yaml:
         # Treat entrypoint as a yaml.
         click.secho('Task from YAML spec: ', fg='yellow', nl=False)
@@ -1165,7 +1174,7 @@ def cli():
               help='Skip setup phase when (re-)launching cluster.')
 @usage_lib.entrypoint
 def launch(
-    entrypoint: str,
+    entrypoint: List[str],
     cluster: Optional[str],
     dryrun: bool,
     detach_setup: bool,
@@ -1182,7 +1191,7 @@ def launch(
     num_nodes: Optional[int],
     use_spot: Optional[bool],
     image_id: Optional[str],
-    env: List[Dict[str, str]],
+    env: List[Tuple[str, str]],
     disk_size: Optional[int],
     idle_minutes_to_autostop: Optional[int],
     down: bool,  # pylint: disable=redefined-outer-name
@@ -1231,6 +1240,7 @@ def launch(
         disk_size=disk_size,
     )
 
+    backend: backends.Backend
     if backend_name == backends.LocalDockerBackend.NAME:
         backend = backends.LocalDockerBackend()
     elif backend_name == backends.CloudVmRayBackend.NAME:
@@ -1239,19 +1249,17 @@ def launch(
         with ux_utils.print_exception_no_traceback():
             raise ValueError(f'{backend_name} backend is not supported.')
 
-    _launch_with_confirm(
-        task,
-        backend,
-        cluster,
-        dryrun=dryrun,
-        detach_setup=detach_setup,
-        detach_run=detach_run,
-        no_confirm=yes,
-        idle_minutes_to_autostop=idle_minutes_to_autostop,
-        down=down,
-        retry_until_up=retry_until_up,
-        no_setup=no_setup,
-        is_local_cloud=onprem_utils.check_if_local_cloud(cluster))
+    _launch_with_confirm(task,
+                         backend,
+                         cluster,
+                         dryrun=dryrun,
+                         detach_setup=detach_setup,
+                         detach_run=detach_run,
+                         no_confirm=yes,
+                         idle_minutes_to_autostop=idle_minutes_to_autostop,
+                         down=down,
+                         retry_until_up=retry_until_up,
+                         no_setup=no_setup)
 
 
 @cli.command(cls=_DocumentedCodeCommand)
@@ -1276,7 +1284,7 @@ def launch(
 # pylint: disable=redefined-builtin
 def exec(
     cluster: str,
-    entrypoint: str,
+    entrypoint: List[str],
     detach_run: bool,
     name: Optional[str],
     cloud: Optional[str],
@@ -1288,7 +1296,7 @@ def exec(
     num_nodes: Optional[int],
     use_spot: Optional[bool],
     image_id: Optional[str],
-    env: List[Dict[str, str]],
+    env: List[Tuple[str, str]],
 ):
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
     """Execute a task or a command on a cluster (skip setup).
@@ -1354,7 +1362,8 @@ def exec(
     if handle is None:
         if onprem_utils.check_if_local_cloud(cluster):
             raise click.BadParameter(
-                constants.UNINITIALIZED_CLUSTER_MESSAGE.format(cluster=cluster))
+                constants.UNINITIALIZED_ONPREM_CLUSTER_MESSAGE.format(
+                    cluster=cluster))
         raise click.BadParameter(f'Cluster {cluster!r} not found. '
                                  'Use `sky launch` to provision first.')
     backend = backend_utils.get_backend_from_handle(handle)
@@ -1450,13 +1459,12 @@ def status(all: bool, refresh: bool, clusters: List[str]):  # pylint: disable=re
       or for autostop-enabled clusters, use ``--refresh`` to query the latest
       cluster statuses from the cloud providers.
     """
+    query_clusters: Optional[List[str]] = None
     if clusters:
-        clusters = _get_glob_clusters(clusters)
-    else:
-        clusters = None
-    cluster_records = core.status(cluster_names=clusters, refresh=refresh)
+        query_clusters = _get_glob_clusters(clusters)
+    cluster_records = core.status(cluster_names=query_clusters, refresh=refresh)
     nonreserved_cluster_records = []
-    reserved_clusters = dict()
+    reserved_clusters = {}
     for cluster_record in cluster_records:
         cluster_name = cluster_record['name']
         if cluster_name in backend_utils.SKY_RESERVED_CLUSTER_NAMES:
@@ -1549,7 +1557,7 @@ def cost_report(all: bool):  # pylint: disable=redefined-builtin
                 nargs=-1,
                 **_get_shell_complete_args(_complete_cluster_name))
 @usage_lib.entrypoint
-def queue(clusters: Tuple[str], skip_finished: bool, all_users: bool):
+def queue(clusters: Sequence[str], skip_finished: bool, all_users: bool):
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
     """Show the job queue for cluster(s)."""
     click.secho('Fetching and parsing job queue...', fg='yellow')
@@ -1968,7 +1976,7 @@ def autostop(
 @usage_lib.entrypoint
 # pylint: disable=redefined-builtin
 def start(
-        clusters: Tuple[str],
+        clusters: Sequence[str],
         all: bool,
         yes: bool,
         idle_minutes_to_autostop: Optional[int],
@@ -2260,7 +2268,7 @@ def _hint_or_raise_for_down_spot_controller(controller_name: str):
 
 
 def _down_or_stop_clusters(
-        names: Tuple[str],
+        names: Sequence[str],
         apply_to_all: Optional[bool],
         down: bool,  # pylint: disable=redefined-outer-name
         no_confirm: bool,
@@ -2899,7 +2907,7 @@ def storage_ls():
               required=False,
               help='Delete all storage objects.')
 @usage_lib.entrypoint
-def storage_delete(names: Tuple[str], all: bool):  # pylint: disable=redefined-builtin
+def storage_delete(names: Sequence[str], all: bool):  # pylint: disable=redefined-builtin
     """Delete storage objects.
 
     Examples:
@@ -2961,6 +2969,7 @@ def admin_deploy(clusterspec_yaml: str):
                                   'Invalid cluster YAML: ')
     if not is_yaml:
         raise ValueError('Must specify cluster config')
+    assert yaml_config is not None, (is_yaml, yaml_config)
 
     auth_config = yaml_config['auth']
     ips = yaml_config['cluster']['ips']
@@ -3059,7 +3068,7 @@ def spot():
 @timeline.event
 @usage_lib.entrypoint
 def spot_launch(
-    entrypoint: str,
+    entrypoint: List[str],
     name: Optional[str],
     workdir: Optional[str],
     cloud: Optional[str],
@@ -3072,7 +3081,7 @@ def spot_launch(
     use_spot: Optional[bool],
     image_id: Optional[str],
     spot_recovery: Optional[str],
-    env: List[Dict[str, str]],
+    env: List[Tuple[str, str]],
     disk_size: Optional[int],
     detach_run: bool,
     retry_until_up: bool,
@@ -3432,7 +3441,7 @@ def benchmark_launch(
     num_nodes: Optional[int],
     use_spot: Optional[bool],
     image_id: Optional[str],
-    env: List[Dict[str, str]],
+    env: List[Tuple[str, str]],
     disk_size: Optional[int],
     idle_minutes_to_autostop: Optional[int],
     yes: bool,
@@ -3459,6 +3468,7 @@ def benchmark_launch(
         raise click.BadParameter(
             'Sky Benchmark does not support command line tasks. '
             'Please provide a YAML file.')
+    assert config is not None, (is_yaml, config)
 
     click.secho('Benchmarking a task from YAML spec: ', fg='yellow', nl=False)
     click.secho(entrypoint, bold=True)
@@ -3494,19 +3504,19 @@ def benchmark_launch(
     # 1. By specifying resources.candidates in the YAML.
     # 2. By specifying gpu types as a command line argument (--gpus).
     if gpus is not None:
-        gpus = gpus.split(',')
-        gpus = [gpu.strip() for gpu in gpus]
+        gpu_list = gpus.split(',')
+        gpu_list = [gpu.strip() for gpu in gpu_list]
         if '' in gpus:
             raise click.BadParameter('Remove blanks in --gpus.')
 
-        if len(gpus) == 1:
-            gpus = gpus[0]
+        if len(gpu_list) == 1:
+            override_gpu = gpu_list[0]
         else:
             # If len(gpus) > 1, gpus is intrepreted
             # as a list of benchmark candidates.
             if candidates is None:
-                candidates = [{'accelerators': gpu} for gpu in gpus]
-                gpus = None
+                candidates = [{'accelerators': gpu} for gpu in gpu_list]
+                override_gpu = None
             else:
                 raise ValueError('Provide benchmark candidates in either '
                                  '--gpus or resources.candidates in the YAML.')
@@ -3527,7 +3537,7 @@ def benchmark_launch(
     override_params = _parse_override_params(cloud=cloud,
                                              region=region,
                                              zone=zone,
-                                             gpus=gpus,
+                                             gpus=override_gpu,
                                              use_spot=use_spot,
                                              image_id=image_id,
                                              disk_size=disk_size)
@@ -3562,14 +3572,14 @@ def benchmark_launch(
         click.confirm(prompt, default=True, abort=True, show_default=True)
 
     # Configs that are only accepted by the CLI.
-    commandline_args = {}
+    commandline_args: Dict[str, Any] = {}
     # Set the default idle minutes to autostop as 5, mimicking
     # the serverless execution.
     if idle_minutes_to_autostop is None:
         idle_minutes_to_autostop = 5
     commandline_args['idle-minutes-to-autostop'] = idle_minutes_to_autostop
     if len(env) > 0:
-        commandline_args['env'] = [f'{k}={v}' for k, v in env.items()]
+        commandline_args['env'] = [f'{k}={v}' for k, v in env]
 
     # Launch the benchmarking clusters in detach mode in parallel.
     benchmark_created = benchmark_utils.launch_benchmark_clusters(
@@ -3811,7 +3821,7 @@ def benchmark_down(
         raise click.BadParameter(f'Benchmark {benchmark} does not exist.')
 
     clusters = benchmark_state.get_benchmark_clusters(benchmark)
-    to_stop = []
+    to_stop: List[str] = []
     for cluster in clusters:
         if cluster in clusters_to_exclude:
             continue
diff --git a/sky/clouds/service_catalog/data_fetchers/fetch_azure.py b/sky/clouds/service_catalog/data_fetchers/fetch_azure.py
index 1b4b6a700..5d79d2c9c 100644
--- a/sky/clouds/service_catalog/data_fetchers/fetch_azure.py
+++ b/sky/clouds/service_catalog/data_fetchers/fetch_azure.py
@@ -231,7 +231,7 @@ def get_all_regions_instance_types_df(region_set: Set[str]):
     before_drop_len = len(df_ret)
     df_ret.dropna(subset=['InstanceType'], inplace=True, how='all')
     after_drop_len = len(df_ret)
-    print('Dropped {} duplicated rows'.format(before_drop_len - after_drop_len))
+    print(f'Dropped {before_drop_len - after_drop_len} duplicated rows')
 
     # Filter out deprecated families
     df_ret = df_ret.loc[~df_ret['family'].isin(DEPRECATED_FAMILIES)]
diff --git a/sky/clouds/service_catalog/data_fetchers/fetch_gcp.py b/sky/clouds/service_catalog/data_fetchers/fetch_gcp.py
index aac5f4a0a..bfe6dbfbb 100644
--- a/sky/clouds/service_catalog/data_fetchers/fetch_gcp.py
+++ b/sky/clouds/service_catalog/data_fetchers/fetch_gcp.py
@@ -87,7 +87,7 @@ def get_regions(doc: 'html.HtmlElement') -> Dict[str, str]:
     # Get the dictionary of regions.
     # E.g., 'kr': 'asia-northeast3'
     region_info = doc.xpath('//md-option')
-    regions = dict()
+    regions = {}
     for region in region_info:
         region_fullname = re.search(r'\((.*?)\)', region.text)
         if region_fullname is None:
diff --git a/sky/core.py b/sky/core.py
index 529469f4d..366c13860 100644
--- a/sky/core.py
+++ b/sky/core.py
@@ -1,8 +1,9 @@
 """SDK functions for cluster/job management."""
-import colorama
 import getpass
 import sys
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+from typing import Any, Dict, List, Optional, Sequence, Union
+
+import colorama
 
 from sky import clouds
 from sky import dag
@@ -162,7 +163,7 @@ def _start(
     retry_until_up: bool = False,
     down: bool = False,  # pylint: disable=redefined-outer-name
     force: bool = False,
-) -> backends.Backend.ResourceHandle:
+) -> backends.CloudVmRayResourceHandle:
 
     cluster_status, handle = backend_utils.refresh_cluster_status_handle(
         cluster_name)
@@ -170,7 +171,7 @@ def _start(
         raise ValueError(f'Cluster {cluster_name!r} does not exist.')
     if not force and cluster_status == global_user_state.ClusterStatus.UP:
         print(f'Cluster {cluster_name!r} is already up.')
-        return
+        return handle
     assert force or cluster_status in (
         global_user_state.ClusterStatus.INIT,
         global_user_state.ClusterStatus.STOPPED), cluster_status
@@ -425,6 +426,11 @@ def autostop(
         cluster_name,
         operation=operation,
     )
+    backend = backend_utils.get_backend_from_handle(handle)
+    if not isinstance(backend, backends.CloudVmRayBackend):
+        raise exceptions.NotSupportedError(
+            f'{operation} cluster {cluster_name!r} with backend '
+            f'{backend.__class__.__name__!r} is not supported.')
 
     if tpu_utils.is_tpu_vm_pod(handle.launched_resources):
         # Reference:
@@ -485,7 +491,7 @@ def queue(cluster_name: str,
         RuntimeError: if failed to get the job queue with ssh.
     """
     all_jobs = not skip_finished
-    username = getpass.getuser()
+    username: Optional[str] = getpass.getuser()
     if all_users:
         username = None
     code = job_lib.JobLibCodeGen.get_job_queue(username, all_jobs)
@@ -528,8 +534,7 @@ def cancel(cluster_name: str,
         sky.exceptions.CloudUserIdentityError: if we fail to get the current
           user identity.
     """
-    job_ids = [] if job_ids is None else job_ids
-    if len(job_ids) == 0 and not all:
+    if not job_ids and not all:
         raise ValueError(
             'sky cancel requires either a job id '
             f'(see `sky queue {cluster_name} -s`) or the --all flag.')
@@ -550,6 +555,7 @@ def cancel(cluster_name: str,
               f'{colorama.Style.RESET_ALL}')
         job_ids = None
     else:
+        assert job_ids is not None, 'job_ids should not be None'
         jobs_str = ', '.join(map(str, job_ids))
         print(f'{colorama.Fore.YELLOW}'
               f'Cancelling jobs ({jobs_str}) on cluster {cluster_name!r}...'
@@ -560,7 +566,7 @@ def cancel(cluster_name: str,
 
 @usage_lib.entrypoint
 def tail_logs(cluster_name: str,
-              job_id: Optional[str],
+              job_id: Optional[int],
               follow: bool = True) -> None:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
     """Tail the logs of a job.
@@ -625,9 +631,10 @@ def download_logs(
         operation='downloading logs',
     )
     backend = backend_utils.get_backend_from_handle(handle)
+    assert isinstance(backend, backends.CloudVmRayBackend), backend
 
     if job_ids is not None and len(job_ids) == 0:
-        return []
+        return {}
 
     usage_lib.record_cluster_name_for_current_operation(cluster_name)
     print(f'{colorama.Fore.YELLOW}'
@@ -638,10 +645,10 @@ def download_logs(
 
 
 @usage_lib.entrypoint
-def job_status(
-        cluster_name: str,
-        job_ids: Optional[List[str]],
-        stream_logs: bool = False) -> Dict[str, Optional[job_lib.JobStatus]]:
+def job_status(cluster_name: str,
+               job_ids: Optional[List[int]],
+               stream_logs: bool = False
+              ) -> Dict[Optional[int], Optional[job_lib.JobStatus]]:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
     """Get the status of jobs.
 
@@ -649,8 +656,8 @@ def job_status(
         cluster_name: (str) name of the cluster.
         job_ids: (List[str]) job ids. If None, get the status of the last job.
     Returns:
-        Dict[str, Optional[job_lib.JobStatus]]: A mapping of job_id to job
-        statuses. The status will be None if the job does not exist.
+        Dict[Optional[int], Optional[job_lib.JobStatus]]: A mapping of job_id to
+        job statuses. The status will be None if the job does not exist.
         If job_ids is None and there is no job on the cluster, it will return
         {None: None}.
     Raises:
@@ -669,9 +676,14 @@ def job_status(
         operation='getting job status',
     )
     backend = backend_utils.get_backend_from_handle(handle)
+    if not isinstance(backend, backends.CloudVmRayBackend):
+        raise exceptions.NotSupportedError(
+            f'Getting job status is not supported for cluster {cluster_name!r} '
+            f'of type {backend.__class__.__name__!r}.')
+    assert isinstance(handle, backends.CloudVmRayResourceHandle), handle
 
     if job_ids is not None and len(job_ids) == 0:
-        return []
+        return {}
 
     print(f'{colorama.Fore.YELLOW}'
           'Getting job status...'
@@ -770,7 +782,7 @@ def spot_queue(refresh: bool,
 @usage_lib.entrypoint
 # pylint: disable=redefined-builtin
 def spot_cancel(name: Optional[str] = None,
-                job_ids: Optional[Tuple[int]] = None,
+                job_ids: Optional[Sequence[int]] = None,
                 all: bool = False) -> None:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
     """Cancel managed spot jobs.
@@ -804,6 +816,7 @@ def spot_cancel(name: Optional[str] = None,
     elif job_ids:
         code = spot.SpotCodeGen.cancel_jobs_by_id(job_ids)
     else:
+        assert name is not None, (job_ids, name, all)
         code = spot.SpotCodeGen.cancel_job_by_name(name)
     # The stderr is redirected to stdout
     returncode, stdout, _ = backend.run_on_head(handle,
@@ -850,6 +863,7 @@ def spot_tail_logs(name: Optional[str], job_id: Optional[int],
     if name is not None and job_id is not None:
         raise ValueError('Cannot specify both name and job_id.')
     backend = backend_utils.get_backend_from_handle(handle)
+    assert isinstance(backend, backends.CloudVmRayBackend), backend
     # Stream the realtime logs
     backend.tail_spot_logs(handle, job_id=job_id, job_name=name, follow=follow)
 
diff --git a/sky/dag.py b/sky/dag.py
index 473b3ed02..d5ef45730 100644
--- a/sky/dag.py
+++ b/sky/dag.py
@@ -1,6 +1,7 @@
 """DAGs: user applications to be run."""
 import pprint
 import threading
+from typing import List
 
 
 class Dag:
@@ -69,7 +70,7 @@ class Dag:
 class _DagContext(threading.local):
     """A thread-local stack of Dags."""
     _current_dag = None
-    _previous_dags = []
+    _previous_dags: List[Dag] = []
 
     def push_dag(self, dag):
         if self._current_dag is not None:
diff --git a/sky/data/storage.py b/sky/data/storage.py
index f16b73282..17e5713be 100644
--- a/sky/data/storage.py
+++ b/sky/data/storage.py
@@ -781,7 +781,7 @@ class Storage(object):
         return storage_obj
 
     def to_yaml_config(self) -> Dict[str, str]:
-        config = dict()
+        config = {}
 
         def add_if_not_none(key: str, value: Optional[Any]):
             if value is not None:
diff --git a/sky/exceptions.py b/sky/exceptions.py
index 18c2da5c1..fbbeb8079 100644
--- a/sky/exceptions.py
+++ b/sky/exceptions.py
@@ -27,7 +27,9 @@ class ResourcesUnavailableError(Exception):
         # Copy the list to avoid modifying from outside.
         self.failover_history: List[Exception] = list(failover_history)
 
-    def with_failover_history(self, failover_history: List[Exception]) -> None:
+    def with_failover_history(
+            self,
+            failover_history: List[Exception]) -> 'ResourcesUnavailableError':
         # Copy the list to avoid modifying from outside.
         self.failover_history = list(failover_history)
         return self
diff --git a/sky/execution.py b/sky/execution.py
index ec2e9c247..8de0c4c9c 100644
--- a/sky/execution.py
+++ b/sky/execution.py
@@ -17,7 +17,7 @@ import enum
 import getpass
 import tempfile
 import os
-from typing import Any, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 import colorama
 
@@ -44,8 +44,6 @@ from sky.utils import ux_utils
 
 logger = sky_logging.init_logger(__name__)
 
-OptimizeTarget = optimizer.OptimizeTarget
-
 # Message thrown when APIs sky.{exec,launch,spot_launch}() received a string
 # instead of a Dag.  CLI (cli.py) is implemented by us so should not trigger
 # this.
@@ -113,7 +111,7 @@ def _execute(
     handle: Any = None,
     backend: Optional[backends.Backend] = None,
     retry_until_up: bool = False,
-    optimize_target: OptimizeTarget = OptimizeTarget.COST,
+    optimize_target: optimizer.OptimizeTarget = optimizer.OptimizeTarget.COST,
     stages: Optional[List[Stage]] = None,
     cluster_name: Optional[str] = None,
     detach_setup: bool = False,
@@ -283,6 +281,7 @@ def _execute(
 
         if Stage.PRE_EXEC in stages:
             if idle_minutes_to_autostop is not None:
+                assert isinstance(backend, backends.CloudVmRayBackend)
                 backend.set_autostop(handle,
                                      idle_minutes_to_autostop,
                                      down=down)
@@ -327,7 +326,7 @@ def launch(
     down: bool = False,
     stream_logs: bool = True,
     backend: Optional[backends.Backend] = None,
-    optimize_target: OptimizeTarget = OptimizeTarget.COST,
+    optimize_target: optimizer.OptimizeTarget = optimizer.OptimizeTarget.COST,
     detach_setup: bool = False,
     detach_run: bool = False,
     no_setup: bool = False,
@@ -547,7 +546,7 @@ def spot_launch(
     assert len(task.resources) == 1, task
     resources = list(task.resources)[0]
 
-    change_default_value = dict()
+    change_default_value: Dict[str, Any] = {}
     if not resources.use_spot_specified:
         change_default_value['use_spot'] = True
     if resources.spot_recovery is None:
@@ -677,7 +676,7 @@ def _maybe_translate_local_file_mounts_and_sync_up(
             f'source paths to SkyPilot Storage...{colorama.Style.RESET_ALL}')
 
     # Step 1: Translate the workdir to SkyPilot storage.
-    new_storage_mounts = dict()
+    new_storage_mounts = {}
     if task.workdir is not None:
         bucket_name = spot.constants.SPOT_WORKDIR_BUCKET_NAME.format(
             username=getpass.getuser(), id=run_id)
@@ -706,7 +705,7 @@ def _maybe_translate_local_file_mounts_and_sync_up(
     # TODO(zhwu): Optimize this by:
     # 1. Use the same bucket for all the mounts.
     # 2. When the src is the same, use the same bucket.
-    copy_mounts_with_file_in_src = dict()
+    copy_mounts_with_file_in_src = {}
     for i, (dst, src) in enumerate(copy_mounts.items()):
         task.file_mounts.pop(dst)
         if os.path.isfile(os.path.abspath(os.path.expanduser(src))):
@@ -735,7 +734,7 @@ def _maybe_translate_local_file_mounts_and_sync_up(
     file_bucket_name = spot.constants.SPOT_FM_FILE_ONLY_BUCKET_NAME.format(
         username=getpass.getuser(), id=run_id)
     if copy_mounts_with_file_in_src:
-        src_to_file_id = dict()
+        src_to_file_id = {}
         for i, src in enumerate(set(copy_mounts_with_file_in_src.values())):
             src_to_file_id[src] = i
             os.link(os.path.abspath(os.path.expanduser(src)),
@@ -775,7 +774,7 @@ def _maybe_translate_local_file_mounts_and_sync_up(
 
     # Step 5: Add the file download into the file mounts, such as
     #  /original-dst: s3://spot-fm-file-only-bucket-name/file-0
-    new_file_mounts = dict()
+    new_file_mounts = {}
     for dst, src in copy_mounts_with_file_in_src.items():
         storage = task.storage_mounts[spot.constants.SPOT_FM_REMOTE_TMP_DIR]
         store_type = list(storage.stores.keys())[0]
diff --git a/sky/global_user_state.py b/sky/global_user_state.py
index 3e98aabf8..ca9dadcc0 100644
--- a/sky/global_user_state.py
+++ b/sky/global_user_state.py
@@ -155,7 +155,7 @@ class StorageStatus(enum.Enum):
 
 
 def add_or_update_cluster(cluster_name: str,
-                          cluster_handle: 'backends.Backend.ResourceHandle',
+                          cluster_handle: 'backends.ResourceHandle',
                           requested_resources: Optional[Set[Any]],
                           ready: bool,
                           is_launch: bool = True):
@@ -163,7 +163,7 @@ def add_or_update_cluster(cluster_name: str,
 
     Args:
         cluster_name: Name of the cluster.
-        cluster_handle: Backend.ResourceHandle of the cluster.
+        cluster_handle: backends.ResourceHandle of the cluster.
         requested_resources: Resources requested for cluster.
         ready: Whether the cluster is ready to use. If False, the cluster will
             be marked as INIT, otherwise it will be marked as UP.
@@ -339,7 +339,7 @@ def remove_cluster(cluster_name: str, terminate: bool) -> float:
 
 
 def get_handle_from_cluster_name(
-        cluster_name: str) -> Optional['backends.Backend.ResourceHandle']:
+        cluster_name: str) -> Optional['backends.ResourceHandle']:
     assert cluster_name is not None, 'cluster_name cannot be None'
     rows = _DB.cursor.execute('SELECT handle FROM clusters WHERE name=(?)',
                               (cluster_name,))
diff --git a/sky/registry.py b/sky/registry.py
deleted file mode 100644
index 553464f21..000000000
--- a/sky/registry.py
+++ /dev/null
@@ -1,69 +0,0 @@
-"""Service registry."""
-import collections
-import typing
-from typing import List, Optional
-
-import sky
-from sky import clouds
-
-if typing.TYPE_CHECKING:
-    from sky import resources as resources_lib
-
-_CLOUDS = [
-    clouds.AWS(),
-    clouds.Azure(),
-    clouds.GCP(),
-    clouds.Lambda(),
-]
-
-
-def _filter_out_blocked_clouds(task: sky.Task):
-    available_clouds = []
-    for cloud in _CLOUDS:
-        for blocked_cloud in task.blocked_clouds:
-            if cloud.is_same_cloud(blocked_cloud):
-                break
-        else:  # non-blocked cloud (no break)
-            available_clouds.append(cloud)
-    return available_clouds
-
-
-def _filter_out_blocked_launchable_resources(
-        launchable_resources: List['resources_lib.Resources'],
-        blocked_launchable_resources: List['resources_lib.Resources']):
-    """Whether the resources are blocked."""
-    available_resources = []
-    for resources in launchable_resources:
-        for blocked_resources in blocked_launchable_resources:
-            if resources.should_be_blocked_by(blocked_resources):
-                break
-        else:  # non-blokced launchable resources. (no break)
-            available_resources.append(resources)
-    return available_resources
-
-
-def fill_in_launchable_resources(
-    task: sky.Task,
-    blocked_launchable_resources: Optional[List['resources_lib.Resources']],
-):
-    if blocked_launchable_resources is None:
-        blocked_launchable_resources = []
-    launchable = collections.defaultdict(list)
-    for resources in task.get_resources():
-        if resources.is_launchable():
-            launchable[resources] = [resources]
-        elif resources.cloud is not None:
-            launchable[
-                resources] = resources.cloud.get_feasible_launchable_resources(
-                    resources)
-        else:
-            # Remove blocked clouds.
-            available_clouds = _filter_out_blocked_clouds(task)
-            for cloud in available_clouds:
-                feasible_resources = cloud.get_feasible_launchable_resources(
-                    resources)
-                launchable[resources].extend(feasible_resources)
-        launchable[resources] = _filter_out_blocked_launchable_resources(
-            launchable[resources], blocked_launchable_resources)
-
-    return launchable
diff --git a/sky/resources.py b/sky/resources.py
index f5b960c74..e4f25c254 100644
--- a/sky/resources.py
+++ b/sky/resources.py
@@ -738,7 +738,7 @@ class Resources:
         backend_utils.validate_schema(config, schemas.get_resources_schema(),
                                       'Invalid resources YAML: ')
 
-        resources_fields = dict()
+        resources_fields = {}
         if config.get('cloud') is not None:
             resources_fields['cloud'] = clouds.CLOUD_REGISTRY.from_str(
                 config.pop('cloud'))
diff --git a/sky/setup_files/setup.py b/sky/setup_files/setup.py
index 32f891f56..08b9ee3af 100644
--- a/sky/setup_files/setup.py
+++ b/sky/setup_files/setup.py
@@ -18,6 +18,7 @@ import os
 import platform
 import re
 import warnings
+from typing import Dict, List
 
 import setuptools
 
@@ -88,6 +89,7 @@ install_requires = [
     'ray[default]>=1.9.0,<=2.2.0',
     'rich',
     'tabulate',
+    'typing-extensions',
     'filelock',  # TODO(mraheja): Enforce >=3.6.0 when python version is >= 3.7
     # This is used by ray. The latest 1.44.0 will generate an error
     # `Fork support is only compatible with the epoll1 and poll
@@ -104,7 +106,7 @@ install_requires = [
 
 # NOTE: Change the templates/spot-controller.yaml.j2 file if any of the following
 # packages dependencies are changed.
-extras_require = {
+extras_require: Dict[str, List[str]] = {
     'aws': [
         # awscli>=1.27.10 is required for SSO support.
         'awscli',
diff --git a/sky/sky_logging.py b/sky/sky_logging.py
index 82992c39f..02bbb59e6 100644
--- a/sky/sky_logging.py
+++ b/sky/sky_logging.py
@@ -27,7 +27,7 @@ class NewLineFormatter(logging.Formatter):
 
 def init_logger(name: str):
     h = logging.StreamHandler(sys.stdout)
-    h.flush = sys.stdout.flush
+    h.flush = sys.stdout.flush  # type: ignore
 
     fmt = NewLineFormatter(FORMAT, datefmt=DATE_FORMAT)
     h.setFormatter(fmt)
diff --git a/sky/skylet/autostop_lib.py b/sky/skylet/autostop_lib.py
index be97e30a2..5358910ba 100644
--- a/sky/skylet/autostop_lib.py
+++ b/sky/skylet/autostop_lib.py
@@ -28,7 +28,7 @@ class AutostopConfig:
 
     def __init__(self,
                  autostop_idle_minutes: int,
-                 boot_time: int,
+                 boot_time: float,
                  backend: Optional[str],
                  down: bool = False):
         assert autostop_idle_minutes < 0 or backend is not None, (
diff --git a/sky/skylet/configs.py b/sky/skylet/configs.py
index b5d177646..ed32f41b9 100644
--- a/sky/skylet/configs.py
+++ b/sky/skylet/configs.py
@@ -4,7 +4,7 @@ import functools
 import os
 import pathlib
 import sqlite3
-from typing import Optional
+from typing import Callable, Optional, Union
 
 _DB_PATH = os.path.expanduser('~/.sky/skylet_config.db')
 os.makedirs(pathlib.Path(_DB_PATH).parents[0], exist_ok=True)
@@ -25,7 +25,7 @@ def _safe_cursor():
         conn.close()
 
 
-def ensure_table(func: callable):
+def ensure_table(func: Callable):
     """Ensure the table exists before calling the function.
 
     Since this module will be imported whenever `sky` is imported (due to
@@ -55,7 +55,7 @@ def ensure_table(func: callable):
 
 
 @ensure_table
-def get_config(key: str) -> Optional[str]:
+def get_config(key: str) -> Optional[bytes]:
     with _safe_cursor() as cursor:
         rows = cursor.execute('SELECT value FROM config WHERE key = ?', (key,))
         for (value,) in rows:
@@ -64,7 +64,7 @@ def get_config(key: str) -> Optional[str]:
 
 
 @ensure_table
-def set_config(key: str, value: str) -> None:
+def set_config(key: str, value: Union[bytes, str]) -> None:
     with _safe_cursor() as cursor:
         cursor.execute(
             """\
diff --git a/sky/skylet/events.py b/sky/skylet/events.py
index b0bdbf1ad..a6f2601bb 100644
--- a/sky/skylet/events.py
+++ b/sky/skylet/events.py
@@ -189,6 +189,6 @@ class AutostopEvent(SkyletEvent):
         # would be launched).
         config['auth'].pop('ssh_proxy_command', None)
         # Empty the file_mounts.
-        config['file_mounts'] = dict()
+        config['file_mounts'] = {}
         common_utils.dump_yaml(yaml_path, config)
         logger.debug('Replaced upscaling speed to 0.')
diff --git a/sky/skylet/job_lib.py b/sky/skylet/job_lib.py
index 360abab05..c52d09181 100644
--- a/sky/skylet/job_lib.py
+++ b/sky/skylet/job_lib.py
@@ -243,7 +243,7 @@ def set_job_started(job_id: int) -> None:
         _CONN.commit()
 
 
-def get_status_no_lock(job_id: int) -> JobStatus:
+def get_status_no_lock(job_id: int) -> Optional[JobStatus]:
     """Get the status of the job with the given id.
 
     This function can return a stale status if there is a concurrent update.
@@ -257,6 +257,7 @@ def get_status_no_lock(job_id: int) -> JobStatus:
         if status is None:
             return None
         return JobStatus(status)
+    return None
 
 
 def get_status(job_id: int) -> Optional[JobStatus]:
@@ -279,7 +280,8 @@ def get_statuses_payload(job_ids: List[Optional[int]]) -> str:
     return common_utils.encode_payload(statuses)
 
 
-def load_statuses_payload(statuses_payload: str) -> Dict[int, JobStatus]:
+def load_statuses_payload(
+        statuses_payload: str) -> Dict[Optional[int], Optional[JobStatus]]:
     statuses = common_utils.decode_payload(statuses_payload)
     for job_id, status in statuses.items():
         if status is not None:
@@ -292,6 +294,7 @@ def get_latest_job_id() -> Optional[int]:
         'SELECT job_id FROM jobs ORDER BY job_id DESC LIMIT 1')
     for (job_id,) in rows:
         return job_id
+    return None
 
 
 def get_job_submitted_or_ended_timestamp_payload(job_id: int,
@@ -403,12 +406,12 @@ def update_job_status(job_owner: str,
     # which contains the job status (str) and submission_id (str).
     job_detail_lists: List['ray_pydantic.JobDetails'] = job_client.list_jobs()
 
-    job_details = dict()
+    job_details = {}
     ray_job_ids_set = set(ray_job_ids)
     for job_detail in job_detail_lists:
         if job_detail.submission_id in ray_job_ids_set:
             job_details[job_detail.submission_id] = job_detail
-    job_statuses: List[JobStatus] = [None] * len(ray_job_ids)
+    job_statuses: List[Optional[JobStatus]] = [None] * len(ray_job_ids)
     for i, ray_job_id in enumerate(ray_job_ids):
         if ray_job_id in job_details:
             ray_status = job_details[ray_job_id].status
@@ -425,8 +428,9 @@ def update_job_status(job_owner: str,
         # updated
         # pylint: disable=abstract-class-instantiated
         with filelock.FileLock(_get_lock_path(job_id)):
+            original_status = get_status_no_lock(job_id)
+            assert original_status is not None, (job_id, status)
             if status is None:
-                original_status = get_status_no_lock(job_id)
                 status = original_status
                 if (original_status is not None and
                         not original_status.is_terminal()):
@@ -439,7 +443,6 @@ def update_job_status(job_owner: str,
                     if not silent:
                         logger.info(f'Updated job {job_id} status to {status}')
             else:
-                original_status = get_status_no_lock(job_id)
                 # Taking max of the status is necessary because:
                 # 1. It avoids race condition, where the original status has
                 # already been set to later state by the job. We skip the
@@ -498,6 +501,7 @@ def is_cluster_idle() -> bool:
         """, in_progress_status)
     for (count,) in rows:
         return count == 0
+    assert False, 'Should not reach here'
 
 
 def format_job_queue(jobs: List[Dict[str, Any]]):
@@ -534,7 +538,9 @@ def dump_job_queue(username: Optional[str], all_jobs: bool) -> str:
         username: The username to show jobs for. Show all the users if None.
         all_jobs: Whether to show all jobs, not just the pending/running ones.
     """
-    status_list = [JobStatus.SETTING_UP, JobStatus.PENDING, JobStatus.RUNNING]
+    status_list: Optional[List[JobStatus]] = [
+        JobStatus.SETTING_UP, JobStatus.PENDING, JobStatus.RUNNING
+    ]
     if all_jobs:
         status_list = None
 
@@ -607,8 +613,7 @@ def get_run_timestamp(job_id: Optional[int]) -> Optional[str]:
     return run_timestamp
 
 
-def run_timestamp_with_globbing_payload(
-        job_ids: List[Optional[str]]) -> Dict[str, str]:
+def run_timestamp_with_globbing_payload(job_ids: List[Optional[str]]) -> str:
     """Returns the relative paths to the log files for job with globbing."""
     query_str = ' OR '.join(['job_id GLOB (?)'] * len(job_ids))
     _CURSOR.execute(
@@ -616,7 +621,7 @@ def run_timestamp_with_globbing_payload(
             SELECT * FROM jobs
             WHERE {query_str}""", job_ids)
     rows = _CURSOR.fetchall()
-    run_timestamps = dict()
+    run_timestamps = {}
     for row in rows:
         job_id = row[JobInfoLoc.JOB_ID.value]
         run_timestamp = row[JobInfoLoc.RUN_TIMESTAMP.value]
diff --git a/sky/skylet/log_lib.py b/sky/skylet/log_lib.py
index afc822093..45e82cc46 100644
--- a/sky/skylet/log_lib.py
+++ b/sky/skylet/log_lib.py
@@ -110,11 +110,12 @@ def process_subprocess_stream(
 def run_with_log(
     cmd: Union[List[str], str],
     log_path: str,
+    *,
+    require_outputs: bool = False,
     stream_logs: bool = False,
     start_streaming_at: str = '',
     end_streaming_at: Optional[str] = None,
     skip_lines: Optional[List[str]] = None,
-    require_outputs: bool = False,
     shell: bool = False,
     with_ray: bool = False,
     process_stream: bool = True,
@@ -161,13 +162,13 @@ def run_with_log(
         os.makedirs(dirname, exist_ok=True)
     # Redirect stderr to stdout when using ray, to preserve the order of
     # stdout and stderr.
-    stdout = stderr = None
+    stdout_arg = stderr_arg = None
     if process_stream:
-        stdout = subprocess.PIPE
-        stderr = subprocess.PIPE if not with_ray else subprocess.STDOUT
+        stdout_arg = subprocess.PIPE
+        stderr_arg = subprocess.PIPE if not with_ray else subprocess.STDOUT
     with subprocess.Popen(cmd,
-                          stdout=stdout,
-                          stderr=stderr,
+                          stdout=stdout_arg,
+                          stderr=stderr_arg,
                           start_new_session=True,
                           shell=shell,
                           **kwargs) as proc:
@@ -302,6 +303,7 @@ def run_bash_command_with_log(bash_command: str,
         # Need this `-i` option to make sure `source ~/.bashrc` work.
         inner_command = f'/bin/bash -i {script_path}'
 
+        subprocess_cmd: Union[str, List[str]]
         if use_sudo:
             subprocess.run(f'chmod a+rwx {script_path}', shell=True, check=True)
             subprocess_cmd = job_lib.make_job_command_with_user_switching(
@@ -359,7 +361,8 @@ def _follow_job_logs(file,
                     time.sleep(1 + _SKY_LOG_TAILING_GAP_SECONDS)
                     wait_last_logs = False
                     continue
-                print(f'INFO: Job finished (status: {status.value}).')
+                status_str = status.value if status is not None else 'None'
+                print(f'INFO: Job finished (status: {status_str}).')
                 return
 
             time.sleep(_SKY_LOG_TAILING_GAP_SECONDS)
diff --git a/sky/skylet/log_lib.pyi b/sky/skylet/log_lib.pyi
new file mode 100644
index 000000000..6a85ba6fb
--- /dev/null
+++ b/sky/skylet/log_lib.pyi
@@ -0,0 +1,113 @@
+"""Type stubs for log_lib.py.
+This file is dynamically generated by stubgen and added with the
+overloaded type hints for run_with_log(), as we need to determine
+the return type based on the value of require_outputs.
+"""
+from sky import sky_logging as sky_logging
+from sky.skylet import constants as constants, job_lib as job_lib
+from sky.utils import log_utils as log_utils
+import typing
+from typing import Dict, List, Optional, Tuple, Union
+from typing_extensions import Literal
+
+
+def process_subprocess_stream(
+        proc,
+        log_path: str,
+        stream_logs: bool,
+        start_streaming_at: str = ...,
+        end_streaming_at: Optional[str] = ...,
+        skip_lines: Optional[List[str]] = ...,
+        replace_crlf: bool = ...,
+        line_processor: Optional[log_utils.LineProcessor] = ...,
+        streaming_prefix: Optional[str] = ...) -> Tuple[str, str]:
+    ...
+
+
+@typing.overload
+def run_with_log(cmd: Union[List[str], str],
+                 log_path: str,
+                 *,
+                 require_outputs: Literal[False] = False,
+                 stream_logs: bool = ...,
+                 start_streaming_at: str = ...,
+                 end_streaming_at: Optional[str] = ...,
+                 skip_lines: Optional[List[str]] = ...,
+                 shell: bool = ...,
+                 with_ray: bool = ...,
+                 process_stream: bool = ...,
+                 line_processor: Optional[log_utils.LineProcessor] = ...,
+                 streaming_prefix: Optional[str] = ...,
+                 ray_job_id: Optional[str] = ...,
+                 use_sudo: bool = ...,
+                 **kwargs) -> int:
+    ...
+
+
+@typing.overload
+def run_with_log(cmd: Union[List[str], str],
+                 log_path: str,
+                 *,
+                 require_outputs: Literal[True],
+                 stream_logs: bool = ...,
+                 start_streaming_at: str = ...,
+                 end_streaming_at: Optional[str] = ...,
+                 skip_lines: Optional[List[str]] = ...,
+                 shell: bool = ...,
+                 with_ray: bool = ...,
+                 process_stream: bool = ...,
+                 line_processor: Optional[log_utils.LineProcessor] = ...,
+                 streaming_prefix: Optional[str] = ...,
+                 ray_job_id: Optional[str] = ...,
+                 use_sudo: bool = ...,
+                 **kwargs) -> Tuple[int, str, str]:
+    ...
+
+
+@typing.overload
+def run_with_log(cmd: Union[List[str], str],
+                 log_path: str,
+                 *,
+                 require_outputs: bool = False,
+                 stream_logs: bool = ...,
+                 start_streaming_at: str = ...,
+                 end_streaming_at: Optional[str] = ...,
+                 skip_lines: Optional[List[str]] = ...,
+                 shell: bool = ...,
+                 with_ray: bool = ...,
+                 process_stream: bool = ...,
+                 line_processor: Optional[log_utils.LineProcessor] = ...,
+                 streaming_prefix: Optional[str] = ...,
+                 ray_job_id: Optional[str] = ...,
+                 use_sudo: bool = ...,
+                 **kwargs) -> Union[int, Tuple[int, str, str]]:
+    ...
+
+
+def make_task_bash_script(codegen: str,
+                          env_vars: Optional[Dict[str, str]] = ...) -> str:
+    ...
+
+
+def add_ray_env_vars(
+        env_vars: Optional[Dict[str, str]] = ...) -> Dict[str, str]:
+    ...
+
+
+def run_bash_command_with_log(bash_command: str,
+                              log_path: str,
+                              job_owner: str,
+                              job_id: int,
+                              env_vars: Optional[Dict[str, str]] = ...,
+                              stream_logs: bool = ...,
+                              with_ray: bool = ...,
+                              use_sudo: bool = ...):
+    ...
+
+
+def tail_logs(job_owner: str,
+              job_id: int,
+              log_dir: Optional[str],
+              spot_job_id: Optional[int] = ...,
+              follow: bool = ...) -> None:
+    ...
diff --git a/sky/skylet/providers/gcp/node_provider.py b/sky/skylet/providers/gcp/node_provider.py
index 0bd987103..1a3443f52 100644
--- a/sky/skylet/providers/gcp/node_provider.py
+++ b/sky/skylet/providers/gcp/node_provider.py
@@ -184,7 +184,7 @@ class GCPNodeProvider(NodeProvider):
         instances.
         """
         with self.lock:
-            result_dict = dict()
+            result_dict = {}
             labels = tags  # gcp uses "labels" instead of aws "tags"
             labels = dict(sorted(copy.deepcopy(labels).items()))
 
diff --git a/sky/skypilot_config.py b/sky/skypilot_config.py
index a9de17651..6967da7c0 100644
--- a/sky/skypilot_config.py
+++ b/sky/skypilot_config.py
@@ -43,7 +43,7 @@ then:
 """
 import copy
 import os
-from typing import Any, Dict, Tuple
+from typing import Any, Dict, Sequence
 
 import yaml
 
@@ -109,7 +109,7 @@ def loaded() -> bool:
     return _dict is not None
 
 
-def get_nested(keys: Tuple[str], default_value: Any) -> Any:
+def get_nested(keys: Sequence[str], default_value: Any) -> Any:
     """Gets a nested key.
 
     If any key is not found, or any intermediate key does not point to a dict
@@ -128,13 +128,14 @@ def get_nested(keys: Tuple[str], default_value: Any) -> Any:
     return curr
 
 
-def pop_nested(keys: Tuple[str]) -> Dict[str, Any]:
+def pop_nested(keys: Sequence[str]) -> Dict[str, Any]:
     """Returns a deep-copied config with the nested key popped.
 
     Like get_nested(), if any key is not found, this will not raise an error.
     """
     _check_loaded_or_die()
     global _dict
+    assert _dict is not None
     curr = copy.deepcopy(_dict)
     to_return = curr
     prev = None
diff --git a/sky/spot/recovery_strategy.py b/sky/spot/recovery_strategy.py
index 538a23f36..de41d685a 100644
--- a/sky/spot/recovery_strategy.py
+++ b/sky/spot/recovery_strategy.py
@@ -7,6 +7,7 @@ import sky
 from sky import exceptions
 from sky import global_user_state
 from sky import sky_logging
+from sky import backends
 from sky.backends import backend_utils
 from sky.skylet import job_lib
 from sky.spot import spot_utils
@@ -15,12 +16,11 @@ from sky.utils import common_utils
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
-    from sky import backends
     from sky import task as task_lib
 
 logger = sky_logging.init_logger(__name__)
 
-SPOT_STRATEGIES = dict()
+SPOT_STRATEGIES = {}
 SPOT_DEFAULT_STRATEGY = None
 
 # Waiting time for job from INIT/PENDING to RUNNING
@@ -43,6 +43,8 @@ class StrategyExecutor:
             task: The task to execute.
             retry_until_up: Whether to retry until the cluster is up.
         """
+        assert isinstance(backend, backends.CloudVmRayBackend), (
+            'Only CloudVMRayBackend is supported.')
         self.dag = sky.Dag()
         self.dag.add(task)
         self.cluster_name = cluster_name
diff --git a/sky/spot/spot_utils.py b/sky/spot/spot_utils.py
index 9b6c39b5e..f2ba6dd19 100644
--- a/sky/spot/spot_utils.py
+++ b/sky/spot/spot_utils.py
@@ -5,7 +5,7 @@ import json
 import pathlib
 import shlex
 import time
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Sequence, Tuple
 
 import colorama
 import filelock
@@ -95,7 +95,7 @@ def update_spot_job_status(job_id: Optional[int] = None):
         job_ids = [job_id]
     for job_id_ in job_ids:
         controller_status = job_lib.get_status(job_id_)
-        if controller_status.is_terminal():
+        if controller_status is None or controller_status.is_terminal():
             logger.error(f'Controller for job {job_id_} has exited abnormally. '
                          'Setting the job status to FAILED_CONTROLLER.')
             task_name = spot_state.get_task_name_by_job_id(job_id_)
@@ -237,20 +237,20 @@ def stream_logs_by_id(job_id: int, follow: bool = True) -> str:
         msg = _JOB_WAITING_STATUS_MESSAGE.format(status_str='')
         status_display.update(msg)
         prev_msg = msg
-        job_status = spot_state.get_status(job_id)
-        while job_status is None:
+        spot_job_status = spot_state.get_status(job_id)
+        while spot_job_status is None:
             time.sleep(1)
-            job_status = spot_state.get_status(job_id)
+            spot_job_status = spot_state.get_status(job_id)
 
-        if job_status.is_terminal():
+        if spot_job_status.is_terminal():
             job_msg = ''
-            if job_status.is_failed():
+            if spot_job_status.is_failed():
                 job_msg = (
                     f'\nFailure reason: {spot_state.get_failure_reason(job_id)}'
                 )
             return (f'{colorama.Fore.YELLOW}'
                     f'Job {job_id} is already in terminal state '
-                    f'{job_status.value}. Logs will not be shown.'
+                    f'{spot_job_status.value}. Logs will not be shown.'
                     f'{colorama.Style.RESET_ALL}{job_msg}')
         task_name = spot_state.get_task_name_by_job_id(job_id)
         cluster_name = generate_spot_cluster_name(task_name, job_id)
@@ -473,7 +473,7 @@ class SpotCodeGen:
         return cls._build(code)
 
     @classmethod
-    def cancel_jobs_by_id(cls, job_ids: Optional[List[int]]) -> str:
+    def cancel_jobs_by_id(cls, job_ids: Optional[Sequence[int]]) -> str:
         code = [
             f'msg = spot_utils.cancel_jobs_by_id({job_ids})',
             'print(msg, end="", flush=True)',
@@ -523,8 +523,15 @@ def dump_job_table_cache(job_table: str):
         json.dump((time.time(), job_table), f)
 
 
-def load_job_table_cache() -> Optional[Tuple[str, str]]:
-    """Load job table cache from file."""
+def load_job_table_cache() -> Optional[Tuple[float, str]]:
+    """Load job table cache from file.
+
+    Returns:
+        A tuple of (timestamp, job_table), where the timestamp is
+        the time when the job table is dumped and the job_table is
+        the dumped job table in string.
+        None if the cache file does not exist.
+    """
     cache_file = pathlib.Path(_SPOT_STATUS_CACHE).expanduser()
     if not cache_file.exists():
         return None
@@ -535,7 +542,7 @@ def load_job_table_cache() -> Optional[Tuple[str, str]]:
 def is_spot_controller_up(
     stopped_message: str,
 ) -> Tuple[Optional[global_user_state.ClusterStatus],
-           Optional[backends.Backend.ResourceHandle]]:
+           Optional['backends.CloudVmRayResourceHandle']]:
     """Check if the spot controller is up.
 
     It can be used to check the actual controller status (since the autostop is
diff --git a/sky/task.py b/sky/task.py
index e4afba182..62c96f3e3 100644
--- a/sky/task.py
+++ b/sky/task.py
@@ -142,7 +142,7 @@ class Task:
         self.storage_mounts = {}
         self.storage_plans = {}
         self.setup = setup
-        self._envs = envs or dict()
+        self._envs = envs or {}
         self.workdir = workdir
         self.docker_image = (docker_image if docker_image else
                              'gpuci/miniforge-cuda:11.4-devel-ubuntu18.04')
@@ -274,7 +274,7 @@ class Task:
         fm_storages = []
         file_mounts = config.pop('file_mounts', None)
         if file_mounts is not None:
-            copy_mounts = dict()
+            copy_mounts = {}
             for dst_path, src in file_mounts.items():
                 # Check if it is str path
                 if isinstance(src, str):
@@ -353,7 +353,7 @@ class Task:
           ValueError: if various invalid inputs errors are detected.
         """
         if envs is None:
-            self._envs = dict()
+            self._envs = {}
             return self
         if isinstance(envs, (list, tuple)):
             keys = set(env[0] for env in envs)
@@ -471,8 +471,8 @@ class Task:
         """
         if self.time_estimator_func is None:
             raise NotImplementedError(
-                'Node [{}] does not have a cost model set; '
-                'call set_time_estimator() first'.format(self))
+                f'Node [{self}] does not have a cost model set; '
+                'call set_time_estimator() first')
         return self.time_estimator_func(resources)
 
     def set_file_mounts(self, file_mounts: Optional[Dict[str, str]]) -> 'Task':
@@ -783,7 +783,7 @@ class Task:
 
         INTERNAL: this method is internal-facing.
         """
-        config = dict()
+        config = {}
 
         def add_if_not_none(key, value, no_empty: bool = False):
             if no_empty and not value:
@@ -812,7 +812,7 @@ class Task:
         add_if_not_none('run', self.run)
         add_if_not_none('envs', self.envs, no_empty=True)
 
-        add_if_not_none('file_mounts', dict())
+        add_if_not_none('file_mounts', {})
 
         if self.file_mounts is not None:
             config['file_mounts'].update(self.file_mounts)
diff --git a/sky/usage/usage_lib.py b/sky/usage/usage_lib.py
index df6842870..90284b364 100644
--- a/sky/usage/usage_lib.py
+++ b/sky/usage/usage_lib.py
@@ -10,7 +10,7 @@ import os
 import time
 import traceback
 import typing
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Union
 
 import requests
 
@@ -47,14 +47,14 @@ class MessageType(enum.Enum):
 class MessageToReport:
     """Abstract class for messages to be sent to Loki."""
 
-    def __init__(self, schema_version: str):
+    def __init__(self, schema_version: int):
         self.schema_version = schema_version
-        self.start_time: int = None
-        self.send_time: int = None
+        self.start_time: Optional[int] = None
+        self.send_time: Optional[int] = None
 
     def start(self):
         if self.start_time is None:
-            self.start_time: int = _get_current_timestamp_ns()
+            self.start_time = _get_current_timestamp_ns()
 
     @property
     def message_sent(self):
@@ -107,8 +107,8 @@ class UsageMessageToReport(MessageToReport):
         self.resources: Optional[Dict[str,
                                       Any]] = None  # update_cluster_resources
         #: Resources of the local cluster.
-        self.local_resources: Optional[Dict[
-            str, Any]] = None  # update_local_cluster_resources
+        self.local_resources: Optional[List[Dict[
+            str, Any]]] = None  # update_local_cluster_resources
         #: The number of nodes in the cluster.
         self.num_nodes: Optional[int] = None  # update_cluster_resources
         #: The status of the cluster.
@@ -148,7 +148,7 @@ class UsageMessageToReport(MessageToReport):
         self.ray_yamls: Optional[List[Dict[str, Any]]] = None
         #: Number of Ray YAML files.
         self.num_tried_regions: Optional[int] = None  # update_ray_yaml
-        self.runtimes: Dict[str, int] = {}  # update_runtime
+        self.runtimes: Dict[str, float] = {}  # update_runtime
         self.stacktrace: Optional[str] = None  # entrypoint_context
 
     def __repr__(self) -> str:
@@ -316,7 +316,7 @@ def _send_to_loki(message_type: MessageType):
     messages.reset(message_type)
 
 
-def _clean_yaml(yaml_info: Dict[str, str]):
+def _clean_yaml(yaml_info: Dict[str, Optional[str]]):
     """Remove sensitive information from user YAML."""
     cleaned_yaml_info = yaml_info.copy()
     for redact_type in constants.USAGE_MESSAGE_REDACT_KEYS:
@@ -424,7 +424,7 @@ def entrypoint_context(name: str, fallback: bool = False):
         _send_local_messages()
 
 
-def entrypoint(name_or_fn: str, fallback: bool = False):
+def entrypoint(name_or_fn: Union[str, Callable], fallback: bool = False):
     return common_utils.make_decorator(entrypoint_context,
                                        name_or_fn,
                                        fallback=fallback)
diff --git a/sky/utils/cli_utils/status_utils.py b/sky/utils/cli_utils/status_utils.py
index c6ff4cc6b..e9956e180 100644
--- a/sky/utils/cli_utils/status_utils.py
+++ b/sky/utils/cli_utils/status_utils.py
@@ -195,7 +195,7 @@ def show_local_status_table(local_clusters: List[str]):
         config = common_utils.read_yaml(config_path)
         username = config['auth']['ssh_user']
 
-        if not isinstance(handle, backends.CloudVmRayBackend.ResourceHandle):
+        if not isinstance(handle, backends.CloudVmRayResourceHandle):
             raise ValueError(f'Unknown handle type {type(handle)} encountered.')
 
         if (handle.launched_nodes is not None and
@@ -285,9 +285,9 @@ def _get_status_for_cost_report(cluster_status):
 def _get_resources(cluster_status):
     handle = cluster_status['handle']
     resources_str = '<initializing>'
-    if isinstance(handle, backends.LocalDockerBackend.ResourceHandle):
+    if isinstance(handle, backends.LocalDockerResourceHandle):
         resources_str = 'docker'
-    elif isinstance(handle, backends.CloudVmRayBackend.ResourceHandle):
+    elif isinstance(handle, backends.CloudVmRayResourceHandle):
         if (handle.launched_nodes is not None and
                 handle.launched_resources is not None):
             launched_resource_str = str(handle.launched_resources)
diff --git a/sky/utils/command_runner.py b/sky/utils/command_runner.py
index 5fd37edf5..410946d04 100644
--- a/sky/utils/command_runner.py
+++ b/sky/utils/command_runner.py
@@ -186,9 +186,9 @@ class SSHCommandRunner:
             self,
             cmd: Union[str, List[str]],
             *,
+            require_outputs: bool = False,
             port_forward: Optional[List[int]] = None,
             # Advanced options.
-            require_outputs: bool = False,
             log_path: str = os.devnull,
             # If False, do not redirect stdout/stderr to optimize performance.
             process_stream: bool = True,
@@ -276,9 +276,9 @@ class SSHCommandRunner:
 
         return log_lib.run_with_log(' '.join(command),
                                     log_path,
-                                    stream_logs,
-                                    process_stream=process_stream,
                                     require_outputs=require_outputs,
+                                    stream_logs=stream_logs,
+                                    process_stream=process_stream,
                                     shell=True,
                                     executable=executable,
                                     **kwargs)
diff --git a/sky/utils/command_runner.pyi b/sky/utils/command_runner.pyi
new file mode 100644
index 000000000..712075544
--- /dev/null
+++ b/sky/utils/command_runner.pyi
@@ -0,0 +1,104 @@
+"""Stub file for sky.utils.command_runner.
+
+This file is dynamically generated by stubgen and added with the
+overloaded type hints for SSHCommandRunner.run(), as we need to
+determine the return type based on the value of require_outputs.
+"""
+import enum
+import typing
+from sky import sky_logging as sky_logging
+from sky.skylet import log_lib as log_lib
+from sky.utils import subprocess_utils as subprocess_utils
+from typing import List, Optional, Tuple, Union
+from typing_extensions import Literal
+
+GIT_EXCLUDE: str
+RSYNC_DISPLAY_OPTION: str
+RSYNC_FILTER_OPTION: str
+RSYNC_EXCLUDE_OPTION: str
+
+
+def ssh_options_list(ssh_private_key: Optional[str],
+                     ssh_control_name: Optional[str],
+                     *,
+                     timeout: int = ...) -> List[str]:
+    ...
+
+
+class SshMode(enum.Enum):
+    NON_INTERACTIVE: int
+    INTERACTIVE: int
+    LOGIN: int
+
+
+class SSHCommandRunner:
+    ip: str
+    ssh_user: str
+    ssh_private_key: str
+    ssh_control_name: Optional[str]
+
+    def __init__(self,
+                 ip: str,
+                 ssh_user: str,
+                 ssh_private_key: str,
+                 ssh_control_name: Optional[str] = ...) -> None:
+        ...
+
+    @staticmethod
+    def make_runner_list(
+            ip_list: List[str],
+            ssh_user: str,
+            ssh_private_key: str,
+            ssh_control_name: Optional[str] = ...) -> List['SSHCommandRunner']:
+        ...
+
+    @typing.overload
+    def run(self,
+            cmd: Union[str, List[str]],
+            *,
+            require_outputs: Literal[False] = ...,
+            port_forward: Optional[List[int]] = ...,
+            log_path: str = ...,
+            process_stream: bool = ...,
+            stream_logs: bool = ...,
+            ssh_mode: SshMode = ...,
+            separate_stderr: bool = ...,
+            **kwargs) -> int:
+        ...
+
+    @typing.overload
+    def run(self,
+            cmd: Union[str, List[str]],
+            *,
+            require_outputs: Literal[True],
+            port_forward: Optional[List[int]] = ...,
+            log_path: str = ...,
+            process_stream: bool = ...,
+            stream_logs: bool = ...,
+            ssh_mode: SshMode = ...,
+            separate_stderr: bool = ...,
+            **kwargs) -> Tuple[int, str, str]:
+        ...
+
+    @typing.overload
+    def run(self,
+            cmd: Union[str, List[str]],
+            *,
+            require_outputs: bool = ...,
+            port_forward: Optional[List[int]] = ...,
+            log_path: str = ...,
+            process_stream: bool = ...,
+            stream_logs: bool = ...,
+            ssh_mode: SshMode = ...,
+            separate_stderr: bool = ...,
+            **kwargs) -> Union[Tuple[int, str, str], int]:
+        ...
+
+    def rsync(self,
+              source: str,
+              target: str,
+              *,
+              up: bool,
+              log_path: str = ...,
+              stream_logs: bool = ...) -> None:
+        ...
diff --git a/sky/utils/common_utils.py b/sky/utils/common_utils.py
index ecb6bfc31..51fcfcd50 100644
--- a/sky/utils/common_utils.py
+++ b/sky/utils/common_utils.py
@@ -12,7 +12,7 @@ import re
 import socket
 import sys
 import time
-from typing import Any, Dict, Optional
+from typing import Any, Callable, Dict, Optional, Union
 import uuid
 import yaml
 
@@ -194,7 +194,7 @@ def dump_yaml_str(config):
                      default_flow_style=False)
 
 
-def make_decorator(cls, name_or_fn, **ctx_kwargs):
+def make_decorator(cls, name_or_fn: Union[str, Callable], **ctx_kwargs):
     """Make the cls a decorator.
 
     class cls:
@@ -318,7 +318,8 @@ def class_fullname(cls):
     return f'{cls.__module__}.{cls.__name__}'
 
 
-def format_exception(e: Exception, use_bracket: bool = False) -> str:
+def format_exception(e: Union[Exception, SystemExit],
+                     use_bracket: bool = False) -> str:
     """Format an exception to a string.
 
     Args:
diff --git a/sky/utils/log_utils.py b/sky/utils/log_utils.py
index 0b2c010f7..9fa5d174c 100644
--- a/sky/utils/log_utils.py
+++ b/sky/utils/log_utils.py
@@ -67,8 +67,8 @@ def create_table(field_names: List[str], **kwargs) -> prettytable.PrettyTable:
     return table
 
 
-def readable_time_duration(start: Optional[int],
-                           end: Optional[int] = None,
+def readable_time_duration(start: Optional[float],
+                           end: Optional[float] = None,
                            absolute: bool = False) -> str:
     """Human readable time duration from timestamps.
 


Test patch:
diff --git a/tests/mypy_files.txt b/tests/mypy_files.txt
index 5c86cd96f..812eb7e5a 100644
--- a/tests/mypy_files.txt
+++ b/tests/mypy_files.txt
@@ -1,5 +1,10 @@
-sky/adaptors
-sky/clouds
-sky/data
-sky/spot
-sky/utils
+sky
+
+--exclude sky/benchmark
+--exclude sky/callbacks
+--exclude sky/skylet/providers
+--exclude sky/resources.py
+--exclude sky/global_user_state.py
+--exclude sky/optimizer.py
+--exclude sky/task.py
+--exclude sky/backends/monkey_patches
diff --git a/tests/test_spot.py b/tests/test_spot.py
index fe61285fb..9c3b89a74 100644
--- a/tests/test_spot.py
+++ b/tests/test_spot.py
@@ -43,7 +43,7 @@ class TestReservedClustersOperations:
     @pytest.fixture
     def _mock_cluster_state(self, _mock_db_conn):
         assert 'state.db' not in global_user_state._DB.db_path
-        handle = backends.CloudVmRayBackend.ResourceHandle(
+        handle = backends.CloudVmRayResourceHandle(
             cluster_name='test-cluster1',
             cluster_yaml='/tmp/cluster1.yaml',
             launched_nodes=2,
@@ -56,7 +56,7 @@ class TestReservedClustersOperations:
             handle,
             requested_resources={handle.launched_resources},
             ready=True)
-        handle = backends.CloudVmRayBackend.ResourceHandle(
+        handle = backends.CloudVmRayResourceHandle(
             cluster_name='test-cluster2',
             cluster_yaml='/tmp/cluster2.yaml',
             launched_nodes=1,
@@ -70,7 +70,7 @@ class TestReservedClustersOperations:
             handle,
             requested_resources={handle.launched_resources},
             ready=True)
-        handle = backends.CloudVmRayBackend.ResourceHandle(
+        handle = backends.CloudVmRayResourceHandle(
             cluster_name='test-cluster3',
             cluster_yaml='/tmp/cluster3.yaml',
             launched_nodes=4,
@@ -83,7 +83,7 @@ class TestReservedClustersOperations:
             handle,
             requested_resources={handle.launched_resources},
             ready=False)
-        handle = backends.CloudVmRayBackend.ResourceHandle(
+        handle = backends.CloudVmRayResourceHandle(
             cluster_name=spot.SPOT_CONTROLLER_NAME,
             cluster_yaml='/tmp/spot_controller.yaml',
             launched_nodes=1,