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

relationalai/semantics/devtools/benchmark_lqp.py

@@ -1,535 +0,0 @@
-# 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()