"""Optimizer: assigns best resources to user tasks."""
import collections
import copy
import json
import typing
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple

import colorama
import numpy as np
import prettytable

from sky import check as sky_check
from sky import clouds
from sky import dag as dag_lib
from sky import exceptions
from sky import resources as resources_lib
from sky import sky_logging
from sky import skypilot_config
from sky import task as task_lib
from sky.adaptors import common as adaptors_common
from sky.clouds import cloud as sky_cloud
from sky.usage import usage_lib
from sky.utils import common
from sky.utils import common_utils
from sky.utils import env_options
from sky.utils import log_utils
from sky.utils import registry
from sky.utils import resources_utils
from sky.utils import rich_utils
from sky.utils import subprocess_utils
from sky.utils import timeline
from sky.utils import ux_utils

if typing.TYPE_CHECKING:
    import networkx as nx
else:
    nx = adaptors_common.LazyImport('networkx')

logger = sky_logging.init_logger(__name__)

_DUMMY_SOURCE_NAME = 'skypilot-dummy-source'
_DUMMY_SINK_NAME = 'skypilot-dummy-sink'

# task -> resources -> estimated cost or time.
_TaskToCostMap = Dict[task_lib.Task, Dict[resources_lib.Resources, float]]
# cloud -> list of resources that have the same accelerators.
_PerCloudCandidates = Dict[clouds.Cloud, List[resources_lib.Resources]]
# task -> per-cloud candidates
_TaskToPerCloudCandidates = Dict[task_lib.Task, _PerCloudCandidates]


# For logging purposes.
def _create_table(field_names: List[str]) -> prettytable.PrettyTable:
    table_kwargs = {
        'hrules': prettytable.FRAME,
        'vrules': prettytable.NONE,
        'border': True,
    }
    return log_utils.create_table(field_names, **table_kwargs)


def _is_dag_resources_ordered(dag: 'dag_lib.Dag') -> bool:
    graph = dag.get_graph()
    topo_order = list(nx.topological_sort(graph))
    for node in topo_order:
        if isinstance(node.resources, list):
            return True
    return False


class Optimizer:
    """Optimizer: assigns best resources to user tasks."""

    @staticmethod
    def _egress_cost(src_cloud: clouds.Cloud, dst_cloud: clouds.Cloud,
                     gigabytes: float) -> float:
        """Returns estimated egress cost."""
        if isinstance(src_cloud, clouds.DummyCloud) or isinstance(
                dst_cloud, clouds.DummyCloud):
            return 0.0

        if not src_cloud.is_same_cloud(dst_cloud):
            egress_cost = src_cloud.get_egress_cost(num_gigabytes=gigabytes)
        else:
            egress_cost = 0.0
        return egress_cost

    @staticmethod
    def _egress_time(src_cloud: clouds.Cloud, dst_cloud: clouds.Cloud,
                     gigabytes: float) -> float:
        """Returns estimated egress time in seconds."""
        # FIXME: estimate bandwidth between each cloud-region pair.
        if isinstance(src_cloud, clouds.DummyCloud) or isinstance(
                dst_cloud, clouds.DummyCloud):
            return 0.0
        if not src_cloud.is_same_cloud(dst_cloud):
            # 10Gbps is close to the average of observed b/w from S3
            # (us-west,east-2) to GCS us-central1, assuming highly sharded files
            # (~128MB per file).
            bandwidth_gbps = 10
            egress_time = gigabytes * 8 / bandwidth_gbps
        else:
            egress_time = 0.0
        return egress_time

    @staticmethod
    @timeline.event
    @usage_lib.entrypoint('sky.optimizer.optimize')
    def optimize(dag: 'dag_lib.Dag',
                 minimize: common.OptimizeTarget = common.OptimizeTarget.COST,
                 blocked_resources: Optional[Iterable[
                     resources_lib.Resources]] = None,
                 quiet: bool = False) -> 'dag_lib.Dag':
        """Find the best execution plan for the given DAG.

        Args:
            dag: the DAG to optimize.
            minimize: whether to minimize cost or time.
            blocked_resources: a list of resources that should not be used.
            quiet: whether to suppress logging.

        Raises:
            exceptions.ResourcesUnavailableError: if no resources are available
                for a task.
            exceptions.NoCloudAccessError: if no public clouds are enabled.
        """
        with rich_utils.safe_status(ux_utils.spinner_message('Optimizing')):
            _check_specified_clouds(dag)
            # This function is effectful: mutates every node in 'dag' by setting
            # node.best_resources if it is None.
            Optimizer._add_dummy_source_sink_nodes(dag)
            try:
                unused_best_plan = Optimizer._optimize_dag(
                    dag=dag,
                    minimize_cost=minimize == common.OptimizeTarget.COST,
                    blocked_resources=blocked_resources,
                    quiet=quiet)
            finally:
                # Make sure to remove the dummy source/sink nodes, even if the
                # optimization fails.
                Optimizer._remove_dummy_source_sink_nodes(dag)
            return dag

    @staticmethod
    def _add_dummy_source_sink_nodes(dag: 'dag_lib.Dag'):
        """Adds special Source and Sink nodes.

        The two special nodes are for conveniently handling cases such as

          { A, B } --> C (multiple sources)

        or

          A --> { B, C } (multiple sinks)

        Adds new edges:
            Source --> for all node with in degree 0
            For all node with out degree 0 --> Sink
        """
        graph = dag.get_graph()
        zero_indegree_nodes = []
        zero_outdegree_nodes = []
        for node, in_degree in graph.in_degree():
            if in_degree == 0:
                zero_indegree_nodes.append(node)
        for node, out_degree in graph.out_degree():
            if out_degree == 0:
                zero_outdegree_nodes.append(node)

        def make_dummy(name):
            dummy = task_lib.Task(name)
            dummy.set_resources({DummyResources(cloud=clouds.DummyCloud())})
            dummy.set_time_estimator(lambda _: 0)
            return dummy

        with dag:
            source = make_dummy(_DUMMY_SOURCE_NAME)
            for real_source_node in zero_indegree_nodes:
                source >> real_source_node  # pylint: disable=pointless-statement
            sink = make_dummy(_DUMMY_SINK_NAME)
            for real_sink_node in zero_outdegree_nodes:
                real_sink_node >> sink  # pylint: disable=pointless-statement

    @staticmethod
    def _remove_dummy_source_sink_nodes(dag: 'dag_lib.Dag'):
        """Removes special Source and Sink nodes."""
        source = [t for t in dag.tasks if t.name == _DUMMY_SOURCE_NAME]
        sink = [t for t in dag.tasks if t.name == _DUMMY_SINK_NAME]
        if not source and not sink:
            return
        assert len(source) == len(sink) == 1, dag.tasks
        dag.remove(source[0])
        dag.remove(sink[0])

    @staticmethod
    def _get_egress_info(
        parent: task_lib.Task,
        parent_resources: resources_lib.Resources,
        node: task_lib.Task,
        resources: resources_lib.Resources,
    ) -> Tuple[Optional[clouds.Cloud], Optional[clouds.Cloud], Optional[float]]:
        if isinstance(parent_resources.cloud, clouds.DummyCloud):
            # Special case.  The current 'node' is a real
            # source node, and its input may be on a different
            # cloud from 'resources'.
            if node.get_inputs() is None:
                # A task_lib.Task may have no inputs specified.
                return None, None, 0
            src_cloud = node.get_inputs_cloud()
            nbytes = node.get_estimated_inputs_size_gigabytes()
        else:
            src_cloud = parent_resources.cloud
            nbytes = parent.get_estimated_outputs_size_gigabytes()
        dst_cloud = resources.cloud
        return src_cloud, dst_cloud, nbytes

    @staticmethod
    def _egress_cost_or_time(minimize_cost: bool, parent: task_lib.Task,
                             parent_resources: resources_lib.Resources,
                             node: task_lib.Task,
                             resources: resources_lib.Resources):
        """Computes the egress cost or time depending on 'minimize_cost'."""
        src_cloud, dst_cloud, nbytes = Optimizer._get_egress_info(
            parent, parent_resources, node, resources)
        if not nbytes:
            # nbytes can be None, if the task has no inputs/outputs.
            return 0
        assert src_cloud is not None and dst_cloud is not None, (src_cloud,
                                                                 dst_cloud,
                                                                 nbytes)

        if minimize_cost:
            fn = Optimizer._egress_cost
        else:
            fn = Optimizer._egress_time
        return fn(src_cloud, dst_cloud, nbytes)

    @staticmethod
    def _estimate_nodes_cost_or_time(
        topo_order: List[task_lib.Task],
        minimize_cost: bool = True,
        blocked_resources: Optional[Iterable[resources_lib.Resources]] = None,
        quiet: bool = False
    ) -> Tuple[_TaskToCostMap, _TaskToPerCloudCandidates]:
        """Estimates the cost/time of each task-resource mapping in the DAG.

        Note that the egress cost/time is not considered in this function.
        The estimated run time of a task running on a resource is given by
        `task.estimate_runtime(resources)` or 1 hour by default.
        The estimated cost is `task.num_nodes * resources.get_cost(runtime)`.
        """
        # Cost/time of running the task on the resources.
        # node -> {resources -> cost/time}
        node_to_cost_map: _TaskToCostMap = collections.defaultdict(dict)

        # node -> cloud -> list of resources that satisfy user's requirements.
        node_to_candidate_map: _TaskToPerCloudCandidates = {}

        def get_available_reservations(
            launchable_resources: Dict[resources_lib.Resources,
                                       List[resources_lib.Resources]]
        ) -> Dict[resources_lib.Resources, int]:
            if not resources_utils.need_to_query_reservations():
                return {}

            num_available_reserved_nodes_per_resource = {}

            def get_reservations_available_resources(
                    resources: resources_lib.Resources):
                num_available_reserved_nodes_per_resource[resources] = sum(
                    resources.get_reservations_available_resources().values())

            launchable_resources_list: List[resources_lib.Resources] = sum(
                launchable_resources.values(), [])
            with rich_utils.safe_status(
                    ux_utils.spinner_message('Checking reserved resources')):
                subprocess_utils.run_in_parallel(
                    get_reservations_available_resources,
                    launchable_resources_list)
            return num_available_reserved_nodes_per_resource

        indent_prefix = ' ' * len('Hint: ')

        # Compute the estimated cost/time for each node.
        for node_i, node in enumerate(topo_order):
            if node_i == 0:
                # Base case: a special source node.
                node_to_cost_map[node][list(node.resources)[0]] = 0
                continue

            # Don't print for the last node, Sink.
            do_print = node_i != len(topo_order) - 1

            fuzzy_candidates: List[str] = []
            if node_i < len(topo_order) - 1:
                # Convert partial resource labels to launchable resources.
                (launchable_resources, cloud_candidates, fuzzy_candidates,
                 resource_hints) = (_fill_in_launchable_resources(
                     task=node,
                     blocked_resources=blocked_resources,
                     quiet=quiet))
                node_to_candidate_map[node] = cloud_candidates
                # Has to call the printing after the launchable resources are
                # computed, because the missing fields of the resources are
                # inferred in the _fill_in_launchable_resources function.
                logger.debug('#### {} ####'.format(node))
            else:
                # Dummy sink node.
                launchable_resources = {
                    list(node.resources)[0]: list(node.resources)
                }

            # Fetch reservations in advance and in parallel to speed up the
            # reservation info fetching.
            num_resources = len(list(node.resources))
            num_available_reserved_nodes_per_resource = (
                get_available_reservations(launchable_resources))

            for orig_resources, launchable_list in launchable_resources.items():
                if num_resources == 1 and node.time_estimator_func is None:
                    logger.debug(
                        'Defaulting the task\'s estimated time to 1 hour.')
                    estimated_runtime = 1 * 3600
                else:
                    # We assume the time estimator takes in a partial resource
                    #    Resources(accelerators='V100')
                    # and treats their launchable versions
                    #    Resources(infra='aws', instance_type='p3.2xlarge'),
                    #    Resources(infra='gcp', accelerators='V100'),
                    #    ...
                    # as having the same run time.
                    # FIXME(zongheng): take 'num_nodes' as an arg/into
                    # account. It may be another reason to treat num_nodes as
                    # part of a Resources.
                    if node.time_estimator_func is None:
                        estimated_runtime = 1 * 3600
                    else:
                        estimated_runtime = node.estimate_runtime(
                            orig_resources)

                for resources in launchable_list:
                    if minimize_cost:
                        cost_per_node = resources.get_cost(estimated_runtime)
                        num_available_reserved_nodes = (
                            num_available_reserved_nodes_per_resource.get(
                                resources, 0))

                        # We consider the cost of the unused reservation
                        # resources to be 0 since we are already paying for
                        # them.
                        # TODO: different policies can be applied here for
                        # whether to choose reserved instances.
                        estimated_cost_or_time = cost_per_node * max(
                            node.num_nodes - num_available_reserved_nodes, 0)
                    else:
                        # Minimize run time.
                        estimated_cost_or_time = estimated_runtime
                    if do_print:
                        debug_msg = (
                            f'resources: {resources}, '
                            f'estimated_runtime: {estimated_runtime} s '
                            f'({estimated_runtime / 3600:.1f} hr)')
                        if minimize_cost:
                            debug_msg += (', estimated_cost: '
                                          f'${estimated_cost_or_time:.1f}')
                        logger.debug(debug_msg)
                    node_to_cost_map[node][resources] = estimated_cost_or_time
            if not node_to_cost_map[node]:
                source_hint = 'catalog'
                # If Kubernetes was included in the search space, then
                # mention "kubernetes cluster" and/instead of "catalog"
                # in the error message.
                enabled_clouds = (
                    sky_check.get_cached_enabled_clouds_or_refresh(
                        sky_cloud.CloudCapability.COMPUTE))
                if clouds.cloud_in_iterable(clouds.Kubernetes(),
                                            enabled_clouds):
                    if any(orig_resources.cloud is None
                           for orig_resources in node.resources):
                        source_hint = 'catalog and kubernetes cluster'
                    elif all(
                            isinstance(orig_resources.cloud, clouds.SSH)
                            for orig_resources in node.resources):
                        source_hint = 'node pool'
                    elif all(
                            isinstance(orig_resources.cloud, clouds.Kubernetes)
                            for orig_resources in node.resources):
                        source_hint = 'kubernetes cluster'

                bold = colorama.Style.BRIGHT
                cyan = colorama.Fore.CYAN
                reset = colorama.Style.RESET_ALL
                fuzzy_candidates_str = ''
                if fuzzy_candidates:
                    fuzzy_candidates_str = (
                        f'\nTry one of these offered accelerators: {cyan}'
                        f'{fuzzy_candidates}{reset}')
                node_resources_reprs = ', '.join(f'{node.num_nodes}x ' +
                                                 r.repr_with_region_zone
                                                 for r in node.resources)
                hints_concat = '\n'.join([
                    f'{bold}Resource: {repr(resource)}{reset}\n' +
                    '\n'.join(hint_list)
                    for resource, hint_list in resource_hints.items()
                    if hint_list
                ])
                hints_formatted = '\n'.join(
                    map(lambda r: f'{indent_prefix}{r}',
                        hints_concat.split('\n')))
                resource_hints_string = (
                    f'Hint: Check Per Resource Hint\n{hints_formatted}'
                    if hints_formatted else '')
                error_msg = (
                    f'{source_hint.capitalize()} does not contain any '
                    f'instances satisfying the request: '
                    f'{node_resources_reprs}.'
                    f'\nTo fix: relax or change the '
                    f'resource requirements.{fuzzy_candidates_str}\n\n'
                    f'Hint: {bold}sky gpus list{reset} '
                    'to list available accelerators.\n'
                    f'{indent_prefix}{bold}sky check{reset} to check the '
                    'enabled clouds.\n'
                    f'{resource_hints_string}')
                with ux_utils.print_exception_no_traceback():
                    raise exceptions.ResourcesUnavailableError(error_msg)
        return node_to_cost_map, node_to_candidate_map

    @staticmethod
    def _optimize_by_dp(
        topo_order: List[task_lib.Task],
        node_to_cost_map: _TaskToCostMap,
        minimize_cost: bool = True,
    ) -> Tuple[Dict[task_lib.Task, resources_lib.Resources], float]:
        """Optimizes a chain DAG using a dynamic programming algorithm."""
        # node -> { resources -> best estimated cost }
        dp_best_objective: Dict[task_lib.Task,
                                Dict[resources_lib.Resources,
                                     float]] = collections.defaultdict(dict)
        # node -> { resources -> best parent resources }
        dp_point_backs: Dict[task_lib.Task, Dict[
            resources_lib.Resources,
            resources_lib.Resources]] = collections.defaultdict(dict)

        # Computes dp_best_objective[node][resources]
        # = my estimated cost + min_phw { dp_best_objective(p, phw) +
        #                                 egress_cost(p, phw, hw) }
        # where p is the parent of the node.
        for node_i, node in enumerate(topo_order):
            if node_i == 0:
                # Base case: a special source node.
                dp_best_objective[node][list(node.resources)[0]] = 0
                continue

            parent = topo_order[node_i - 1]
            # FIXME: Account for egress costs for multi-node clusters
            for resources, execution_cost in node_to_cost_map[node].items():
                min_pred_cost_plus_egress = np.inf
                for parent_resources, parent_cost in \
                    dp_best_objective[parent].items():
                    egress_cost = Optimizer._egress_cost_or_time(
                        minimize_cost, parent, parent_resources, node,
                        resources)

                    if parent_cost + egress_cost < min_pred_cost_plus_egress:
                        min_pred_cost_plus_egress = parent_cost + egress_cost
                        best_parent_hardware = parent_resources

                dp_point_backs[node][resources] = best_parent_hardware
                dp_best_objective[node][resources] = \
                    execution_cost + min_pred_cost_plus_egress

        # Compute the total objective value of the DAG.
        sink_node = topo_order[-1]
        total_objective = dp_best_objective[sink_node]
        assert len(total_objective) == 1, \
            f'Should be DummyCloud: {total_objective}'
        best_resources, best_total_objective = list(total_objective.items())[0]

        # Find the best plan for the DAG.
        # node -> best resources
        best_plan = {}
        for node in reversed(topo_order):
            best_plan[node] = best_resources
            node.best_resources = best_resources
            if node.name != _DUMMY_SOURCE_NAME:
                best_resources = dp_point_backs[node][best_resources]
        return best_plan, best_total_objective

    @staticmethod
    def _optimize_by_ilp(
        graph: 'nx.DiGraph',
        topo_order: List[task_lib.Task],
        node_to_cost_map: _TaskToCostMap,
        minimize_cost: bool = True,
    ) -> Tuple[Dict[task_lib.Task, resources_lib.Resources], float]:
        """Optimizes a general DAG using an ILP solver.

        Notations:
            V: the set of nodes (tasks).
            E: the set of edges (dependencies).
            k: node -> [r.cost for r in node.resources].
            F: (node i, node j) -> the egress cost/time between node i and j.
            c: node -> one-hot decision vector. c[node][i] = 1 means
                the node is assigned to the i-th resource.
            e: (node i, node j) -> linearization of c[node i] x c[node j].
              e[node i][node j][a][b] = 1 means node i and node j are assigned
              to the a-th and the b-th resources, respectively.

        Objective:
            For cost optimization,
                minimize_{c} sum(c[v]^T @ k[v] for each v in V) +
                             sum(c[u]^T @ F[u][v] @ c[v] for each u, v in E)
                s.t. sum(c[v] == 1) for each v in V
            which is equivalent (linearized) to,
                minimize_{c, e} sum(c[v]^T @ k[v] for each v in V) +
                                sum(e[u][v]^T @ F[u][v] for each u, v in E)
                s.t. sum(c[v] == 1) for each v in V (i.e., c is one-hot)
                     sum(e[u][v] == 1) for each u, v in E (i.e., e is one-hot)
                     e[u][v] = flatten(c[u] @ c[v]^T) for each u, v in E
            The first term of the objective indicates the execution cost
            of the task v, and the second term indicates the egress cost
            of the parent task u to the task v.

            For time optimization,
                minimize_{c} finish_time[sink_node]
                s.t. finish_time[v] >= c[v]^T @ k[v] + finish_time[u] +
                                       c[u]^T @ F[u][v] @ c[v]
                     for each u, v in E
                     sum(c[v] == 1) for each v in V
            which is equivalent (linearized) to,
                minimize_{c, e} finish_time[sink_node]
                s.t. finish_time[v] >= c[v]^T @ k[v] + finish_time[u] +
                                       e[u][v]^T @ F[u][v]
                     for each u, v in E
                     sum(c[v] == 1) for each v in V (i.e., c is one-hot)
                     sum(e[u][v] == 1) for each u, v in E (i.e., e is one-hot)
                     e[u][v] = flatten(c[u] @ c[v]^T) for each u, v in E
            The first term of the objective indicates the execution time
            of the task v, and the other two terms indicate that the task v
            starts executing no sooner than its parent tasks are finished and
            the output data from the parents has arrived to the task v.
        """
        import pulp  # pylint: disable=import-outside-toplevel

        if minimize_cost:
            prob = pulp.LpProblem('Cost-Optimization', pulp.LpMinimize)
        else:
            prob = pulp.LpProblem('Runtime-Optimization', pulp.LpMinimize)

        # Prepare the constants.
        V = topo_order  # pylint: disable=invalid-name
        E = graph.edges()  # pylint: disable=invalid-name
        k = {
            node: list(resource_cost_map.values())
            for node, resource_cost_map in node_to_cost_map.items()
        }
        F: Dict[Any, Dict[Any, List[float]]] = collections.defaultdict(dict)  # pylint: disable=invalid-name
        for u, v in E:
            F[u][v] = []
            for r_u in node_to_cost_map[u].keys():
                for r_v in node_to_cost_map[v].keys():
                    F[u][v].append(
                        Optimizer._egress_cost_or_time(minimize_cost, u, r_u, v,
                                                       r_v))

        # Define the decision variables.
        c = {
            v: pulp.LpVariable.matrix(v.name, (range(len(k[v])),), cat='Binary')
            for v in V
        }

        e: Dict[Any,
                Dict[Any,
                     List[pulp.LpVariable]]] = collections.defaultdict(dict)
        for u, v in E:
            num_vars = len(c[u]) * len(c[v])
            e[u][v] = pulp.LpVariable.matrix(f'({u.name}->{v.name})',
                                             (range(num_vars),),
                                             cat='Binary')

        # Formulate the constraints.
        # 1. c[v] is an one-hot vector.
        for v in V:
            prob += pulp.lpSum(c[v]) == 1

        # 2. e[u][v] is an one-hot vector.
        for u, v in E:
            prob += pulp.lpSum(e[u][v]) == 1

        # 3. e[u][v] linearizes c[u] x c[v].
        for u, v in E:
            e_uv = e[u][v]  # 1-d one-hot vector
            N_u = len(c[u])  # pylint: disable=invalid-name
            N_v = len(c[v])  # pylint: disable=invalid-name

            for row in range(N_u):
                prob += pulp.lpSum(
                    e_uv[N_v * row + col] for col in range(N_v)) == c[u][row]

            for col in range(N_v):
                prob += pulp.lpSum(
                    e_uv[N_v * row + col] for row in range(N_u)) == c[v][col]

        # Formulate the objective.
        if minimize_cost:
            objective = 0
            for v in V:
                objective += pulp.lpDot(c[v], k[v])
            for u, v in E:
                objective += pulp.lpDot(e[u][v], F[u][v])
        else:
            # We need additional decision variables.
            finish_time = {
                v: pulp.LpVariable(f'lat({v})', lowBound=0) for v in V
            }
            for u, v in E:
                prob += finish_time[v] >= (pulp.lpDot(
                    c[v], k[v]) + finish_time[u] + pulp.lpDot(e[u][v], F[u][v]))
            sink_node = V[-1]
            objective = finish_time[sink_node]
        prob += objective

        # Solve the optimization problem.
        prob.solve(solver=pulp.PULP_CBC_CMD(msg=False))
        assert prob.status != pulp.LpStatusInfeasible, \
            'Cannot solve the optimization problem'
        best_total_objective = prob.objective.value()

        # Find the best plan for the DAG.
        # node -> best resources
        best_plan = {}
        for node, variables in c.items():
            selected = [variable.value() for variable in variables].index(1)
            best_resources = list(node_to_cost_map[node].keys())[selected]
            node.best_resources = best_resources
            best_plan[node] = best_resources
        return best_plan, best_total_objective

    @staticmethod
    def _compute_total_time(
        graph,
        topo_order: List[task_lib.Task],
        plan: Dict[task_lib.Task, resources_lib.Resources],
    ) -> float:
        """Estimates the total time of running the DAG by the plan."""
        cache_finish_time: Dict[task_lib.Task, float] = {}

        def finish_time(node):
            if node in cache_finish_time:
                return cache_finish_time[node]

            resources = plan[node]
            if node.time_estimator_func is None:
                execution_time = 1 * 3600
            else:
                # The execution time of dummy nodes is always 0,
                # as they have a time estimator lambda _: 0.
                execution_time = node.estimate_runtime(resources)

            pred_finish_times = [0]
            for pred in graph.predecessors(node):
                # FIXME: Account for egress time for multi-node clusters
                egress_time = Optimizer._egress_cost_or_time(
                    False, pred, plan[pred], node, resources)
                pred_finish_times.append(finish_time(pred) + egress_time)

            cache_finish_time[node] = execution_time + max(pred_finish_times)
            return cache_finish_time[node]

        sink_node = topo_order[-1]
        return finish_time(sink_node)

    @staticmethod
    def _compute_total_cost(
        graph,
        topo_order: List[task_lib.Task],
        plan: Dict[task_lib.Task, resources_lib.Resources],
    ) -> float:
        """Estimates the total cost of running the DAG by the plan."""
        total_cost = 0.
        for node in topo_order:
            resources = plan[node]
            if node.time_estimator_func is None:
                execution_time = 1 * 3600
            else:
                # The execution time of dummy nodes is always 0,
                # as they have a time estimator lambda _: 0.
                execution_time = node.estimate_runtime(resources)

            cost_per_node = resources.get_cost(execution_time)
            total_cost += cost_per_node * node.num_nodes

            for pred in graph.predecessors(node):
                # FIXME: Account for egress costs for multi-node clusters
                egress_cost = Optimizer._egress_cost_or_time(
                    True, pred, plan[pred], node, resources)
                total_cost += egress_cost
        return total_cost

    @staticmethod
    def _print_egress_plan(graph, plan, minimize_cost):
        message_data = []
        for parent, child in graph.edges():
            src_cloud, dst_cloud, nbytes = Optimizer._get_egress_info(
                parent, plan[parent], child, plan[child])
            if not nbytes:
                # nbytes can be None, if the task has no inputs/outputs.
                return 0
            assert src_cloud is not None and dst_cloud is not None, (src_cloud,
                                                                     dst_cloud,
                                                                     nbytes)

            if minimize_cost:
                fn = Optimizer._egress_cost
            else:
                fn = Optimizer._egress_time
            cost_or_time = fn(src_cloud, dst_cloud, nbytes)

            if cost_or_time > 0:
                if parent.name == _DUMMY_SOURCE_NAME:
                    egress = [
                        f'{child.get_inputs()} ({src_cloud})',
                        f'{child} ({dst_cloud})'
                    ]
                else:
                    egress = [
                        f'{parent} ({src_cloud})', f'{child} ({dst_cloud})'
                    ]
                message_data.append((*egress, nbytes, cost_or_time))

        if message_data:
            metric = 'COST ($)' if minimize_cost else 'TIME (s)'
            table = _create_table(['SOURCE', 'TARGET', 'SIZE (GB)', metric])
            table.add_rows(list(reversed(message_data)))
            logger.info(f'Egress plan:\n{table}\n')

    @staticmethod
    def print_optimized_plan(
        graph,
        topo_order: List[task_lib.Task],
        best_plan: Dict[task_lib.Task, resources_lib.Resources],
        total_time: float,
        total_cost: float,
        node_to_cost_map: _TaskToCostMap,
        minimize_cost: bool,
    ):
        ordered_node_to_cost_map = collections.OrderedDict()
        ordered_best_plan = collections.OrderedDict()
        for node in topo_order:
            if node.name not in (_DUMMY_SOURCE_NAME, _DUMMY_SINK_NAME):
                ordered_node_to_cost_map[node] = node_to_cost_map[node]
                ordered_best_plan[node] = best_plan[node]

        is_trivial = all(len(v) == 1 for v in node_to_cost_map.values())
        if not is_trivial and not env_options.Options.MINIMIZE_LOGGING.get():
            metric_str = 'cost' if minimize_cost else 'run time'
            logger.info(
                f'{colorama.Style.BRIGHT}Target:{colorama.Style.RESET_ALL}'
                f' minimizing {metric_str}')

        print_hourly_cost = False
        if len(ordered_node_to_cost_map) == 1:
            node = list(ordered_node_to_cost_map.keys())[0]
            if (node.time_estimator_func is None and
                    node.get_inputs() is None and node.get_outputs() is None):
                print_hourly_cost = True

        if not env_options.Options.MINIMIZE_LOGGING.get():
            if print_hourly_cost:
                logger.info(
                    f'{colorama.Style.BRIGHT}Estimated cost: '
                    f'{colorama.Style.RESET_ALL}${total_cost:.1f} / hour\n')
            else:
                logger.info(
                    f'{colorama.Style.BRIGHT}Estimated total runtime: '
                    f'{colorama.Style.RESET_ALL}{total_time / 3600:.1f} '
                    'hours\n'
                    f'{colorama.Style.BRIGHT}Estimated total cost: '
                    f'{colorama.Style.RESET_ALL}${total_cost:.1f}\n')

        def _instance_type_str(resources: 'resources_lib.Resources') -> str:
            instance_type = resources.instance_type
            assert instance_type is not None, 'Instance type must be specified'
            if isinstance(resources.cloud, (clouds.Kubernetes, clouds.Slurm)):
                instance_type = '-'
                if resources.use_spot:
                    instance_type = ''
            return instance_type

        def _get_resources_element_list(
                resources: 'resources_lib.Resources') -> List[str]:
            accelerators = resources.get_accelerators_str()
            spot = resources.get_spot_str()
            cloud = resources.cloud
            assert cloud is not None, 'Cloud must be specified'
            assert (resources.instance_type is not None), \
                'Instance type must be specified'
            vcpus_, mem_ = cloud.get_vcpus_mem_from_instance_type(
                resources.instance_type)

            def format_number(x: Optional[float]) -> str:
                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_)

            # Format infra as CLOUD (REGION/ZONE)
            infra = resources.infra.formatted_str()

            return [
                infra,
                _instance_type_str(resources) + spot,
                vcpus,
                mem,
                str(accelerators),
            ]

        Row = collections.namedtuple('Row', [
            'infra', 'instance', 'vcpus', 'mem', 'accelerators', 'cost_str',
            'chosen_str'
        ])

        def _get_resources_named_tuple(resources: 'resources_lib.Resources',
                                       cost_str: str, chosen: bool) -> Row:

            accelerators = resources.get_accelerators_str()
            spot = resources.get_spot_str()
            resources = resources.assert_launchable()
            cloud = resources.cloud
            vcpus_, mem_ = cloud.get_vcpus_mem_from_instance_type(
                resources.instance_type)

            def format_number(x: Optional[float]) -> str:
                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_)

            infra = resources.infra.formatted_str()

            chosen_str = ''
            if chosen:
                chosen_str = (colorama.Fore.GREEN + '   ' + '\u2714' +
                              colorama.Style.RESET_ALL)
            row = Row(infra,
                      _instance_type_str(resources) + spot, vcpus, mem,
                      str(accelerators), cost_str, chosen_str)

            return row

        def _get_resource_group_hash(resources: 'resources_lib.Resources'):
            resource_key_dict = {
                'cloud': f'{resources.cloud}',
                'accelerators': f'{resources.accelerators}',
                'use_spot': resources.use_spot
            }

            # Handle special case for Kubernetes, SSH, and SLURM clouds
            if isinstance(resources.cloud, (clouds.Kubernetes, clouds.Slurm)):
                # Region for Kubernetes-like clouds (SSH, Kubernetes) is the
                # context name, i.e. different Kubernetes clusters.
                # Region for SLURM is the cluster name.
                # We add region to the key to show all the clusters in the
                # optimizer table for better UX.

                if resources.cloud.__class__.__name__ == 'SSH':
                    resource_key_dict[
                        'cloud'] = 'SSH'  # Force the cloud name to be SSH
                resource_key_dict['region'] = resources.region

            return json.dumps(resource_key_dict, sort_keys=True)

        # Print the list of resouces that the optimizer considered.
        resource_fields = ['INFRA', 'INSTANCE', 'vCPUs', 'Mem(GB)', 'GPUS']
        if len(ordered_best_plan) > 1:
            best_plan_rows = []
            for t, r in ordered_best_plan.items():
                assert t.name is not None, t
                best_plan_rows.append([t.name, str(t.num_nodes)] +
                                      _get_resources_element_list(r))
            logger.info(
                f'{colorama.Style.BRIGHT}Best plan: {colorama.Style.RESET_ALL}')
            best_plan_table = _create_table(['TASK', '#NODES'] +
                                            resource_fields)
            best_plan_table.add_rows(best_plan_rows)
            logger.info(f'{best_plan_table}')

        # Print the egress plan if any data egress is scheduled.
        Optimizer._print_egress_plan(graph, best_plan, minimize_cost)

        metric = 'COST ($)' if minimize_cost else 'TIME (hr)'
        field_names = resource_fields + [metric, 'CHOSEN']

        num_tasks = len(ordered_node_to_cost_map)
        for task, v in ordered_node_to_cost_map.items():
            # Hack: convert the dictionary values
            # (resources) to their yaml config
            # For dictionary comparison later.
            v_yaml = {
                json.dumps(resource.to_yaml_config()): cost
                for resource, cost in v.items()
            }
            task_str = (f'for task {task.name!r} ' if num_tasks > 1 else '')
            plural = 's' if task.num_nodes > 1 else ''
            if num_tasks > 1:
                # Add a new line for better readability, when there are multiple
                # tasks.
                logger.info('')
            logger.info(f'Considered resources {task_str}'
                        f'({task.num_nodes} node{plural}):')

            # Only print 1 row per cloud.
            # The following code is to generate the table
            # of optimizer table for display purpose.
            best_per_resource_group: Dict[str, Tuple[resources_lib.Resources,
                                                     float]] = {}
            for resources, cost in v.items():
                resource_table_key = _get_resource_group_hash(resources)
                if resource_table_key in best_per_resource_group:
                    if cost < best_per_resource_group[resource_table_key][1]:
                        best_per_resource_group[resource_table_key] = (
                            resources, cost)
                else:
                    best_per_resource_group[resource_table_key] = (resources,
                                                                   cost)

            # If the DAG has multiple tasks, the chosen resources may not be
            # the best resources for the task.
            chosen_resources = best_plan[task]
            resource_table_key = _get_resource_group_hash(chosen_resources)
            best_per_resource_group[resource_table_key] = (
                chosen_resources,
                v_yaml[json.dumps(chosen_resources.to_yaml_config())])
            rows = []
            for resources, cost in best_per_resource_group.values():
                if minimize_cost:
                    cost_str = f'{cost:.2f}'
                else:
                    cost_str = f'{cost / 3600:.2f}'

                row = _get_resources_named_tuple(resources, cost_str,
                                                 resources == best_plan[task])
                rows.append(row)

            # NOTE: we've converted the cost to a string above, so we should
            # convert it back to float for sorting.
            if isinstance(task.resources, list):
                accelerator_spot_list = [
                    r.get_accelerators_str() + r.get_spot_str()
                    for r in list(task.resources)
                ]

                def sort_key(row, accelerator_spot_list=accelerator_spot_list):
                    accelerator_index = accelerator_spot_list.index(
                        row.accelerators +
                        ('[Spot]' if '[Spot]' in row.instance else ''))
                    cost = float(row.cost_str)
                    return (accelerator_index, cost)

                rows = sorted(rows, key=sort_key)
            else:
                rows = sorted(rows, key=lambda row: float(row.cost_str))

            row_list = []
            for row in rows:
                row_in_list = []
                if row.chosen_str != '':
                    for _, cell in enumerate(row):
                        row_in_list.append((f'{colorama.Style.BRIGHT}{cell}'
                                            f'{colorama.Style.RESET_ALL}'))
                else:
                    row_in_list = list(row)
                row_list.append(row_in_list)

            table = _create_table(field_names)
            table.add_rows(rows)
            logger.info(f'{table}')

            # Warning message for using disk_tier=ultra
            # TODO(yi): Consider price of disks in optimizer and
            # move this warning there.
            if chosen_resources.disk_tier == resources_utils.DiskTier.ULTRA:
                logger.warning(
                    'Using disk_tier=ultra will utilize more advanced disks '
                    '(io2 Block Express on AWS and extreme persistent disk on '
                    'GCP), which can lead to significant higher costs (~$2/h).')

    @staticmethod
    def _print_candidates(node_to_candidate_map: _TaskToPerCloudCandidates):
        for node, candidate_set in node_to_candidate_map.items():
            best_resources = node.best_resources
            if best_resources is None:
                best_resources = list(node.resources)[0]
            is_multi_instances = False
            if best_resources.accelerators:
                acc_name, acc_count = list(
                    best_resources.accelerators.items())[0]
                for cloud, candidate_list in candidate_set.items():
                    # Filter only the candidates matching the best
                    # resources chosen by the optimizer.
                    best_resources_candidates = [
                        res for res in candidate_list if
                        res.get_accelerators_str() == f'{acc_name}:{acc_count}'
                    ]
                    if len(best_resources_candidates) > 1:
                        is_multi_instances = True
                        instance_list = set([
                            res.instance_type
                            for res in best_resources_candidates
                            if res.instance_type is not None
                        ])
                        candidate_str = resources_utils.format_resource(
                            best_resources, simplified_only=True)[0]

                        logger.info(
                            f'{colorama.Style.DIM}🔍 Multiple {cloud} instances '
                            f'satisfy {acc_name}:{int(acc_count)}. '
                            f'The cheapest {candidate_str} is considered '
                            f'among: {", ".join(instance_list)}.'
                            f'{colorama.Style.RESET_ALL}')
            if is_multi_instances:
                logger.info(
                    f'To list more details, run: sky gpus list {acc_name}\n')

    @staticmethod
    @timeline.event
    def optimize_job_group(
            dag: 'dag_lib.Dag',
            minimize: common.OptimizeTarget = common.OptimizeTarget.COST,
            blocked_resources: Optional[Iterable[
                resources_lib.Resources]] = None,
            quiet: bool = False) -> 'dag_lib.Dag':
        """Optimize a JobGroup DAG.

        This method optimizes all tasks in a JobGroup to run on the same
        infrastructure (cloud + region/zone or K8s cluster).

        Args:
            dag: The JobGroup DAG to optimize.
            minimize: Whether to minimize cost or time.
            blocked_resources: Resources that should not be used.
            quiet: Whether to suppress logging.

        Returns:
            The optimized DAG with best_resources set for each task.

        Raises:
            exceptions.ResourcesUnavailableError: If no common infrastructure
                can satisfy all tasks.
        """
        if not dag.is_job_group():
            # Fall back to normal optimization for non-JobGroup DAGs
            return Optimizer.optimize(dag, minimize, blocked_resources, quiet)

        tasks = dag.tasks

        if not quiet:
            logger.info(
                f'Optimizing JobGroup "{dag.name}" with {len(tasks)} jobs')

        # Find common infrastructure for all tasks in the JobGroup
        return Optimizer._optimize_same_infra(dag, minimize, blocked_resources,
                                              quiet)

    @staticmethod
    def _optimize_independent(dag: 'dag_lib.Dag',
                              minimize: common.OptimizeTarget,
                              blocked_resources: Optional[Iterable[
                                  resources_lib.Resources]],
                              quiet: bool) -> 'dag_lib.Dag':
        """Optimize each task in the JobGroup independently."""
        for task in dag.tasks:
            # Create a temporary single-task DAG for optimization
            temp_dag = dag_lib.Dag()
            temp_dag.add(task)
            temp_dag.name = task.name

            # Optimize this task
            Optimizer.optimize(temp_dag, minimize, blocked_resources, quiet)

        return dag

    @staticmethod
    def _optimize_same_infra(dag: 'dag_lib.Dag',
                             minimize: common.OptimizeTarget,
                             blocked_resources: Optional[Iterable[
                                 resources_lib.Resources]],
                             quiet: bool) -> 'dag_lib.Dag':
        """Optimize JobGroup with SAME_INFRA constraint.

        Find a common cloud+region that can satisfy all tasks.
        """
        tasks = dag.tasks

        # Step 1: Get feasible resources for each task
        # Use launchable_resources (with regions) instead of cloud_candidates
        # (without regions) to ensure proper region constraints.
        task_launchables: Dict[task_lib.Task, _PerCloudCandidates] = {}

        for task in tasks:
            # Get launchable resources for this task
            launchable_resources, _, _, _ = (_fill_in_launchable_resources(
                task, blocked_resources=blocked_resources, quiet=quiet))

            if not any(launchable_resources.values()):
                with ux_utils.print_exception_no_traceback():
                    raise exceptions.ResourcesUnavailableError(
                        f'No resources available for job "{task.name}" '
                        f'in JobGroup "{dag.name}"')

            # Build cloud -> launchable resources mapping (with regions)
            cloud_launchables: _PerCloudCandidates = collections.defaultdict(
                list)
            for launchable_list in launchable_resources.values():
                for res in launchable_list:
                    if res.cloud is not None:
                        cloud_launchables[res.cloud].append(res)
            task_launchables[task] = cloud_launchables

        # Step 2: Find common cloud+region combinations
        common_infras = Optimizer._find_common_infras(task_launchables)

        if not common_infras:
            # If no common infra, fallback to independent optimization
            if not quiet:
                logger.warning('No common infrastructure found for all jobs. '
                               'Falling back to independent optimization.')
            return Optimizer._optimize_independent(dag, minimize,
                                                   blocked_resources, quiet)

        # Step 3: Select best infra based on minimize target
        best_infra = Optimizer._select_best_infra(
            common_infras, task_launchables, tasks,
            minimize == common.OptimizeTarget.COST)

        if not quiet:
            cloud, region = best_infra
            # Format infra as lowercase cloud/region
            infra_str = f'{str(cloud).lower()}/{region}'
            logger.info(f'Selected infrastructure: {infra_str}')
            # Hint user about other available infras
            other_infras = [(c, r)
                            for c, r in common_infras
                            if not (str(c) == str(cloud) and r == region)]
            if other_infras:
                other_infras_str = ', '.join(
                    [f'{str(c).lower()}/{r}' for c, r in other_infras[:3]])
                if len(other_infras) > 3:
                    other_infras_str += f', ... ({len(other_infras) - 3} more)'
                logger.info(
                    f'Other available common infras: {other_infras_str}')

        # Step 4: Assign resources for each task on the selected infra
        cloud, region = best_infra
        cloud_name = str(cloud)
        for task in tasks:
            candidates = task_launchables[task]
            # Find the cloud object in candidates that matches (by name)
            # since different tasks may have different Cloud instances
            matching_cloud = None
            for cand_cloud in candidates.keys():
                if str(cand_cloud) == cloud_name:
                    matching_cloud = cand_cloud
                    break
            if matching_cloud is None:
                continue

            # Find resources in this cloud+region
            for resources in candidates[matching_cloud]:
                if resources.region == region:
                    # Set best_resources on the task
                    task.best_resources = resources
                    # Also set resources override to ensure the constraint
                    # persists through YAML serialization to the controller.
                    # Without this, the controller would re-optimize each
                    # task independently, placing them on different infras.
                    override_params: Dict[str, Any] = {}
                    if resources.cloud is not None:
                        override_params['cloud'] = resources.cloud
                    if resources.region is not None:
                        override_params['region'] = resources.region
                    if override_params:
                        task.set_resources_override(override_params)
                    break

        # Step 5: Print optimizer table for job groups
        if not quiet and len(tasks) > 1:
            Optimizer._print_job_group_plan(tasks)

        return dag

    @staticmethod
    def _print_job_group_plan(tasks: List[task_lib.Task]) -> None:
        """Print the optimizer table for a job group."""
        resource_fields = ['INFRA', 'INSTANCE', 'vCPUs', 'Mem(GB)', 'GPUS']
        table = _create_table(['TASK', '#NODES'] + resource_fields)

        rows = []
        for task in tasks:
            best_resources = task.best_resources
            if best_resources is None:
                continue

            # Get instance type string (display '-' for K8s/Slurm in table)
            instance_type = best_resources.instance_type
            if instance_type is None:
                display_instance_type = '-'
            elif isinstance(best_resources.cloud,
                            (clouds.Kubernetes, clouds.Slurm)):
                display_instance_type = '-'
            else:
                display_instance_type = instance_type

            # Get vCPUs and memory
            vcpus = '-'
            mem = '-'
            if best_resources.cloud is not None and instance_type is not None:
                cloud = best_resources.cloud
                vcpus_, mem_ = cloud.get_vcpus_mem_from_instance_type(
                    instance_type)
                if vcpus_ is not None:
                    vcpus = (str(int(vcpus_))
                             if vcpus_.is_integer() else f'{vcpus_:.1f}')
                if mem_ is not None:
                    mem = (str(int(mem_))
                           if mem_.is_integer() else f'{mem_:.1f}')

            # Get accelerators
            accelerators = best_resources.get_accelerators_str()

            # Get spot string
            spot = best_resources.get_spot_str()

            # Get infra string
            infra = best_resources.infra.formatted_str()

            row = [
                task.name,
                str(task.num_nodes), infra, display_instance_type + spot, vcpus,
                mem,
                str(accelerators)
            ]
            rows.append(row)

        if rows:
            table.add_rows(rows)
            logger.info(f'{colorama.Style.BRIGHT}Best plan: '
                        f'{colorama.Style.RESET_ALL}')
            logger.info(f'{table}')

    @staticmethod
    def _find_common_infras(
        task_candidates: Dict[task_lib.Task, _PerCloudCandidates]
    ) -> List[Tuple[clouds.Cloud, Optional[str]]]:
        """Find cloud+region combinations that satisfy all tasks.

        Returns:
            List of (cloud, region) tuples that can run all tasks.
        """
        if not task_candidates:
            return []

        # Collect infras per task
        infras_per_task: List[Set[Tuple[str, Optional[str]]]] = []

        for _, cloud_candidates in task_candidates.items():
            task_infras: Set[Tuple[str, Optional[str]]] = set()
            for cloud, resources_list in cloud_candidates.items():
                cloud_name = str(cloud)
                for resources in resources_list:
                    # Use (cloud_name, region) as infra identifier
                    region = resources.region
                    task_infras.add((cloud_name, region))
            infras_per_task.append(task_infras)

        if not infras_per_task:
            return []

        # Find intersection of all task infras
        common_infras = infras_per_task[0]
        for infra_set in infras_per_task[1:]:
            common_infras = common_infras & infra_set

        # Convert back to (Cloud, region) tuples
        result: List[Tuple[clouds.Cloud, Optional[str]]] = []
        for cloud_name, region in common_infras:
            # Get Cloud object from name
            cloud_obj = None
            for task_cloud_candidates in task_candidates.values():
                for cloud in task_cloud_candidates.keys():
                    if str(cloud) == cloud_name:
                        cloud_obj = cloud
                        break
                if cloud_obj:
                    break
            if cloud_obj:
                result.append((cloud_obj, region))

        return result

    @staticmethod
    def _select_best_infra(
            common_infras: List[Tuple[clouds.Cloud, Optional[str]]],
            task_candidates: Dict[task_lib.Task, _PerCloudCandidates],
            tasks: List[task_lib.Task],
            minimize_cost: bool) -> Tuple[clouds.Cloud, Optional[str]]:
        """Select the best infrastructure from common options.

        Args:
            common_infras: List of (cloud, region) that can run all tasks.
            task_candidates: Per-task candidate resources.
            tasks: List of tasks.
            minimize_cost: If True, minimize cost; else minimize time.

        Returns:
            Best (cloud, region) tuple.
        """
        if len(common_infras) == 1:
            return common_infras[0]

        # Estimate total cost/time for each infra
        best_infra: Optional[Tuple[clouds.Cloud, Optional[str]]] = None
        best_score = float('inf')

        for cloud, region in common_infras:
            total_score = 0.0
            all_tasks_valid = True

            for task in tasks:
                candidates = task_candidates.get(task, {})
                if cloud not in candidates:
                    # Task cannot run on this cloud - skip this infra
                    all_tasks_valid = False
                    break

                # Find cheapest/fastest resources in this infra
                best_task_score = float('inf')
                for resources in candidates[cloud]:
                    if resources.region == region or resources.region is None:
                        # Estimate score for this resource
                        if task.time_estimator_func is None:
                            runtime = 3600  # Default 1 hour
                        else:
                            runtime = task.estimate_runtime(resources)
                            if runtime is None:
                                runtime = 3600  # Default 1 hour

                        if minimize_cost:
                            hourly_cost = resources.get_cost(runtime)
                            score = hourly_cost * task.num_nodes
                        else:
                            score = runtime

                        if score < best_task_score:
                            best_task_score = score

                if best_task_score < float('inf'):
                    total_score += best_task_score
                else:
                    # No valid resources found for this task on this infra
                    all_tasks_valid = False
                    break

            if all_tasks_valid and total_score < best_score:
                best_score = total_score
                best_infra = (cloud, region)

        return best_infra if best_infra else common_infras[0]

    @staticmethod
    def _optimize_dag(
        dag: 'dag_lib.Dag',
        minimize_cost: bool = True,
        blocked_resources: Optional[Iterable[resources_lib.Resources]] = None,
        quiet: bool = False,
    ) -> Dict[task_lib.Task, resources_lib.Resources]:
        """Finds the optimal task-resource mapping for the entire DAG.

        The optimal mapping should consider the egress cost/time so that
        the total estimated cost/time of the DAG becomes the minimum.
        """

        # TODO: The output of this function is useful. Should generate a
        # text plan and print to both console and a log file.

        def ordinal_number(n):
            if 10 <= n % 100 <= 20:
                suffix = 'th'
            else:
                suffix = {1: 'st', 2: 'nd', 3: 'rd'}.get(n % 10, 'th')
            return str(n) + suffix

        has_resources_ordered = _is_dag_resources_ordered(dag)
        if has_resources_ordered:
            # Honor the user's choice.
            # The actual task dag can store dummy tasks.
            for task_id, task in enumerate(dag.tasks):
                if isinstance(task.resources, list):
                    resources_list = task.resources
                    accelerators_str = ', '.join(
                        [f'{r}' for r in resources_list])
                    task_id_str = ordinal_number(task_id + 1)
                    if len(dag.tasks) > 3:
                        # User-provided dag has more than one task.
                        # Comparing with 3,
                        # as there are two dummy tasks added by the optimizer.
                        logger.info(f'{colorama.Fore.YELLOW}{task_id_str} '
                                    'task is using user-specified '
                                    'accelerators list '
                                    f'{colorama.Style.RESET_ALL}'
                                    '(will be tried in the listed order): '
                                    f'{accelerators_str}')
                    else:
                        logger.info(
                            f'{colorama.Fore.YELLOW}Using user-specified '
                            'accelerators list '
                            f'{colorama.Style.RESET_ALL}'
                            '(will be tried in the listed order): '
                            f'{accelerators_str}')

        graph = dag.get_graph()
        local_dag = copy.deepcopy(dag) if has_resources_ordered else dag
        for task_id in range(len(dag.tasks)):
            task = dag.tasks[task_id]
            if isinstance(task.resources, list):
                # For ordered resources, we try the resources in the order
                # specified by the user.
                local_task = local_dag.tasks[task_id]
                for resources in task.resources:
                    # Check if there exists launchable resources
                    local_task.set_resources(resources)
                    launchable_resources_map, _, _, _ = (
                        _fill_in_launchable_resources(
                            task=local_task,
                            blocked_resources=blocked_resources,
                            quiet=False))
                    if launchable_resources_map.get(resources, []):
                        break

        local_graph = local_dag.get_graph()
        local_topo_order = list(nx.topological_sort(local_graph))
        local_node_to_cost_map, local_node_to_candidate_map = (
            Optimizer._estimate_nodes_cost_or_time(local_topo_order,
                                                   minimize_cost,
                                                   blocked_resources))
        if local_dag.is_chain():
            local_best_plan, best_total_objective = Optimizer._optimize_by_dp(
                local_topo_order, local_node_to_cost_map, minimize_cost)
        else:
            local_best_plan, best_total_objective = Optimizer._optimize_by_ilp(
                local_graph, local_topo_order, local_node_to_cost_map,
                minimize_cost)

        if minimize_cost:
            total_time = Optimizer._compute_total_time(local_graph,
                                                       local_topo_order,
                                                       local_best_plan)
            total_cost = best_total_objective
        else:
            total_time = best_total_objective
            total_cost = Optimizer._compute_total_cost(local_graph,
                                                       local_topo_order,
                                                       local_best_plan)

        if local_best_plan is None:
            error_msg = (f'No launchable resource found for task {task}. '
                         'To fix: relax its resource requirements.\n'
                         'Hint: \'sky gpus list --all\' '
                         'to list available accelerators.\n'
                         '      \'sky check\' to check the enabled clouds.')
            with ux_utils.print_exception_no_traceback():
                raise exceptions.ResourcesUnavailableError(error_msg)

        if has_resources_ordered:
            best_plan = {}
            # We have to manually set the best_resources for the tasks in the
            # original dag, to pass the optimization results to the caller, as
            # we deep copied the dag when the dag has nodes with ordered
            # resources.
            for task, resources in local_best_plan.items():
                task_idx = local_dag.tasks.index(task)
                dag.tasks[task_idx].best_resources = resources
                best_plan[dag.tasks[task_idx]] = resources

            topo_order = list(nx.topological_sort(graph))
            # Get the cost of each specified resources for display purpose.
            node_to_cost_map, _ = Optimizer._estimate_nodes_cost_or_time(
                topo_order=topo_order,
                minimize_cost=minimize_cost,
                blocked_resources=blocked_resources,
                quiet=True)
        else:
            best_plan = local_best_plan
            topo_order = local_topo_order
            node_to_cost_map = local_node_to_cost_map

        if not quiet:
            Optimizer.print_optimized_plan(graph, topo_order, best_plan,
                                           total_time, total_cost,
                                           node_to_cost_map, minimize_cost)
            Optimizer._print_candidates(local_node_to_candidate_map)
        return best_plan


class DummyResources(resources_lib.Resources):
    """A dummy Resources that has zero egress cost from/to."""

    _REPR = 'DummyResources'

    def __repr__(self) -> str:
        return DummyResources._REPR

    def get_cost(self, seconds):
        return 0


def _filter_out_blocked_launchable_resources(
        launchable_resources: Iterable[resources_lib.Resources],
        blocked_resources: Iterable[resources_lib.Resources]):
    """Whether the resources are blocked."""
    available_resources = []
    for resources in launchable_resources:
        for blocked in blocked_resources:
            if resources.should_be_blocked_by(blocked):
                break
        else:  # non-blocked launchable resources. (no break)
            available_resources.append(resources)
    return available_resources


def _check_specified_clouds(dag: 'dag_lib.Dag') -> None:
    """Check if specified clouds are enabled in cache and refresh if needed.

    Our enabled cloud list is cached in a local database, and if a user
    specified a cloud that is not enabled, we should refresh the cache for that
    cloud in case the cloud access has been enabled since the last cache update.

    Args:
        dag: The DAG specified by a user.
    """
    enabled_clouds = sky_check.get_cached_enabled_clouds_or_refresh(
        capability=sky_cloud.CloudCapability.COMPUTE,
        raise_if_no_cloud_access=True)

    global_disabled_clouds: Set[str] = set()
    for task in dag.tasks:
        # Recheck the enabled clouds if the task's requested resources are on a
        # cloud that is not enabled in the cached enabled_clouds.
        all_clouds_specified: Set[str] = set()
        clouds_need_recheck: Set[str] = set()
        for resources in task.resources:
            cloud_str = str(resources.cloud)
            if (resources.cloud is not None and not clouds.cloud_in_iterable(
                    resources.cloud, enabled_clouds)):
                # Explicitly check again to update the enabled cloud list.
                clouds_need_recheck.add(cloud_str)
            all_clouds_specified.add(cloud_str)

        # Explicitly check again to update the enabled cloud list.
        clouds_to_check_again = list(clouds_need_recheck -
                                     global_disabled_clouds)
        if len(clouds_to_check_again) > 0:
            sky_check.check_capability(
                sky_cloud.CloudCapability.COMPUTE,
                quiet=True,
                clouds=clouds_to_check_again,
                workspace=skypilot_config.get_active_workspace())
        enabled_clouds = sky_check.get_cached_enabled_clouds_or_refresh(
            capability=sky_cloud.CloudCapability.COMPUTE,
            raise_if_no_cloud_access=True)
        disabled_clouds = (clouds_need_recheck -
                           {str(c) for c in enabled_clouds})
        global_disabled_clouds.update(disabled_clouds)
        if disabled_clouds:
            is_or_are = 'is' if len(disabled_clouds) == 1 else 'are'
            task_name = f' {task.name!r}' if task.name is not None else ''
            disabled_display_names = []
            for c in disabled_clouds:
                cloud_obj_one = registry.CLOUD_REGISTRY.from_str(c)
                if cloud_obj_one is not None:
                    disabled_display_names.append(cloud_obj_one.display_name())
            cloud_names = ', '.join(disabled_display_names)
            msg = (f'Task{task_name} requires {cloud_names} '
                   f'which {is_or_are} not enabled. To enable access, change '
                   f'the task cloud requirement or run: {colorama.Style.BRIGHT}'
                   f'sky check {" ".join(c.lower() for c in disabled_clouds)}'
                   f'{colorama.Style.RESET_ALL}')
            if all_clouds_specified == disabled_clouds:
                # If all resources are specified with a disabled cloud, we
                # should raise an error as no resource can satisfy the
                # requirement. Otherwise, we should just skip the resource.
                with ux_utils.print_exception_no_traceback():
                    raise exceptions.ResourcesUnavailableError(msg)
            logger.warning(
                f'{colorama.Fore.YELLOW}{msg}{colorama.Style.RESET_ALL}')

        _check_specified_regions(task)


def _check_specified_regions(task: task_lib.Task) -> None:
    """Check if specified regions (Kubernetes/SSH contexts) are enabled.

    Args:
        task: The task to check.
    """
    # Only check for Kubernetes/SSH for now
    # Below check works because SSH inherits Kubernetes cloud.
    if not all(
            isinstance(resources.cloud, clouds.Kubernetes)
            for resources in task.resources):
        return
    # Kubernetes region is a context if set
    for resources in task.resources:
        if resources.region is None:
            continue

        is_ssh = isinstance(resources.cloud, clouds.SSH)
        if is_ssh:
            existing_contexts = clouds.SSH.existing_allowed_contexts()
        else:
            existing_contexts = clouds.Kubernetes.existing_allowed_contexts()

        region = resources.region
        task_name = f' {task.name!r}' if task.name is not None else ''
        msg = f'Task{task_name} requires '
        if region not in existing_contexts:
            if is_ssh:
                infra_str = f'SSH/{common_utils.removeprefix(region, "ssh-")}'
            else:
                infra_str = f'Kubernetes/{region}'
            logger.warning(f'{infra_str} is not enabled.')
            volume_mounts_str = ''
            if task.volume_mounts:
                if len(task.volume_mounts) > 1:
                    volume_mounts_str += 'volumes '
                else:
                    volume_mounts_str += 'volume '
                volume_mounts_str += ', '.join(
                    [f'{v.volume_name}' for v in task.volume_mounts])
                volume_mounts_str += f' with infra {infra_str}'
            if volume_mounts_str:
                msg += volume_mounts_str
            else:
                msg += f'infra {infra_str}'
            msg += (
                f' which is not enabled. To enable access, change '
                f'the task infra requirement or run: {colorama.Style.BRIGHT}'
                f'sky check {colorama.Style.RESET_ALL}'
                f'to ensure the infra is enabled.')
            with ux_utils.print_exception_no_traceback():
                raise exceptions.ResourcesUnavailableError(msg)


def _fill_in_launchable_resources(
    task: task_lib.Task,
    blocked_resources: Optional[Iterable[resources_lib.Resources]],
    quiet: bool = False
) -> Tuple[Dict[resources_lib.Resources, List[resources_lib.Resources]],
           _PerCloudCandidates, List[str], Dict[resources_lib.Resources,
                                                List[str]]]:
    """Fills in the launchable resources for the task.

    Returns:
      A tuple of:
        Dict mapping the task's requested Resources to a list of launchable
          Resources,
        Dict mapping Cloud to a list of feasible Resources (for printing),
        Sorted list of fuzzy candidates (alternative GPU names).
        Dict mapping requested Resources and a list of hints for why the
          resource is unavailable if so.
    Raises:
      ResourcesUnavailableError: if all resources required by the task are on
        a cloud that is not enabled.
    """
    enabled_clouds = sky_check.get_cached_enabled_clouds_or_refresh(
        capability=sky_cloud.CloudCapability.COMPUTE,
        raise_if_no_cloud_access=True)

    launchable: Dict[resources_lib.Resources, List[resources_lib.Resources]] = (
        collections.defaultdict(list))
    all_fuzzy_candidates = set()
    cloud_candidates: _PerCloudCandidates = collections.defaultdict(list)
    resource_hints: Dict[resources_lib.Resources,
                         List[str]] = collections.defaultdict(list)
    if blocked_resources is None:
        blocked_resources = []
    for resources in task.resources:
        # Validate the resources first which may fill in missing fields
        # automatically for the resources.
        resources.validate()
        if (resources.cloud is not None and
                not clouds.cloud_in_iterable(resources.cloud, enabled_clouds)):
            # Skip the resources that are on a cloud that is not enabled. The
            # hint has been printed in _check_specified_clouds.
            launchable[resources] = []
            continue
        clouds_list = ([resources.cloud]
                       if resources.cloud is not None else enabled_clouds)
        # If clouds provide hints, store them for later printing.
        hints: Dict[clouds.Cloud, str] = {}

        feasible_list = subprocess_utils.run_in_parallel(
            lambda cloud, r=resources, n=task.num_nodes:
            (cloud, cloud.get_feasible_launchable_resources(r, n)),
            clouds_list)
        for cloud, feasible_resources in feasible_list:
            if feasible_resources.hint is not None:
                hints[cloud] = feasible_resources.hint
                resource_hints[resources].append(feasible_resources.hint)
            if feasible_resources.resources_list:
                # Assume feasible_resources is sorted by prices. Guaranteed by
                # the implementation of get_feasible_launchable_resources and
                # the underlying catalog filtering
                cheapest = feasible_resources.resources_list[0]
                # Generate region/zone-specified resources.
                launchable[resources].extend(
                    resources_utils.make_launchables_for_valid_region_zones(
                        cheapest))
                # Each cloud can occur multiple times in feasible_list,
                # for different region/zone.
                cloud_candidates[cloud].extend(
                    feasible_resources.resources_list)
            else:
                all_fuzzy_candidates.update(
                    feasible_resources.fuzzy_candidate_list)
        if not launchable[resources]:
            clouds_str = str(clouds_list) if len(clouds_list) > 1 else str(
                clouds_list[0])
            num_node_str = ''
            if task.num_nodes > 1:
                num_node_str = f'{task.num_nodes}x '
            if not (quiet or resources.no_missing_accel_warnings):
                logger.info(
                    f'No resource satisfying {num_node_str}'
                    f'{resources.repr_with_region_zone} on {clouds_str}.')
                if all_fuzzy_candidates:
                    logger.info('Did you mean: '
                                f'{colorama.Fore.CYAN}'
                                f'{sorted(all_fuzzy_candidates)}'
                                f'{colorama.Style.RESET_ALL}')
                else:
                    if resources.cpus is not None:
                        logger.info(f'{colorama.Fore.LIGHTBLACK_EX}'
                                    '- Try specifying a different CPU count, '
                                    'or add "+" to the end of the CPU count '
                                    'to allow for larger instances.'
                                    f'{colorama.Style.RESET_ALL}')
                    if resources.memory is not None:
                        logger.info(f'{colorama.Fore.LIGHTBLACK_EX}'
                                    '- Try specifying a different memory size, '
                                    'or add "+" to the end of the memory size '
                                    'to allow for larger instances.'
                                    f'{colorama.Style.RESET_ALL}')
                    if resources.local_disk is not None:
                        logger.info(
                            f'{colorama.Fore.LIGHTBLACK_EX}'
                            '- Try using "+" suffix for at-least matching '
                            '(e.g., "nvme:500+"), or reduce the size '
                            f'requirement. {colorama.Style.RESET_ALL}')
                for cloud, hint in hints.items():
                    logger.info(f'{colorama.Fore.LIGHTBLACK_EX}'
                                f'{repr(cloud)}: {hint}'
                                f'{colorama.Style.RESET_ALL}')

        launchable[resources] = _filter_out_blocked_launchable_resources(
            launchable[resources], blocked_resources)
    return launchable, cloud_candidates, list(
        sorted(all_fuzzy_candidates)), resource_hints


# Expose optimize_job_group as a module-level function
def optimize_job_group(
        dag: 'dag_lib.Dag',
        minimize: common.OptimizeTarget = common.OptimizeTarget.COST,
        blocked_resources: Optional[Iterable[resources_lib.Resources]] = None,
        quiet: bool = False) -> 'dag_lib.Dag':
    """Optimize a JobGroup DAG with SAME_INFRA constraint.

    This is a convenience wrapper around Optimizer.optimize_job_group().
    """
    return Optimizer.optimize_job_group(dag, minimize, blocked_resources, quiet)
