diff --git a/docs/source/contributing/profiling/profiling_index.md b/docs/source/contributing/profiling/profiling_index.md index d6e597ea..ce25daa3 100644 --- a/docs/source/contributing/profiling/profiling_index.md +++ b/docs/source/contributing/profiling/profiling_index.md @@ -124,3 +124,52 @@ nsys stats report1.nsys-rep GUI example: Screenshot 2025-03-05 at 11 48 42 AM + +## Profiling vLLM Python Code + +The Python standard library includes +[cProfile](https://docs.python.org/3/library/profile.html) for profiling Python +code. vLLM includes a couple of helpers that make it easy to apply it to a section of vLLM. +Both the `vllm.utils.cprofile` and `vllm.utils.cprofile_context` functions can be +used to profile a section of code. + +### Example usage - decorator + +The first helper is a Python decorator that can be used to profile a function. +If a filename is specified, the profile will be saved to that file. If no filename is +specified, profile data will be printed to stdout. + +```python +import vllm.utils + +@vllm.utils.cprofile("expensive_function.prof") +def expensive_function(): + # some expensive code + pass +``` + +### Example Usage - context manager + +The second helper is a context manager that can be used to profile a block of +code. Similar to the decorator, the filename is optional. + +```python +import vllm.utils + +def another_function(): + # more expensive code + pass + +with vllm.utils.cprofile_context("another_function.prof"): + another_function() +``` + +### Analyzing Profile Results + +There are multiple tools available that can help analyze the profile results. +One example is [snakeviz](https://jiffyclub.github.io/snakeviz/). + +```bash +pip install snakeviz +snakeviz expensive_function.prof +``` diff --git a/vllm/utils.py b/vllm/utils.py index 55ee044b..64d9faeb 100644 --- a/vllm/utils.py +++ b/vllm/utils.py @@ -2405,3 +2405,51 @@ def swap_dict_values(obj: dict[_K, _V], key1: _K, key2: _K) -> None: obj[key1] = v2 else: obj.pop(key1, None) + + +@contextlib.contextmanager +def cprofile_context(save_file: Optional[str] = None): + """Run a cprofile + + Args: + save_file: path to save the profile result. "1" or + None will result in printing to stdout. + """ + import cProfile + + prof = cProfile.Profile() + prof.enable() + + try: + yield + finally: + prof.disable() + if save_file and save_file != "1": + prof.dump_stats(save_file) + else: + prof.print_stats(sort="cumtime") + + +def cprofile(save_file: Optional[str] = None, enabled: bool = True): + """Decorator to profile a Python method using cProfile. + + Args: + save_file: Path to save the profile result. + If "1", None, or "", results will be printed to stdout. + enabled: Set to false to turn this into a no-op + """ + + def decorator(func: Callable): + + @wraps(func) + def wrapper(*args, **kwargs): + if not enabled: + # If profiling is disabled, just call the function directly. + return func(*args, **kwargs) + + with cprofile_context(save_file): + return func(*args, **kwargs) + + return wrapper + + return decorator