|
2 | 2 |
|
3 | 3 | from __future__ import annotations
|
4 | 4 |
|
| 5 | +import gc |
5 | 6 | import os
|
6 | 7 | import json
|
7 | 8 | import asyncio
|
8 | 9 | import inspect
|
| 10 | +import tracemalloc |
9 | 11 | from typing import Any, Union, cast
|
10 | 12 | from unittest import mock
|
11 | 13 |
|
@@ -195,6 +197,67 @@ def test_copy_signature(self) -> None:
|
195 | 197 | copy_param = copy_signature.parameters.get(name)
|
196 | 198 | assert copy_param is not None, f"copy() signature is missing the {name} param"
|
197 | 199 |
|
| 200 | + def test_copy_build_request(self) -> None: |
| 201 | + options = FinalRequestOptions(method="get", url="/foo") |
| 202 | + |
| 203 | + def build_request(options: FinalRequestOptions) -> None: |
| 204 | + client = self.client.copy() |
| 205 | + client._build_request(options) |
| 206 | + |
| 207 | + # ensure that the machinery is warmed up before tracing starts. |
| 208 | + build_request(options) |
| 209 | + gc.collect() |
| 210 | + |
| 211 | + tracemalloc.start(1000) |
| 212 | + |
| 213 | + snapshot_before = tracemalloc.take_snapshot() |
| 214 | + |
| 215 | + ITERATIONS = 10 |
| 216 | + for _ in range(ITERATIONS): |
| 217 | + build_request(options) |
| 218 | + gc.collect() |
| 219 | + |
| 220 | + snapshot_after = tracemalloc.take_snapshot() |
| 221 | + |
| 222 | + tracemalloc.stop() |
| 223 | + |
| 224 | + def add_leak(leaks: list[tracemalloc.StatisticDiff], diff: tracemalloc.StatisticDiff) -> None: |
| 225 | + if diff.count == 0: |
| 226 | + # Avoid false positives by considering only leaks (i.e. allocations that persist). |
| 227 | + return |
| 228 | + |
| 229 | + if diff.count % ITERATIONS != 0: |
| 230 | + # Avoid false positives by considering only leaks that appear per iteration. |
| 231 | + return |
| 232 | + |
| 233 | + for frame in diff.traceback: |
| 234 | + if any( |
| 235 | + frame.filename.endswith(fragment) |
| 236 | + for fragment in [ |
| 237 | + # to_raw_response_wrapper leaks through the @functools.wraps() decorator. |
| 238 | + # |
| 239 | + # removing the decorator fixes the leak for reasons we don't understand. |
| 240 | + "openai/_response.py", |
| 241 | + # pydantic.BaseModel.model_dump || pydantic.BaseModel.dict leak memory for some reason. |
| 242 | + "openai/_compat.py", |
| 243 | + # Standard library leaks we don't care about. |
| 244 | + "/logging/__init__.py", |
| 245 | + ] |
| 246 | + ): |
| 247 | + return |
| 248 | + |
| 249 | + leaks.append(diff) |
| 250 | + |
| 251 | + leaks: list[tracemalloc.StatisticDiff] = [] |
| 252 | + for diff in snapshot_after.compare_to(snapshot_before, "traceback"): |
| 253 | + add_leak(leaks, diff) |
| 254 | + if leaks: |
| 255 | + for leak in leaks: |
| 256 | + print("MEMORY LEAK:", leak) |
| 257 | + for frame in leak.traceback: |
| 258 | + print(frame) |
| 259 | + raise AssertionError() |
| 260 | + |
198 | 261 | def test_request_timeout(self) -> None:
|
199 | 262 | request = self.client._build_request(FinalRequestOptions(method="get", url="/foo"))
|
200 | 263 | timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
|
@@ -858,6 +921,67 @@ def test_copy_signature(self) -> None:
|
858 | 921 | copy_param = copy_signature.parameters.get(name)
|
859 | 922 | assert copy_param is not None, f"copy() signature is missing the {name} param"
|
860 | 923 |
|
| 924 | + def test_copy_build_request(self) -> None: |
| 925 | + options = FinalRequestOptions(method="get", url="/foo") |
| 926 | + |
| 927 | + def build_request(options: FinalRequestOptions) -> None: |
| 928 | + client = self.client.copy() |
| 929 | + client._build_request(options) |
| 930 | + |
| 931 | + # ensure that the machinery is warmed up before tracing starts. |
| 932 | + build_request(options) |
| 933 | + gc.collect() |
| 934 | + |
| 935 | + tracemalloc.start(1000) |
| 936 | + |
| 937 | + snapshot_before = tracemalloc.take_snapshot() |
| 938 | + |
| 939 | + ITERATIONS = 10 |
| 940 | + for _ in range(ITERATIONS): |
| 941 | + build_request(options) |
| 942 | + gc.collect() |
| 943 | + |
| 944 | + snapshot_after = tracemalloc.take_snapshot() |
| 945 | + |
| 946 | + tracemalloc.stop() |
| 947 | + |
| 948 | + def add_leak(leaks: list[tracemalloc.StatisticDiff], diff: tracemalloc.StatisticDiff) -> None: |
| 949 | + if diff.count == 0: |
| 950 | + # Avoid false positives by considering only leaks (i.e. allocations that persist). |
| 951 | + return |
| 952 | + |
| 953 | + if diff.count % ITERATIONS != 0: |
| 954 | + # Avoid false positives by considering only leaks that appear per iteration. |
| 955 | + return |
| 956 | + |
| 957 | + for frame in diff.traceback: |
| 958 | + if any( |
| 959 | + frame.filename.endswith(fragment) |
| 960 | + for fragment in [ |
| 961 | + # to_raw_response_wrapper leaks through the @functools.wraps() decorator. |
| 962 | + # |
| 963 | + # removing the decorator fixes the leak for reasons we don't understand. |
| 964 | + "openai/_response.py", |
| 965 | + # pydantic.BaseModel.model_dump || pydantic.BaseModel.dict leak memory for some reason. |
| 966 | + "openai/_compat.py", |
| 967 | + # Standard library leaks we don't care about. |
| 968 | + "/logging/__init__.py", |
| 969 | + ] |
| 970 | + ): |
| 971 | + return |
| 972 | + |
| 973 | + leaks.append(diff) |
| 974 | + |
| 975 | + leaks: list[tracemalloc.StatisticDiff] = [] |
| 976 | + for diff in snapshot_after.compare_to(snapshot_before, "traceback"): |
| 977 | + add_leak(leaks, diff) |
| 978 | + if leaks: |
| 979 | + for leak in leaks: |
| 980 | + print("MEMORY LEAK:", leak) |
| 981 | + for frame in leak.traceback: |
| 982 | + print(frame) |
| 983 | + raise AssertionError() |
| 984 | + |
861 | 985 | async def test_request_timeout(self) -> None:
|
862 | 986 | request = self.client._build_request(FinalRequestOptions(method="get", url="/foo"))
|
863 | 987 | timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
|
|
0 commit comments