ILP-based DAG Optimizer (skypilot-org#637)

WoosukKwon · web-flow · commit 5da7c4dd9cf3 · 2022-04-02T11:50:22.000-07:00
* Refactor optimizer

* Remove unnecessary import

* yapf

* Minor fix

* Add NotImplementedError

* ILP-based optimization

* yapf

* Add pulp in setup.py

* Minor

* Rename vars &amp; Annotate types

* Minor fix

* Minor

* Minor fix

* yapf

* Fix type annotation

* yapf

* [Minor] Address comment

* Add type alias &amp; enhance comments

* yapf

* Fix minor error in dag_lib.Dag

* Add is_chain to Dag

* Address comments

* yapf

* yapf

* Address comments

* Add total in optimizer msg

* Add a comment in is_chain

* Address reviews &amp; Fix egress msg

* yapf

* Minor fix

* Fix egress msg

* yapf

* obj -&gt; objective

* pass yapf

* cost -&gt; cost/time

* Add random DAG generator

* Add random DAG generator

* Change variable names

* Minor fix

* yapf on test_random_dag.py

* Add docstring

* Rename

* _optimize_cost -&gt; _optimize_objective

* Minor

* Default num_tasks to 10

* Add docstrings &amp; Fix variable names

* yapf

* Minor

* Improve test_optimizer_random_dag

* yapf

* Fix optimizer

* Add docstring about ILP objective

* fix typo

* yapf

* Minor

* Add monkeyptach

* Fix docstring

* yapf
diff --git a/sky/optimizer.py b/sky/optimizer.py
@@ -85,7 +85,7 @@ def optimize(dag: 'dag_lib.Dag',
         # This function is effectful: mutates every node in 'dag' by setting
         # node.best_resources if it is None.
         dag = Optimizer._add_dummy_source_sink_nodes(dag)
-        optimized_dag, unused_best_plan = Optimizer._optimize_cost(
+        optimized_dag, unused_best_plan = Optimizer._optimize_objective(
             dag,
             minimize_cost=minimize == OptimizeTarget.COST,
             blocked_launchable_resources=blocked_launchable_resources,
@@ -339,6 +339,154 @@ def _optimize_by_dp(
                 best_resources = dp_point_backs[node][best_resources]
         return best_plan, best_total_objective
 
+    @staticmethod
+    def _optimize_by_ilp(
+        graph,
+        topo_order: List[Task],
+        node_to_cost_map: _TaskToCostMap,
+        minimize_cost: bool = True,
+    ) -> Tuple[Dict[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('Sky-Cost-Optimization', pulp.LpMinimize)
+        else:
+            prob = pulp.LpProblem('Sky-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 = 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 = 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,
@@ -510,7 +658,7 @@ def _print_candidates(node_to_candidate_map: _TaskToPerCloudCandidates):
                     f'To list more details, run \'sky show-gpus {acc_name}\'.')
 
     @staticmethod
-    def _optimize_cost(
+    def _optimize_objective(
         dag: 'dag_lib.Dag',
         minimize_cost: bool = True,
         blocked_launchable_resources: Optional[List[
@@ -540,7 +688,8 @@ def _optimize_cost(
             best_plan, best_total_objective = Optimizer._optimize_by_dp(
                 topo_order, node_to_cost_map, minimize_cost)
         else:
-            raise NotImplementedError('Currently Sky only supports chain DAGs.')
+            best_plan, best_total_objective = Optimizer._optimize_by_ilp(
+                graph, topo_order, node_to_cost_map, minimize_cost)
 
         if minimize_cost:
             total_time = Optimizer._compute_total_time(graph, topo_order,
diff --git a/sky/setup_files/setup.py b/sky/setup_files/setup.py
@@ -43,7 +43,8 @@
     # This is used by ray. The latest 1.44.0 will generate an error
     # `Fork support is only compatible with the epoll1 and poll
     # polling strategies`
-    'grpcio<=1.43.0'
+    'grpcio<=1.43.0',
+    'pulp',
 ]
 
 extras_require = {
diff --git a/tests/test_optimizer_random_dag.py b/tests/test_optimizer_random_dag.py
@@ -0,0 +1,133 @@
+import copy
+import random
+
+import numpy as np
+import sky
+
+CLOUDS = {
+    'AWS': sky.AWS(),
+    'GCP': sky.GCP(),
+    'Azure': sky.Azure(),
+}
+ALL_INSTANCE_TYPES = sum(sky.list_accelerators(gpus_only=True).values(), [])
+GCP_INSTANCE_TYPES = list(sky.GCP._ON_DEMAND_PRICES.keys())
+
+DUMMY_NODES = [
+    sky.optimizer._DUMMY_SOURCE_NAME,
+    sky.optimizer._DUMMY_SINK_NAME,
+]
+
+
+def generate_random_dag(
+    num_tasks: int,
+    seed: int = 0,
+    max_num_nodes: int = 10,
+    max_num_parents: int = 5,
+    max_num_candidate_resources: int = 5,
+    max_task_runtime: int = 3600,
+    max_data_size: int = 1000,
+) -> sky.Dag:
+    """Generates a random Sky DAG to test Sky optimizer."""
+    random.seed(seed)
+    with sky.Dag() as dag:
+        for i in range(num_tasks):
+            op = sky.Task(name=f'task{i}')
+            task_runtime = random.random() * max_task_runtime
+            op.set_time_estimator(lambda _: task_runtime)
+            op.num_nodes = random.randint(1, max_num_nodes)
+
+            if i == 0:
+                num_parents = 0
+            else:
+                num_parents = random.randint(0, min(i, max_num_parents))
+
+            if num_parents == 0:
+                src_cloud = random.choice(['s3:', 'gs:', None])
+                src_volume = random.randint(0, max_data_size)
+            else:
+                parents = random.choices(dag.tasks[:-1], k=num_parents)
+                for parent in parents:
+                    parent >> op
+                # NOTE: Sky only takes single input data source
+                src_cloud = parents[0]
+                # Sky uses the parent's output data size
+                src_volume = None
+
+            if src_cloud is not None:
+                op.set_inputs(src_cloud, src_volume)
+            op.set_outputs('CLOUD', random.randint(0, max_data_size))
+
+            num_candidates = random.randint(1, max_num_candidate_resources)
+            candidate_instance_types = random.choices(ALL_INSTANCE_TYPES,
+                                                      k=num_candidates)
+            op.set_resources({
+                sky.Resources(
+                    cloud=CLOUDS[candidate.cloud],
+                    instance_type=candidate.instance_type \
+                        if candidate.cloud != 'GCP' \
+                        else random.choice(GCP_INSTANCE_TYPES),
+                    accelerators={
+                        candidate.accelerator_name: candidate.accelerator_count},
+                )
+                for candidate in candidate_instance_types
+            })
+    return dag
+
+
+def find_min_objective(dag: sky.Dag, minimize_cost: bool) -> float:
+    """Manually finds the minimum objective value."""
+    graph = dag.get_graph()
+    topo_order = dag.tasks
+
+    def _optimize_by_brute_force(tasks, plan):
+        """Optimizes a Sky DAG in a brute-force manner."""
+        # NOTE: Here we assume that the Sky DAG is topologically sorted.
+        task = tasks[0]
+        min_objective = np.inf
+        for resources in task.get_resources():
+            assert task.name in DUMMY_NODES or resources.is_launchable()
+            plan[task] = resources
+            if len(tasks) == 1:
+                if minimize_cost:
+                    objective = sky.Optimizer._compute_total_cost(
+                        graph, topo_order, plan)
+                else:
+                    objective = sky.Optimizer._compute_total_time(
+                        graph, topo_order, plan)
+            else:
+                objective = _optimize_by_brute_force(tasks[1:], plan)
+            if objective < min_objective:
+                min_objective = objective
+        return min_objective
+
+    return _optimize_by_brute_force(topo_order, {})
+
+
+def compare_optimization_results(dag: sky.Dag, minimize_cost: bool):
+    copy_dag = copy.deepcopy(dag)
+
+    _, optimizer_plan = sky.Optimizer._optimize_objective(dag, minimize_cost)
+    if minimize_cost:
+        objective = sky.Optimizer._compute_total_cost(dag.get_graph(),
+                                                      dag.tasks, optimizer_plan)
+    else:
+        objective = sky.Optimizer._compute_total_time(dag.get_graph(),
+                                                      dag.tasks, optimizer_plan)
+
+    min_objective = find_min_objective(copy_dag, minimize_cost)
+    assert objective == min_objective
+
+
+def test_optimizer(monkeypatch):
+    enabled_clouds = list(CLOUDS.values())
+    monkeypatch.setattr(
+        'sky.global_user_state.get_enabled_clouds',
+        lambda: enabled_clouds,
+    )
+    monkeypatch.setattr('sky.check.check', lambda *_args, **_kwargs: None)
+
+    dag = generate_random_dag(num_tasks=5, seed=0)
+    dag = sky.Optimizer._add_dummy_source_sink_nodes(dag)
+
+    compare_optimization_results(dag, minimize_cost=True)
+    compare_optimization_results(dag, minimize_cost=False)

Original file line number	Diff line number	Diff line change
`@@ -43,7 +43,8 @@`
`43`	`43`	`# This is used by ray. The latest 1.44.0 will generate an error`
`44`	`44`	# `Fork support is only compatible with the epoll1 and poll
`45`	`45`	# polling strategies`
`46`		`- 'grpcio<=1.43.0'`
	`46`	`+ 'grpcio<=1.43.0',`
	`47`	`+ 'pulp',`
`47`	`48`	`]`
`48`	`49`
`49`	`50`	`extras_require = {`