aiqtoolkit 1.2.0rc4__py3-none-any.whl → 1.2.0rc5__py3-none-any.whl

aiq/profiler/inference_optimization/bottleneck_analysis/simple_stack_analysis.py

@@ -1,258 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-Advanced bottleneck analysis for a dataframe that contains:
-- event_type in {LLM_START, LLM_END, TOOL_START, TOOL_END, ...}
-- llm_name
-- tool_name
-- UUID
-- event_timestamp (float or datetime)
-- other metadata...
-
-We pair start/end events by UUID, compute operation durations,
-then analyze concurrency and produce a summary report.
-"""
-
-import numpy as np
-import pandas as pd
-
-from aiq.data_models.intermediate_step import IntermediateStep
-from aiq.profiler.inference_optimization.data_models import SimpleBottleneckReport
-from aiq.profiler.inference_optimization.data_models import SimpleOperationStats
-from aiq.profiler.utils import create_standardized_dataframe
-
-
-# ----------------------------------------------------------------------
-# Main Function
-# ----------------------------------------------------------------------
-def profile_workflow_bottlenecks(all_steps: list[list[IntermediateStep]]) -> SimpleBottleneckReport:
-    """
-    Perform advanced bottleneck profiling on a workflow dataframe.
-
-    1) Pair LLM_START/LLM_END and TOOL_START/TOOL_END by UUID.
-    2) Compute operation durations.
-    3) Analyze concurrency (max concurrent usage).
-    4) Summarize as SimpleOperationStats and produce a final SimpleBottleneckReport.
-
-    Parameters
-    ----------
-    all_steps : Intermediate Steps
-
-    Returns
-    -------
-    SimpleBottleneckReport
-        Contains detailed stats per operation and a textual summary of top bottlenecks.
-    """
-    df = create_standardized_dataframe(all_steps)
-    # -------------------------------------------------------------
-    # 1) Separate events by operation type and match start/end
-    # -------------------------------------------------------------
-    required_cols = {"event_type", "UUID", "event_timestamp"}
-    missing_cols = required_cols - set(df.columns)
-    if missing_cols:
-        raise ValueError(f"DataFrame missing required columns: {missing_cols}")
-
-    # We'll unify LLM and TOOL operations into a single set, with:
-    #   operation_type = 'LLM' or 'TOOL'
-    #   operation_name = llm_name/tool_name
-    #   start_time
-    #   end_time
-    #   duration = end_time - start_time
-    # We'll store them in a list of dicts, then convert to DataFrame.
-    operations_records = []
-
-    # We'll create a copy to avoid mutating user data
-    dfc = df.copy()
-
-    # We define a small helper to map event_type -> (operation_type, which_name_field)
-    def get_operation_info(event_type: str) -> str | None:
-        """
-        Return 'LLM' if event_type starts with 'LLM_', 'TOOL' if event_type starts with 'TOOL_',
-        else None (unknown).
-        """
-        if event_type.startswith("LLM_"):
-            return "LLM"
-        if event_type.startswith("TOOL_"):
-            return "TOOL"
-        return None
-
-    # Group by UUID so we can pair each START with the corresponding END
-    grouped = dfc.groupby("UUID", as_index=False, group_keys=True)
-
-    for uuid_val, group_df in grouped:
-        if len(group_df) < 2:
-            # Possibly incomplete or single event, skip
-            continue
-
-        # We might have multiple events with the same UUID, but typically we expect:
-        #   LLM_START, LLM_END (or TOOL_START, TOOL_END).
-        # Sort by timestamp
-        group_df = group_df.sort_values("event_timestamp")
-
-        # Identify operation_type from the first row's event_type
-        first_event_type = group_df["event_type"].iloc[0]
-        operation_type = get_operation_info(first_event_type)
-        if not operation_type:
-            # unknown or not LLM_/TOOL_
-            continue
-
-        # We'll attempt to find the start row and the end row
-        # Usually there's exactly 1 start, 1 end
-        start_rows = group_df[group_df["event_type"] == f"{operation_type}_START"]
-        end_rows = group_df[group_df["event_type"] == f"{operation_type}_END"]
-
-        if len(start_rows) == 0 or len(end_rows) == 0:
-            # No matching start/end
-            continue
-
-        # We'll just take the earliest start and the latest end for the entire group.
-        start_time = start_rows["event_timestamp"].min()
-        end_time = end_rows["event_timestamp"].max()
-        duration = end_time - start_time
-
-        # For the name, we pick 'llm_name' or 'tool_name' depending on operation_type
-        if operation_type == "LLM":
-            # Among the rows, pick a non-null llm_name if present
-            op_names = group_df["llm_name"].dropna().unique()
-            # fallback to a default if none
-            operation_name = op_names[0] if len(op_names) else "unknown_llm"
-        else:
-            op_names = group_df["tool_name"].dropna().unique()
-            operation_name = op_names[0] if len(op_names) else "unknown_tool"
-
-        operations_records.append({
-            "operation_type": operation_type,
-            "operation_name": operation_name,
-            "start_time": start_time,
-            "end_time": end_time,
-            "duration": duration,
-            "UUID": uuid_val
-        })
-
-    if not operations_records:
-        # No valid operations found
-        return SimpleBottleneckReport(stats={}, summary="No operations found to profile.")
-
-    operations_df = pd.DataFrame(operations_records)
-
-    # -------------------------------------------------------------
-    # 2) Concurrency Analysis
-    # -------------------------------------------------------------
-    # We want to find the maximum concurrency for each operation_name.
-    # We'll do a timeline-based approach: for each operation we have a start_time, end_time
-    # We'll create +1 event at start_time, -1 event at end_time, then do a running sum.
-    # Then we can measure concurrency across the entire timeline. However, we want concurrency
-    # specifically *by operation_name* as well as overall.
-    #
-    # We'll do it in two passes:
-    #   A) Overall concurrency ignoring operation_name
-    #   B) concurrency per (operation_type, operation_name)
-    # Then we can combine them for a "peak concurrency" measure.
-
-    # A) Overall concurrency (not always essential, but might be interesting)
-    timeline_events = []
-    for row in operations_df.itertuples(index=False):
-        timeline_events.append((row.start_time, +1))
-        timeline_events.append((row.end_time, -1))
-
-    timeline_events.sort(key=lambda x: x[0])  # sort by time
-    current_concurrency = 0
-    concurrency_trace = []
-    for ts, delta in timeline_events:
-        current_concurrency += delta
-        concurrency_trace.append((ts, current_concurrency))
-    overall_max_concurrency = max(c[1] for c in concurrency_trace) if concurrency_trace else 0
-
-    # B) concurrency by operation_name
-    # We'll generate timeline events per operation_name
-    # Then compute the max concurrency for that subset
-    operation_names = operations_df["operation_name"].unique()
-    max_concurrency_by_name = {}
-
-    for op_name in operation_names:
-        sub = operations_df[operations_df["operation_name"] == op_name]
-        events_sub = []
-        for row in sub.itertuples(index=False):
-            events_sub.append((row.start_time, +1))
-            events_sub.append((row.end_time, -1))
-        if not events_sub:
-            max_concurrency_by_name[op_name] = 0
-            continue
-        events_sub.sort(key=lambda x: x[0])
-        c_curr = 0
-        c_max = 0
-        for ts, delta in events_sub:
-            c_curr += delta
-            if c_curr > c_max:  # pylint: disable=consider-using-max-builtin
-                c_max = c_curr
-        max_concurrency_by_name[op_name] = c_max
-
-    # -------------------------------------------------------------
-    # 3) Compute summary stats per (operation_type, operation_name)
-    # -------------------------------------------------------------
-    # We'll gather durations in a list, compute average, p95, p99, etc.
-
-    stats_dict = {}
-    grouped_ops = operations_df.groupby(["operation_type", "operation_name"])
-    for (op_type, op_name), grp in grouped_ops:
-        durations = grp["duration"].values
-        usage_count = len(durations)
-        avg_duration = durations.mean()
-        p95_duration = np.percentile(durations, 95)
-        p99_duration = np.percentile(durations, 99)
-
-        # concurrency
-        max_concur = max_concurrency_by_name.get(op_name, 0)
-
-        # define a custom "bottleneck_score":
-        # We say score = avg_duration * max_concurrency,
-        bottleneck_score = float(avg_duration * max_concur)
-
-        # store in dictionary
-        key = f"{op_type}:{op_name}"
-        stats_dict[key] = SimpleOperationStats(operation_type=op_type,
-                                               operation_name=op_name,
-                                               usage_count=usage_count,
-                                               avg_duration=float(avg_duration),
-                                               p95_duration=float(p95_duration),
-                                               p99_duration=float(p99_duration),
-                                               max_concurrency=int(max_concur),
-                                               bottleneck_score=bottleneck_score)
-
-    # -------------------------------------------------------------
-    # 4) Produce a textual summary highlighting top bottlenecks
-    # -------------------------------------------------------------
-    # We'll rank by bottleneck_score descending and show top 3.
-    if not stats_dict:
-        return SimpleBottleneckReport(stats={}, summary="No stats to report.")
-
-    top_items = sorted(stats_dict.values(), key=lambda x: x.bottleneck_score, reverse=True)
-    top_3 = top_items[:3]
-
-    # Build a simple textual summary
-    lines = []
-    lines.append("---- BOTTLENECK REPORT ----")
-    lines.append(f"Total distinct operations found: {len(stats_dict)}")
-    lines.append(f"Overall max concurrency (all ops): {overall_max_concurrency}")
-    lines.append("Top 3 Bottlenecks by bottleneck_score (avg_duration * max_concurrency):")
-    for i, item in enumerate(top_3, start=1):
-        lines.append(f"{i}) {item.operation_type} '{item.operation_name}': "
-                     f"score={item.bottleneck_score:.2f}, "
-                     f"avg_dur={item.avg_duration:.2f}, "
-                     f"max_concurrency={item.max_concurrency}")
-    summary_report = "\n".join(lines)
-
-    # Construct a final Pydantic model
-    return SimpleBottleneckReport(stats=stats_dict, summary=summary_report)

aiq/profiler/inference_optimization/data_models.py

@@ -1,386 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from typing import Any
-
-from pydantic import BaseModel
-from pydantic import ConfigDict
-from pydantic import Field
-from pydantic import RootModel
-
-# -----------------------------------------------------------
-# Prompt Caching Data Models
-# -----------------------------------------------------------
-
-
-class PrefixInfo(BaseModel):
-    """
-    Stores metadata about a particular prefix observed in the LLM text input.
-    """
-    prefix: str
-    prefix_length: int
-    calls_count: int
-    calls_percentage: float = Field(..., ge=0.0, le=1.0)
-
-
-class FrameworkLLMPrefixData(BaseModel):
-    """
-    Metadata for a single (framework, llm_name) group,
-    including total calls and all prefix statistics.
-    """
-    total_calls: int
-    prefix_info: list[PrefixInfo]
-
-
-class CommonPrefixesOutput(RootModel[dict[str, FrameworkLLMPrefixData]]):
-    """
-    A root model storing a dictionary keyed by '<framework>-<llm>',
-    each value is a FrameworkLLMPrefixData instance.
-    """
-
-    def to_dict(self) -> dict[str, FrameworkLLMPrefixData]:
-        """
-        Return the raw dictionary of data, discarding the 'root' wrapper.
-        """
-        return self.root
-
-
-# ----------------------------------------------------------------
-# Token Uniqueness Models
-# ----------------------------------------------------------------
-
-
-class LLMUniquenessMetrics(BaseModel):
-    """
-    Stores p90, p95, and p99 for the 'new words' metric.
-    """
-    p90: float
-    p95: float
-    p99: float
-
-
-class LLMUniquenessMetricsByLLM(RootModel[dict[str, LLMUniquenessMetrics]]):
-    """
-    A RootModel containing a dictionary where each key is an LLM name
-    and each value is the LLMUniquenessMetrics for that LLM.
-    """
-
-    def to_dict(self) -> dict[str, Any]:
-        # Return the raw dictionary for convenience
-        return self.root
-
-
-# ----------------------------------------------------------------
-# Workflow Runtime Models
-# ----------------------------------------------------------------
-
-
-class WorkflowRuntimeMetrics(BaseModel):
-    """
-    Stores p90, p95, and p99 for workflow runtimes across all examples.
-    """
-    p90: float
-    p95: float
-    p99: float
-
-
-# ----------------------------------------------------------------------
-# Simple Bottleneck Detection Models
-# ----------------------------------------------------------------------
-
-
-class SimpleOperationStats(BaseModel):
-    """
-    Statistics for a particular operation name (LLM or tool),
-    capturing concurrency, duration, usage, etc.
-    """
-    operation_type: str  # 'LLM' or 'TOOL'
-    operation_name: str  # e.g., "llama-3" or "serpapi"
-    usage_count: int  # how many times it appears
-    avg_duration: float  # average duration
-    p95_duration: float
-    p99_duration: float
-    max_concurrency: int  # maximum number of concurrent operations
-    bottleneck_score: float = Field(..., description="Custom metric to rank bottlenecks.")
-
-
-class SimpleBottleneckReport(BaseModel):
-    """
-    A container for all operation stats keyed by 'operation_type:operation_name',
-    plus a textual summary that highlights top bottlenecks.
-    """
-    stats: dict[str, SimpleOperationStats]
-    summary: str
-
-
-# ----------------------------------------------------------------------
-# Nested Bottleneck Models
-# ----------------------------------------------------------------------
-
-
-class CallNode(BaseModel):
-    """
-    A single call (LLM or TOOL) in a nested call tree.
-
-    Attributes
-    ----------
-    uuid: str
-        Unique ID tying together START/END events.
-    operation_type: str
-        e.g. 'LLM' or 'TOOL'.
-    operation_name: str
-        e.g. 'llama-3', 'bing-search', ...
-    start_time: float
-        Time when the call started.
-    end_time: float
-        Time when the call ended.
-    duration: float
-        end_time - start_time
-    children: list["CallNode"]
-        List of nested calls inside this call's time window.
-    parent: "CallNode" | None
-        Reference to the parent call in the tree (None if top-level).
-    """
-    model_config = ConfigDict(arbitrary_types_allowed=True)
-
-    uuid: str
-    operation_type: str
-    operation_name: str
-    start_time: float
-    end_time: float
-    duration: float = Field(..., description="end_time - start_time")
-    children: list["CallNode"] = Field(default_factory=list)
-    parent: "CallNode | None" = None
-
-    def compute_self_time(self) -> float:
-        """
-        'Self time' = duration minus the union of child intervals.
-        Overlapping child intervals are merged so we don't double-count them.
-        """
-        if not self.children:
-            return self.duration
-
-        intervals = [(c.start_time, c.end_time) for c in self.children]  # pylint: disable=not-an-iterable
-        # Sort by start time
-        intervals.sort(key=lambda x: x[0])
-
-        merged = []
-        cur_start, cur_end = intervals[0]
-        for i in range(1, len(intervals)):
-            s, e = intervals[i]
-            if s <= cur_end:
-                # Overlap
-                cur_end = max(cur_end, e)
-            else:
-                merged.append((cur_start, cur_end))
-                cur_start, cur_end = s, e
-        merged.append((cur_start, cur_end))
-
-        # Sum coverage, clamped to [start_time, end_time]
-        covered = 0.0
-        for (s, e) in merged:
-            s_clamped = max(s, self.start_time)
-            e_clamped = min(e, self.end_time)
-            if e_clamped > s_clamped:
-                covered += (e_clamped - s_clamped)
-
-        return max(0.0, self.duration - covered)
-
-    def compute_subtree_time(self) -> float:
-        """
-        Recursively compute the sum of self_time + children's subtree_time.
-        This ensures no overlap double-counting among children.
-        """
-        total = self.compute_self_time()
-        for c in self.children:  # pylint: disable=not-an-iterable
-            total += c.compute_subtree_time()
-        return total
-
-    def __str__(self) -> str:
-        return self._repr(0)
-
-    def _repr(self, level: int) -> str:
-        indent = "  " * level
-        info = (f"{indent}- {self.operation_type} '{self.operation_name}' "
-                f"(uuid={self.uuid}, start={self.start_time:.2f}, "
-                f"end={self.end_time:.2f}, dur={self.duration:.2f})")
-        child_strs = [child._repr(level + 1) for child in self.children]  # pylint: disable=not-an-iterable
-        return "\n".join([info] + child_strs)
-
-
-CallNode.model_rebuild()
-
-
-class NodeMetrics(BaseModel):
-    """
-    Metrics for a single node:
-      - self_time
-      - subtree_time
-      - concurrency_midpoint (optional)
-      - bottleneck_score (example: subtree_time)
-    """
-    uuid: str
-    operation_type: str
-    operation_name: str
-    start_time: float
-    end_time: float
-    duration: float
-    self_time: float
-    subtree_time: float
-    concurrency_midpoint: float | None = None
-    bottleneck_score: float
-
-
-class ConcurrencyDistribution(BaseModel):
-    """
-    Overall concurrency distribution info:
-      - timeline_segments: List of (start, end, concurrency)
-      - p50, p90, p95, p99 concurrency
-    """
-    timeline_segments: list[tuple[float, float, int]]
-    p50: float
-    p90: float
-    p95: float
-    p99: float
-
-
-class NestedCallProfilingResult(BaseModel):
-    """
-    The final Pydantic model returned by 'multi_example_call_profiling'.
-
-    Contains:
-      - concurrency: ConcurrencyDistribution
-      - node_metrics: dict[uuid, NodeMetrics]
-      - top_bottlenecks: The top calls by bottleneck_score
-      - textual_report: A multiline string summarizing everything
-    """
-    concurrency: ConcurrencyDistribution
-    node_metrics: dict[str, NodeMetrics]
-    top_bottlenecks: list[NodeMetrics]
-    textual_report: str
-
-
-# ----------------------------------------------------------------------
-# Concurrency Spike Analysis Models
-# ----------------------------------------------------------------------
-
-
-class ConcurrencyCallNode(CallNode):
-    """
-    A single call in the nested call tree for one example.
-    Each call is matched by a UUID with a `*_START` and `*_END` event.
-
-    Because fields like prompt_tokens, completion_tokens, total_tokens
-    may only exist at the END event, we store them only after seeing `*_END`".
-    """
-
-    example_number: int
-
-    # Additional fields from END events
-    prompt_tokens: int | None = None
-    completion_tokens: int | None = None
-    total_tokens: int | None = None
-    tool_outputs: str | None = None
-    llm_text_output: str | None = None
-
-
-ConcurrencyCallNode.model_rebuild()
-
-
-class ConcurrencySpikeInfo(BaseModel):
-    """
-    Info about one concurrency spike interval:
-    - start, end of the spike
-    - concurrency level
-    - list of calls that overlap
-    """
-    start_time: float
-    end_time: float
-    concurrency: int
-    active_uuids: list[str] = Field(default_factory=list)
-
-
-class ConcurrencyCorrelationStats(BaseModel):
-    """
-    Simple container for correlation / summarized stats of calls overlapping concurrency spikes.
-    """
-    avg_prompt_tokens: float
-    avg_total_tokens: float
-
-
-class ConcurrencyAnalysisResult(BaseModel):
-    """
-    The final Pydantic model returned by concurrency_spike_analysis(...).
-    Contains:
-    - concurrency_distribution: concurrency_level => total_time
-    - p50_concurrency, p90_concurrency, p95_concurrency, p99_concurrency
-    - spike_threshold, spike_intervals
-    - correlation_stats
-    - textual_report
-    """
-    concurrency_distribution: dict[int, float]
-    p50_concurrency: float
-    p90_concurrency: float
-    p95_concurrency: float
-    p99_concurrency: float
-
-    spike_threshold: int
-    spike_intervals: list[ConcurrencySpikeInfo]
-    correlation_stats: ConcurrencyCorrelationStats
-
-    average_latency_by_concurrency: dict[int, float]
-
-    textual_report: str
-
-
-# ----------------------------------------------------------------------
-# PrefixSpan Analysis Models
-# ----------------------------------------------------------------------
-
-
-class PrefixCallNode(BaseModel):
-    """
-    Represents a single call in an example's workflow.
-    - For LLM calls, we also store llm_text_input if available so we can incorporate it into the token.
-    """
-    uuid: str
-    example_number: int
-    operation_type: str  # "LLM" or "TOOL"
-    operation_name: str  # e.g. "llama-3", "internet-search"
-    start_time: float
-    end_time: float
-    duration: float
-    llm_text_input: str | None = None
-
-
-class FrequentPattern(BaseModel):
-    """
-    Frequent sub-sequence discovered by PrefixSpan, with coverage and average duration data.
-    """
-    pattern: list[str]  # e.g. ["LLM:llama-3|Hello world", "TOOL:internet-search"]
-    frequency: int  # total occurrences across all examples
-    coverage: float  # fraction of distinct examples that contain this pattern
-    average_duration: float  # average sum of call durations for calls in that sub-sequence
-    examples_containing: list[int]  # which examples have at least one occurrence
-
-
-class PrefixSpanSubworkflowResult(BaseModel):
-    """
-    Pydantic model for the final outcome:
-    - A list of frequent patterns
-    - A textual summary
-    """
-    patterns: list[FrequentPattern]
-    textual_report: str