|
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 |
|
@@ -224,6 +226,67 @@ def test_copy_signature(self) -> None:
|
224 | 226 | copy_param = copy_signature.parameters.get(name)
|
225 | 227 | assert copy_param is not None, f"copy() signature is missing the {name} param"
|
226 | 228 |
|
| 229 | + def test_copy_build_request(self) -> None: |
| 230 | + options = FinalRequestOptions(method="get", url="/foo") |
| 231 | + |
| 232 | + def build_request(options: FinalRequestOptions) -> None: |
| 233 | + client = self.client.copy() |
| 234 | + client._build_request(options) |
| 235 | + |
| 236 | + # ensure that the machinery is warmed up before tracing starts. |
| 237 | + build_request(options) |
| 238 | + gc.collect() |
| 239 | + |
| 240 | + tracemalloc.start(1000) |
| 241 | + |
| 242 | + snapshot_before = tracemalloc.take_snapshot() |
| 243 | + |
| 244 | + ITERATIONS = 10 |
| 245 | + for _ in range(ITERATIONS): |
| 246 | + build_request(options) |
| 247 | + gc.collect() |
| 248 | + |
| 249 | + snapshot_after = tracemalloc.take_snapshot() |
| 250 | + |
| 251 | + tracemalloc.stop() |
| 252 | + |
| 253 | + def add_leak(leaks: list[tracemalloc.StatisticDiff], diff: tracemalloc.StatisticDiff) -> None: |
| 254 | + if diff.count == 0: |
| 255 | + # Avoid false positives by considering only leaks (i.e. allocations that persist). |
| 256 | + return |
| 257 | + |
| 258 | + if diff.count % ITERATIONS != 0: |
| 259 | + # Avoid false positives by considering only leaks that appear per iteration. |
| 260 | + return |
| 261 | + |
| 262 | + for frame in diff.traceback: |
| 263 | + if any( |
| 264 | + frame.filename.endswith(fragment) |
| 265 | + for fragment in [ |
| 266 | + # to_raw_response_wrapper leaks through the @functools.wraps() decorator. |
| 267 | + # |
| 268 | + # removing the decorator fixes the leak for reasons we don't understand. |
| 269 | + "anthropic_bedrock/_response.py", |
| 270 | + # pydantic.BaseModel.model_dump || pydantic.BaseModel.dict leak memory for some reason. |
| 271 | + "anthropic_bedrock/_compat.py", |
| 272 | + # Standard library leaks we don't care about. |
| 273 | + "/logging/__init__.py", |
| 274 | + ] |
| 275 | + ): |
| 276 | + return |
| 277 | + |
| 278 | + leaks.append(diff) |
| 279 | + |
| 280 | + leaks: list[tracemalloc.StatisticDiff] = [] |
| 281 | + for diff in snapshot_after.compare_to(snapshot_before, "traceback"): |
| 282 | + add_leak(leaks, diff) |
| 283 | + if leaks: |
| 284 | + for leak in leaks: |
| 285 | + print("MEMORY LEAK:", leak) |
| 286 | + for frame in leak.traceback: |
| 287 | + print(frame) |
| 288 | + raise AssertionError() |
| 289 | + |
227 | 290 | def test_request_timeout(self) -> None:
|
228 | 291 | request = self.client._build_request(FinalRequestOptions(method="get", url="/foo"))
|
229 | 292 | timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
|
@@ -997,6 +1060,67 @@ def test_copy_signature(self) -> None:
|
997 | 1060 | copy_param = copy_signature.parameters.get(name)
|
998 | 1061 | assert copy_param is not None, f"copy() signature is missing the {name} param"
|
999 | 1062 |
|
| 1063 | + def test_copy_build_request(self) -> None: |
| 1064 | + options = FinalRequestOptions(method="get", url="/foo") |
| 1065 | + |
| 1066 | + def build_request(options: FinalRequestOptions) -> None: |
| 1067 | + client = self.client.copy() |
| 1068 | + client._build_request(options) |
| 1069 | + |
| 1070 | + # ensure that the machinery is warmed up before tracing starts. |
| 1071 | + build_request(options) |
| 1072 | + gc.collect() |
| 1073 | + |
| 1074 | + tracemalloc.start(1000) |
| 1075 | + |
| 1076 | + snapshot_before = tracemalloc.take_snapshot() |
| 1077 | + |
| 1078 | + ITERATIONS = 10 |
| 1079 | + for _ in range(ITERATIONS): |
| 1080 | + build_request(options) |
| 1081 | + gc.collect() |
| 1082 | + |
| 1083 | + snapshot_after = tracemalloc.take_snapshot() |
| 1084 | + |
| 1085 | + tracemalloc.stop() |
| 1086 | + |
| 1087 | + def add_leak(leaks: list[tracemalloc.StatisticDiff], diff: tracemalloc.StatisticDiff) -> None: |
| 1088 | + if diff.count == 0: |
| 1089 | + # Avoid false positives by considering only leaks (i.e. allocations that persist). |
| 1090 | + return |
| 1091 | + |
| 1092 | + if diff.count % ITERATIONS != 0: |
| 1093 | + # Avoid false positives by considering only leaks that appear per iteration. |
| 1094 | + return |
| 1095 | + |
| 1096 | + for frame in diff.traceback: |
| 1097 | + if any( |
| 1098 | + frame.filename.endswith(fragment) |
| 1099 | + for fragment in [ |
| 1100 | + # to_raw_response_wrapper leaks through the @functools.wraps() decorator. |
| 1101 | + # |
| 1102 | + # removing the decorator fixes the leak for reasons we don't understand. |
| 1103 | + "anthropic_bedrock/_response.py", |
| 1104 | + # pydantic.BaseModel.model_dump || pydantic.BaseModel.dict leak memory for some reason. |
| 1105 | + "anthropic_bedrock/_compat.py", |
| 1106 | + # Standard library leaks we don't care about. |
| 1107 | + "/logging/__init__.py", |
| 1108 | + ] |
| 1109 | + ): |
| 1110 | + return |
| 1111 | + |
| 1112 | + leaks.append(diff) |
| 1113 | + |
| 1114 | + leaks: list[tracemalloc.StatisticDiff] = [] |
| 1115 | + for diff in snapshot_after.compare_to(snapshot_before, "traceback"): |
| 1116 | + add_leak(leaks, diff) |
| 1117 | + if leaks: |
| 1118 | + for leak in leaks: |
| 1119 | + print("MEMORY LEAK:", leak) |
| 1120 | + for frame in leak.traceback: |
| 1121 | + print(frame) |
| 1122 | + raise AssertionError() |
| 1123 | + |
1000 | 1124 | async def test_request_timeout(self) -> None:
|
1001 | 1125 | request = self.client._build_request(FinalRequestOptions(method="get", url="/foo"))
|
1002 | 1126 | timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
|
|
0 commit comments