Repository: skypilot-org/skypilot
Pull request #1717 (2023-02-23 22:20:55)

Fix spot tpu bug
<!-- Describe the changes in this PR -->

Fixes
- https://github.com/skypilot-org/skypilot/issues/1694
- misclassification of user program error and preeemption error

<!-- 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):

- [ ] Any manual or new tests for this PR (please specify below)
- [ ] All smoke tests: `pytest tests/test_smoke.py` 
- [ ] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_fill_in_the_name` 
- [ ] Backward compatibility tests: `bash tests/backward_comaptibility_tests.sh`



Patch:
diff --git a/pyproject.toml b/pyproject.toml
index 53a88edda..36e03fb38 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -2,6 +2,8 @@
 requires = ["setuptools>=58.0"]
 build-backend = "setuptools.build_meta"
 
+[project]
+name = "skypilot"
 
 [tool.yapf]
 based_on_style = "google"
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index 060397956..869032c1c 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -142,10 +142,8 @@ def _get_yaml_path_from_cluster_name(cluster_name: str,
     return str(output_path)
 
 
-def fill_template(template_name: str,
-                  variables: Dict,
-                  output_path: Optional[str] = None,
-                  output_prefix: str = SKY_USER_FILE_PATH) -> str:
+def fill_template(template_name: str, variables: Dict,
+                  output_path: str) -> None:
     """Create a file from a Jinja template and return the filename."""
     assert template_name.endswith('.j2'), template_name
     template_path = os.path.join(sky.__root_dir__, 'templates', template_name)
@@ -153,19 +151,13 @@ def fill_template(template_name: str,
         raise FileNotFoundError(f'Template "{template_name}" does not exist.')
     with open(template_path) as fin:
         template = fin.read()
-    if output_path is None:
-        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)
+    output_path = os.path.abspath(os.path.expanduser(output_path))
 
     # Write out yaml config.
     j2_template = jinja2.Template(template)
     content = j2_template.render(**variables)
     with open(output_path, 'w') as fout:
         fout.write(content)
-    return output_path
 
 
 def _optimize_file_mounts(yaml_path: str) -> None:
@@ -434,7 +426,7 @@ class SSHConfigHelper(object):
             # If an existing config with `cluster_name` exists, raise a warning.
             for i, line in enumerate(config):
                 if line.strip() == f'Host {cluster_name}':
-                    prev_line = config[i - 1] if i - 1 > 0 else ''
+                    prev_line = config[i - 1] if i - 1 >= 0 else ''
                     if prev_line.strip().startswith(sky_autogen_comment):
                         overwrite = True
                         overwrite_begin_idx = i - 1
@@ -445,7 +437,7 @@ class SSHConfigHelper(object):
                         logger.warning(f'Using {ip} to identify host instead.')
 
                 if line.strip() == f'Host {ip}':
-                    prev_line = config[i - 1] if i - 1 > 0 else ''
+                    prev_line = config[i - 1] if i - 1 >= 0 else ''
                     if prev_line.strip().startswith(sky_autogen_comment):
                         overwrite = True
                         overwrite_begin_idx = i - 1
@@ -837,7 +829,7 @@ def write_cluster_config(
     # Use a tmp file path to avoid incomplete YAML file being re-used in the
     # future.
     tmp_yaml_path = yaml_path + '.tmp'
-    tmp_yaml_path = fill_template(
+    fill_template(
         cluster_config_template,
         dict(
             resources_vars,
@@ -946,7 +938,10 @@ def write_cluster_config(
         config = common_utils.read_yaml(os.path.expanduser(config_dict['ray']))
         vpc_name = gcp_config.get_usable_vpc(config)
 
-        scripts = tuple(
+        scripts = []
+        for template_name in ('gcp-tpu-create.sh.j2', 'gcp-tpu-delete.sh.j2'):
+            script_path = os.path.join(user_file_dir, template_name).replace(
+                '.sh.j2', f'.{cluster_name}.sh')
             fill_template(
                 template_name,
                 dict(
@@ -958,10 +953,10 @@ def write_cluster_config(
                 # Use new names for TPU scripts so that different runs can use
                 # different TPUs.  Put in SKY_USER_FILE_PATH to be consistent
                 # with cluster yamls.
-                output_path=os.path.join(user_file_dir, template_name).replace(
-                    '.sh.j2', f'.{cluster_name}.sh'),
-            ) for template_name in
-            ['gcp-tpu-create.sh.j2', 'gcp-tpu-delete.sh.j2'])
+                output_path=script_path,
+            )
+            scripts.append(script_path)
+
         config_dict['tpu-create-script'] = scripts[0]
         config_dict['tpu-delete-script'] = scripts[1]
         config_dict['tpu_name'] = tpu_name
@@ -1001,6 +996,37 @@ def get_timestamp_from_run_timestamp(run_timestamp: str) -> float:
         run_timestamp.partition('-')[2], '%Y-%m-%d-%H-%M-%S-%f').timestamp()
 
 
+def _count_healthy_nodes_from_ray(output: str,
+                                  is_local_cloud: bool = False
+                                 ) -> Tuple[int, int]:
+    """Count the number of healthy nodes from the output of `ray status`."""
+
+    def get_ready_nodes(pattern, output):
+        result = pattern.findall(output)
+        # On-prem/local case is handled differently.
+        # `ray status` produces different output for local case, and
+        # we poll for number of nodes launched instead of counting for
+        # head and number of worker nodes separately (it is impossible
+        # to distinguish between head and worker node for local case).
+        if is_local_cloud:
+            # In the local case, ready_workers mean the total number
+            # of nodes launched, including head.
+            return len(result)
+        if len(result) == 0:
+            return 0
+        assert len(result) == 1, result
+        return int(result[0])
+
+    if is_local_cloud:
+        ready_head = 0
+        ready_workers = get_ready_nodes(_LAUNCHED_LOCAL_WORKER_PATTERN, output)
+    else:
+        ready_head = get_ready_nodes(_LAUNCHED_HEAD_PATTERN, output)
+        ready_workers = get_ready_nodes(_LAUNCHED_WORKER_PATTERN, output)
+    assert ready_head <= 1, f'#head node should be <=1 (Got {ready_head}).'
+    return ready_head, ready_workers
+
+
 @timeline.event
 def wait_until_ray_cluster_ready(
     cluster_config_file: str,
@@ -1040,32 +1066,8 @@ def wait_until_ray_cluster_ready(
                 stderr)
             logger.debug(output)
 
-            # Workers that are ready
-            ready_workers = 0
-            # On-prem/local case is handled differently.
-            # `ray status` produces different output for local case, and
-            # we poll for number of nodes launched instead of counting for
-            # head and number of worker nodes separately (it is impossible
-            # to distinguish between head and worker node for local case).
-            if is_local_cloud:
-                result = _LAUNCHED_LOCAL_WORKER_PATTERN.findall(output)
-                # In the local case, ready_workers mean the total number
-                # of nodes launched, including head.
-                ready_workers = len(result)
-            else:
-                result = _LAUNCHED_WORKER_PATTERN.findall(output)
-                if len(result) == 0:
-                    ready_workers = 0
-                else:
-                    assert len(result) == 1, result
-                    ready_workers = int(result[0])
-
-            result = _LAUNCHED_HEAD_PATTERN.findall(output)
-            ready_head = 0
-            if result:
-                assert len(result) == 1, result
-                ready_head = int(result[0])
-                assert ready_head <= 1, ready_head
+            ready_head, ready_workers = _count_healthy_nodes_from_ray(
+                output, is_local_cloud=is_local_cloud)
 
             worker_status.update('[bold cyan]'
                                  f'{ready_workers} out of {num_nodes - 1} '
@@ -1450,16 +1452,16 @@ def _get_tpu_vm_pod_ips(ray_config: Dict[str, Any],
                          'Hint: make sure it is not leaked.')
             continue
 
-        if not get_internal_ips:
-            ips = [
-                endpoint['accessConfig']['externalIp']
-                for endpoint in tpuvm_json['networkEndpoints']
-            ]
-        else:
-            ips = [
-                endpoint['ipAddress']
-                for endpoint in tpuvm_json['networkEndpoints']
-            ]
+        ips = []
+        for endpoint in tpuvm_json['networkEndpoints']:
+            # Note: if TPU VM is being preempted, its IP field may not exist.
+            # We use get() to avoid KeyError.
+            if get_internal_ips:
+                ip = endpoint.get('ipAddress', None)
+            else:
+                ip = endpoint['accessConfig'].get('externalIp', None)
+            if ip is not None:
+                ips.append(ip)
         all_ips.extend(ips)
 
     return all_ips
@@ -1767,6 +1769,8 @@ def _update_cluster_status_no_lock(
         return record
 
     cluster_name = handle.cluster_name
+    use_spot = handle.launched_resources.use_spot
+    ray_cluster_up = False
     try:
         # TODO(zhwu): This function cannot distinguish transient network error
         # in ray's get IPs vs. ray runtime failing.
@@ -1775,20 +1779,50 @@ def _update_cluster_status_no_lock(
         if external_ips is None or len(external_ips) == 0:
             raise exceptions.FetchIPError(
                 reason=exceptions.FetchIPError.Reason.HEAD)
-        if handle.launched_nodes == 1:
-            # Check the ray cluster status. We have to check it for single node
-            # case, since the get_node_ips() does not require ray cluster to be
-            # running.
-            ssh_credentials = ssh_credential_from_yaml(handle.cluster_yaml)
-            runner = command_runner.SSHCommandRunner(external_ips[0],
-                                                     **ssh_credentials)
-            returncode = runner.run('ray status', stream_logs=False)
-            if returncode:
-                raise exceptions.FetchIPError(
-                    reason=exceptions.FetchIPError.Reason.HEAD)
-        # If we get node ips correctly, the cluster is UP. It is safe to
-        # set the status to UP, as the `handle.external_ips` function uses ray
-        # to fetch IPs and starting ray is the final step of sky launch.
+        # Check if ray cluster status is healthy.
+        ssh_credentials = ssh_credential_from_yaml(handle.cluster_yaml)
+        runner = command_runner.SSHCommandRunner(external_ips[0],
+                                                 **ssh_credentials)
+        rc, output, _ = runner.run('ray status',
+                                   stream_logs=False,
+                                   require_outputs=True,
+                                   separate_stderr=True)
+        if rc:
+            raise exceptions.FetchIPError(
+                reason=exceptions.FetchIPError.Reason.HEAD)
+
+        ready_head, ready_workers = _count_healthy_nodes_from_ray(output)
+
+        if ready_head + ready_workers == handle.launched_nodes:
+            ray_cluster_up = True
+
+        # For non-spot clusters:
+        # If ray status shows all nodes are healthy, it is safe to set
+        # the status to UP as starting ray is the final step of sky launch.
+        # For spot clusters, the above can be unsafe because the Ray cluster
+        # may remain healthy for a while before the cloud completely
+        # preempts the VMs.
+        # Additionally, we query the VM state from the cloud provider.
+        if ray_cluster_up and not use_spot:
+            record['status'] = global_user_state.ClusterStatus.UP
+            global_user_state.add_or_update_cluster(cluster_name,
+                                                    handle,
+                                                    requested_resources=None,
+                                                    ready=True,
+                                                    is_launch=False)
+            return record
+    except exceptions.FetchIPError:
+        logger.debug('Refreshing status: Failed to get IPs from cluster '
+                     f'{cluster_name!r}, trying to fetch from provider.')
+    # For all code below, we query cluster status by cloud CLI for two cases:
+    # 1) ray fails to get IPs for the cluster.
+    # 2) the cluster is a spot cluster.
+    node_statuses = _get_cluster_status_via_cloud_cli(handle)
+
+    all_nodes_up = (all(status == global_user_state.ClusterStatus.UP
+                        for status in node_statuses) and
+                    len(node_statuses) == handle.launched_nodes)
+    if ray_cluster_up and all_nodes_up:
         record['status'] = global_user_state.ClusterStatus.UP
         global_user_state.add_or_update_cluster(cluster_name,
                                                 handle,
@@ -1796,11 +1830,6 @@ def _update_cluster_status_no_lock(
                                                 ready=True,
                                                 is_launch=False)
         return record
-    except exceptions.FetchIPError:
-        logger.debug('Refreshing status: Failed to get IPs from cluster '
-                     f'{cluster_name!r}, trying to fetch from provider.')
-    # For all code below, ray fails to get IPs for the cluster.
-    node_statuses = _get_cluster_status_via_cloud_cli(handle)
 
     if len(node_statuses) > handle.launched_nodes:
         # Unexpected: in the queried region more than 1 cluster with the same
diff --git a/sky/benchmark/benchmark_utils.py b/sky/benchmark/benchmark_utils.py
index f2b71a62f..2522ea88c 100644
--- a/sky/benchmark/benchmark_utils.py
+++ b/sky/benchmark/benchmark_utils.py
@@ -119,6 +119,7 @@ def _print_candidate_resources(
         '# NODES',
         'INSTANCE',
         'vCPUs',
+        'Mem(GB)',
         'ACCELERATORS',
         'PRICE ($/hr)',
     ]
@@ -136,18 +137,24 @@ def _print_candidate_resources(
             accelerator, count = list(resources.accelerators.items())[0]
             accelerators = f'{accelerator}:{count}'
         cloud = resources.cloud
-        vcpus = cloud.get_vcpus_from_instance_type(resources.instance_type)
-        if vcpus is None:
-            vcpus = '-'
-        elif vcpus.is_integer():
-            vcpus = str(int(vcpus))
-        else:
-            vcpus = f'{vcpus:.1f}'
+        vcpus, mem = cloud.get_vcpus_mem_from_instance_type(
+            resources.instance_type)
+
+        def format_number(x):
+            if x is None:
+                return '-'
+            elif x.is_integer():
+                return str(int(x))
+            else:
+                return f'{x:.1f}'
+
+        vcpus = format_number(vcpus)
+        mem = format_number(mem)
         cost = num_nodes * resources.get_cost(3600)
         spot = '[Spot]' if resources.use_spot else ''
         row = [
             cluster, cloud, num_nodes, resources.instance_type + spot, vcpus,
-            accelerators, f'{cost:.2f}'
+            mem, accelerators, f'{cost:.2f}'
         ]
         candidate_table.add_row(row)
     logger.info(f'{candidate_table}\n')
diff --git a/sky/cli.py b/sky/cli.py
index 0c8c0d586..3f5b73862 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -186,6 +186,14 @@ def _interactive_node_cli_command(cli_func):
         help=('Number of vCPUs each instance must have '
               '(e.g., ``--cpus=4`` (exactly 4) or ``--cpus=4+`` (at least 4)). '
               'This is used to automatically select the instance type.'))
+    memory = click.option(
+        '--memory',
+        default=None,
+        type=str,
+        required=False,
+        help=('Amount of memory each instance must have in GB (e.g., '
+              '``--memory=16`` (exactly 16GB), ``--memory=16+`` (at least '
+              '16GB))'))
     gpus = click.option('--gpus',
                         default=None,
                         type=str,
@@ -281,6 +289,7 @@ def _interactive_node_cli_command(cli_func):
         zone_option,
         instance_type_option,
         cpus,
+        memory,
         *([gpus] if cli_func.__name__ == 'gpunode' else []),
         *([tpus] if cli_func.__name__ == 'tpunode' else []),
         spot_option,
@@ -575,6 +584,7 @@ def _parse_override_params(cloud: Optional[str] = None,
                            zone: Optional[str] = None,
                            gpus: Optional[str] = None,
                            cpus: Optional[str] = None,
+                           memory: Optional[str] = None,
                            instance_type: Optional[str] = None,
                            use_spot: Optional[bool] = None,
                            image_id: Optional[str] = None,
@@ -606,6 +616,11 @@ def _parse_override_params(cloud: Optional[str] = None,
             override_params['cpus'] = None
         else:
             override_params['cpus'] = cpus
+    if memory is not None:
+        if memory.lower() == 'none':
+            override_params['memory'] = None
+        else:
+            override_params['memory'] = memory
     if instance_type is not None:
         if instance_type.lower() == 'none':
             override_params['instance_type'] = None
@@ -936,6 +951,7 @@ def _make_task_from_entrypoint_with_overrides(
     zone: Optional[str] = None,
     gpus: Optional[str] = None,
     cpus: Optional[str] = None,
+    memory: Optional[str] = None,
     instance_type: Optional[str] = None,
     num_nodes: Optional[int] = None,
     use_spot: Optional[bool] = None,
@@ -979,6 +995,7 @@ def _make_task_from_entrypoint_with_overrides(
                                              zone=zone,
                                              gpus=gpus,
                                              cpus=cpus,
+                                             memory=memory,
                                              instance_type=instance_type,
                                              use_spot=use_spot,
                                              image_id=image_id,
@@ -1127,6 +1144,13 @@ def cli():
               help=('Number of vCPUs each instance must have (e.g., '
                     '``--cpus=4`` (exactly 4) or ``--cpus=4+`` (at least 4)). '
                     'This is used to automatically select the instance type.'))
+@click.option(
+    '--memory',
+    default=None,
+    type=str,
+    required=False,
+    help=('Amount of memory each instance must have in GB (e.g., '
+          '``--memory=16`` (exactly 16GB), ``--memory=16+`` (at least 16GB))'))
 @click.option('--disk-size',
               default=None,
               type=int,
@@ -1192,6 +1216,7 @@ def launch(
     zone: Optional[str],
     gpus: Optional[str],
     cpus: Optional[str],
+    memory: Optional[str],
     instance_type: Optional[str],
     num_nodes: Optional[int],
     use_spot: Optional[bool],
@@ -1237,6 +1262,7 @@ def launch(
         zone=zone,
         gpus=gpus,
         cpus=cpus,
+        memory=memory,
         instance_type=instance_type,
         num_nodes=num_nodes,
         use_spot=use_spot,
@@ -1383,6 +1409,7 @@ def exec(
         zone=zone,
         gpus=gpus,
         cpus=None,
+        memory=None,
         instance_type=instance_type,
         use_spot=use_spot,
         image_id=image_id,
@@ -1662,6 +1689,11 @@ def cost_report(all: bool):  # pylint: disable=redefined-builtin
     for cluster_group_name, cluster_record in reserved_clusters.items():
         status_utils.show_cost_report_table(
             [cluster_record], all, reserved_group_name=cluster_group_name)
+    click.secho(
+        'NOTE: This feature is experimental. '
+        'Costs for clusters with auto{stop,down} '
+        'scheduled may not be accurate.',
+        fg='yellow')
 
 
 @cli.command()
@@ -2592,10 +2624,11 @@ def _down_or_stop_clusters(
 def gpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
             cloud: Optional[str], region: Optional[str], zone: Optional[str],
             instance_type: Optional[str], cpus: Optional[str],
-            gpus: Optional[str], use_spot: Optional[bool],
-            screen: Optional[bool], tmux: Optional[bool],
-            disk_size: Optional[int], idle_minutes_to_autostop: Optional[int],
-            down: bool, retry_until_up: bool):
+            memory: Optional[str], gpus: Optional[str],
+            use_spot: Optional[bool], screen: Optional[bool],
+            tmux: Optional[bool], disk_size: Optional[int],
+            idle_minutes_to_autostop: Optional[int], down: bool,
+            retry_until_up: bool):
     """Launch or attach to an interactive GPU node.
 
     Examples:
@@ -2634,8 +2667,8 @@ def gpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
 
     user_requested_resources = not (cloud is None and region is None and
                                     zone is None and instance_type is None and
-                                    cpus is None and gpus is None and
-                                    use_spot is None)
+                                    cpus is None and memory is None and
+                                    gpus is None and use_spot is None)
     default_resources = _INTERACTIVE_NODE_DEFAULT_RESOURCES['gpunode']
     cloud_provider = clouds.CLOUD_REGISTRY.from_str(cloud)
     if gpus is None and instance_type is None:
@@ -2649,6 +2682,7 @@ def gpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
                               zone=zone,
                               instance_type=instance_type,
                               cpus=cpus,
+                              memory=memory,
                               accelerators=gpus,
                               use_spot=use_spot,
                               disk_size=disk_size)
@@ -2673,10 +2707,10 @@ def gpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
 def cpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
             cloud: Optional[str], region: Optional[str], zone: Optional[str],
             instance_type: Optional[str], cpus: Optional[str],
-            use_spot: Optional[bool], screen: Optional[bool],
-            tmux: Optional[bool], disk_size: Optional[int],
-            idle_minutes_to_autostop: Optional[int], down: bool,
-            retry_until_up: bool):
+            memory: Optional[str], use_spot: Optional[bool],
+            screen: Optional[bool], tmux: Optional[bool],
+            disk_size: Optional[int], idle_minutes_to_autostop: Optional[int],
+            down: bool, retry_until_up: bool):
     """Launch or attach to an interactive CPU node.
 
     Examples:
@@ -2714,7 +2748,8 @@ def cpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
 
     user_requested_resources = not (cloud is None and region is None and
                                     zone is None and instance_type is None and
-                                    cpus is None and use_spot is None)
+                                    cpus is None and memory is None and
+                                    use_spot is None)
     default_resources = _INTERACTIVE_NODE_DEFAULT_RESOURCES['cpunode']
     cloud_provider = clouds.CLOUD_REGISTRY.from_str(cloud)
     if instance_type is None:
@@ -2726,6 +2761,7 @@ def cpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
                               zone=zone,
                               instance_type=instance_type,
                               cpus=cpus,
+                              memory=memory,
                               use_spot=use_spot,
                               disk_size=disk_size)
 
@@ -2749,11 +2785,11 @@ def cpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
 def tpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
             region: Optional[str], zone: Optional[str],
             instance_type: Optional[str], cpus: Optional[str],
-            tpus: Optional[str], use_spot: Optional[bool],
-            tpu_vm: Optional[bool], screen: Optional[bool],
-            tmux: Optional[bool], disk_size: Optional[int],
-            idle_minutes_to_autostop: Optional[int], down: bool,
-            retry_until_up: bool):
+            memory: Optional[str], tpus: Optional[str],
+            use_spot: Optional[bool], tpu_vm: Optional[bool],
+            screen: Optional[bool], tmux: Optional[bool],
+            disk_size: Optional[int], idle_minutes_to_autostop: Optional[int],
+            down: bool, retry_until_up: bool):
     """Launch or attach to an interactive TPU node.
 
     Examples:
@@ -2791,7 +2827,8 @@ def tpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
 
     user_requested_resources = not (region is None and zone is None and
                                     instance_type is None and cpus is None and
-                                    tpus is None and use_spot is None)
+                                    memory is None and tpus is None and
+                                    use_spot is None)
     default_resources = _INTERACTIVE_NODE_DEFAULT_RESOURCES['tpunode']
     accelerator_args = default_resources.accelerator_args
     if tpu_vm:
@@ -2808,6 +2845,7 @@ def tpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
                               zone=zone,
                               instance_type=instance_type,
                               cpus=cpus,
+                              memory=memory,
                               accelerators=tpus,
                               accelerator_args=accelerator_args,
                               use_spot=use_spot,
@@ -3159,6 +3197,13 @@ def spot():
               help=('Number of vCPUs each instance must have (e.g., '
                     '``--cpus=4`` (exactly 4) or ``--cpus=4+`` (at least 4)). '
                     'This is used to automatically select the instance type.'))
+@click.option(
+    '--memory',
+    default=None,
+    type=str,
+    required=False,
+    help=('Amount of memory each instance must have in GB (e.g., '
+          '``--memory=16`` (exactly 16GB), ``--memory=16+`` (at least 16GB))'))
 @click.option('--spot-recovery',
               default=None,
               type=str,
@@ -3202,6 +3247,7 @@ def spot_launch(
     zone: Optional[str],
     gpus: Optional[str],
     cpus: Optional[str],
+    memory: Optional[str],
     instance_type: Optional[str],
     num_nodes: Optional[int],
     use_spot: Optional[bool],
@@ -3239,6 +3285,7 @@ def spot_launch(
         zone=zone,
         gpus=gpus,
         cpus=cpus,
+        memory=memory,
         instance_type=instance_type,
         num_nodes=num_nodes,
         use_spot=use_spot,
@@ -3599,6 +3646,7 @@ def benchmark_launch(
     # The user can specify the benchmark candidates in either of the two ways:
     # 1. By specifying resources.candidates in the YAML.
     # 2. By specifying gpu types as a command line argument (--gpus).
+    override_gpu = None
     if gpus is not None:
         gpu_list = gpus.split(',')
         gpu_list = [gpu.strip() for gpu in gpu_list]
diff --git a/sky/clouds/aws.py b/sky/clouds/aws.py
index 348792e16..b29d1380e 100644
--- a/sky/clouds/aws.py
+++ b/sky/clouds/aws.py
@@ -266,9 +266,12 @@ class AWS(clouds.Cloud):
         return isinstance(other, AWS)
 
     @classmethod
-    def get_default_instance_type(cls,
-                                  cpus: Optional[str] = None) -> Optional[str]:
+    def get_default_instance_type(
+            cls,
+            cpus: Optional[str] = None,
+            memory: Optional[str] = None) -> Optional[str]:
         return service_catalog.get_default_instance_type(cpus=cpus,
+                                                         memory=memory,
                                                          clouds='aws')
 
     # TODO: factor the following three methods, as they are the same logic
@@ -282,12 +285,12 @@ class AWS(clouds.Cloud):
             instance_type, clouds='aws')
 
     @classmethod
-    def get_vcpus_from_instance_type(
+    def get_vcpus_mem_from_instance_type(
         cls,
         instance_type: str,
-    ) -> Optional[float]:
-        return service_catalog.get_vcpus_from_instance_type(instance_type,
-                                                            clouds='aws')
+    ) -> Tuple[Optional[float], Optional[float]]:
+        return service_catalog.get_vcpus_mem_from_instance_type(instance_type,
+                                                                clouds='aws')
 
     def make_deploy_resources_variables(
             self, resources: 'resources_lib.Resources', region: 'clouds.Region',
@@ -334,6 +337,7 @@ class AWS(clouds.Cloud):
                     # attach the accelerators.  Billed as part of the VM type.
                     accelerators=None,
                     cpus=None,
+                    memory=None,
                 )
                 resource_list.append(r)
             return resource_list
@@ -343,7 +347,7 @@ class AWS(clouds.Cloud):
         if accelerators is None:
             # Return a default instance type with the given number of vCPUs.
             default_instance_type = AWS.get_default_instance_type(
-                cpus=resources.cpus)
+                cpus=resources.cpus, memory=resources.memory)
             if default_instance_type is None:
                 return ([], [])
             else:
@@ -357,6 +361,7 @@ class AWS(clouds.Cloud):
             acc_count,
             use_spot=resources.use_spot,
             cpus=resources.cpus,
+            memory=resources.memory,
             region=resources.region,
             zone=resources.zone,
             clouds='aws')
diff --git a/sky/clouds/azure.py b/sky/clouds/azure.py
index 90eaf379d..e2e69768e 100644
--- a/sky/clouds/azure.py
+++ b/sky/clouds/azure.py
@@ -6,6 +6,8 @@ import subprocess
 import typing
 from typing import Dict, Iterator, List, Optional, Tuple
 
+import colorama
+
 from sky import clouds
 from sky import exceptions
 from sky import sky_logging
@@ -109,9 +111,12 @@ class Azure(clouds.Cloud):
         return isinstance(other, Azure)
 
     @classmethod
-    def get_default_instance_type(cls,
-                                  cpus: Optional[str] = None) -> Optional[str]:
+    def get_default_instance_type(
+            cls,
+            cpus: Optional[str] = None,
+            memory: Optional[str] = None) -> Optional[str]:
         return service_catalog.get_default_instance_type(cpus=cpus,
+                                                         memory=memory,
                                                          clouds='azure')
 
     def _get_image_config(self, gen_version, instance_type):
@@ -194,12 +199,12 @@ class Azure(clouds.Cloud):
             instance_type, clouds='azure')
 
     @classmethod
-    def get_vcpus_from_instance_type(
+    def get_vcpus_mem_from_instance_type(
         cls,
         instance_type: str,
-    ) -> Optional[float]:
-        return service_catalog.get_vcpus_from_instance_type(instance_type,
-                                                            clouds='azure')
+    ) -> Tuple[Optional[float], Optional[float]]:
+        return service_catalog.get_vcpus_mem_from_instance_type(instance_type,
+                                                                clouds='azure')
 
     @classmethod
     def get_zone_shell_cmd(cls) -> Optional[str]:
@@ -256,6 +261,7 @@ class Azure(clouds.Cloud):
                     # attach the accelerators.  Billed as part of the VM type.
                     accelerators=None,
                     cpus=None,
+                    memory=None,
                 )
                 resource_list.append(r)
             return resource_list
@@ -265,7 +271,7 @@ class Azure(clouds.Cloud):
         if accelerators is None:
             # Return a default instance type with the given number of vCPUs.
             default_instance_type = Azure.get_default_instance_type(
-                cpus=resources.cpus)
+                cpus=resources.cpus, memory=resources.memory)
             if default_instance_type is None:
                 return ([], [])
             else:
@@ -278,6 +284,7 @@ class Azure(clouds.Cloud):
             acc,
             acc_count,
             cpus=resources.cpus,
+            memory=resources.memory,
             use_spot=resources.use_spot,
             region=resources.region,
             zone=resources.zone,
@@ -349,6 +356,14 @@ class Azure(clouds.Cloud):
             retry_cnt += 1
             try:
                 import knack  # pylint: disable=import-outside-toplevel
+            except ModuleNotFoundError as e:
+                with ux_utils.print_exception_no_traceback():
+                    raise exceptions.CloudUserIdentityError(
+                        'Failed to import \'knack\'. To install the dependencies for Azure, '
+                        'Please install SkyPilot with: '
+                        f'{colorama.Style.BRIGHT}pip install skypilot[azure]'
+                        f'{colorama.Style.RESET_ALL}') from e
+            try:
                 account_email = azure.get_current_account_user()
                 break
             except (FileNotFoundError, knack.util.CLIError) as e:
diff --git a/sky/clouds/cloud.py b/sky/clouds/cloud.py
index a70a5ad84..54ae403db 100644
--- a/sky/clouds/cloud.py
+++ b/sky/clouds/cloud.py
@@ -234,9 +234,9 @@ class Cloud:
         raise NotImplementedError
 
     @classmethod
-    def get_vcpus_from_instance_type(cls,
-                                     instance_type: str) -> Optional[float]:
-        """Returns the number of virtual CPUs that the instance type offers."""
+    def get_vcpus_mem_from_instance_type(
+            cls, instance_type: str) -> Tuple[Optional[float], Optional[float]]:
+        """Returns the #vCPUs and memory that the instance type offers."""
         raise NotImplementedError
 
     @classmethod
@@ -248,15 +248,21 @@ class Cloud:
         raise NotImplementedError
 
     @classmethod
-    def get_default_instance_type(cls,
-                                  cpus: Optional[str] = None) -> Optional[str]:
-        """Returns the default instance type with the given number of vCPUs.
+    def get_default_instance_type(
+            cls,
+            cpus: Optional[str] = None,
+            memory: Optional[str] = None) -> Optional[str]:
+        """Returns the default instance type with the given #vCPUs and memory.
 
         For example, if cpus='4', this method returns the default instance type
         with 4 vCPUs.  If cpus='4+', this method returns the default instance
         type with 4 or more vCPUs.
 
-        When cpus is None, this method will never return None.
+        If 'memory=4', this method returns the default instance type with 4GB
+        memory.  If 'memory=4+', this method returns the default instance
+        type with 4GB or more memory.
+
+        When cpus is None or memory is None, this method will never return None.
         This method may return None if the cloud's default instance family
         does not have a VM with the given number of vCPUs (e.g., when cpus='7').
         """
diff --git a/sky/clouds/gcp.py b/sky/clouds/gcp.py
index 63738fcd6..87bea4b3d 100644
--- a/sky/clouds/gcp.py
+++ b/sky/clouds/gcp.py
@@ -31,6 +31,14 @@ GCP_CONFIG_PATH = '~/.config/gcloud/configurations/config_default'
 # autoscaler can overwrite that directory on the remote nodes.
 GCP_CONFIG_SKY_BACKUP_PATH = '~/.sky/.sky_gcp_config_default'
 
+# A list of permissions required to run SkyPilot on GCP.
+# This is not a complete list but still useful to check first
+# and hint users if not sufficient during sky check.
+GCP_PREMISSION_CHECK_LIST = [
+    'compute.projects.get',
+    'iam.serviceAccounts.actAs',
+]
+
 # Minimum set of files under ~/.config/gcloud that grant GCP access.
 _CREDENTIAL_FILES = [
     'credentials.db',
@@ -258,9 +266,12 @@ class GCP(clouds.Cloud):
             raise
 
     @classmethod
-    def get_default_instance_type(cls,
-                                  cpus: Optional[str] = None) -> Optional[str]:
+    def get_default_instance_type(
+            cls,
+            cpus: Optional[str] = None,
+            memory: Optional[str] = None) -> Optional[str]:
         return service_catalog.get_default_instance_type(cpus=cpus,
+                                                         memory=memory,
                                                          clouds='gcp')
 
     def make_deploy_resources_variables(
@@ -350,7 +361,8 @@ class GCP(clouds.Cloud):
 
         if resources.accelerators is None:
             # Return a default instance type with the given number of vCPUs.
-            host_vm_type = GCP.get_default_instance_type(cpus=resources.cpus)
+            host_vm_type = GCP.get_default_instance_type(
+                cpus=resources.cpus, memory=resources.memory)
             if host_vm_type is None:
                 return ([], [])
             else:
@@ -359,6 +371,7 @@ class GCP(clouds.Cloud):
                     instance_type=host_vm_type,
                     accelerators=None,
                     cpus=None,
+                    memory=None,
                 )
                 return ([r], [])
 
@@ -378,6 +391,7 @@ class GCP(clouds.Cloud):
             acc,
             acc_count,
             cpus=resources.cpus if not use_tpu_vm else None,
+            memory=resources.memory if not use_tpu_vm else None,
             use_spot=resources.use_spot,
             region=resources.region,
             zone=resources.zone,
@@ -391,8 +405,10 @@ class GCP(clouds.Cloud):
 
         if use_tpu_vm:
             host_vm_type = 'TPU-VM'
-            # FIXME(woosuk): This leverages the fact that TPU VMs have 96 vCPUs.
-            num_cpus_in_tpu_vm = 96
+            # FIXME(woosuk, wei-lin): This leverages the fact that TPU VMs
+            # have 96 vCPUs, and 240 vCPUs for tpu-v4. We need to move
+            # this to service catalog, instead.
+            num_cpus_in_tpu_vm = 240 if 'v4' in acc else 96
             if resources.cpus is not None:
                 if resources.cpus.endswith('+'):
                     cpus = float(resources.cpus[:-1])
@@ -402,6 +418,19 @@ class GCP(clouds.Cloud):
                     cpus = float(resources.cpus)
                     if cpus != num_cpus_in_tpu_vm:
                         return ([], fuzzy_candidate_list)
+            # FIXME(woosuk, wei-lin): This leverages the fact that TPU VMs
+            # have 334 GB RAM, and 400 GB RAM for tpu-v4. We need to move
+            # this to service catalog, instead.
+            memory_in_tpu_vm = 400 if 'v4' in acc else 334
+            if resources.memory is not None:
+                if resources.memory.endswith('+'):
+                    memory = float(resources.memory[:-1])
+                    if memory > memory_in_tpu_vm:
+                        return ([], fuzzy_candidate_list)
+                else:
+                    memory = float(resources.memory)
+                    if memory != memory_in_tpu_vm:
+                        return ([], fuzzy_candidate_list)
         else:
             host_vm_type = instance_list[0]
 
@@ -411,6 +440,7 @@ class GCP(clouds.Cloud):
             instance_type=host_vm_type,
             accelerators=acc_dict,
             cpus=None,
+            memory=None,
         )
         return ([r], fuzzy_candidate_list)
 
@@ -424,12 +454,12 @@ class GCP(clouds.Cloud):
         return None
 
     @classmethod
-    def get_vcpus_from_instance_type(
+    def get_vcpus_mem_from_instance_type(
         cls,
         instance_type: str,
-    ) -> Optional[float]:
-        return service_catalog.get_vcpus_from_instance_type(instance_type,
-                                                            clouds='gcp')
+    ) -> Tuple[Optional[float], Optional[float]]:
+        return service_catalog.get_vcpus_mem_from_instance_type(instance_type,
+                                                                clouds='gcp')
 
     @classmethod
     def check_credentials(cls) -> Tuple[bool, Optional[str]]:
@@ -457,7 +487,7 @@ class GCP(clouds.Cloud):
             project_id = cls.get_project_id()
 
             # Check if the user is activated.
-            cls.get_current_user_identity()
+            identity = cls.get_current_user_identity()
         except (auth.exceptions.DefaultCredentialsError,
                 subprocess.CalledProcessError,
                 exceptions.CloudUserIdentityError, FileNotFoundError,
@@ -525,6 +555,27 @@ class GCP(clouds.Cloud):
                   'effect. If any SkyPilot commands/calls failed, retry after '
                   'some time.')
 
+        # pylint: disable=import-outside-toplevel,unused-import
+        import googleapiclient.discovery
+        import google.auth
+
+        # This takes user's credential info from "~/.config/gcloud/application_default_credentials.json".  # pylint: disable=line-too-long
+        credentials, project = google.auth.default()
+        service = googleapiclient.discovery.build('cloudresourcemanager',
+                                                  'v1',
+                                                  credentials=credentials)
+        permissions = {'permissions': GCP_PREMISSION_CHECK_LIST}
+        request = service.projects().testIamPermissions(resource=project,
+                                                        body=permissions)
+        ret_permissions = request.execute().get('permissions', [])
+        if len(ret_permissions) < len(GCP_PREMISSION_CHECK_LIST):
+            diffs = set(GCP_PREMISSION_CHECK_LIST).difference(
+                set(ret_permissions))
+            return False, (
+                'The following permissions are not enabled for the current '
+                f'GCP identity ({identity}):\n    '
+                f'{diffs}\n    '
+                'For more details, visit: https://skypilot.readthedocs.io/en/latest/reference/faq.html#what-are-the-required-iam-permissons-on-gcp-for-skypilot')  # pylint: disable=line-too-long
         return True, None
 
     def get_credential_file_mounts(self) -> Dict[str, str]:
diff --git a/sky/clouds/lambda_cloud.py b/sky/clouds/lambda_cloud.py
index 60151dcd8..b8c96b663 100644
--- a/sky/clouds/lambda_cloud.py
+++ b/sky/clouds/lambda_cloud.py
@@ -112,9 +112,12 @@ class Lambda(clouds.Cloud):
         return isinstance(other, Lambda)
 
     @classmethod
-    def get_default_instance_type(cls,
-                                  cpus: Optional[str] = None) -> Optional[str]:
+    def get_default_instance_type(
+            cls,
+            cpus: Optional[str] = None,
+            memory: Optional[str] = None) -> Optional[str]:
         return service_catalog.get_default_instance_type(cpus=cpus,
+                                                         memory=memory,
                                                          clouds='lambda')
 
     @classmethod
@@ -126,12 +129,12 @@ class Lambda(clouds.Cloud):
             instance_type, clouds='lambda')
 
     @classmethod
-    def get_vcpus_from_instance_type(
+    def get_vcpus_mem_from_instance_type(
         cls,
         instance_type: str,
-    ) -> Optional[float]:
-        return service_catalog.get_vcpus_from_instance_type(instance_type,
-                                                            clouds='lambda')
+    ) -> Tuple[Optional[float], Optional[float]]:
+        return service_catalog.get_vcpus_mem_from_instance_type(instance_type,
+                                                                clouds='lambda')
 
     @classmethod
     def get_zone_shell_cmd(cls) -> Optional[str]:
@@ -175,6 +178,7 @@ class Lambda(clouds.Cloud):
                     # attach the accelerators.  Billed as part of the VM type.
                     accelerators=None,
                     cpus=None,
+                    memory=None,
                 )
                 resource_list.append(r)
             return resource_list
@@ -184,7 +188,7 @@ class Lambda(clouds.Cloud):
         if accelerators is None:
             # Return a default instance type with the given number of vCPUs.
             default_instance_type = Lambda.get_default_instance_type(
-                cpus=resources.cpus)
+                cpus=resources.cpus, memory=resources.memory)
             if default_instance_type is None:
                 return ([], [])
             else:
@@ -198,6 +202,7 @@ class Lambda(clouds.Cloud):
             acc_count,
             use_spot=resources.use_spot,
             cpus=resources.cpus,
+            memory=resources.memory,
             region=resources.region,
             zone=resources.zone,
             clouds='lambda')
diff --git a/sky/clouds/local.py b/sky/clouds/local.py
index 12d96a3ef..2adcc8f78 100644
--- a/sky/clouds/local.py
+++ b/sky/clouds/local.py
@@ -108,15 +108,17 @@ class Local(clouds.Cloud):
         return isinstance(other, Local)
 
     @classmethod
-    def get_default_instance_type(cls, cpus: Optional[str] = None) -> str:
+    def get_default_instance_type(cls,
+                                  cpus: Optional[str] = None,
+                                  memory: Optional[str] = None) -> str:
         # There is only "1" instance type for local cloud: on-prem
-        del cpus  # Unused.
+        del cpus, memory  # Unused.
         return Local._DEFAULT_INSTANCE_TYPE
 
     @classmethod
-    def get_vcpus_from_instance_type(cls,
-                                     instance_type: str) -> Optional[float]:
-        return None
+    def get_vcpus_mem_from_instance_type(
+            cls, instance_type: str) -> Tuple[Optional[float], Optional[float]]:
+        return None, None
 
     @classmethod
     def get_accelerators_from_instance_type(
diff --git a/sky/clouds/service_catalog/__init__.py b/sky/clouds/service_catalog/__init__.py
index b84580e8a..15289adbc 100644
--- a/sky/clouds/service_catalog/__init__.py
+++ b/sky/clouds/service_catalog/__init__.py
@@ -164,22 +164,30 @@ def get_hourly_cost(instance_type: str,
                                use_spot, region, zone)
 
 
-def get_vcpus_from_instance_type(instance_type: str,
-                                 clouds: CloudFilter = None) -> Optional[float]:
+def get_vcpus_mem_from_instance_type(
+        instance_type: str,
+        clouds: CloudFilter = None) -> Tuple[Optional[float], Optional[float]]:
     """Returns the number of virtual CPUs from a instance type."""
-    return _map_clouds_catalog(clouds, 'get_vcpus_from_instance_type',
+    return _map_clouds_catalog(clouds, 'get_vcpus_mem_from_instance_type',
                                instance_type)
 
 
 def get_default_instance_type(cpus: Optional[str] = None,
+                              memory: Optional[str] = None,
                               clouds: CloudFilter = None) -> Optional[str]:
-    """Returns the cloud's default instance type for the given number of vCPUs.
+    """Returns the cloud's default instance type for given #vCPUs and memory.
 
     For example, if cpus='4', this method returns the default instance type
     with 4 vCPUs.  If cpus='4+', this method returns the default instance
     type with 4 or more vCPUs.
+
+    If memory_gb_or_ratio is not specified, this method returns the General
+    Purpose instance type with the given number of vCPUs. If memory_gb_or_ratio
+    is specified, this method returns the cheapest instance type that meets
+    the given CPU and memory requirement.
     """
-    return _map_clouds_catalog(clouds, 'get_default_instance_type', cpus)
+    return _map_clouds_catalog(clouds, 'get_default_instance_type', cpus,
+                               memory)
 
 
 def get_accelerators_from_instance_type(
@@ -194,6 +202,7 @@ def get_instance_type_for_accelerator(
     acc_name: str,
     acc_count: int,
     cpus: Optional[str] = None,
+    memory: Optional[str] = None,
     use_spot: bool = False,
     region: Optional[str] = None,
     zone: Optional[str] = None,
@@ -204,8 +213,8 @@ def get_instance_type_for_accelerator(
     accelerators with sorted prices and a list of candidates with fuzzy search.
     """
     return _map_clouds_catalog(clouds, 'get_instance_type_for_accelerator',
-                               acc_name, acc_count, cpus, use_spot, region,
-                               zone)
+                               acc_name, acc_count, cpus, memory, use_spot,
+                               region, zone)
 
 
 def get_accelerator_hourly_cost(
diff --git a/sky/clouds/service_catalog/aws_catalog.py b/sky/clouds/service_catalog/aws_catalog.py
index 4b0b45329..0639a584a 100644
--- a/sky/clouds/service_catalog/aws_catalog.py
+++ b/sky/clouds/service_catalog/aws_catalog.py
@@ -25,11 +25,23 @@ if typing.TYPE_CHECKING:
 
 logger = sky_logging.init_logger(__name__)
 
-# This is the latest general-purpose instance family as of Jan 2023.
-# CPU: Intel Ice Lake 8375C.
-# Memory: 4 GiB RAM per 1 vCPU.
-_DEFAULT_INSTANCE_FAMILY = 'm6i'
+# We will select from the following three instance families:
+_DEFAULT_INSTANCE_FAMILY = [
+    # This is the latest general-purpose instance family as of Mar 2023.
+    # CPU: Intel Ice Lake 8375C.
+    # Memory: 4 GiB RAM per 1 vCPU;
+    'm6i',
+    # This is the latest memory-optimized instance family as of Mar 2023.
+    # CPU: Intel Ice Lake 8375C
+    # Memory: 8 GiB RAM per 1 vCPU;
+    'r6i',
+    # This is the latest compute-optimized instance family as of Mar 2023.
+    # CPU: Intel Ice Lake 8375C
+    # Memory: 2 GiB RAM per 1 vCPU;
+    'c6i',
+]
 _DEFAULT_NUM_VCPUS = 8
+_DEFAULT_MEMORY_CPU_RATIO = 4
 
 # Keep it synced with the frequency in
 # skypilot-catalog/.github/workflows/update-aws-catalog.yml
@@ -161,17 +173,27 @@ def get_hourly_cost(instance_type: str,
                                        region, zone)
 
 
-def get_vcpus_from_instance_type(instance_type: str) -> Optional[float]:
-    return common.get_vcpus_from_instance_type_impl(_get_df(), instance_type)
+def get_vcpus_mem_from_instance_type(
+        instance_type: str) -> Tuple[Optional[float], Optional[float]]:
+    return common.get_vcpus_mem_from_instance_type_impl(_get_df(),
+                                                        instance_type)
 
 
-def get_default_instance_type(cpus: Optional[str] = None) -> Optional[str]:
-    if cpus is None:
-        cpus = str(_DEFAULT_NUM_VCPUS)
-    instance_type_prefix = f'{_DEFAULT_INSTANCE_FAMILY}.'
+def get_default_instance_type(cpus: Optional[str] = None,
+                              memory: Optional[str] = None) -> Optional[str]:
+    if cpus is None and memory is None:
+        cpus = f'{_DEFAULT_NUM_VCPUS}+'
+
+    if memory is None:
+        memory_gb_or_ratio = f'{_DEFAULT_MEMORY_CPU_RATIO}x'
+    else:
+        memory_gb_or_ratio = memory
+    instance_type_prefix = tuple(
+        f'{family}.' for family in _DEFAULT_INSTANCE_FAMILY)
     df = _get_df()
     df = df[df['InstanceType'].str.startswith(instance_type_prefix)]
-    return common.get_instance_type_for_cpus_impl(df, cpus)
+    return common.get_instance_type_for_cpus_mem_impl(df, cpus,
+                                                      memory_gb_or_ratio)
 
 
 def get_accelerators_from_instance_type(
@@ -184,6 +206,7 @@ def get_instance_type_for_accelerator(
     acc_name: str,
     acc_count: int,
     cpus: Optional[str] = None,
+    memory: Optional[str] = None,
     use_spot: bool = False,
     region: Optional[str] = None,
     zone: Optional[str] = None,
@@ -196,6 +219,7 @@ def get_instance_type_for_accelerator(
                                                          acc_name=acc_name,
                                                          acc_count=acc_count,
                                                          cpus=cpus,
+                                                         memory=memory,
                                                          use_spot=use_spot,
                                                          region=region,
                                                          zone=zone)
diff --git a/sky/clouds/service_catalog/azure_catalog.py b/sky/clouds/service_catalog/azure_catalog.py
index a3428c679..061c8948b 100644
--- a/sky/clouds/service_catalog/azure_catalog.py
+++ b/sky/clouds/service_catalog/azure_catalog.py
@@ -12,11 +12,23 @@ from sky.utils import ux_utils
 
 _df = common.read_catalog('azure/vms.csv')
 
-# This is the latest general-purpose instance family as of Jan 2023.
-# CPU: Intel Ice Lake 8370C.
-# Memory: 4 GiB RAM per 1 vCPU.
-_DEFAULT_INSTANCE_FAMILY = 'D_v5'
+# We will select from the following three instance families:
+_DEFAULT_INSTANCE_FAMILY = [
+    # The latest general-purpose instance family as of Mar. 2023.
+    # CPU: Intel Ice Lake 8370C.
+    # Memory: 4 GiB RAM per 1 vCPU;
+    'D_v5',
+    # The latest memory-optimized instance family as of Mar. 2023.
+    # CPU: Intel Ice Lake 8370C.
+    # Memory: 8 GiB RAM per 1 vCPU.
+    'E_v5',
+    # The latest compute-optimized instance family as of Mar 2023.
+    # CPU: Intel Ice Lake 8370C, Cascade Lake 8272CL, or Skylake 8168.
+    # Memory: 2 GiB RAM per 1 vCPU.
+    'Fs_v2'
+]
 _DEFAULT_NUM_VCPUS = 8
+_DEFAULT_MEMORY_CPU_RATIO = 4
 
 
 def instance_type_exists(instance_type: str) -> bool:
@@ -56,8 +68,9 @@ def get_hourly_cost(instance_type: str,
                                        zone)
 
 
-def get_vcpus_from_instance_type(instance_type: str) -> Optional[float]:
-    return common.get_vcpus_from_instance_type_impl(_df, instance_type)
+def get_vcpus_mem_from_instance_type(
+        instance_type: str) -> Tuple[Optional[float], Optional[float]]:
+    return common.get_vcpus_mem_from_instance_type_impl(_df, instance_type)
 
 
 def _get_instance_family(instance_type: str) -> str:
@@ -83,12 +96,18 @@ def _get_instance_family(instance_type: str) -> str:
     return instance_family
 
 
-def get_default_instance_type(cpus: Optional[str] = None) -> Optional[str]:
-    if cpus is None:
-        cpus = str(_DEFAULT_NUM_VCPUS)
-    df = _df[_df['InstanceType'].apply(_get_instance_family) ==
-             _DEFAULT_INSTANCE_FAMILY]
-    return common.get_instance_type_for_cpus_impl(df, cpus)
+def get_default_instance_type(cpus: Optional[str] = None,
+                              memory: Optional[str] = None) -> Optional[str]:
+    if cpus is None and memory is None:
+        cpus = f'{_DEFAULT_NUM_VCPUS}+'
+    if memory is None:
+        memory_gb_or_ratio = f'{_DEFAULT_MEMORY_CPU_RATIO}x'
+    else:
+        memory_gb_or_ratio = memory
+    df = _df[_df['InstanceType'].apply(_get_instance_family).isin(
+        _DEFAULT_INSTANCE_FAMILY)]
+    return common.get_instance_type_for_cpus_mem_impl(df, cpus,
+                                                      memory_gb_or_ratio)
 
 
 def get_accelerators_from_instance_type(
@@ -100,6 +119,7 @@ def get_instance_type_for_accelerator(
         acc_name: str,
         acc_count: int,
         cpus: Optional[str] = None,
+        memory: Optional[str] = None,
         use_spot: bool = False,
         region: Optional[str] = None,
         zone: Optional[str] = None) -> Tuple[Optional[List[str]], List[str]]:
@@ -114,6 +134,7 @@ def get_instance_type_for_accelerator(
                                                          acc_name=acc_name,
                                                          acc_count=acc_count,
                                                          cpus=cpus,
+                                                         memory=memory,
                                                          use_spot=use_spot,
                                                          region=region,
                                                          zone=zone)
diff --git a/sky/clouds/service_catalog/common.py b/sky/clouds/service_catalog/common.py
index 22588dcef..9f1a69bb1 100644
--- a/sky/clouds/service_catalog/common.py
+++ b/sky/clouds/service_catalog/common.py
@@ -253,10 +253,16 @@ def get_hourly_cost_impl(
     return cheapest[price_str]
 
 
-def get_vcpus_from_instance_type_impl(
+def _get_value(value):
+    if pd.isna(value):
+        return None
+    return float(value)
+
+
+def get_vcpus_mem_from_instance_type_impl(
     df: pd.DataFrame,
     instance_type: str,
-) -> Optional[float]:
+) -> Tuple[Optional[float], Optional[float]]:
     df = _get_instance_type(df, instance_type, None)
     if len(df) == 0:
         with ux_utils.print_exception_no_traceback():
@@ -264,10 +270,15 @@ def get_vcpus_from_instance_type_impl(
     assert len(set(df['vCPUs'])) == 1, ('Cannot determine the number of vCPUs '
                                         f'of the instance type {instance_type}.'
                                         f'\n{df}')
+    assert len(set(
+        df['MemoryGiB'])) == 1, ('Cannot determine the memory size '
+                                 f'of the instance type {instance_type}.'
+                                 f'\n{df}')
+
     vcpus = df['vCPUs'].iloc[0]
-    if pd.isna(vcpus):
-        return None
-    return float(vcpus)
+    mem = df['MemoryGiB'].iloc[0]
+
+    return _get_value(vcpus), _get_value(mem)
 
 
 def _filter_with_cpus(df: pd.DataFrame, cpus: Optional[str]) -> pd.DataFrame:
@@ -293,9 +304,50 @@ def _filter_with_cpus(df: pd.DataFrame, cpus: Optional[str]) -> pd.DataFrame:
         return df[df['vCPUs'] == num_cpus]
 
 
-def get_instance_type_for_cpus_impl(
-        df: pd.DataFrame, cpus: Optional[str] = None) -> Optional[str]:
+def _filter_with_mem(df: pd.DataFrame,
+                     memory_gb_or_ratio: Optional[str]) -> pd.DataFrame:
+    if memory_gb_or_ratio is None:
+        return df
+
+    # The following code is partially redundant with the code in
+    # resources.py::_set_memory() but we add it here for safety.
+    if memory_gb_or_ratio.endswith(('+', 'x')):
+        memory_gb_str = memory_gb_or_ratio[:-1]
+    else:
+        memory_gb_str = memory_gb_or_ratio
+    try:
+        memory = float(memory_gb_str)
+    except ValueError:
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(f'The "memory" field should be either a number or '
+                             'a string "<number>+" or "<number>x". Found: '
+                             f'{memory_gb_or_ratio!r}') from None
+    if memory_gb_or_ratio.endswith('+'):
+        return df[df['MemoryGiB'] >= memory]
+    elif memory_gb_or_ratio.endswith('x'):
+        return df[df['MemoryGiB'] >= df['vCPUs'] * memory]
+    else:
+        return df[df['MemoryGiB'] == memory]
+
+
+def get_instance_type_for_cpus_mem_impl(
+        df: pd.DataFrame, cpus: Optional[str],
+        memory_gb_or_ratio: Optional[str]) -> Optional[str]:
+    """Returns the cheapest instance type that satisfies the requirements.
+
+    Args:
+        df: The catalog cloud catalog data frame.
+        cpus: The number of vCPUs. Can be a number or a string "<number>+". If
+            the string ends with "+", then the returned instance type should
+            have at least the given number of vCPUs.
+        memory_gb_or_ratio: The memory size in GB. Can be a number or a string
+            "<number>+" or "<number>x". If the string ends with "+", then the
+            returned instance type should have at least the given memory size.
+            If the string ends with "x", then the returned instance type should
+            have at least the given number of vCPUs times the given ratio.
+    """
     df = _filter_with_cpus(df, cpus)
+    df = _filter_with_mem(df, memory_gb_or_ratio)
     if df.empty:
         return None
     # Sort by the number of vCPUs and then by the price.
@@ -323,6 +375,7 @@ def get_instance_type_for_accelerator_impl(
     acc_name: str,
     acc_count: int,
     cpus: Optional[str] = None,
+    memory: Optional[str] = None,
     use_spot: bool = False,
     region: Optional[str] = None,
     zone: Optional[str] = None,
@@ -348,6 +401,7 @@ def get_instance_type_for_accelerator_impl(
         return (None, fuzzy_candidate_list)
 
     result = _filter_with_cpus(result, cpus)
+    result = _filter_with_mem(result, memory)
     if region is not None:
         result = result[result['Region'] == region]
     if zone is not None:
diff --git a/sky/clouds/service_catalog/data_fetchers/fetch_azure.py b/sky/clouds/service_catalog/data_fetchers/fetch_azure.py
index 6260c13d4..184e759f0 100644
--- a/sky/clouds/service_catalog/data_fetchers/fetch_azure.py
+++ b/sky/clouds/service_catalog/data_fetchers/fetch_azure.py
@@ -189,7 +189,7 @@ def get_all_regions_instance_types_df(region_set: Set[str]):
         gpu_name = None
         gpu_count = np.nan
         vcpus = np.nan
-        memory_gb = np.nan
+        memory = np.nan
         gen_version = None
         caps = row['capabilities']
         for item in caps:
@@ -201,19 +201,18 @@ def get_all_regions_instance_types_df(region_set: Set[str]):
             elif item['name'] == 'vCPUs':
                 vcpus = float(item['value'])
             elif item['name'] == 'MemoryGB':
-                memory_gb = item['value']
+                memory = item['value']
             elif item['name'] == 'HyperVGenerations':
                 gen_version = item['value']
-        return gpu_name, gpu_count, vcpus, memory_gb, gen_version
+        return gpu_name, gpu_count, vcpus, memory, gen_version
 
     def get_additional_columns(row):
-        gpu_name, gpu_count, vcpus, memory_gb, gen_version = get_capabilities(
-            row)
+        gpu_name, gpu_count, vcpus, memory, gen_version = get_capabilities(row)
         return pd.Series({
             'AcceleratorName': gpu_name,
             'AcceleratorCount': gpu_count,
             'vCPUs': vcpus,
-            'MemoryGiB': memory_gb,
+            'MemoryGiB': memory,
             'GpuInfo': gpu_name,
             'Generation': gen_version,
         })
diff --git a/sky/clouds/service_catalog/gcp_catalog.py b/sky/clouds/service_catalog/gcp_catalog.py
index 01fbea392..7962bcac9 100644
--- a/sky/clouds/service_catalog/gcp_catalog.py
+++ b/sky/clouds/service_catalog/gcp_catalog.py
@@ -25,16 +25,28 @@ _TPU_REGIONS = [
     'asia-east1',
 ]
 
-# Default instance family for CPU-only VMs.
-# This is the latest general-purpose instance family as of Jan 2023.
-# CPU: Intel Ice Lake 8373C or Cascade Lake 6268CL.
-# Memory: 4 GiB RAM per 1 vCPU.
-_DEFAULT_INSTANCE_FAMILY = 'n2-standard'
+# We will select from the following three CPU instance families:
+_DEFAULT_INSTANCE_FAMILY = [
+    # This is the latest general-purpose instance family as of Mar 2023.
+    # CPU: Intel Ice Lake 8373C or Cascade Lake 6268CL.
+    # Memory: 4 GiB RAM per 1 vCPU;
+    'n2-standard',
+    # This is the latest memory-optimized instance family as of Mar 2023.
+    # CPU: Intel Ice Lake 8373C or Cascade Lake 6268CL.
+    # Memory: 8 GiB RAM per 1 vCPU;
+    'n2-highmem',
+    # This is the latest compute-optimized instance family as of Mar 2023.
+    # CPU: Intel Ice Lake 8373C or Cascade Lake 6268CL.
+    # Memory: 1 GiB RAM per 1 vCPU;
+    'n2-highcpu',
+]
 _DEFAULT_NUM_VCPUS = 8
+_DEFAULT_MEMORY_CPU_RATIO = 4
 
 # This can be switched between n1 and n2.
 # n2 is not allowed for launching GPUs.
 _DEFAULT_HOST_VM_FAMILY = 'n1'
+_DEFAULT_GPU_MEMORY_CPU_RATIO = 4
 
 # TODO(zongheng): fix A100 info directly in catalog.
 # https://cloud.google.com/blog/products/compute/a2-vms-with-nvidia-a100-gpus-are-ga
@@ -164,26 +176,36 @@ def get_hourly_cost(
                                        zone)
 
 
-def get_vcpus_from_instance_type(instance_type: str) -> Optional[float]:
-    # The number of vCPUs provided with a TPU VM is not officially documented.
+def get_vcpus_mem_from_instance_type(
+        instance_type: str) -> Tuple[Optional[float], Optional[float]]:
+    # The number of vCPUs and memory size provided with a TPU VM is not
+    # officially documented.
     if instance_type == 'TPU-VM':
-        return None
-    return common.get_vcpus_from_instance_type_impl(_df, instance_type)
+        return None, None
+    return common.get_vcpus_mem_from_instance_type_impl(_df, instance_type)
 
 
-def get_default_instance_type(cpus: Optional[str] = None) -> Optional[str]:
-    if cpus is None:
-        cpus = str(_DEFAULT_NUM_VCPUS)
-    instance_type_prefix = f'{_DEFAULT_INSTANCE_FAMILY}-'
+def get_default_instance_type(cpus: Optional[str] = None,
+                              memory: Optional[str] = None) -> Optional[str]:
+    if cpus is None and memory is None:
+        cpus = f'{_DEFAULT_NUM_VCPUS}+'
+    if memory is None:
+        memory_gb_or_ratio = f'{_DEFAULT_MEMORY_CPU_RATIO}x'
+    else:
+        memory_gb_or_ratio = memory
+    instance_type_prefix = tuple(
+        f'{family}-' for family in _DEFAULT_INSTANCE_FAMILY)
     df = _df[_df['InstanceType'].notna()]
     df = df[df['InstanceType'].str.startswith(instance_type_prefix)]
-    return common.get_instance_type_for_cpus_impl(df, cpus)
+    return common.get_instance_type_for_cpus_mem_impl(df, cpus,
+                                                      memory_gb_or_ratio)
 
 
 def get_instance_type_for_accelerator(
         acc_name: str,
         acc_count: int,
         cpus: Optional[str] = None,
+        memory: Optional[str] = None,
         use_spot: bool = False,
         region: Optional[str] = None,
         zone: Optional[str] = None) -> Tuple[Optional[List[str]], List[str]]:
@@ -195,7 +217,7 @@ def get_instance_type_for_accelerator(
     """
     (instance_list,
      fuzzy_candidate_list) = common.get_instance_type_for_accelerator_impl(
-         _df, acc_name, acc_count, cpus, use_spot, region, zone)
+         _df, acc_name, acc_count, cpus, memory, use_spot, region, zone)
     if instance_list is None:
         return None, fuzzy_candidate_list
 
@@ -203,53 +225,34 @@ def get_instance_type_for_accelerator(
         # If A100 is used, host VM type must be A2.
         # https://cloud.google.com/compute/docs/gpus#a100-gpus
 
-        # FIXME(woosuk): This uses the knowledge that the A2 machines provide
-        # 12 vCPUs per GPU, except for a2-megagpu-16g which has 16 GPUs.
-        if cpus is not None:
-            num_a2_cpus = min(12 * acc_count, 96)
-            if cpus.endswith('+'):
-                if num_a2_cpus < float(cpus[:-1]):
-                    return None, []
-            else:
-                if num_a2_cpus != float(cpus):
-                    return None, []
-        return [_A100_INSTANCE_TYPE_DICTS[acc_name][acc_count]], []
+        df = _df[_df['InstanceType'].notna()]
+        instance_type = _A100_INSTANCE_TYPE_DICTS[acc_name][acc_count]
+        df = df[df['InstanceType'] == instance_type]
+
+        # Check the cpus and memory specified by the user.
+        instance_type = common.get_instance_type_for_cpus_mem_impl(
+            df, cpus, memory)
+        if instance_type is None:
+            return None, []
+        return [instance_type], []
 
     if acc_name not in _NUM_ACC_TO_NUM_CPU:
         acc_name = 'DEFAULT'
 
     assert _DEFAULT_HOST_VM_FAMILY == 'n1'
-    num_cpus = None
-    if cpus is None:
-        num_cpus = _NUM_ACC_TO_NUM_CPU[acc_name].get(acc_count, None)
-    else:
-        # FIXME(woosuk): This uses the knowledge that the N1-highmem machines
-        # have 2, 4, 8, 16, 32, 64, or 96 vCPUs.
-        for num_n1_cpus in [2, 4, 8, 16, 32, 64, 96]:
-            if cpus.endswith('+'):
-                if num_n1_cpus >= float(cpus[:-1]):
-                    num_cpus = num_n1_cpus
-                    break
-            else:
-                if num_n1_cpus == float(cpus):
-                    num_cpus = num_n1_cpus
-                    break
-    if num_cpus is None:
-        return None, []
 
-    mem_type = 'highmem'
-    # patches for the number of cores per GPU, as some of the combinations
-    # are not supported by GCP.
-    if _DEFAULT_HOST_VM_FAMILY == 'n1':
-        if num_cpus < 96:
-            num_cpus = _closest_power_of_two(num_cpus)
-        else:
-            num_cpus = 96
-    else:
-        if num_cpus > 80:
-            mem_type = 'standard'
+    if cpus is None and memory is None:
+        cpus = f'{_NUM_ACC_TO_NUM_CPU[acc_name].get(acc_count, None)}+'
+    if memory is None:
+        memory = f'{_DEFAULT_GPU_MEMORY_CPU_RATIO}x'
+    df = _df[_df['InstanceType'].notna()]
+    df = df[df['InstanceType'].str.startswith(f'{_DEFAULT_HOST_VM_FAMILY}-')]
+
+    instance_type = common.get_instance_type_for_cpus_mem_impl(df, cpus, memory)
     # The fuzzy candidate should have already been fetched in the caller.
-    return [f'{_DEFAULT_HOST_VM_FAMILY}-{mem_type}-{num_cpus}'], []
+    if instance_type is None:
+        return None, []
+    return [instance_type], []
 
 
 def validate_region_zone(
diff --git a/sky/clouds/service_catalog/lambda_catalog.py b/sky/clouds/service_catalog/lambda_catalog.py
index ea04877f5..c1a09b7ee 100644
--- a/sky/clouds/service_catalog/lambda_catalog.py
+++ b/sky/clouds/service_catalog/lambda_catalog.py
@@ -16,6 +16,7 @@ _df = common.read_catalog('lambda/vms.csv')
 
 # Number of vCPUS for gpu_1x_a100_sxm4
 _DEFAULT_NUM_VCPUS = 30
+_DEFAULT_MEMORY_CPU_RATIO = 4
 
 
 def instance_type_exists(instance_type: str) -> bool:
@@ -55,19 +56,28 @@ def get_hourly_cost(instance_type: str,
                                        zone)
 
 
-def get_vcpus_from_instance_type(instance_type: str) -> Optional[float]:
-    return common.get_vcpus_from_instance_type_impl(_df, instance_type)
+def get_vcpus_mem_from_instance_type(
+        instance_type: str) -> Tuple[Optional[float], Optional[float]]:
+    return common.get_vcpus_mem_from_instance_type_impl(_df, instance_type)
 
 
-def get_default_instance_type(cpus: Optional[str] = None) -> Optional[str]:
-    if cpus is None:
-        cpus = str(_DEFAULT_NUM_VCPUS)
+def get_default_instance_type(cpus: Optional[str] = None,
+                              memory: Optional[str] = None) -> Optional[str]:
+    if cpus is None and memory is None:
+        cpus = f'{_DEFAULT_NUM_VCPUS}+'
+    if memory is None:
+        memory_gb_or_ratio = f'{_DEFAULT_MEMORY_CPU_RATIO}x'
+    else:
+        memory_gb_or_ratio = memory
+
     # Set to gpu_1x_a100_sxm4 to be the default instance type if match vCPU
     # requirement.
     df = _df[_df['InstanceType'].eq('gpu_1x_a100_sxm4')]
-    instance = common.get_instance_type_for_cpus_impl(df, cpus)
+    instance = common.get_instance_type_for_cpus_mem_impl(
+        df, cpus, memory_gb_or_ratio)
     if not instance:
-        instance = common.get_instance_type_for_cpus_impl(_df, cpus)
+        instance = common.get_instance_type_for_cpus_mem_impl(
+            _df, cpus, memory_gb_or_ratio)
     return instance
 
 
@@ -80,6 +90,7 @@ def get_instance_type_for_accelerator(
         acc_name: str,
         acc_count: int,
         cpus: Optional[str] = None,
+        memory: Optional[str] = None,
         use_spot: bool = False,
         region: Optional[str] = None,
         zone: Optional[str] = None) -> Tuple[Optional[List[str]], List[str]]:
@@ -94,6 +105,7 @@ def get_instance_type_for_accelerator(
                                                          acc_name=acc_name,
                                                          acc_count=acc_count,
                                                          cpus=cpus,
+                                                         memory=memory,
                                                          use_spot=use_spot,
                                                          region=region,
                                                          zone=zone)
diff --git a/sky/execution.py b/sky/execution.py
index 9bc64d89a..599ea10b7 100644
--- a/sky/execution.py
+++ b/sky/execution.py
@@ -17,6 +17,7 @@ import enum
 import getpass
 import tempfile
 import os
+import uuid
 from typing import Any, Dict, List, Optional, Union
 
 import colorama
@@ -176,7 +177,8 @@ def _execute(
             cluster_name)
         cluster_exists = existing_handle is not None
         # TODO(woosuk): If the cluster exists, print a warning that
-        # `cpus` is not used as a job scheduling constraint, unlike `gpus`.
+        # `cpus` and `memory` are not used as a job scheduling constraint,
+        # unlike `gpus`.
 
     stages = stages if stages is not None else list(Stage)
 
@@ -542,8 +544,7 @@ def spot_launch(
         sky.exceptions.NotSupportedError: the feature is not supported.
     """
     entrypoint = task
-    if name is None:
-        name = backend_utils.generate_cluster_name()
+    task_uuid = str(uuid.uuid4().hex[:4])
 
     dag = _convert_to_dag(entrypoint)
     assert len(dag.tasks) == 1, ('Only one task is allowed in a spot launch.',
@@ -569,6 +570,15 @@ def spot_launch(
 
     task = _maybe_translate_local_file_mounts_and_sync_up(task)
 
+    if name is None:
+        if task.name is not None:
+            name = task.name
+        else:
+            name = backend_utils.generate_cluster_name()
+    # Override the task name with the specified name or generated name, so that
+    # the controller process can retrieve the task name from the task config.
+    task.name = name
+
     with tempfile.NamedTemporaryFile(prefix=f'spot-task-{name}-',
                                      mode='w') as f:
         task_config = task.to_yaml_config()
@@ -580,7 +590,9 @@ def spot_launch(
             'user_yaml_path': f.name,
             'user_config_path': None,
             'spot_controller': controller_name,
-            'cluster_name': name,
+            # Note: actual spot cluster name will be <task_name>-<spot job ID>
+            'task_name': name,
+            'uuid': task_uuid,
             'gcloud_installation_commands': gcp.GCLOUD_INSTALLATION_COMMAND,
             'is_dev': env_options.Options.IS_DEVELOPER.get(),
             'disable_logging': env_options.Options.DISABLE_LOGGING.get(),
@@ -629,10 +641,11 @@ def spot_launch(
                         skypilot_config.ENV_VAR_SKYPILOT_CONFIG,
                 })
 
-        yaml_path = backend_utils.fill_template(
-            spot.SPOT_CONTROLLER_TEMPLATE,
-            vars_to_fill,
-            output_prefix=spot.SPOT_CONTROLLER_YAML_PREFIX)
+        yaml_path = os.path.join(spot.SPOT_CONTROLLER_YAML_PREFIX,
+                                 f'{name}-{task_uuid}.yaml')
+        backend_utils.fill_template(spot.SPOT_CONTROLLER_TEMPLATE,
+                                    vars_to_fill,
+                                    output_path=yaml_path)
         controller_task = task_lib.Task.from_yaml(yaml_path)
         controller_task.spot_task = task
         assert len(controller_task.resources) == 1
diff --git a/sky/optimizer.py b/sky/optimizer.py
index d786e7649..af5bf3195 100644
--- a/sky/optimizer.py
+++ b/sky/optimizer.py
@@ -656,13 +656,20 @@ class Optimizer:
                 accelerators = f'{accelerators}:{count}'
             spot = '[Spot]' if resources.use_spot else ''
             cloud = resources.cloud
-            vcpus = cloud.get_vcpus_from_instance_type(resources.instance_type)
-            if vcpus is None:
-                vcpus = '-'
-            elif vcpus.is_integer():
-                vcpus = str(int(vcpus))
-            else:
-                vcpus = f'{vcpus:.1f}'
+            vcpus, mem = cloud.get_vcpus_mem_from_instance_type(
+                resources.instance_type)
+
+            def format_number(x):
+                if x is None:
+                    return '-'
+                elif x.is_integer():
+                    return str(int(x))
+                else:
+                    return f'{x:.1f}'
+
+            vcpus = format_number(vcpus)
+            mem = format_number(mem)
+
             if resources.zone is None:
                 region_or_zone = resources.region
             else:
@@ -671,13 +678,15 @@ class Optimizer:
                 str(cloud),
                 resources.instance_type + spot,
                 vcpus,
+                mem,
                 str(accelerators),
                 str(region_or_zone),
             ]
 
         # Print the list of resouces that the optimizer considered.
         resource_fields = [
-            'CLOUD', 'INSTANCE', 'vCPUs', 'ACCELERATORS', 'REGION/ZONE'
+            'CLOUD', 'INSTANCE', 'vCPUs', 'Mem(GB)', 'ACCELERATORS',
+            'REGION/ZONE'
         ]
         # Do not print Source or Sink.
         best_plan_rows = [[t, t.num_nodes] + _get_resources_element_list(r)
@@ -972,10 +981,15 @@ def _fill_in_launchable_resources(
                                 f'{colorama.Fore.CYAN}'
                                 f'{sorted(all_fuzzy_candidates)}'
                                 f'{colorama.Style.RESET_ALL}')
-                elif resources.cpus is not None:
-                    logger.info('Try specifying a different CPU count, '
-                                'or add "+" to the end of the CPU count '
-                                'to allow for larger instances.')
+                else:
+                    if resources.cpus is not None:
+                        logger.info('Try specifying a different CPU count, '
+                                    'or add "+" to the end of the CPU count '
+                                    'to allow for larger instances.')
+                    if resources.memory is not None:
+                        logger.info('Try specifying a different memory size, '
+                                    'or add "+" to the end of the memory size '
+                                    'to allow for larger instances.')
 
         launchable[resources] = _filter_out_blocked_launchable_resources(
             launchable[resources], blocked_resources)
diff --git a/sky/resources.py b/sky/resources.py
index e5e2f9451..9be99f7a6 100644
--- a/sky/resources.py
+++ b/sky/resources.py
@@ -44,13 +44,14 @@ class Resources:
     """
     # If any fields changed, increment the version. For backward compatibility,
     # modify the __setstate__ method to handle the old version.
-    _VERSION = 7
+    _VERSION = 8
 
     def __init__(
         self,
         cloud: Optional[clouds.Cloud] = None,
         instance_type: Optional[str] = None,
         cpus: Union[None, int, float, str] = None,
+        memory: Union[None, int, float, str] = None,
         accelerators: Union[None, str, Dict[str, int]] = None,
         accelerator_args: Optional[Dict[str, str]] = None,
         use_spot: Optional[bool] = None,
@@ -97,11 +98,12 @@ class Resources:
                 }
 
         self._set_cpus(cpus)
+        self._set_memory(memory)
         self._set_accelerators(accelerators, accelerator_args)
 
         self._try_validate_local()
         self._try_validate_instance_type()
-        self._try_validate_cpus()
+        self._try_validate_cpus_mem()
         self._try_validate_accelerators()
         self._try_validate_spot()
         self._try_validate_image_id()
@@ -141,6 +143,10 @@ class Resources:
         if self.cpus is not None:
             cpus = f', cpus={self.cpus}'
 
+        memory = ''
+        if self.memory is not None:
+            memory = f', mem={self.memory}'
+
         if isinstance(self.cloud, clouds.Local):
             return f'{self.cloud}({self.accelerators})'
 
@@ -166,7 +172,8 @@ class Resources:
 
         hardware_str = (
             f'{instance_type}{use_spot}'
-            f'{cpus}{accelerators}{accelerator_args}{image_id}{disk_size}')
+            f'{cpus}{memory}{accelerators}{accelerator_args}{image_id}'
+            f'{disk_size}')
         # It may have leading ',' (for example, instance_type not set) or empty
         # spaces.  Remove them.
         while hardware_str and hardware_str[0] in (',', ' '):
@@ -207,6 +214,20 @@ class Resources:
         """
         return self._cpus
 
+    @property
+    def memory(self) -> Optional[str]:
+        """Returns the memory that each instance must have in GB.
+
+        For example, memory='16' means each instance must have exactly 16GB
+        memory; memory='16+' means each instance must have at least 16GB
+        memory.
+
+        (Developer note: The memory field is only used to select the instance
+        type at launch time. Thus, Resources in the backend's ResourceHandle
+        will always have the memory field set to None.)
+        """
+        return self._memory
+
     @property
     def accelerators(self) -> Optional[Dict[str, int]]:
         """Returns the accelerators field directly or by inferring.
@@ -276,6 +297,36 @@ class Resources:
                 raise ValueError(
                     f'The "cpus" field should be positive. Found: {cpus!r}')
 
+    def _set_memory(
+        self,
+        memory: Union[None, int, float, str],
+    ) -> None:
+        if memory is None:
+            self._memory = None
+            return
+
+        self._memory = str(memory)
+        if isinstance(memory, str):
+            if memory.endswith('+'):
+                num_memory_gb = memory[:-1]
+            else:
+                num_memory_gb = memory
+
+            try:
+                memory_gb = float(num_memory_gb)
+            except ValueError:
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError(
+                        f'The "memory" field should be either a number or '
+                        f'a string "<number>+". Found: {memory!r}') from None
+        else:
+            memory_gb = float(memory)
+
+        if memory_gb <= 0:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    f'The "cpus" field should be positive. Found: {memory!r}')
+
     def _set_accelerators(
         self,
         accelerators: Union[None, str, Dict[str, int]],
@@ -447,8 +498,8 @@ class Resources:
                 f'inferred from the instance_type {self.instance_type!r}.')
             self._cloud = valid_clouds[0]
 
-    def _try_validate_cpus(self) -> None:
-        if self.cpus is None:
+    def _try_validate_cpus_mem(self) -> None:
+        if self.cpus is None and self.memory is None:
             return
         if self.instance_type is not None:
             # The assertion should be true because we have already executed
@@ -456,20 +507,36 @@ class Resources:
             # The _try_validate_instance_type() method infers and sets
             # self.cloud if self.instance_type is not None.
             assert self.cloud is not None
-            cpus = self.cloud.get_vcpus_from_instance_type(self.instance_type)
-            if self.cpus.endswith('+'):
-                if cpus < float(self.cpus[:-1]):
+            cpus, mem = self.cloud.get_vcpus_mem_from_instance_type(
+                self.instance_type)
+            if self.cpus is not None:
+                if self.cpus.endswith('+'):
+                    if cpus < float(self.cpus[:-1]):
+                        with ux_utils.print_exception_no_traceback():
+                            raise ValueError(
+                                f'{self.instance_type} does not have enough '
+                                f'vCPUs. {self.instance_type} has {cpus} '
+                                f'vCPUs, but {self.cpus} is requested.')
+                elif cpus != float(self.cpus):
                     with ux_utils.print_exception_no_traceback():
                         raise ValueError(
-                            f'{self.instance_type} does not have enough vCPUs. '
-                            f'{self.instance_type} has {cpus} vCPUs, '
-                            f'but {self.cpus} is requested.')
-            elif cpus != float(self.cpus):
-                with ux_utils.print_exception_no_traceback():
-                    raise ValueError(
-                        f'{self.instance_type} does not have the requested '
-                        f'number of vCPUs. {self.instance_type} has {cpus} '
-                        f'vCPUs, but {self.cpus} is requested.')
+                            f'{self.instance_type} does not have the requested '
+                            f'number of vCPUs. {self.instance_type} has {cpus} '
+                            f'vCPUs, but {self.cpus} is requested.')
+            if self.memory is not None:
+                if self.memory.endswith('+'):
+                    if mem < float(self.memory[:-1]):
+                        with ux_utils.print_exception_no_traceback():
+                            raise ValueError(
+                                f'{self.instance_type} does not have enough '
+                                f'memory. {self.instance_type} has {mem} GB '
+                                f'memory, but {self.memory} is requested.')
+                elif mem != float(self.memory):
+                    with ux_utils.print_exception_no_traceback():
+                        raise ValueError(
+                            f'{self.instance_type} does not have the requested '
+                            f'memory. {self.instance_type} has {mem} GB '
+                            f'memory, but {self.memory} is requested.')
 
     def _try_validate_accelerators(self) -> None:
         """Validate accelerators against the instance type and region/zone."""
@@ -718,6 +785,7 @@ class Resources:
             self.cloud is None,
             self._instance_type is None,
             self.cpus is None,
+            self.memory is None,
             self.accelerators is None,
             self.accelerator_args is None,
             not self._use_spot_specified,
@@ -730,6 +798,7 @@ class Resources:
             cloud=override.pop('cloud', self.cloud),
             instance_type=override.pop('instance_type', self.instance_type),
             cpus=override.pop('cpus', self.cpus),
+            memory=override.pop('memory', self.memory),
             accelerators=override.pop('accelerators', self.accelerators),
             accelerator_args=override.pop('accelerator_args',
                                           self.accelerator_args),
@@ -770,6 +839,8 @@ class Resources:
             resources_fields['instance_type'] = config.pop('instance_type')
         if config.get('cpus') is not None:
             resources_fields['cpus'] = str(config.pop('cpus'))
+        if config.get('memory') is not None:
+            resources_fields['memory'] = str(config.pop('memory'))
         if config.get('accelerators') is not None:
             resources_fields['accelerators'] = config.pop('accelerators')
         if config.get('accelerator_args') is not None:
@@ -804,6 +875,7 @@ class Resources:
         add_if_not_none('cloud', str(self.cloud))
         add_if_not_none('instance_type', self.instance_type)
         add_if_not_none('cpus', self.cpus)
+        add_if_not_none('memory', self.memory)
         add_if_not_none('accelerators', self.accelerators)
         add_if_not_none('accelerator_args', self.accelerator_args)
 
@@ -866,6 +938,9 @@ class Resources:
         if version < 7:
             self._cpus = None
 
+        if version < 8:
+            self._memory = None
+
         image_id = state.get('_image_id', None)
         if isinstance(image_id, str):
             state['_image_id'] = {state.get('_region', None): image_id}
diff --git a/sky/spot/controller.py b/sky/spot/controller.py
index 1f1870f89..773dd4b0a 100644
--- a/sky/spot/controller.py
+++ b/sky/spot/controller.py
@@ -31,8 +31,8 @@ class SpotController:
     def __init__(self, job_id: int, task_yaml: str,
                  retry_until_up: bool) -> None:
         self._job_id = job_id
-        self._task_name = pathlib.Path(task_yaml).stem
         self._task = sky.Task.from_yaml(task_yaml)
+        self._task_name = self._task.name
 
         self._retry_until_up = retry_until_up
         # TODO(zhwu): this assumes the specific backend.
@@ -330,7 +330,6 @@ if __name__ == '__main__':
                         help='Retry until the spot cluster is up.')
     parser.add_argument('task_yaml',
                         type=str,
-                        help='The path to the user spot task yaml file. '
-                        'The file name is the spot task name.')
+                        help='The path to the user spot task yaml file.')
     args = parser.parse_args()
     start(args.job_id, args.task_yaml, args.retry_until_up)
diff --git a/sky/spot/recovery_strategy.py b/sky/spot/recovery_strategy.py
index de41d685a..3d0980ae0 100644
--- a/sky/spot/recovery_strategy.py
+++ b/sky/spot/recovery_strategy.py
@@ -1,5 +1,6 @@
 """The strategy to handle launching/recovery/termination of spot clusters."""
 import time
+import traceback
 import typing
 from typing import Optional, Tuple
 
@@ -117,7 +118,9 @@ class StrategyExecutor:
                     raise RuntimeError('Failed to terminate the spot cluster '
                                        f'{self.cluster_name}.') from e
                 logger.error('Failed to terminate the spot cluster '
-                             f'{self.cluster_name}. Retrying.')
+                             f'{self.cluster_name}. Retrying.'
+                             f'Details: {common_utils.format_exception(e)}')
+                logger.error(f'  Traceback: {traceback.format_exc()}')
 
     def _try_cancel_all_jobs(self):
         handle = global_user_state.get_handle_from_cluster_name(
@@ -288,7 +291,6 @@ class StrategyExecutor:
                 # code.
                 logger.info('Failed to launch the spot cluster with error: '
                             f'{common_utils.format_exception(e)})')
-                import traceback  # pylint: disable=import-outside-toplevel
                 logger.info(f'  Traceback: {traceback.format_exc()}')
             else:  # No exception, the launch succeeds.
                 # At this point, a sky.launch() has succeeded. Cluster may be
diff --git a/sky/templates/spot-controller.yaml.j2 b/sky/templates/spot-controller.yaml.j2
index 919818a21..7931706ce 100644
--- a/sky/templates/spot-controller.yaml.j2
+++ b/sky/templates/spot-controller.yaml.j2
@@ -1,6 +1,6 @@
 # The template for the spot controller
 
-name: {{cluster_name}}
+name: {{task_name}}
 
 resources:
   disk_size: 50
@@ -11,9 +11,9 @@ resources:
 #  instance_type: t3.xlarge
 
 file_mounts:
-  {{remote_user_yaml_prefix}}/{{cluster_name}}.yaml: {{user_yaml_path}}
+  {{remote_user_yaml_prefix}}/{{task_name}}-{{uuid}}.yaml: {{user_yaml_path}}
 {% if user_config_path is not none %}
-  {{remote_user_yaml_prefix}}/{{cluster_name}}.config_yaml: {{user_config_path}}
+  {{remote_user_yaml_prefix}}/{{task_name}}-{{uuid}}.config_yaml: {{user_config_path}}
 {% endif %}
 
 setup: |
@@ -37,7 +37,7 @@ setup: |
 
 run: |
   python -u -m sky.spot.controller \
-    {{remote_user_yaml_prefix}}/{{cluster_name}}.yaml \
+    {{remote_user_yaml_prefix}}/{{task_name}}-{{uuid}}.yaml \
     --job-id $SKYPILOT_INTERNAL_JOB_ID {% if retry_until_up %}--retry-until-up{% endif %}
 
 envs:
@@ -46,7 +46,7 @@ envs:
   SKYPILOT_SKIP_CLOUD_IDENTITY_CHECK: 1
   SKYPILOT_USER: {{user}}
 {% if user_config_path is not none %}
-  {{env_var_skypilot_config}}: {{remote_user_yaml_prefix}}/{{cluster_name}}.config_yaml
+  {{env_var_skypilot_config}}: {{remote_user_yaml_prefix}}/{{task_name}}-{{uuid}}.config_yaml
 {% endif %}
 {% if is_dev %}
   SKYPILOT_DEV: 1
diff --git a/sky/utils/schemas.py b/sky/utils/schemas.py
index c71c83b0e..4fb582b9b 100644
--- a/sky/utils/schemas.py
+++ b/sky/utils/schemas.py
@@ -32,6 +32,13 @@ def get_resources_schema():
                     'type': 'number',
                 }],
             },
+            'memory': {
+                'anyOf': [{
+                    'type': 'string',
+                }, {
+                    'type': 'number',
+                }],
+            },
             'accelerators': {
                 'anyOf': [{
                     'type': 'string',
diff --git a/sky/utils/tpu_utils.py b/sky/utils/tpu_utils.py
index bfa9955e5..fe5a4c918 100644
--- a/sky/utils/tpu_utils.py
+++ b/sky/utils/tpu_utils.py
@@ -27,7 +27,7 @@ def is_tpu_vm_pod(resources: Optional[resources_lib.Resources]) -> bool:
     if resources is None or not is_tpu_vm(resources):
         return False
     acc, _ = list(resources.accelerators.items())[0]
-    return acc not in ['tpu-v2-8', 'tpu-v3-8']
+    return acc not in ['tpu-v2-8', 'tpu-v3-8', 'tpu-v4-8']
 
 
 def get_num_tpu_devices(


Test patch:
diff --git a/tests/test_optimizer_dryruns.py b/tests/test_optimizer_dryruns.py
index b30a7ae55..8b66283a6 100644
--- a/tests/test_optimizer_dryruns.py
+++ b/tests/test_optimizer_dryruns.py
@@ -18,6 +18,15 @@ def _test_parse_cpus(spec, expected_cpus):
             assert list(task.resources)[0].cpus == expected_cpus
 
 
+def _test_parse_memory(spec, expected_memory):
+    with tempfile.NamedTemporaryFile('w') as f:
+        f.write(spec)
+        f.flush()
+        with sky.Dag():
+            task = sky.Task.from_yaml(f.name)
+            assert list(task.resources)[0].memory == expected_memory
+
+
 def _test_parse_accelerators(spec, expected_accelerators):
     with tempfile.NamedTemporaryFile('w') as f:
         f.write(spec)
@@ -105,6 +114,12 @@ def test_partial_cpus(monkeypatch):
     _test_resources_launch(monkeypatch, cpus='7+')
 
 
+def test_partial_memory(monkeypatch):
+    _test_resources_launch(monkeypatch, memory=32)
+    _test_resources_launch(monkeypatch, memory='32')
+    _test_resources_launch(monkeypatch, memory='32+')
+
+
 def test_partial_k80(monkeypatch):
     _test_resources_launch(monkeypatch, accelerators='K80')
 
@@ -179,6 +194,22 @@ def test_instance_type_mismatches_cpus(monkeypatch):
         assert 'does not have the requested number of vCPUs' in str(e.value)
 
 
+def test_instance_type_mismatches_memory(monkeypatch):
+    bad_instance_and_memory = [
+        # Actual: 32
+        ('m6i.2xlarge', 4),
+        # Actual: 4
+        ('c6i.large', 2),
+    ]
+    for instance, memory in bad_instance_and_memory:
+        with pytest.raises(ValueError) as e:
+            _test_resources_launch(monkeypatch,
+                                   sky.AWS(),
+                                   instance_type=instance,
+                                   memory=memory)
+        assert 'does not have the requested memory' in str(e.value)
+
+
 def test_instance_type_matches_cpus(monkeypatch):
     _test_resources_launch(monkeypatch,
                            sky.AWS(),
@@ -198,6 +229,84 @@ def test_instance_type_matches_cpus(monkeypatch):
                            cpus=8.0)
 
 
+def test_instance_type_matches_memory(monkeypatch):
+    _test_resources_launch(monkeypatch,
+                           sky.AWS(),
+                           instance_type='c6i.8xlarge',
+                           memory=64)
+    _test_resources_launch(monkeypatch,
+                           sky.Azure(),
+                           instance_type='Standard_E8s_v5',
+                           memory='64')
+    _test_resources_launch(monkeypatch,
+                           sky.GCP(),
+                           instance_type='n1-standard-8',
+                           memory='30+')
+    _test_resources_launch(monkeypatch,
+                           sky.AWS(),
+                           instance_type='g4dn.2xlarge',
+                           memory=32)
+
+
+def test_instance_type_from_cpu_memory(monkeypatch, capfd):
+    _test_resources_launch(monkeypatch, cpus=8)
+    stdout, _ = capfd.readouterr()
+    # Choose General Purpose instance types
+    assert 'm6i.2xlarge' in stdout  # AWS, 8 vCPUs, 32 GB memory
+    assert 'Standard_D8_v5' in stdout  # Azure, 8 vCPUs, 32 GB memory
+    assert 'n2-standard-8' in stdout  # GCP, 8 vCPUs, 32 GB memory
+
+    _test_resources_launch(monkeypatch, memory=32)
+    stdout, _ = capfd.readouterr()
+    # Choose memory-optimized instance types, when the memory
+    # is specified
+    assert 'r6i.xlarge' in stdout  # AWS, 4 vCPUs, 32 GB memory
+    assert 'Standard_E4_v5' in stdout  # Azure, 4 vCPUs, 32 GB memory
+    assert 'n2-highmem-4' in stdout  # GCP, 4 vCPUs, 32 GB memory
+
+    _test_resources_launch(monkeypatch, memory='64+')
+    stdout, _ = capfd.readouterr()
+    # Choose memory-optimized instance types
+    assert 'r6i.2xlarge' in stdout  # AWS, 8 vCPUs, 64 GB memory
+    assert 'Standard_E8_v5' in stdout  # Azure, 8 vCPUs, 64 GB memory
+    assert 'n2-highmem-8' in stdout  # GCP, 8 vCPUs, 64 GB memory
+    assert 'gpu_1x_a100_sxm4' in stdout  # Lambda, 30 vCPUs, 200 GB memory
+
+    _test_resources_launch(monkeypatch, cpus='4+', memory='4+')
+    stdout, _ = capfd.readouterr()
+    # Choose compute-optimized instance types, when the memory
+    # requirement is less than the memory of General Purpose
+    # instance types.
+    assert 'n2-highcpu-4' in stdout  # GCP, 4 vCPUs, 4 GB memory
+    assert 'c6i.xlarge' in stdout  # AWS, 4 vCPUs, 8 GB memory
+    assert 'Standard_F4s_v2' in stdout  # Azure, 4 vCPUs, 8 GB memory
+    assert 'gpu_1x_a100_sxm4' in stdout  # Lambda, 30 vCPUs, 200 GB memory
+
+    _test_resources_launch(monkeypatch, accelerators='T4')
+    stdout, _ = capfd.readouterr()
+    # Choose cheapest T4 instance type
+    assert 'g4dn.xlarge' in stdout  # AWS, 4 vCPUs, 16 GB memory, 1 T4 GPU
+    assert 'Standard_NC4as_T4_v3' in stdout  # Azure, 4 vCPUs, 28 GB memory, 1 T4 GPU
+    assert 'n1-highmem-4' in stdout  # GCP, 4 vCPUs, 26 GB memory, 1 T4 GPU
+
+    _test_resources_launch(monkeypatch,
+                           cpus='16+',
+                           memory='32+',
+                           accelerators='T4')
+    stdout, _ = capfd.readouterr()
+    # Choose cheapest T4 instance type that satisfies the requirement
+    assert 'n1-standard-16' in stdout  # GCP, 16 vCPUs, 60 GB memory, 1 T4 GPU
+    assert 'g4dn.4xlarge' in stdout  # AWS, 16 vCPUs, 64 GB memory, 1 T4 GPU
+    assert 'Standard_NC16as_T4_v3' in stdout  # Azure, 16 vCPUs, 110 GB memory, 1 T4 GPU
+
+    _test_resources_launch(monkeypatch, memory='200+', accelerators='T4')
+    stdout, _ = capfd.readouterr()
+    # Choose cheapest T4 instance type that satisfies the requirement
+    assert 'n1-highmem-32' in stdout  # GCP, 32 vCPUs, 208 GB memory, 1 T4 GPU
+    assert 'g4dn.16xlarge' in stdout  # AWS, 64 vCPUs, 256 GB memory, 1 T4 GPU
+    assert 'Azure' not in stdout  # Azure does not have a 1 T4 GPU instance type with 200+ GB memory
+
+
 def test_instance_type_mistmatches_accelerators(monkeypatch):
     bad_instance_and_accs = [
         # Actual: V100
@@ -266,6 +375,13 @@ def test_invalid_cpus(monkeypatch):
         assert '"cpus" field should be' in str(e.value)
 
 
+def test_invalid_memory(monkeypatch):
+    for cloud in [sky.AWS(), sky.Azure(), sky.GCP(), None]:
+        with pytest.raises(ValueError) as e:
+            _test_resources(monkeypatch, cloud, memory='invalid')
+        assert '"memory" field should be' in str(e.value)
+
+
 def test_invalid_region(monkeypatch):
     for cloud in [sky.AWS(), sky.Azure(), sky.GCP()]:
         with pytest.raises(ValueError) as e:
@@ -360,6 +476,33 @@ def test_parse_cpus_from_yaml():
             cpus: '3+' """)
     _test_parse_cpus(spec, '3+')
 
+    spec = textwrap.dedent("""\
+        resources:
+            cpus: 3+ """)
+    _test_parse_cpus(spec, '3+')
+
+
+def test_parse_memory_from_yaml():
+    spec = textwrap.dedent("""\
+        resources:
+            memory: 32""")
+    _test_parse_memory(spec, '32')
+
+    spec = textwrap.dedent("""\
+        resources:
+            memory: 1.5""")
+    _test_parse_memory(spec, '1.5')
+
+    spec = textwrap.dedent("""\
+        resources:
+            memory: '3+' """)
+    _test_parse_memory(spec, '3+')
+
+    spec = textwrap.dedent("""\
+        resources:
+            memory: 3+ """)
+    _test_parse_memory(spec, '3+')
+
 
 def test_parse_accelerators_from_yaml():
     spec = textwrap.dedent("""\
diff --git a/tests/test_optimizer_random_dag.py b/tests/test_optimizer_random_dag.py
index c4f3c551c..175363a46 100644
--- a/tests/test_optimizer_random_dag.py
+++ b/tests/test_optimizer_random_dag.py
@@ -116,7 +116,7 @@ def compare_optimization_results(dag: sky.Dag, minimize_cost: bool):
                                                       dag.tasks, optimizer_plan)
 
     min_objective = find_min_objective(copy_dag, minimize_cost)
-    assert objective == min_objective
+    assert abs(objective - min_objective) < 1e-3
 
 
 def test_optimizer(enable_all_clouds):