Skip to content

Commit f23a78f

Browse files
committed
fix: avoid leaking memory when Client.with_options is used
Fixes openai/openai-python#865.
1 parent 5edeac0 commit f23a78f

File tree

4 files changed

+141
-17
lines changed

4 files changed

+141
-17
lines changed

Diff for: 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
]

Diff for: src/lithic/_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

Diff for: src/lithic/_client.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -282,7 +282,7 @@ def copy(
282282
return self.__class__(
283283
api_key=api_key or self.api_key,
284284
webhook_secret=webhook_secret or self.webhook_secret,
285-
base_url=base_url or str(self.base_url),
285+
base_url=base_url or self.base_url,
286286
environment=environment or self._environment,
287287
timeout=self.timeout if isinstance(timeout, NotGiven) else timeout,
288288
http_client=http_client,
@@ -584,7 +584,7 @@ def copy(
584584
return self.__class__(
585585
api_key=api_key or self.api_key,
586586
webhook_secret=webhook_secret or self.webhook_secret,
587-
base_url=base_url or str(self.base_url),
587+
base_url=base_url or self.base_url,
588588
environment=environment or self._environment,
589589
timeout=self.timeout if isinstance(timeout, NotGiven) else timeout,
590590
http_client=http_client,

Diff for: 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

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

199+
def test_copy_build_request(self) -> None:
200+
options = FinalRequestOptions(method="get", url="/foo")
201+
202+
def build_request(options: FinalRequestOptions) -> None:
203+
client = self.client.copy()
204+
client._build_request(options)
205+
206+
# ensure that the machinery is warmed up before tracing starts.
207+
build_request(options)
208+
gc.collect()
209+
210+
tracemalloc.start(1000)
211+
212+
snapshot_before = tracemalloc.take_snapshot()
213+
214+
ITERATIONS = 10
215+
for _ in range(ITERATIONS):
216+
build_request(options)
217+
gc.collect()
218+
219+
snapshot_after = tracemalloc.take_snapshot()
220+
221+
tracemalloc.stop()
222+
223+
def add_leak(leaks: list[tracemalloc.StatisticDiff], diff: tracemalloc.StatisticDiff) -> None:
224+
if diff.count == 0:
225+
# Avoid false positives by considering only leaks (i.e. allocations that persist).
226+
return
227+
228+
if diff.count % ITERATIONS != 0:
229+
# Avoid false positives by considering only leaks that appear per iteration.
230+
return
231+
232+
for frame in diff.traceback:
233+
if any(
234+
frame.filename.endswith(fragment)
235+
for fragment in [
236+
# to_raw_response_wrapper leaks through the @functools.wraps() decorator.
237+
#
238+
# removing the decorator fixes the leak for reasons we don't understand.
239+
"lithic/_response.py",
240+
# pydantic.BaseModel.model_dump || pydantic.BaseModel.dict leak memory for some reason.
241+
"lithic/_compat.py",
242+
# Standard library leaks we don't care about.
243+
"/logging/__init__.py",
244+
]
245+
):
246+
return
247+
248+
leaks.append(diff)
249+
250+
leaks: list[tracemalloc.StatisticDiff] = []
251+
for diff in snapshot_after.compare_to(snapshot_before, "traceback"):
252+
add_leak(leaks, diff)
253+
if leaks:
254+
for leak in leaks:
255+
print("MEMORY LEAK:", leak)
256+
for frame in leak.traceback:
257+
print(frame)
258+
raise AssertionError()
259+
197260
def test_request_timeout(self) -> None:
198261
request = self.client._build_request(FinalRequestOptions(method="get", url="/foo"))
199262
timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
@@ -953,6 +1016,67 @@ def test_copy_signature(self) -> None:
9531016
copy_param = copy_signature.parameters.get(name)
9541017
assert copy_param is not None, f"copy() signature is missing the {name} param"
9551018

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

0 commit comments

Comments
 (0)