[RPC] Fetch TPCH profiled info for application at runtime

profiles/workload/tpch_decima/query_dag.json

@@ -0,0 +1,27 @@
+{
+    "workload_type": "tpch",
+    "query_number": {
+        "1": "[(0, 1), (1, 2)]",
+        "2": "[(0, 2), (0, 9), (1, 2), (2, 4), (3, 4), (4, 7), (5, 7), (5, 13), (6, 15), (7, 15), (8, 9), (9, 11), (10, 11), (11, 13), (12, 14), (13, 14), (14, 16), (15, 16)]",
+        "3": "[(0, 2), (1, 2), (2, 4), (3, 4)]",
+        "4": "[(0, 2), (1, 2), (2, 3), (3, 4)]",
+        "5": "[(0, 2), (1, 2), (2, 4), (3, 4), (4, 6), (5, 6), (6, 8), (7, 8), (8, 10), (9, 10), (10, 11), (11, 12)]",
+        "6": "[(0, 1)]",
+        "7": "[(0, 2), (0, 5), (1, 2), (2, 8), (3, 8), (4, 5), (5, 7), (6, 7), (7, 9), (8, 9), (9, 10), (10, 11)]",
+        "8": "[(0, 5), (1, 5), (2, 4), (3, 4), (4, 13), (5, 13), (6, 8), (7, 8), (8, 10), (9, 10), (10, 12), (11, 12), (12, 14), (13, 14), (14, 15), (15, 16)]",
+        "9": "[(0, 5), (1, 5), (2, 4), (3, 4), (4, 6), (5, 6), (6, 8), (7, 8), (8, 10), (9, 10), (10, 11), (11, 12)]",
+        "10": "[(0, 2), (1, 2), (2, 4), (3, 4), (4, 6), (5, 6), (6, 7)]",
+        "11": "",
+        "12": "[(0, 2), (1, 2), (2, 3), (3, 4)]",
+        "13": "[(0, 2), (1, 2), (2, 3), (3, 4), (4, 5)]",
+        "14": "[(0, 2), (1, 2)]",
+        "15": "",
+        "16": "[(0, 2), (1, 2), (2, 4), (3, 4), (4, 5), (5, 6), (6, 7)]",
+        "17": "",
+        "18": "",
+        "19": "[(0, 2), (1, 2)]",
+        "20": "",
+        "21": "",
+        "22": ""
+    }
+}

rpc/service.py

@@ -16,6 +16,7 @@
 import erdos_scheduler_pb2_grpc
 import grpc
 from absl import app, flags
+from tpch_utils import get_all_stage_info_for_query, verify_and_relable_tpch_app_graph
 
 from schedulers import EDFScheduler, FIFOScheduler
 from utils import EventTime, setup_logging
@@ -367,10 +368,62 @@ async def RegisterTaskGraph(self, request, context):
                 num_executors=0,
             )
 
+        self._logger.info(
+            "Attempting to register application ID %s with name %s",
+            request.id,
+            request.name,
+        )
+        # Check if query is from TPC-H workload.
+        # If yes, retrieve profiled slots and runtime info. If no, use default values
+        is_tpch_query = False
+        tpch_query_all_stage_info = None
+        if request.name.startswith("TPCH_"):
+            is_tpch_query = True
+            # retrieve tasks-per-stage and runtime info based on query number
+            tpch_query_num = request.name.split("TPCH_Q", 1)[1]
+            tpch_query_all_stage_info = get_all_stage_info_for_query(tpch_query_num)
+            same_structure, stage_id_mapping = verify_and_relable_tpch_app_graph(
+                query_num=tpch_query_num, dependencies=request.dependencies
+            )
+
+            # return failure message if not tpch app isnt of same DAG structure
+            if not same_structure:
+                self._logger.warning(
+                    "TPCH application with ID %s and name %s couldn't be registered."
+                    "DAG structure mismatch!",
+                    request.id,
+                    request.name,
+                )
+                return erdos_scheduler_pb2.RegisterTaskGraphResponse(
+                    success=False,
+                    message=f"TPCH application ID {request.id} with name {request.name}"
+                    f" couldn't be registered. DAG structure mismatch!",
+                    num_executors=0,
+                )
+
         # Construct all the Tasks for the TaskGraph.
         task_ids_to_task: Mapping[int, Task] = {}
+        default_resource = Resources(
+            resource_vector={Resource(name="Slot_CPU", _id="any"): 30}
+        )
+        default_runtime = EventTime(20, EventTime.Unit.US)
+
         for task_dependency in request.dependencies:
             framework_task = task_dependency.key
+            if is_tpch_query:
+                mapped_stage_id = stage_id_mapping[framework_task.id]
+                task_slots = tpch_query_all_stage_info[mapped_stage_id]["num_tasks"]
+                task_runtime = tpch_query_all_stage_info[mapped_stage_id][
+                    "avg_task_duration"
+                ]
+                self._logger.info(
+                    "Creating Task for given app TPCH stage: %s, mapped to "
+                    "original stage id %s, with tasks: %s and avg runtime: %s",
+                    framework_task.id,
+                    mapped_stage_id,
+                    task_slots,
+                    task_runtime,
+                )
             task_ids_to_task[framework_task.id] = Task(
                 name=framework_task.name,
                 task_graph=request.id,
@@ -380,17 +433,24 @@ async def RegisterTaskGraph(self, request, context):
                         name=f"WorkProfile_{framework_task.name}",
                         execution_strategies=ExecutionStrategies(
                             [
-                                # TODO (Sukrit): Find the actual resource requirements
-                                # for the particular TaskGraph, along with the expected
-                                # runtime and set it here.
                                 ExecutionStrategy(
-                                    resources=Resources(
-                                        resource_vector={
-                                            Resource(name="Slot_CPU", _id="any"): 30
-                                        }
+                                    resources=(
+                                        default_resource
+                                        if not is_tpch_query
+                                        else Resources(
+                                            resource_vector={
+                                                Resource(
+                                                    name="Slot_CPU", _id="any"
+                                                ): task_slots
+                                            }
+                                        )
                                     ),
                                     batch_size=1,
-                                    runtime=EventTime(20, EventTime.Unit.US),
+                                    runtime=(
+                                        default_runtime
+                                        if not is_tpch_query
+                                        else EventTime(task_runtime, EventTime.Unit.US)
+                                    ),
                                 )
                             ]
                         ),

rpc/tpch_utils.py

@@ -0,0 +1,170 @@
+# Code adopted from decima-sim to use their profiles of TPC-H queries
+
+import ast
+import json
+import os
+from typing import Mapping, Sequence
+
+import networkx as nx
+import numpy as np
+
+HOME_TPCH_DIR = "../profiles/workload/tpch_decima/"
+TPCH_SUBDIR = "2g/"
+
+
+class SetWithCount(object):
+    """
+    allow duplication in set
+    """
+
+    def __init__(self):
+        self.set = {}
+
+    def __contains__(self, item):
+        return item in self.set
+
+    def add(self, item):
+        if item in self.set:
+            self.set[item] += 1
+        else:
+            self.set[item] = 1
+
+    def clear(self):
+        self.set.clear()
+
+    def remove(self, item):
+        self.set[item] -= 1
+        if self.set[item] == 0:
+            del self.set[item]
+
+
+def pre_process_task_duration(task_duration):
+    # remove fresh durations from first wave
+    clean_first_wave = {}
+    for e in task_duration["first_wave"]:
+        clean_first_wave[e] = []
+        fresh_durations = SetWithCount()
+        # O(1) access
+        for d in task_duration["fresh_durations"][e]:
+            fresh_durations.add(d)
+        for d in task_duration["first_wave"][e]:
+            if d not in fresh_durations:
+                clean_first_wave[e].append(d)
+            else:
+                # prevent duplicated fresh duration blocking first wave
+                fresh_durations.remove(d)
+
+
+def get_all_stage_info_for_query(query_num):
+    task_durations = np.load(
+        os.path.join(
+            HOME_TPCH_DIR, TPCH_SUBDIR, "task_duration_" + str(query_num) + ".npy"
+        ),
+        allow_pickle=True,
+    ).item()
+
+    num_nodes = len(task_durations)
+
+    stage_info = {}
+
+    for n in range(num_nodes):
+        task_duration = task_durations[n]
+        e = next(iter(task_duration["first_wave"]))
+        # NOTE: somehow only picks the first element {2: [n_tasks_in_ms]}
+
+        num_tasks = len(task_duration["first_wave"][e]) + len(
+            task_duration["rest_wave"][e]
+        )
+
+        # remove fresh duration from first wave duration
+        # drag nearest neighbor first wave duration to empty spots
+        pre_process_task_duration(task_duration)
+        rough_duration = np.mean(
+            [i for t in task_duration["first_wave"].values() for i in t]
+            + [i for t in task_duration["rest_wave"].values() for i in t]
+            + [i for t in task_duration["fresh_durations"].values() for i in t]
+        )
+
+        curr_stage = {
+            "stage_id": n,
+            "num_tasks": num_tasks,
+            "avg_task_duration": round(rough_duration),
+        }
+        stage_info[n] = curr_stage
+
+    return stage_info
+
+
+def get_base_tpch_graph_structure(query_num):
+    # use query_num to read string from file
+    with open(os.path.join(HOME_TPCH_DIR, "query_dag.json")) as f:
+        tpch_query_json = json.load(f)
+
+    # get query dependency from file
+    query_dependency = ast.literal_eval(tpch_query_json["query_number"][str(query_num)])
+
+    # convert job structure into a nx graph
+    base_tpch_graph = nx.DiGraph(query_dependency)
+
+    # return the job nx graph for query
+    return base_tpch_graph
+
+
+def get_graph_from_deps(dependencies):
+    # parse dependencies to get it in list of tuples
+    # construct the TaskGraph from the dependencies.
+    task_graph_structure: Mapping[int, Sequence[int]] = {}
+    for task_dependency in dependencies:
+        task_graph_structure[task_dependency.key.id] = [
+            task_id for task_id in task_dependency.children_ids
+        ]
+
+    # convert our TaskGraph into a nx friendly notation
+    all_edges_in_app = []
+    for parent in task_graph_structure.keys():
+        for child in task_graph_structure[parent]:
+            all_edges_in_app.append((parent, child))
+
+    # construct nx graph
+    given_app_tpch_graph = nx.DiGraph(all_edges_in_app)
+
+    # return the graph
+    return given_app_tpch_graph
+
+
+def are_structurally_same(graph1, graph2):
+    # Step 1: Check if both graphs have the same number of vertices
+    if len(graph1.nodes) != len(graph2.nodes):
+        print(
+            f"DAG structure mismatch! Graph1 has "
+            f"{graph1.nodes} while Graph2 has {graph2.nodes}"
+        )
+        return False, None
+
+    # Step 2: Check if there exists a bijection between the vertices
+    #         of the two graphs such that their adjacency relationships match
+    for mapping in nx.isomorphism.GraphMatcher(graph1, graph2).isomorphisms_iter():
+        # Check if the adjacency relationships match
+        if all(v in mapping for u, v in graph1.edges):
+            # graph structures match
+            # mapping is a dict {key=original-stage-id, val=app-stage-id}
+            # we reverse reversed mapping from app-stage-id to orig-stage-id
+            reversed_mapping = {v: k for k, v in mapping.items()}
+
+            return True, reversed_mapping
+
+    print("DAG structure mismatch! No mapping could be found")
+    return False, None
+
+
+def verify_and_relable_tpch_app_graph(query_num, dependencies):
+    # get nx graphs from base tpch file and application dependencies
+    base_tpch_graph = get_base_tpch_graph_structure(query_num)
+    app_graph = get_graph_from_deps(dependencies)
+
+    # verify that the graphs are isomorphic
+    # returns true if same_structure, along with stage_id_mapping wrt base tpch file
+    same_structure, stage_id_mapping = are_structurally_same(base_tpch_graph, app_graph)
+
+    # return with stage_id_mapping back to assign correct runtime and num_executors
+    return same_structure, stage_id_mapping