vllm_pdb_helper.py

import pdb
import traceback
import inspect


def function_in_stack(function_name):
    stack = traceback.extract_stack() 
    for frame in stack:
        if frame.name == function_name:
            return True
    return False

# Using this to insert breakpoints when running on actual data.
# Useful when e.g. comparing hidden states for correctness. 
# If profile_run is in our stack trace, we're profiling for KV cache usage
# If initialize_cache is in our stack trace, we're creating CUDA graphs.
if not function_in_stack("profile_run") or function_in_stack("initialize_cache"):
    pdb.set_trace()