Skip to content

Commit 9b98da6

Browse files
committed
fix: avoid leaking memory when Client.with_options is used
Fixes #865.
1 parent c70e194 commit 9b98da6

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
@@ -152,8 +152,6 @@ select = [
152152
"T203",
153153
]
154154
ignore = [
155-
# lru_cache in methods, will be fixed separately
156-
"B019",
157155
# mutable defaults
158156
"B006",
159157
]

Diff for: src/openai/_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/openai/_client.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -192,7 +192,7 @@ def copy(
192192
return self.__class__(
193193
api_key=api_key or self.api_key,
194194
organization=organization or self.organization,
195-
base_url=base_url or str(self.base_url),
195+
base_url=base_url or self.base_url,
196196
timeout=self.timeout if isinstance(timeout, NotGiven) else timeout,
197197
http_client=http_client,
198198
max_retries=max_retries if is_given(max_retries) else self.max_retries,
@@ -402,7 +402,7 @@ def copy(
402402
return self.__class__(
403403
api_key=api_key or self.api_key,
404404
organization=organization or self.organization,
405-
base_url=base_url or str(self.base_url),
405+
base_url=base_url or self.base_url,
406406
timeout=self.timeout if isinstance(timeout, NotGiven) else timeout,
407407
http_client=http_client,
408408
max_retries=max_retries if is_given(max_retries) else self.max_retries,

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

@@ -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+
"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+
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
@@ -858,6 +921,67 @@ def test_copy_signature(self) -> None:
858921
copy_param = copy_signature.parameters.get(name)
859922
assert copy_param is not None, f"copy() signature is missing the {name} param"
860923

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+
861985
async def test_request_timeout(self) -> None:
862986
request = self.client._build_request(FinalRequestOptions(method="get", url="/foo"))
863987
timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore

0 commit comments

Comments
 (0)