Repository: skypilot-org/skypilot
Pull request: Refactor DAG optimizer
Closes #545.

This PR refactors the Sky DAG optimizer. This PR includes:
* breaking up the `_optimize_cost` function, which was quite long (~170 lines) and a bit hard to follow. Now the optimization process is split into multiple smaller stages/functions.
* simplifying the DP algorithm by restricting the input graphs to chain DAGs. As reported in #556, the DP algorithm is not applicable to non-chain DAGs. So I pruned out the code for non-chain DAGs.
* more clearly printing the estimated run time and estimated cost. Now the optimizer prints both estimated run time and estimated total cost.
 
  * Before:
```bash
$ python examples/example_app.py
Optimizer - plan minimizing cost (~$57.1):
```
*
  * After: 
```bash
$ python examples/example_app.py
Optimizer - plan minimizing cost
Estimated Run time: ~21.7 hr, Cost: ~$57.1
```
* fixing a bug in printing the DAG optimization result.

NOTE: the problem in #556 is not fixed by this PR. The problem will fixed by #637, which should be merged after this PR.

Tested:
* `examples/example_app.py`
* smoke tests.

Problem statement:
Optimizer's message should be clarified
Kevin's question: what does $0.9 mean?
```
(sky) weichiang@blaze:~/repos/sky$ sky gpunode --cloud aws
I 03-11 13:58:37 optimizer.py:355] Optimizer - plan minimizing cost (~$0.9):
I 03-11 13:58:37 optimizer.py:370] 
I 03-11 13:58:37 optimizer.py:370] TASK     BEST_RESOURCE
I 03-11 13:58:37 optimizer.py:370] gpunode  AWS(p2.xlarge, {'K80': 1})
I 03-11 13:58:37 optimizer.py:370] 
```

Hints from issue discussion:
I assume this is because our default time estimator is set to 1 hour... since we don't really seem to use those we might as well update to $0.9/hour
Yes, we should clarify.  Currently, that number is `sum(estimated cost of node n) for n in DAG`.  The estimated cost of each node is based on its estimated run time, which is either
- (if time estimator specified) some number from the estimator [see `example_app.py`]; or
- (otherwise) 1 hour

So, it's not sufficient to just display an extra "per hour" string there.

How about doing the math correctly, such as dividing that total cost by the total estimated run time?   Then we can display "per hour" next to the price.
@WoosukKwon fyi - may be a good first issue to get to know the optimizer code.
@concretevitamin Yeah I'll take a look at it.

Commit message: yapf

Patch to review:
diff --git a/examples/example_app.py b/examples/example_app.py
index add2891f0e5..d09b2c9c7aa 100644
--- a/examples/example_app.py
+++ b/examples/example_app.py
@@ -76,4 +76,4 @@ def make_application():
 
 dag = make_application()
 sky.optimize(dag, minimize=sky.OptimizeTarget.COST)
-# sky.optimize(dag, minimize=OptimizeTarget.TIME)
+# sky.optimize(dag, minimize=sky.OptimizeTarget.TIME)
diff --git a/sky/optimizer.py b/sky/optimizer.py
index 43e2eda6f1b..f0133945918 100644
--- a/sky/optimizer.py
+++ b/sky/optimizer.py
@@ -165,35 +165,34 @@ def _egress_cost_or_time(minimize_cost: bool, parent: Task,
         return fn(src_cloud, dst_cloud, nbytes)
 
     @staticmethod
-    def _optimize_cost(
-        dag: Dag,
+    def _per_node_cost_or_time(
+        topo_order: List[Task],
         minimize_cost: bool = True,
         blocked_launchable_resources: Optional[List[Resources]] = None,
         raise_error: bool = False,
-    ):
-        import networkx as nx  # pylint: disable=import-outside-toplevel
-        # TODO: The output of this function is useful. Should generate a
-        # text plan and print to both console and a log file.
-        graph = dag.get_graph()
-        topo_order = list(nx.topological_sort(graph))
-
-        # FIXME: write to (node, cloud) as egress cost depends only on clouds.
-        # node -> {resources -> best estimated cost}
-        dp_best_cost = collections.defaultdict(dict)
-        # d[node][resources][parent] = (best parent resources, best parent cost)
-        dp_point_backs = collections.defaultdict(
-            lambda: collections.defaultdict(dict))
+    ) -> Tuple[Dict[Task, Dict[Resources, float]], \
+        Dict[Task, Dict[str, List[Resources]]]]:
+        """Estimates the compute cost of feasible task-resource mappings."""
+        # Cost of running the task on the resources
+        # node -> {resources -> cost}
+        node_to_cost_map = collections.defaultdict(dict)
+
+        # If a cloud has multiple instance types with the same accelerators,
+        # Sky informs users of the candidates and its decisions.
+        # node -> cloud candidates
         node_to_candidates = collections.defaultdict(dict)
 
+        # Compute the estimated cost/time for each node
         for node_i, node in enumerate(topo_order):
-            # Base case: a special source node.
             if node_i == 0:
-                dp_best_cost[node][list(node.get_resources())[0]] = 0
+                # Base case: a special source node.
+                node_to_cost_map[node][list(node.get_resources())[0]] = 0
                 continue
-            # Don't print for the last node, Sink.
+
             do_print = node_i != len(topo_order) - 1
             if do_print:
                 logger.debug('#### {} ####'.format(node))
+
             if node_i < len(topo_order) - 1:
                 # Convert partial resource labels to launchable resources.
                 launchable_resources, cloud_candidates = \
@@ -201,15 +200,15 @@ def _optimize_cost(
                         node,
                         blocked_launchable_resources
                     )
-                node_to_candidates[node_i] = cloud_candidates
+                node_to_candidates[node] = cloud_candidates
             else:
                 # Dummy sink node.
                 launchable_resources = node.get_resources()
                 launchable_resources = {
                     list(node.get_resources())[0]: launchable_resources
                 }
+
             num_resources = len(node.get_resources())
-            parents = list(graph.predecessors(node))
             for orig_resources, launchable_list in launchable_resources.items():
                 if not launchable_list:
                     error_msg = (
@@ -240,12 +239,6 @@ def _optimize_cost(
                     # part of a Resources.
                     estimated_runtime = node.estimate_runtime(orig_resources)
                 for resources in launchable_list:
-                    # Computes dp_best_cost[node][resources]
-                    # = my estimated cost +
-                    #      sum_p min_phw { dp_best_cost(p, phw) +
-                    #                      egress_cost(p, phw, hw) }
-                    # where p in Parents(node).
-                    assert resources not in dp_best_cost[node]
                     if do_print:
                         logger.debug(f'resources: {resources}')
 
@@ -263,57 +256,175 @@ def _optimize_cost(
                             logger.debug(
                                 '  estimated_cost (not incl. egress): ${:.1f}'.
                                 format(estimated_cost))
+                    node_to_cost_map[node][resources] = estimated_cost
+        return node_to_cost_map, node_to_candidates
+
+    @staticmethod
+    def _optimize_by_dp(
+        topo_order: List[Task],
+        node_to_cost_map: Dict[Task, Dict[Resources, float]],
+        minimize_cost: bool = True,
+    ) -> Tuple[Dict[Task, Resources], float]:
+        """Optimizes a chain DAG using a dynamic programming algorithm."""
+        # node -> { resources -> best estimated cost }
+        dp_best_cost = collections.defaultdict(dict)
+        # node -> { resources -> best parent resources }
+        dp_point_backs = collections.defaultdict(dict)
 
-                    # FIXME: Account for egress costs for multi-node clusters
-                    sum_parent_cost_and_egress = 0
-                    for parent in parents:
-                        min_pred_cost_plus_egress = np.inf
-                        for parent_resources, parent_cost in dp_best_cost[
-                                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
-                        sum_parent_cost_and_egress += min_pred_cost_plus_egress
-                        dp_point_backs[node][resources][parent] = (
-                            best_parent_hardware, min_pred_cost_plus_egress)
-                    dp_best_cost[node][
-                        resources] = estimated_cost + sum_parent_cost_and_egress
-
-        # Dict: node -> (resources, cost).
-        best_plan = Optimizer.read_optimized_plan(dp_best_cost, topo_order,
-                                                  dp_point_backs, minimize_cost)
-
-        # If it's 1 resource choice, the info is already printed.
-        should_print = any(len(v) > 1 for v in dp_best_cost.values())
+        # Computes dp_best_cost[node][resources]
+        # = my estimated cost + min_phw { dp_best_cost(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_cost[node][list(node.get_resources())[0]] = 0
+                continue
+
+            parent = topo_order[node_i - 1]
+            # FIXME: Account for egress costs for multi-node clusters
+            for resources, compute_cost in node_to_cost_map[node].items():
+                min_pred_cost_plus_egress = np.inf
+                for parent_resources, parent_cost in \
+                    dp_best_cost[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_cost[node][resources] = \
+                    compute_cost + min_pred_cost_plus_egress
+
+        # Compute the total cost of the DAG.
+        sink_node = topo_order[-1]
+        total_costs = dp_best_cost[sink_node]
+        assert len(total_costs) == 1, f'Should be DummyCloud: {total_costs}'
+        best_resources, best_total_cost = list(total_costs.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_cost
+
+    @staticmethod
+    def _compute_total_time(
+        graph,
+        topo_order: List[Task],
+        plan: Dict[Task, Resources],
+    ) -> float:
+        cache_finish_time = {}
+
+        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:
+                compute_time = 1 * 3600
+            else:
+                compute_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] = compute_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],
+        plan: Dict[Task, Resources],
+    ) -> float:
+        total_cost = 0
+        for node in topo_order:
+            resources = plan[node]
+            if node.time_estimator_func is None:
+                compute_time = 1 * 3600
+            else:
+                compute_time = node.estimate_runtime(resources)
+
+            cost_per_node = resources.get_cost(compute_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_optimized_plan(
+        best_plan: Dict[Task, Resources],
+        total_time: float,
+        total_cost: float,
+        node_to_cost_map: Dict[Task, Dict[Resources, float]],
+        minimize_cost: bool,
+    ):
+        if minimize_cost:
+            logger.info('Optimizer - plan minimizing cost')
+        else:
+            logger.info('Optimizer - plan minimizing run time')
+        logger.info(f'Estimated Run time: ~{total_time / 3600:.1f} hr, '
+                    f'Cost: ~${total_cost:.1f}')
+
+        # Do not print Source or Sink.
+        message_data = [
+            (t, f'{t.num_nodes}x {repr(r)}' if t.num_nodes > 1 else repr(r))
+            for t, r in best_plan.items()
+            if t.name not in (_DUMMY_SOURCE_NAME, _DUMMY_SINK_NAME)
+        ]
+        message = tabulate.tabulate(reversed(message_data),
+                                    headers=['TASK', 'BEST_RESOURCE'],
+                                    tablefmt='plain')
+        logger.info(f'\n{message}\n')
+
+        # Print the list of resouces that the optimizer considered.
+        should_print = any(len(v) > 1 for v in node_to_cost_map.values())
         if should_print:
-            dp_best_cost = {
+            node_to_cost_map = {
                 k: v
-                for k, v in dp_best_cost.items()
+                for k, v in node_to_cost_map.items()
                 if k.name not in (_DUMMY_SOURCE_NAME, _DUMMY_SINK_NAME)
             }
             metric = 'cost ($)' if minimize_cost else 'time (hr)'
-            for k, v in dp_best_cost.items():
-                dp_best_cost[k] = {
+            for k, v in node_to_cost_map.items():
+                node_to_cost_map[k] = {
                     resources: round(cost, 2) if minimize_cost \
                         else round(cost / 3600, 2)
                     for resources, cost in v.items()
                 }
 
-            if len(dp_best_cost) > 1:
+            num_tasks = len(node_to_cost_map)
+            if num_tasks > 1:
                 logger.info(f'Details: task -> {{resources -> {metric}}}')
-                logger.info('%s\n', pprint.pformat(dp_best_cost))
-            elif len(dp_best_cost) == 1:
+                logger.info('%s\n', pprint.pformat(node_to_cost_map))
+            elif num_tasks == 1:
                 logger.info(f'Considered resources -> {metric}')
                 logger.info('%s\n',
-                            pprint.pformat(list(dp_best_cost.values())[0]))
+                            pprint.pformat(list(node_to_cost_map.values())[0]))
 
-        for node_i, candidate_set in node_to_candidates.items():
-            node = topo_order[node_i]
+    @staticmethod
+    def _print_candidates(node_to_candidates: Dict[Task,
+                                                   Dict[str, List[Resources]]]):
+        for node, candidate_set in node_to_candidates.items():
             accelerator = list(node.get_resources())[0].accelerators
             is_multi_instances = False
             if accelerator:
@@ -332,51 +443,52 @@ def _optimize_cost(
             if is_multi_instances:
                 logger.info(
                     f'To list more details, run \'sky show-gpus {acc_name}\'.')
-        return dag, best_plan
 
     @staticmethod
-    def read_optimized_plan(dp_best_cost, topo_order, dp_point_backs,
-                            minimize_cost):
-        message_data = []
-        best_plan = {}
+    def _optimize_cost(
+        dag: Dag,
+        minimize_cost: bool = True,
+        blocked_launchable_resources: Optional[List[Resources]] = None,
+        raise_error: bool = False,
+    ) -> Tuple[Dag, Dict[Task, Resources]]:
+        import networkx as nx  # pylint: disable=import-outside-toplevel
+        # TODO: The output of this function is useful. Should generate a
+        # text plan and print to both console and a log file.
+
+        graph = dag.get_graph()
+        topo_order = list(nx.topological_sort(graph))
 
-        def _walk(node, best_hardware, best_cost):
-            if node.best_resources is None:
-                # Record the best decision for 'node'.
-                message_data.append((node, best_hardware))
-                best_plan[node] = (best_hardware, best_cost)
-                node.best_resources = best_hardware
-            # Recurse back to parent(s).
-            for tup in dp_point_backs[node][best_hardware].items():
-                parent, (parent_best_hardware, parent_best_cost) = tup
-                _walk(parent, parent_best_hardware, parent_best_cost)
-
-        # Start at Sink, the last node in the topo order.
-        node = topo_order[-1]
-        # Find the best (hardware, cost) for node.
-        best_costs = dp_best_cost[node]
-        assert len(best_costs) == 1, f'Should be DummyCloud: {best_costs}'
-        h, overall_best = list(best_costs.items())[0]
-        _walk(node, h, overall_best)
+        node_to_cost_map, node_to_candidates = Optimizer._per_node_cost_or_time(
+            topo_order,
+            minimize_cost,
+            blocked_launchable_resources,
+            raise_error,
+        )
+
+        is_chain = True
+        for node in topo_order[:-1]:
+            if len(list(graph.successors(node))) != 1:
+                is_chain = False
+
+        if is_chain:
+            best_plan, best_total_cost = Optimizer._optimize_by_dp(
+                topo_order, node_to_cost_map, minimize_cost)
+        else:
+            raise NotImplementedError('Currently Sky only supports chain DAGs.')
 
         if minimize_cost:
-            logger.info('Optimizer - plan minimizing cost (~${:.1f}):'.format(
-                overall_best))
+            total_time = Optimizer._compute_total_time(graph, topo_order,
+                                                       best_plan)
+            total_cost = best_total_cost
         else:
-            logger.info(
-                'Optimizer - plan minimizing run time (~{:.1f} hr):'.format(
-                    overall_best / 3600))
-        # Do not print Source or Sink.
-        message_data = [
-            (t, f'{t.num_nodes}x {repr(r)}' if t.num_nodes > 1 else repr(r))
-            for (t, r) in message_data
-            if t.name not in (_DUMMY_SOURCE_NAME, _DUMMY_SINK_NAME)
-        ]
-        message = tabulate.tabulate(reversed(message_data),
-                                    headers=['TASK', 'BEST_RESOURCE'],
-                                    tablefmt='plain')
-        logger.info(f'\n{message}\n')
-        return best_plan
+            total_time = best_total_cost
+            total_cost = Optimizer._compute_total_cost(graph, topo_order,
+                                                       best_plan)
+
+        Optimizer.print_optimized_plan(best_plan, total_time, total_cost,
+                                       node_to_cost_map, minimize_cost)
+        Optimizer._print_candidates(node_to_candidates)
+        return dag, best_plan
 
 
 class DummyResources(Resources):
@@ -418,7 +530,7 @@ def _fill_in_launchable_resources(
     task: Task,
     blocked_launchable_resources: Optional[List[Resources]],
     try_fix_with_sky_check: bool = True,
-) -> Tuple[Dict[Resources, List[Resources]], Dict[str, set]]:
+) -> Tuple[Dict[Resources, List[Resources]], Dict[str, List[Resources]]]:
     enabled_clouds = global_user_state.get_enabled_clouds()
     if len(enabled_clouds) == 0 and try_fix_with_sky_check:
         check.check(quiet=True)


Code review comments:
diff_hunk: @@ -332,51 +443,52 @@ def _optimize_cost(
             if is_multi_instances:
                 logger.info(
                     f'To list more details, run \'sky show-gpus {acc_name}\'.')
-        return dag, best_plan
 
     @staticmethod
-    def read_optimized_plan(dp_best_cost, topo_order, dp_point_backs,
-                            minimize_cost):
-        message_data = []
-        best_plan = {}
+    def _optimize_cost(
+        dag: Dag,
+        minimize_cost: bool = True,
+        blocked_launchable_resources: Optional[List[Resources]] = None,
+        raise_error: bool = False,
+    ) -> Tuple[Dag, Dict[Task, Resources]]:
+        import networkx as nx  # pylint: disable=import-outside-toplevel
+        # TODO: The output of this function is useful. Should generate a
+        # text plan and print to both console and a log file.
+
+        graph = dag.get_graph()
+        topo_order = list(nx.topological_sort(graph))
 
-        def _walk(node, best_hardware, best_cost):
-            if node.best_resources is None:
-                # Record the best decision for 'node'.
-                message_data.append((node, best_hardware))
-                best_plan[node] = (best_hardware, best_cost)
-                node.best_resources = best_hardware
-            # Recurse back to parent(s).
-            for tup in dp_point_backs[node][best_hardware].items():
-                parent, (parent_best_hardware, parent_best_cost) = tup
-                _walk(parent, parent_best_hardware, parent_best_cost)
-
-        # Start at Sink, the last node in the topo order.
-        node = topo_order[-1]
-        # Find the best (hardware, cost) for node.
-        best_costs = dp_best_cost[node]
-        assert len(best_costs) == 1, f'Should be DummyCloud: {best_costs}'
-        h, overall_best = list(best_costs.items())[0]
-        _walk(node, h, overall_best)
+        node_to_cost_map, node_to_candidates = Optimizer._per_node_cost_or_time(
+            topo_order,
+            minimize_cost,
+            blocked_launchable_resources,
+            raise_error,
+        )
+
+        is_chain = True
+        for node in topo_order[:-1]:
+            if len(list(graph.successors(node))) != 1:
path: sky/optimizer.py
text: @user1:
1. Does `len(graph.successors(node))` work?  Skipping a list copy.

2. Once `is_chain = False` we can `break`.

@user1:
Minor - Can also try https://networkx.org/documentation/networkx-1.10/reference/generated/networkx.DiGraph.out_degree.html

@author:
Thanks for the suggestion!

> Does len(graph.successors(node)) work? Skipping a list copy.

No it doesn't. Error msg: `TypeError: object of type 'dict_keyiterator' has no len()`.

> Once is_chain = False we can break.

Sure. Will do that.

> Minor - Can also try https://networkx.org/documentation/networkx-1.10/reference/generated/networkx.DiGraph.out_degree.html

Yeah I'll change this line to `if graph.out_degree(node) != 1`.

diff_hunk: @@ -165,51 +165,50 @@ def _egress_cost_or_time(minimize_cost: bool, parent: Task,
         return fn(src_cloud, dst_cloud, nbytes)
 
     @staticmethod
-    def _optimize_cost(
-        dag: Dag,
+    def _per_node_cost_or_time(
+        topo_order: List[Task],
         minimize_cost: bool = True,
         blocked_launchable_resources: Optional[List[Resources]] = None,
         raise_error: bool = False,
-    ):
-        import networkx as nx  # pylint: disable=import-outside-toplevel
-        # TODO: The output of this function is useful. Should generate a
-        # text plan and print to both console and a log file.
-        graph = dag.get_graph()
-        topo_order = list(nx.topological_sort(graph))
-
-        # FIXME: write to (node, cloud) as egress cost depends only on clouds.
-        # node -> {resources -> best estimated cost}
-        dp_best_cost = collections.defaultdict(dict)
-        # d[node][resources][parent] = (best parent resources, best parent cost)
-        dp_point_backs = collections.defaultdict(
-            lambda: collections.defaultdict(dict))
+    ) -> Tuple[Dict[Task, Dict[Resources, float]], \
+        Dict[Task, Dict[str, List[Resources]]]]:
+        """Estimates the compute cost of feasible task-resource mappings."""
+        # Cost of running the task on the resources
+        # node -> {resources -> cost}
+        node_to_cost_map = collections.defaultdict(dict)
+
+        # If a cloud has multiple instance types with the same accelerators,
+        # Sky informs users of the candidates and its decisions.
+        # node -> cloud candidates
path: sky/optimizer.py
text: @user1:
Slightly inaccurate because the candidates can span clouds: e.g., AWS(p3.2xlarge), GCP(n1-..., {V100:1}), Azure(...), etc.  Update the comment's text and add this example?

@author:
Oh I think the variable name `node_to_candidates` is misleading. Its type is `task -> cloud -> List[Resources]`. That is, for each cloud, it stores a list of resources which have the same accelerators (e.g., `{AWS: ['g4dn.xlarge', 'g4dn.2xlarge', ...], Azure: ['Standard_NC4as_T4_v3', 'Standard_NC8as_T4_v3', 'Standard_NC16as_T4_v3']}` when the requested accelerator is `T4:1`).

Will clarify this.

diff_hunk: @@ -165,51 +165,50 @@ def _egress_cost_or_time(minimize_cost: bool, parent: Task,
         return fn(src_cloud, dst_cloud, nbytes)
 
     @staticmethod
-    def _optimize_cost(
-        dag: Dag,
+    def _per_node_cost_or_time(
+        topo_order: List[Task],
         minimize_cost: bool = True,
         blocked_launchable_resources: Optional[List[Resources]] = None,
         raise_error: bool = False,
-    ):
-        import networkx as nx  # pylint: disable=import-outside-toplevel
-        # TODO: The output of this function is useful. Should generate a
-        # text plan and print to both console and a log file.
-        graph = dag.get_graph()
-        topo_order = list(nx.topological_sort(graph))
-
-        # FIXME: write to (node, cloud) as egress cost depends only on clouds.
-        # node -> {resources -> best estimated cost}
-        dp_best_cost = collections.defaultdict(dict)
-        # d[node][resources][parent] = (best parent resources, best parent cost)
-        dp_point_backs = collections.defaultdict(
-            lambda: collections.defaultdict(dict))
+    ) -> Tuple[Dict[Task, Dict[Resources, float]], \
+        Dict[Task, Dict[str, List[Resources]]]]:
path: sky/optimizer.py
text: @user1:
style: 
If a single name and type is too long, consider using an [alias](https://google.github.io/styleguide/pyguide.html#typing-aliases) for the type.  https://google.github.io/styleguide/pyguide.html#3196-type-aliases

Alternatively, since it's pretty complex here, maybe create a container class / namedtuple to hold these fields. Example: 
- https://github.com/sky-proj/sky/blob/master/sky/backends/cloud_vm_ray_backend.py#L1104-L1145
- https://github.com/sky-proj/sky/blob/master/sky/backends/cloud_vm_ray_backend.py#L408-L422

@author:
Thanks for the suggestion. Yeah The type is very complicated.

As the type length is 2, I think creating a new container class is not useful here. Will try type alias instead.