nvidia-nat 1.2.0__py3-none-any.whl

nat/profiler/inference_optimization/bottleneck_analysis/simple_stack_analysis.py

@@ -0,0 +1,258 @@
+# 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 nat.data_models.intermediate_step import IntermediateStep
+from nat.profiler.inference_optimization.data_models import SimpleBottleneckReport
+from nat.profiler.inference_optimization.data_models import SimpleOperationStats
+from nat.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)

nat/profiler/inference_optimization/data_models.py

@@ -0,0 +1,386 @@
+# 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