Skip to content

Commit 42a94b3

Browse files
committed
fix: avoid leaking memory when Client.with_options is used (#275)
Fixes openai/openai-python#865.
1 parent 78cdc53 commit 42a94b3

File tree

4 files changed

+141
-17
lines changed

4 files changed

+141
-17
lines changed

pyproject.toml

-2
Original file line numberDiff line numberDiff line change
@@ -149,8 +149,6 @@ select = [
149149
"T203",
150150
]
151151
ignore = [
152-
# lru_cache in methods, will be fixed separately
153-
"B019",
154152
# mutable defaults
155153
"B006",
156154
]

src/anthropic/_base_client.py

+15-13
Original file line numberDiff line numberDiff line change
@@ -403,14 +403,12 @@ def _build_headers(self, options: FinalRequestOptions) -> httpx.Headers:
403403
headers_dict = _merge_mappings(self.default_headers, custom_headers)
404404
self._validate_headers(headers_dict, custom_headers)
405405

406+
# headers are case-insensitive while dictionaries are not.
406407
headers = httpx.Headers(headers_dict)
407408

408409
idempotency_header = self._idempotency_header
409410
if idempotency_header and options.method.lower() != "get" and idempotency_header not in headers:
410-
if not options.idempotency_key:
411-
options.idempotency_key = self._idempotency_key()
412-
413-
headers[idempotency_header] = options.idempotency_key
411+
headers[idempotency_header] = options.idempotency_key or self._idempotency_key()
414412

415413
return headers
416414

@@ -594,16 +592,8 @@ def base_url(self) -> URL:
594592
def base_url(self, url: URL | str) -> None:
595593
self._base_url = self._enforce_trailing_slash(url if isinstance(url, URL) else URL(url))
596594

597-
@lru_cache(maxsize=None)
598595
def platform_headers(self) -> Dict[str, str]:
599-
return {
600-
"X-Stainless-Lang": "python",
601-
"X-Stainless-Package-Version": self._version,
602-
"X-Stainless-OS": str(get_platform()),
603-
"X-Stainless-Arch": str(get_architecture()),
604-
"X-Stainless-Runtime": platform.python_implementation(),
605-
"X-Stainless-Runtime-Version": platform.python_version(),
606-
}
596+
return platform_headers(self._version)
607597

608598
def _calculate_retry_timeout(
609599
self,
@@ -1691,6 +1681,18 @@ def get_platform() -> Platform:
16911681
return "Unknown"
16921682

16931683

1684+
@lru_cache(maxsize=None)
1685+
def platform_headers(version: str) -> Dict[str, str]:
1686+
return {
1687+
"X-Stainless-Lang": "python",
1688+
"X-Stainless-Package-Version": version,
1689+
"X-Stainless-OS": str(get_platform()),
1690+
"X-Stainless-Arch": str(get_architecture()),
1691+
"X-Stainless-Runtime": platform.python_implementation(),
1692+
"X-Stainless-Runtime-Version": platform.python_version(),
1693+
}
1694+
1695+
16941696
class OtherArch:
16951697
def __init__(self, name: str) -> None:
16961698
self.name = name

src/anthropic/_client.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -239,7 +239,7 @@ def copy(
239239
return self.__class__(
240240
api_key=api_key or self.api_key,
241241
auth_token=auth_token or self.auth_token,
242-
base_url=base_url or str(self.base_url),
242+
base_url=base_url or self.base_url,
243243
timeout=self.timeout if isinstance(timeout, NotGiven) else timeout,
244244
http_client=http_client,
245245
connection_pool_limits=connection_pool_limits,
@@ -500,7 +500,7 @@ def copy(
500500
return self.__class__(
501501
api_key=api_key or self.api_key,
502502
auth_token=auth_token or self.auth_token,
503-
base_url=base_url or str(self.base_url),
503+
base_url=base_url or self.base_url,
504504
timeout=self.timeout if isinstance(timeout, NotGiven) else timeout,
505505
http_client=http_client,
506506
connection_pool_limits=connection_pool_limits,

tests/test_client.py

+124
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,12 @@
22

33
from __future__ import annotations
44

5+
import gc
56
import os
67
import json
78
import asyncio
89
import inspect
10+
import tracemalloc
911
from typing import Any, Union, cast
1012
from unittest import mock
1113

@@ -195,6 +197,67 @@ def test_copy_signature(self) -> None:
195197
copy_param = copy_signature.parameters.get(name)
196198
assert copy_param is not None, f"copy() signature is missing the {name} param"
197199

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+
"anthropic/_response.py",
241+
# pydantic.BaseModel.model_dump || pydantic.BaseModel.dict leak memory for some reason.
242+
"anthropic/_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+
198261
def test_request_timeout(self) -> None:
199262
request = self.client._build_request(FinalRequestOptions(method="get", url="/foo"))
200263
timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
@@ -947,6 +1010,67 @@ def test_copy_signature(self) -> None:
9471010
copy_param = copy_signature.parameters.get(name)
9481011
assert copy_param is not None, f"copy() signature is missing the {name} param"
9491012

1013+
def test_copy_build_request(self) -> None:
1014+
options = FinalRequestOptions(method="get", url="/foo")
1015+
1016+
def build_request(options: FinalRequestOptions) -> None:
1017+
client = self.client.copy()
1018+
client._build_request(options)
1019+
1020+
# ensure that the machinery is warmed up before tracing starts.
1021+
build_request(options)
1022+
gc.collect()
1023+
1024+
tracemalloc.start(1000)
1025+
1026+
snapshot_before = tracemalloc.take_snapshot()
1027+
1028+
ITERATIONS = 10
1029+
for _ in range(ITERATIONS):
1030+
build_request(options)
1031+
gc.collect()
1032+
1033+
snapshot_after = tracemalloc.take_snapshot()
1034+
1035+
tracemalloc.stop()
1036+
1037+
def add_leak(leaks: list[tracemalloc.StatisticDiff], diff: tracemalloc.StatisticDiff) -> None:
1038+
if diff.count == 0:
1039+
# Avoid false positives by considering only leaks (i.e. allocations that persist).
1040+
return
1041+
1042+
if diff.count % ITERATIONS != 0:
1043+
# Avoid false positives by considering only leaks that appear per iteration.
1044+
return
1045+
1046+
for frame in diff.traceback:
1047+
if any(
1048+
frame.filename.endswith(fragment)
1049+
for fragment in [
1050+
# to_raw_response_wrapper leaks through the @functools.wraps() decorator.
1051+
#
1052+
# removing the decorator fixes the leak for reasons we don't understand.
1053+
"anthropic/_response.py",
1054+
# pydantic.BaseModel.model_dump || pydantic.BaseModel.dict leak memory for some reason.
1055+
"anthropic/_compat.py",
1056+
# Standard library leaks we don't care about.
1057+
"/logging/__init__.py",
1058+
]
1059+
):
1060+
return
1061+
1062+
leaks.append(diff)
1063+
1064+
leaks: list[tracemalloc.StatisticDiff] = []
1065+
for diff in snapshot_after.compare_to(snapshot_before, "traceback"):
1066+
add_leak(leaks, diff)
1067+
if leaks:
1068+
for leak in leaks:
1069+
print("MEMORY LEAK:", leak)
1070+
for frame in leak.traceback:
1071+
print(frame)
1072+
raise AssertionError()
1073+
9501074
async def test_request_timeout(self) -> None:
9511075
request = self.client._build_request(FinalRequestOptions(method="get", url="/foo"))
9521076
timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore

0 commit comments

Comments
 (0)