A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 18:30 UTC
[Misc] Add cProfile helpers (#15074)
Signed-off-by: Russell Bryant <rbryant@redhat.com>
Russell Bryant committed
Mar 21, 2025 at 07:56 UTC
61e8c183505355e2005314e7004a76049439e165
2 files changed
+97
docs/source/contributing/profiling/profiling_index.md
+49
@@ -124,3 +124,52 @@ nsys stats report1.nsys-rep
124
GUI example:
125
126
<img width="1799" alt="Screenshot 2025-03-05 at 11 48 42 AM" src="https://github.com/user-attachments/assets/c7cff1ae-6d6f-477d-a342-bd13c4fc424c" />
127
+
128
+## Profiling vLLM Python Code
129
+
130
+The Python standard library includes
131
+[cProfile](https://docs.python.org/3/library/profile.html) for profiling Python
132
+code. vLLM includes a couple of helpers that make it easy to apply it to a section of vLLM.
133
+Both the `vllm.utils.cprofile` and `vllm.utils.cprofile_context` functions can be
134
+used to profile a section of code.
135
+
136
+### Example usage - decorator
137
+
138
+The first helper is a Python decorator that can be used to profile a function.
139
+If a filename is specified, the profile will be saved to that file. If no filename is
140
+specified, profile data will be printed to stdout.
141
+
142
+```python
143
+import vllm.utils
144
+
145
+@vllm.utils.cprofile("expensive_function.prof")
146
+def expensive_function():
147
+ # some expensive code
148
+ pass
149
+```
150
+
151
+### Example Usage - context manager
152
+
153
+The second helper is a context manager that can be used to profile a block of
154
+code. Similar to the decorator, the filename is optional.
155
+
156
+```python
157
+import vllm.utils
158
+
159
+def another_function():
160
+ # more expensive code
161
+ pass
162
+
163
+with vllm.utils.cprofile_context("another_function.prof"):
164
+ another_function()
165
+```
166
+
167
+### Analyzing Profile Results
168
+
169
+There are multiple tools available that can help analyze the profile results.
170
+One example is [snakeviz](https://jiffyclub.github.io/snakeviz/).
171
+
172
+```bash
173
+pip install snakeviz
174
+snakeviz expensive_function.prof
175
+```
vllm/utils.py
+48
@@ -2405,3 +2405,51 @@ def swap_dict_values(obj: dict[_K, _V], key1: _K, key2: _K) -> None:
2405
obj[key1] = v2
2406
else:
2407
obj.pop(key1, None)
2408
+
2409
+
2410
+@contextlib.contextmanager
2411
+def cprofile_context(save_file: Optional[str] = None):
2412
+ """Run a cprofile
2413
+
2414
+ Args:
2415
+ save_file: path to save the profile result. "1" or
2416
+ None will result in printing to stdout.
2417
+ """
2418
+ import cProfile
2419
+
2420
+ prof = cProfile.Profile()
2421
+ prof.enable()
2422
+
2423
+ try:
2424
+ yield
2425
+ finally:
2426
+ prof.disable()
2427
+ if save_file and save_file != "1":
2428
+ prof.dump_stats(save_file)
2429
+ else:
2430
+ prof.print_stats(sort="cumtime")
2431
+
2432
+
2433
+def cprofile(save_file: Optional[str] = None, enabled: bool = True):
2434
+ """Decorator to profile a Python method using cProfile.
2435
+
2436
+ Args:
2437
+ save_file: Path to save the profile result.
2438
+ If "1", None, or "", results will be printed to stdout.
2439
+ enabled: Set to false to turn this into a no-op
2440
+ """
2441
+
2442
+ def decorator(func: Callable):
2443
+
2444
+ @wraps(func)
2445
+ def wrapper(*args, **kwargs):
2446
+ if not enabled:
2447
+ # If profiling is disabled, just call the function directly.
2448
+ return func(*args, **kwargs)
2449
+
2450
+ with cprofile_context(save_file):
2451
+ return func(*args, **kwargs)
2452
+
2453
+ return wrapper
2454
+
2455
+ return decorator