Skip to content

Commit c634771

Browse files
fix: avoid leaking memory when Client.with_options is used (#220)
Fixes openai/openai-python#865.
1 parent ee8fb39 commit c634771

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/finch/_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/finch/_client.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -236,7 +236,7 @@ def copy(
236236
client_id=client_id or self.client_id,
237237
client_secret=client_secret or self.client_secret,
238238
webhook_secret=webhook_secret or self.webhook_secret,
239-
base_url=base_url or str(self.base_url),
239+
base_url=base_url or self.base_url,
240240
timeout=self.timeout if isinstance(timeout, NotGiven) else timeout,
241241
http_client=http_client,
242242
connection_pool_limits=connection_pool_limits,
@@ -542,7 +542,7 @@ def copy(
542542
client_id=client_id or self.client_id,
543543
client_secret=client_secret or self.client_secret,
544544
webhook_secret=webhook_secret or self.webhook_secret,
545-
base_url=base_url or str(self.base_url),
545+
base_url=base_url or self.base_url,
546546
timeout=self.timeout if isinstance(timeout, NotGiven) else timeout,
547547
http_client=http_client,
548548
connection_pool_limits=connection_pool_limits,

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

@@ -174,6 +176,67 @@ def test_copy_signature(self) -> None:
174176
copy_param = copy_signature.parameters.get(name)
175177
assert copy_param is not None, f"copy() signature is missing the {name} param"
176178

179+
def test_copy_build_request(self) -> None:
180+
options = FinalRequestOptions(method="get", url="/foo")
181+
182+
def build_request(options: FinalRequestOptions) -> None:
183+
client = self.client.copy()
184+
client._build_request(options)
185+
186+
# ensure that the machinery is warmed up before tracing starts.
187+
build_request(options)
188+
gc.collect()
189+
190+
tracemalloc.start(1000)
191+
192+
snapshot_before = tracemalloc.take_snapshot()
193+
194+
ITERATIONS = 10
195+
for _ in range(ITERATIONS):
196+
build_request(options)
197+
gc.collect()
198+
199+
snapshot_after = tracemalloc.take_snapshot()
200+
201+
tracemalloc.stop()
202+
203+
def add_leak(leaks: list[tracemalloc.StatisticDiff], diff: tracemalloc.StatisticDiff) -> None:
204+
if diff.count == 0:
205+
# Avoid false positives by considering only leaks (i.e. allocations that persist).
206+
return
207+
208+
if diff.count % ITERATIONS != 0:
209+
# Avoid false positives by considering only leaks that appear per iteration.
210+
return
211+
212+
for frame in diff.traceback:
213+
if any(
214+
frame.filename.endswith(fragment)
215+
for fragment in [
216+
# to_raw_response_wrapper leaks through the @functools.wraps() decorator.
217+
#
218+
# removing the decorator fixes the leak for reasons we don't understand.
219+
"finch/_response.py",
220+
# pydantic.BaseModel.model_dump || pydantic.BaseModel.dict leak memory for some reason.
221+
"finch/_compat.py",
222+
# Standard library leaks we don't care about.
223+
"/logging/__init__.py",
224+
]
225+
):
226+
return
227+
228+
leaks.append(diff)
229+
230+
leaks: list[tracemalloc.StatisticDiff] = []
231+
for diff in snapshot_after.compare_to(snapshot_before, "traceback"):
232+
add_leak(leaks, diff)
233+
if leaks:
234+
for leak in leaks:
235+
print("MEMORY LEAK:", leak)
236+
for frame in leak.traceback:
237+
print(frame)
238+
raise AssertionError()
239+
177240
def test_request_timeout(self) -> None:
178241
request = self.client._build_request(FinalRequestOptions(method="get", url="/foo"))
179242
timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore
@@ -882,6 +945,67 @@ def test_copy_signature(self) -> None:
882945
copy_param = copy_signature.parameters.get(name)
883946
assert copy_param is not None, f"copy() signature is missing the {name} param"
884947

948+
def test_copy_build_request(self) -> None:
949+
options = FinalRequestOptions(method="get", url="/foo")
950+
951+
def build_request(options: FinalRequestOptions) -> None:
952+
client = self.client.copy()
953+
client._build_request(options)
954+
955+
# ensure that the machinery is warmed up before tracing starts.
956+
build_request(options)
957+
gc.collect()
958+
959+
tracemalloc.start(1000)
960+
961+
snapshot_before = tracemalloc.take_snapshot()
962+
963+
ITERATIONS = 10
964+
for _ in range(ITERATIONS):
965+
build_request(options)
966+
gc.collect()
967+
968+
snapshot_after = tracemalloc.take_snapshot()
969+
970+
tracemalloc.stop()
971+
972+
def add_leak(leaks: list[tracemalloc.StatisticDiff], diff: tracemalloc.StatisticDiff) -> None:
973+
if diff.count == 0:
974+
# Avoid false positives by considering only leaks (i.e. allocations that persist).
975+
return
976+
977+
if diff.count % ITERATIONS != 0:
978+
# Avoid false positives by considering only leaks that appear per iteration.
979+
return
980+
981+
for frame in diff.traceback:
982+
if any(
983+
frame.filename.endswith(fragment)
984+
for fragment in [
985+
# to_raw_response_wrapper leaks through the @functools.wraps() decorator.
986+
#
987+
# removing the decorator fixes the leak for reasons we don't understand.
988+
"finch/_response.py",
989+
# pydantic.BaseModel.model_dump || pydantic.BaseModel.dict leak memory for some reason.
990+
"finch/_compat.py",
991+
# Standard library leaks we don't care about.
992+
"/logging/__init__.py",
993+
]
994+
):
995+
return
996+
997+
leaks.append(diff)
998+
999+
leaks: list[tracemalloc.StatisticDiff] = []
1000+
for diff in snapshot_after.compare_to(snapshot_before, "traceback"):
1001+
add_leak(leaks, diff)
1002+
if leaks:
1003+
for leak in leaks:
1004+
print("MEMORY LEAK:", leak)
1005+
for frame in leak.traceback:
1006+
print(frame)
1007+
raise AssertionError()
1008+
8851009
async def test_request_timeout(self) -> None:
8861010
request = self.client._build_request(FinalRequestOptions(method="get", url="/foo"))
8871011
timeout = httpx.Timeout(**request.extensions["timeout"]) # type: ignore

0 commit comments

Comments
 (0)