relationalai 0.13.0.dev0__py3-none-any.whl → 0.13.2__py3-none-any.whl

relationalai/semantics/devtools/benchmark_lqp.py

@@ -0,0 +1,535 @@
+# This script runs a QB program and captures high-level timing information about different parts of the system.
+# The output is a JSON file containing timing information that can be used to analyze performance.
+
+import argparse
+import os
+import json
+
+from relationalai.clients.resources.snowflake import Resources as snowflake_api, APP_NAME
+from relationalai.semantics.lqp.executor import LQPExecutor
+from relationalai.semantics.internal import internal
+from relationalai.clients.resources.snowflake.use_index_poller import UseIndexPoller as index_poller
+from snowflake.connector.cursor import DictCursor
+
+from enum import Enum
+
+import time
+
+# Categories that spans can belong to
+class SpanCategory(Enum):
+    DATA_PREP = "Data Preparation"
+    QB_COMPILE = "QB Compilation"
+    EXEC = "RAICode Execution"
+    SNOWFLAKE = "Snowflake"
+    RESULTS = "Result Processing"
+    OTHER = "Other"
+    UNKNOWN = "Unknown"
+
+# Contextual information that describes an instrumented method. These are used to create spans.
+#
+# To better distinguish important spans, we have an `is_core` flag to indicate "core methods".
+# Core methods represent high-level concepts that make sense to be reasoned about in isolation.
+#
+# Core methods must be mutually exclusive. This means that they should not overlap in time with one another,
+# and should represent distinct concepts.
+# For example, `exec_rai_app` and `process_results` are core methods since they are distinct high-level concepts
+# that aid with understanding the overall time breakdown and do not overlap with one another.
+# On the other hand, `exec_snowflake` is not a core method since it overlaps with other core methods and
+# is tricky to reason about in isolation without knowing /what/ was being executed.
+#
+# Ideally, the sum of all core methods should equal the total time of the root span.
+#
+# Additionally, we have a `count_for_cat` flag to indicate whether this method should count towards the time of
+# its category. This prevents duplicate counting for overlapping methods in the same category.
+class TimeMarker:
+    def __init__(self, name, category: SpanCategory, is_core=False, count_for_cat=False):
+        self.name = name
+        self.category = category
+        self.is_core = is_core
+        self.count_for_cat = count_for_cat
+
+# Time span for a specific operation.
+class TimeSpan:
+    def __init__(self, id: str, marker: TimeMarker, parent_id, tags={}):
+        self.id = id
+        self.marker = marker
+        self.parent_id = parent_id
+        self.tags = tags
+
+        self.start_time = None
+        self.end_time = None
+
+    def start(self):
+        self.start_time = time.time()
+
+    def end(self):
+        self.end_time = time.time()
+
+    def duration(self):
+        assert self.start_time is not None, "Span has not started yet"
+        assert self.end_time is not None, "Span has not ended yet"
+        return self.end_time - self.start_time
+
+# Context for the benchmarking. Keeps track of active spans and finished spans.
+class BenchmarkContext:
+    def __init__(self):
+        self.current_spans = []
+        self.finished_spans = []
+        self.span_id_counter = 0
+        self.print_logs = True
+        self.gap_size = 0.0
+        self.last_end_time = 0
+
+        self.snowflake_calls = []
+        self.last_qid = None
+        self.qid_to_txnid = dict()
+
+    def start_span(self, marker: TimeMarker, tags={}):
+        parent_id = self.current_spans[-1].id if self.current_spans else None
+        span_id = f"span_{self.span_id_counter}"
+        self.span_id_counter += 1
+
+        new_span = TimeSpan(span_id, marker, parent_id, tags=tags)
+        new_span.start()
+        self.current_spans.append(new_span)
+
+        if new_span.marker.is_core:
+            assert new_span.start_time is not None, "Span has not started yet"
+            last_end_diff = new_span.start_time - self.last_end_time
+            self.gap_size += last_end_diff
+            if last_end_diff > 0.01:
+                print(f"\tWarning: Gap of {last_end_diff:.4f} seconds before start of span `{new_span.id}` ({new_span.marker.name})")
+
+        return new_span
+
+    def end_span(self):
+        assert self.current_spans, "No active span to end"
+        span = self.current_spans.pop()
+        span.end()
+        self.finished_spans.append(span)
+        if span.marker.is_core:
+            self.last_end_time = span.end_time
+
+        if self.print_logs:
+            duration = span.duration()
+            print(f"Span {span.id} ({span.marker.name}) ended, duration: {duration:.4f} seconds")
+        return span
+
+    def add_snowflake_call(self, qid: str, code: str):
+        self.snowflake_calls.append({
+            "sf_query_id": qid,
+            "code": code
+        })
+        self.last_qid = qid
+        if self.print_logs:
+            print(f"Snowflake call [{qid}]: {code}")
+
+    def add_txn_id(self, txn_id: str):
+        assert self.last_qid is not None, "No Snowflake query ID recorded yet"
+        assert self.last_qid not in self.qid_to_txnid, f"Snowflake query ID {self.last_qid} already has a transaction ID recorded"
+        self.qid_to_txnid[self.last_qid] = txn_id
+        self.last_qid = None
+
+BENCHMARK_CTX = BenchmarkContext()
+
+# Wrapper for instrumenting a method.
+# If more fine-grained control is needed, consider monkey-patching manually with spans instead.
+# E.g. if you want to add tags or have more complex logic.
+class InstrumentedMethod:
+    def __init__(self, module, method_name, category, is_core=False, count_for_cat=False, alias=None, override=None):
+        self.module = module
+        self.method_name = method_name
+
+        marker_name = alias if alias else method_name
+        self.marker = TimeMarker(marker_name, category, is_core, count_for_cat)
+
+        self.original_method = getattr(module, method_name)
+        if override:
+            self.new_method = override(BENCHMARK_CTX, self.original_method, self.marker)
+        else:
+            self.new_method = self._create_instrumented_method()
+
+    def _create_instrumented_method(self):
+        def instrumented_method(*args, **kwargs):
+            global BENCHMARK_CTX
+            BENCHMARK_CTX.start_span(self.marker)
+            result = self.original_method(*args, **kwargs)
+            BENCHMARK_CTX.end_span()
+            return result
+        return instrumented_method
+
+    def apply(self):
+        setattr(self.module, self.method_name, self.new_method)
+
+    def restore(self):
+        setattr(self.module, self.method_name, self.original_method)
+
+SF_COMPILATION_MARKER = TimeMarker("sf_compilation", SpanCategory.SNOWFLAKE)
+SF_EXECUTION_MARKER = TimeMarker("sf_execution", SpanCategory.SNOWFLAKE)
+SF_REQ_OVERHEAD_MARKER = TimeMarker("sf_request_overhead", SpanCategory.SNOWFLAKE)
+def _exec_snowflake_override(bench_ctx, old_func, marker):
+    def new_func(self, code, params, raw=False):
+        cur = self._session.connection.cursor(DictCursor)
+        try:
+            cur.execute(code.replace(APP_NAME, self.get_app_name()), params)
+            rows = cur.fetchall()
+            qid = str(getattr(cur, "sfqid", None))
+            assert qid is not None, "Snowflake query ID was not available"
+        finally:
+            cur.close()
+        bench_ctx.add_snowflake_call(qid, code)
+
+        # ===
+        # Add dummy spans that will be populated later
+        bench_ctx.start_span(SF_REQ_OVERHEAD_MARKER, tags={"sf_query_id": qid})
+        bench_ctx.end_span()
+        bench_ctx.start_span(SF_COMPILATION_MARKER, tags={"sf_query_id": qid})
+        bench_ctx.end_span()
+        bench_ctx.start_span(SF_EXECUTION_MARKER, tags={"sf_query_id": qid})
+        bench_ctx.end_span()
+        # ===
+
+        return rows
+    return new_func
+
+def _is_dummy_sf_span(span):
+    return span.marker in {SF_COMPILATION_MARKER, SF_EXECUTION_MARKER, SF_REQ_OVERHEAD_MARKER}
+
+def _download_results_override(bench_ctx, old_func, marker):
+    def new_func(self, artifact_info, txn_id: str, state):
+        result = old_func(self, artifact_info, txn_id, state)
+        bench_ctx.add_txn_id(txn_id)
+        print("txn id: ", txn_id)
+        return result
+    return new_func
+
+# List of all methods to track.
+INSTRUMENTED_METHODS = [
+    # Core methods. They must be mutually exclusive, high-level concepts, and ideally sum up to the total time.
+    InstrumentedMethod(LQPExecutor, "prepare_data", SpanCategory.DATA_PREP, is_core=True, count_for_cat=True),
+    InstrumentedMethod(LQPExecutor, "compile_lqp", SpanCategory.QB_COMPILE, is_core=True, count_for_cat=True),
+    InstrumentedMethod(snowflake_api, "_exec_rai_app", SpanCategory.EXEC, is_core=True, count_for_cat=True, alias="exec_rai_app"),
+    InstrumentedMethod(snowflake_api, "_fetch_exec_async_artifacts", SpanCategory.RESULTS, is_core=True, count_for_cat=True, alias="fetch_artifacts"),
+    InstrumentedMethod(LQPExecutor, "_process_results", SpanCategory.RESULTS, is_core=True, count_for_cat=True, alias="process_results"),
+
+    # Other trackers, these should also be on the high-level side but can overlap
+    InstrumentedMethod(LQPExecutor, "execute", SpanCategory.EXEC, alias="execute_lqp"),
+    InstrumentedMethod(snowflake_api, "_exec", SpanCategory.SNOWFLAKE, count_for_cat=True, alias="exec_snowflake"),
+    InstrumentedMethod(index_poller, "_poll_loop", SpanCategory.DATA_PREP, alias="poll_once"),
+    InstrumentedMethod(index_poller, "poll", SpanCategory.DATA_PREP, alias="poll_main"),
+    InstrumentedMethod(snowflake_api, "_exec_sql", SpanCategory.SNOWFLAKE, alias="exec_sql", override=_exec_snowflake_override),
+    InstrumentedMethod(snowflake_api, "_build_snowflake_session", SpanCategory.SNOWFLAKE),
+    InstrumentedMethod(snowflake_api, "_download_results", SpanCategory.RESULTS, override=_download_results_override),
+]
+
+
+def _make_unknown_span(parent_id, parent_duration, known_children_duration):
+    unknown_duration = parent_duration - known_children_duration
+    if unknown_duration <= 0.01:
+        return None
+    return {
+        "name": "unknown",
+        "duration": unknown_duration,
+        "category": SpanCategory.UNKNOWN.value,
+        "parent": parent_id,
+        "tags": {},
+        "is_core": True,
+        "children": [],
+        "total_percent": (unknown_duration / parent_duration) * 100,
+        "parent_percent": (unknown_duration / parent_duration) * 100,
+    }
+
+def _make_task_times_dict(spans, root_span):
+    total_duration = root_span.duration()
+    task_times = dict()
+    for span in spans:
+        if span.marker.name not in task_times:
+            task_times[span.marker.name] = {"duration": 0.0, "category": span.marker.category.value}
+        task_times[span.marker.name]["duration"] += span.duration()
+
+    # Add total percent to each
+    for v in task_times.values():
+        v["percent"] = (v["duration"] / total_duration) * 100
+    return task_times
+
+def _make_spans_dict(spans, root_span):
+    spans_dict = dict()
+    for span in spans:
+        spans_dict[span.id] = {
+            "name": span.marker.name,
+            "duration": span.duration(),
+            "category": span.marker.category.value,
+            "parent": span.parent_id,
+            "tags": span.tags,
+            "is_core": span.marker.is_core,
+
+            # To be filled later
+            "children": [],
+        }
+
+    # Fill in children
+    for span in spans:
+        if span.parent_id:
+            spans_dict[span.parent_id]["children"].append(span.id)
+
+    # Add total_percent and parent_percent
+    total_duration = root_span.duration()
+    for span in spans:
+        span_info = spans_dict[span.id]
+        span_info["total_percent"] = (span_info["duration"] / total_duration) * 100
+        if span.parent_id:
+            parent_duration = spans_dict[span.parent_id]["duration"]
+            span_info["parent_percent"] = (span_info["duration"] / parent_duration) * 100
+        else:
+            # Root span has no parent
+            span_info["parent_percent"] = 100.0
+
+    # Add "Unknown" spans for when sum of children < parent
+    for span in spans:
+        children = spans_dict[span.id]["children"]
+        if not children:
+            continue
+        known_children_duration = sum(spans_dict[child_id]["duration"] for child_id in children)
+        unknown_span = _make_unknown_span(span.id, spans_dict[span.id]["duration"], known_children_duration)
+        if unknown_span:
+            unknown_id = f"{span.id}_unknown"
+            spans_dict[unknown_id] = unknown_span
+            spans_dict[span.id]["children"].append(unknown_id)
+
+    return spans_dict
+
+def _make_core_times_dict(spans, root_span):
+    core_times = dict()
+    total_duration = root_span.duration()
+    duration_accounted = 0.0
+    for span in spans:
+        if span.marker.is_core:
+            if span.marker.name not in core_times:
+                core_times[span.marker.name] = {
+                    "duration": 0.0,
+                }
+            core_times[span.marker.name]["duration"] += span.duration()
+            duration_accounted += span.duration()
+    for v in core_times.values():
+        v["total_percent"] = (v["duration"] / total_duration) * 100
+
+    if duration_accounted < total_duration - 0.01:
+        core_times["unknown"] = {
+            "duration": total_duration - duration_accounted,
+            "total_percent": ((total_duration - duration_accounted) / total_duration) * 100
+        }
+    return core_times
+
+def _make_category_times_dict(spans, root_span):
+    category_times = dict()
+    for span in spans:
+        if span.marker.count_for_cat:
+            cat_name = span.marker.category.value
+            if cat_name not in category_times:
+                category_times[cat_name] = {
+                    "duration": 0.0,
+                }
+            category_times[cat_name]["duration"] += span.duration()
+    return category_times
+
+def main():
+    parser = argparse.ArgumentParser(description="Extract LQP requests to run locally")
+    parser.add_argument("--program", type=str, help="Path to the QB program to run", required=True)
+    parser.add_argument("--outdir", type=str, help="Path to the output directory", required=True)
+
+    args = parser.parse_args()
+    program = args.program
+    outdir = args.outdir
+
+    if not os.path.isfile(program):
+        raise ValueError(f"Program path {program} is not a file")
+    if not os.path.isdir(outdir):
+        raise ValueError(f"Output path {outdir} is not a directory")
+
+    benchmark_lqp(program, outdir)
+
+def _run_instrumented_qb(qb_prog_path: str):
+    global BENCHMARK_CTX
+    global INSTRUMENTED_METHODS
+
+    # Run the QB program with instrumentation
+    for method in INSTRUMENTED_METHODS:
+        method.apply()
+
+    # Special marker for the root span
+    e2e_marker = TimeMarker("e2e", SpanCategory.OTHER)
+    root_span = BENCHMARK_CTX.start_span(e2e_marker)
+    BENCHMARK_CTX.last_end_time = root_span.start_time
+    with internal.with_overrides(**{
+        'dry_run': False,
+        'reasoner.rule.use_lqp': True,
+    }):
+        with open(qb_prog_path) as f:
+            # Compile with the correct filename
+            code = compile(f.read(), qb_prog_path, "exec")
+            exec(code)
+    root_span = BENCHMARK_CTX.end_span()
+    assert root_span.marker == e2e_marker, "Mismatched root span"
+    for method in INSTRUMENTED_METHODS:
+        method.restore()
+    spans = BENCHMARK_CTX.finished_spans
+    return spans, root_span
+
+def get_sf_query_info(bench_ctx):
+    snowflake_calls = bench_ctx.snowflake_calls
+    qids = []
+    for call in snowflake_calls:
+        qids.append(call["sf_query_id"])
+    assert len(qids) == len(snowflake_calls), "Duplicate Snowflake query IDs found"
+
+    result = dict()
+
+    qids_to_info = _get_query_info(qids)
+    for call in snowflake_calls:
+        qid = call["sf_query_id"]
+        assert qid in qids_to_info, f"Missing info for Snowflake query ID {qid}"
+        qid_info = qids_to_info[qid]
+        call["sf_compilation_time"] = qid_info["compilation_time"]
+        call["sf_execution_time"] = qid_info["execution_time"]
+        if qid in bench_ctx.qid_to_txnid:
+            call["rai_txn_id"] = bench_ctx.qid_to_txnid[qid]
+        result[qid] = call
+    return result
+
+def _get_query_info(qids):
+    from relationalai.clients.resources.snowflake import Resources as snowflake_client
+    client = snowflake_client()
+
+    qids_str = "','".join(qids)
+    timing_query = f"""
+        SELECT
+            query_id,
+            start_time,
+            end_time,
+            total_elapsed_time,
+            compilation_time,
+            execution_time,
+            queued_overload_time,
+            queued_provisioning_time,
+            queued_repair_time
+        FROM table(snowflake.information_schema.query_history())
+        WHERE query_id IN ('{qids_str}')
+        ORDER BY start_time;
+        """
+    results = client._exec(timing_query)
+    if len(results) != len(qids):
+        print(f"Expected {len(qids)} results, but got {len(results)}")
+        print("Trying again after a few more seconds...")
+        time.sleep(5)
+        results = client._exec(timing_query)
+    assert len(results) == len(qids), f"Expected {len(qids)} results, but got {len(results)}"
+
+    qid_to_info = dict()
+    for row in results:
+        qid = row["QUERY_ID"]
+        qid_to_info[qid] = {
+            "compilation_time": row["COMPILATION_TIME"]/1000.0,
+            "execution_time": row["EXECUTION_TIME"]/1000.0,
+
+        }
+    return qid_to_info
+
+def populate_dummy_spans(spans, sf_query_info):
+    parentid_to_qid = dict()
+    for span in spans:
+        if not _is_dummy_sf_span(span):
+            continue
+        parentid_to_qid[span.parent_id] = span.tags["sf_query_id"]
+    assert len(parentid_to_qid) == len(sf_query_info), "Mismatch between dummy spans and Snowflake query info"
+
+    parentid_to_times = dict()
+    for span in spans:
+        if span.id in parentid_to_qid:
+            parentid_to_times[span.id] = span.duration()
+
+    for span in spans:
+        if not _is_dummy_sf_span(span):
+            continue
+
+        qid = span.tags.get("sf_query_id", None)
+        assert qid is not None, f"Span {span.id} ({span.marker.name}) is missing sf_query_id tag"
+
+        info = sf_query_info.get(qid, None)
+        assert info is not None, f"Missing Snowflake info for query ID {qid}"
+
+        comp_time = info["sf_compilation_time"]
+        exec_time = info["sf_execution_time"]
+
+        total_parent_time = parentid_to_times[span.parent_id]
+        # We order it like: compilation -> execution -> overhead
+        # The overhead is the remaining time after accounting for the snowflake-reported times
+        if span.marker.name == "sf_compilation":
+            span.start_time = span.start_time
+            span.end_time = span.start_time + comp_time
+        elif span.marker.name == "sf_execution":
+            offset = span.start_time + comp_time
+            span.start_time = offset
+            span.end_time = offset + exec_time
+        elif span.marker.name == "sf_request_overhead":
+            offset = span.start_time + comp_time + exec_time
+            duration = total_parent_time - (comp_time + exec_time)
+            span.start_time = offset
+            span.end_time = offset + duration
+        else:
+            raise ValueError(f"Unknown marker name {span.marker.name}")
+        span.tags["rai_txn_id"] = info.get("rai_txn_id", None)
+        assert span.end_time >= span.start_time, f"Span {span.id} ({span.marker.name}) has negative duration after population"
+
+def _make_snowflake_calls_dict(sf_query_info):
+    result = []
+    for qid, info in sf_query_info.items():
+        result.append(info)
+    return result
+
+def _make_metadata_dict(bench_ctx, root_span, prog_path):
+    metadata = []
+    metadata.append({
+        "key": "total_time",
+        "alias": "Total Time",
+        "value": root_span.duration(),
+    })
+    metadata.append({
+        "key": "gap_time",
+        "alias": "Untracked Time",
+        "value": bench_ctx.gap_size,
+    })
+    metadata.append({
+        "key": "program",
+        "alias": "QB Program",
+        "value": prog_path,
+    })
+    return metadata
+
+def benchmark_lqp(qb_prog_path: str, outdir: str):
+    global BENCHMARK_CTX
+    spans, root_span = _run_instrumented_qb(qb_prog_path)
+
+    print(f"Total benchmark time: {root_span.duration():.4f} seconds")
+    print(f"\tMissing time not tracked by core methods: {BENCHMARK_CTX.gap_size:.4f} seconds")
+
+    sf_query_info = get_sf_query_info(BENCHMARK_CTX)
+    populate_dummy_spans(spans, sf_query_info)
+
+    # Write results to file
+    final_json_dict = dict()
+    final_json_dict["spans"] = _make_spans_dict(spans, root_span)
+    final_json_dict["task_times"] = _make_task_times_dict(spans, root_span)
+    final_json_dict["core_times"] = _make_core_times_dict(spans, root_span)
+    final_json_dict["category_times"] = _make_category_times_dict(spans, root_span)
+    final_json_dict["snowflake_calls"] = _make_snowflake_calls_dict(sf_query_info)
+    final_json_dict["metadata"] = _make_metadata_dict(BENCHMARK_CTX, root_span, qb_prog_path)
+
+    final_json = json.dumps(final_json_dict, default=str)
+    out_path = os.path.join(outdir, "benchmark_results.json")
+    with open(out_path, "w") as f:
+        f.write(final_json)
+
+    print(f"Wrote benchmark results to {out_path}")
+
+if __name__ == "__main__":
+    main()