From 7010686b45dc3c8d3075a9030a77ab7a8a9f37f3 Mon Sep 17 00:00:00 2001 From: Yao You Date: Tue, 28 Jul 2026 19:37:46 -0500 Subject: [PATCH 1/3] experiment: NDJSON elements-file mode to bound split-PDF recombination memory EXPERIMENT - for local editable-install testing, not for release. The split-PDF path already streams each chunk response to a temp file without reading it into memory (call_api_async), but recombination then made four full copies of the document: 1. json.load per chunk (load_elements_from_response) 2. all chunk lists held at once (_elements_from_task_responses) 3. flattened list (same) 4. json.dumps(elements).encode() (request_utils.create_response) and general.partition then re-parsed that blob into PartitionResponse.elements for a fifth. On a document with large image_base64 payloads this makes the caller's peak worse than the server's. This adds an opt-in path that never builds a list: - PartitionAcceptEnum.APPLICATION_X_NDJSON. Previously the enum held only json/csv AND the response matcher rejected anything else, so NDJSON was not expressible at all. - PartitionResponse.elements_file: path to an NDJSON file, set instead of `elements`. The CALLER owns deleting it. - request_utils.combine_chunk_files_to_ndjson: concatenates the per-chunk temp files on disk. Detects each chunk's shape from its first non-space byte, so a json-array chunk is converted (bounded by one chunk, 20 pages by default) and an NDJSON chunk is copied through with no parsing at all. - split_pdf_hook: records ndjson mode from the Accept header in before_request, collects chunk paths instead of parsing, and returns the combined path via an x-unstructured-elements-file header, following the existing convention where a cached chunk response carries its temp-file path as the body. Elements-file mode requires cache_tmp_data (it needs the chunk files) and falls back to the existing in-memory recombination otherwise. Default behavior is unchanged: without the NDJSON accept header nothing here activates. Note the combined file is deliberately written to the cache PARENT dir, not the operation's TemporaryDirectory, which _clear_operation deletes immediately after after_success returns. Tests: 10 new unit tests for the combine helper (array/ndjson/mixed/empty chunks, order, exact payload round-trip, non-ascii). Existing suite unaffected: 224 passed, 1 xfailed. Co-Authored-By: Claude Opus 5 (1M context) --- .../unit/test_ndjson_combine.py | 159 ++++++++++++++++++ .../_hooks/custom/request_utils.py | 76 +++++++++ .../_hooks/custom/split_pdf_hook.py | 55 +++++- src/unstructured_client/general.py | 61 +++++++ .../models/operations/partition.py | 9 + 5 files changed, 359 insertions(+), 1 deletion(-) create mode 100644 _test_unstructured_client/unit/test_ndjson_combine.py diff --git a/_test_unstructured_client/unit/test_ndjson_combine.py b/_test_unstructured_client/unit/test_ndjson_combine.py new file mode 100644 index 00000000..64d40caa --- /dev/null +++ b/_test_unstructured_client/unit/test_ndjson_combine.py @@ -0,0 +1,159 @@ +"""Unit tests for the on-disk NDJSON recombination used by elements-file mode. + +This helper is what removes the four in-memory copies the split-PDF recombination used to +make (per-chunk list, flattened list, json.dumps blob, SDK re-parse). It must: + - handle chunk files that are JSON arrays (server returned application/json) + - handle chunk files that are already NDJSON (server honored application/x-ndjson) + - preserve element order across chunks + - round-trip payload strings byte-for-byte +""" + +import json + +import pytest + +from unstructured_client._hooks.custom.request_utils import ( + ELEMENTS_FILE_HEADER, + combine_chunk_files_to_ndjson, + create_elements_file_response, +) + + +def _elements(prefix, count): + return [ + { + "type": "Table" if i % 2 == 0 else "NarrativeText", + "text": f"{prefix}-{i}", + "metadata": {"page_number": i + 1, "image_base64": f"PAYLOAD{prefix}{i}" * 4}, + } + for i in range(count) + ] + + +def _write_json_array(path, elements): + path.write_text(json.dumps(elements), encoding="utf-8") + + +def _write_ndjson(path, elements): + with path.open("w", encoding="utf-8") as f: + for element in elements: + f.write(json.dumps(element)) + f.write("\n") + + +def _read_ndjson(path): + with open(path, encoding="utf-8") as f: + return [json.loads(line) for line in f if line.strip()] + + +def test_combines_json_array_chunks_in_order(tmp_path): + chunk_a = tmp_path / "a.json" + chunk_b = tmp_path / "b.json" + elements_a = _elements("a", 3) + elements_b = _elements("b", 2) + _write_json_array(chunk_a, elements_a) + _write_json_array(chunk_b, elements_b) + + out = tmp_path / "combined.ndjson" + written = combine_chunk_files_to_ndjson([str(chunk_a), str(chunk_b)], str(out)) + + assert written == 5 + assert _read_ndjson(out) == elements_a + elements_b + + +def test_combines_ndjson_chunks_without_parsing(tmp_path): + """The zero-parse path: both ends speak NDJSON.""" + chunk_a = tmp_path / "a.ndjson" + chunk_b = tmp_path / "b.ndjson" + elements_a = _elements("a", 4) + elements_b = _elements("b", 1) + _write_ndjson(chunk_a, elements_a) + _write_ndjson(chunk_b, elements_b) + + out = tmp_path / "combined.ndjson" + written = combine_chunk_files_to_ndjson([str(chunk_a), str(chunk_b)], str(out)) + + assert written == 5 + assert _read_ndjson(out) == elements_a + elements_b + + +def test_mixed_chunk_formats(tmp_path): + """A server upgraded mid-flight, or a retry served by an older pod.""" + chunk_a = tmp_path / "a.json" + chunk_b = tmp_path / "b.ndjson" + elements_a = _elements("a", 2) + elements_b = _elements("b", 2) + _write_json_array(chunk_a, elements_a) + _write_ndjson(chunk_b, elements_b) + + out = tmp_path / "combined.ndjson" + written = combine_chunk_files_to_ndjson([str(chunk_a), str(chunk_b)], str(out)) + + assert written == 4 + assert _read_ndjson(out) == elements_a + elements_b + + +@pytest.mark.parametrize("body", ["", " ", "\n\n"]) +def test_empty_chunk_files_are_skipped(tmp_path, body): + chunk_a = tmp_path / "a.json" + chunk_empty = tmp_path / "empty.json" + elements_a = _elements("a", 2) + _write_json_array(chunk_a, elements_a) + chunk_empty.write_text(body, encoding="utf-8") + + out = tmp_path / "combined.ndjson" + written = combine_chunk_files_to_ndjson([str(chunk_a), str(chunk_empty)], str(out)) + + assert written == 2 + assert _read_ndjson(out) == elements_a + + +def test_empty_array_chunk_contributes_nothing(tmp_path): + """A chunk that legitimately produced no elements (e.g. blank pages).""" + chunk_a = tmp_path / "a.json" + chunk_b = tmp_path / "b.json" + _write_json_array(chunk_a, []) + elements_b = _elements("b", 3) + _write_json_array(chunk_b, elements_b) + + out = tmp_path / "combined.ndjson" + written = combine_chunk_files_to_ndjson([str(chunk_a), str(chunk_b)], str(out)) + + assert written == 3 + assert _read_ndjson(out) == elements_b + + +def test_payload_round_trips_exactly(tmp_path): + """Base64 payloads are the reason this path exists; they must survive unchanged.""" + payload = "A" * 100_000 + element = {"type": "Table", "text": "t", "metadata": {"image_base64": payload}} + chunk = tmp_path / "a.json" + _write_json_array(chunk, [element]) + + out = tmp_path / "combined.ndjson" + combine_chunk_files_to_ndjson([str(chunk)], str(out)) + + result = _read_ndjson(out) + assert result[0]["metadata"]["image_base64"] == payload + + +def test_non_ascii_is_preserved(tmp_path): + element = {"type": "NarrativeText", "text": "日本語 café ✓", "metadata": {}} + chunk = tmp_path / "a.json" + _write_json_array(chunk, [element]) + + out = tmp_path / "combined.ndjson" + combine_chunk_files_to_ndjson([str(chunk)], str(out)) + + assert _read_ndjson(out)[0]["text"] == "日本語 café ✓" + + +def test_elements_file_response_carries_path_in_header_and_body(tmp_path): + path = str(tmp_path / "combined.ndjson") + response = create_elements_file_response(path) + + assert response.status_code == 200 + assert response.headers[ELEMENTS_FILE_HEADER] == path + assert response.headers["Content-Type"] == "application/x-ndjson" + # Body-as-path mirrors the existing cached-chunk convention in the split hook. + assert response.text == path diff --git a/src/unstructured_client/_hooks/custom/request_utils.py b/src/unstructured_client/_hooks/custom/request_utils.py index bfc9cb0f..993b2bcf 100644 --- a/src/unstructured_client/_hooks/custom/request_utils.py +++ b/src/unstructured_client/_hooks/custom/request_utils.py @@ -277,6 +277,82 @@ def create_response(elements: list) -> httpx.Response: return response +ELEMENTS_FILE_HEADER = "x-unstructured-elements-file" +NDJSON_MEDIA_TYPE = "application/x-ndjson" + + +def combine_chunk_files_to_ndjson(chunk_paths: list[str], out_path: str) -> int: + """Combine per-chunk split-PDF response files into one NDJSON file on disk. + + This is the whole point of the elements-file path: the split-PDF hook already streams + each chunk response to a temp file without reading it into memory + (`call_api_async`), but recombining them used to build a list per chunk, a flattened + list, a `json.dumps` blob, and then the SDK re-parsed that blob -- four full copies of + the document. Concatenating on disk keeps peak memory at ~one chunk. + + Each chunk file is either a JSON array (a server that returns + `application/json`) or already NDJSON (a server that honors + `application/x-ndjson`). The first non-whitespace byte tells us which, so the + fast path -- byte-level concatenation with no parsing at all -- is taken whenever both + ends speak NDJSON, and the JSON-array path still works against an older server. + + Returns the number of elements written. + """ + total = 0 + with open(out_path, "w", encoding="utf-8") as out: + for chunk_path in chunk_paths: + with open(chunk_path, "r", encoding="utf-8") as chunk: + first_char = "" + while True: + char = chunk.read(1) + if char == "": + break + if not char.isspace(): + first_char = char + break + if first_char == "": + # Empty chunk file; nothing to append. + continue + chunk.seek(0) + + if first_char == "[": + # JSON array: bounded by one chunk (20 pages by default), so a plain + # load is fine and avoids taking on a streaming-parser dependency. + for element in json.load(chunk): + out.write(json.dumps(element, ensure_ascii=False)) + out.write("\n") + total += 1 + else: + # Already NDJSON: copy through, no parsing. + for line in chunk: + line = line.strip() + if line: + out.write(line) + out.write("\n") + total += 1 + return total + + +def create_elements_file_response(elements_file: str) -> httpx.Response: + """Synthetic 200 whose payload is a path to an NDJSON file of elements. + + Mirrors the existing convention in the split hook, where a cached chunk response + carries its temp-file path as the body. The path is also exposed as a header so the + SDK can distinguish this from a real NDJSON body streamed from the server. + """ + content = elements_file.encode() + response = httpx.Response( + status_code=200, + headers={ + "Content-Type": NDJSON_MEDIA_TYPE, + "Content-Length": str(len(content)), + ELEMENTS_FILE_HEADER: elements_file, + }, + ) + setattr(response, "_content", content) + return response + + def get_base_url(url: str | URL) -> str: """Extracts the base URL from the given URL. diff --git a/src/unstructured_client/_hooks/custom/split_pdf_hook.py b/src/unstructured_client/_hooks/custom/split_pdf_hook.py index 912da332..c14e9760 100644 --- a/src/unstructured_client/_hooks/custom/split_pdf_hook.py +++ b/src/unstructured_client/_hooks/custom/split_pdf_hook.py @@ -524,6 +524,11 @@ def __init__(self) -> None: self.allow_failed: dict[str, bool] = {} self.cache_tmp_data_feature: dict[str, bool] = {} self.cache_tmp_data_dir: dict[str, str] = {} + # NDJSON elements-file mode: when the caller asked for application/x-ndjson, the + # per-chunk temp files are concatenated on disk instead of being parsed and + # re-serialized, and the combined path is handed back via the response header. + self.ndjson_mode: dict[str, bool] = {} + self.ndjson_output_path: dict[str, str] = {} @staticmethod def _get_operation_id_from_request(request: Optional[httpx.Request]) -> Optional[str]: @@ -791,6 +796,12 @@ def _before_request_unlocked( self.cache_tmp_data_feature[operation_id] = cache_tmp_data_feature self.cache_tmp_data_dir[operation_id] = cache_tmp_data_dir self.concurrency_level[operation_id] = concurrency_level + # Elements-file mode requires the per-chunk temp files, so it is only available + # when chunk caching is on; otherwise fall back to the in-memory recombination. + self.ndjson_mode[operation_id] = ( + request_utils.NDJSON_MEDIA_TYPE in request.headers.get("Accept", "") + and cache_tmp_data_feature + ) timeout_seconds = _get_request_timeout_seconds(request) if timeout_seconds is None and hook_ctx.config.timeout_ms is not None: @@ -1382,6 +1393,8 @@ def _elements_from_task_responses( failed_responses: list[tuple[int, httpx.Response]] = [] transport_failure_count = 0 elements = [] + ndjson_mode = self.ndjson_mode.get(operation_id, False) + chunk_paths: list[str] = [] for response_number, res in task_responses: if res.status_code == 200: logger.debug( @@ -1390,7 +1403,11 @@ def _elements_from_task_responses( response_number, ) successful_responses.append(res) - if self.cache_tmp_data_feature.get(operation_id, DEFAULT_CACHE_TMP_DATA): + if ndjson_mode: + # Record the cached temp-file path only. Nothing is parsed here -- + # that is what keeps peak memory at ~one chunk during recombination. + chunk_paths.append(res.text) + elif self.cache_tmp_data_feature.get(operation_id, DEFAULT_CACHE_TMP_DATA): elements.append(load_elements_from_response(res)) else: elements.append(res.json()) @@ -1433,6 +1450,27 @@ def _elements_from_task_responses( total_chunks=len(task_responses), response=response, ) + if ndjson_mode: + # Combine on disk and stash the path; `_build_after_success_response` turns it + # into the response. Deliberately written to the cache *parent* dir, not the + # operation's TemporaryDirectory: that directory is cleaned up by + # `_clear_operation` -> `_finalize_operation_resources` immediately after + # after_success returns, which would delete this file before the caller could + # read it. The combined file therefore outlives the operation and the CALLER + # owns deleting it (see `PartitionResponse.elements_file`). + temp_dir_path = self.cache_tmp_data_dir.get(operation_id) or tempfile.gettempdir() + out_path = f"{temp_dir_path}/{uuid.uuid4()}.ndjson" + written = request_utils.combine_chunk_files_to_ndjson(chunk_paths, out_path) + self.ndjson_output_path[operation_id] = out_path + logger.info( + "split_pdf event=ndjson_combined operation_id=%s chunk_count=%d element_count=%d out_file=%s", + operation_id, + len(chunk_paths), + written, + Path(out_path).name, + ) + return [] + flattened_elements = [element for sublist in elements for element in sublist] return flattened_elements @@ -1464,6 +1502,19 @@ def _build_after_success_response( ) return self.api_failed_responses[operation_id][0] + # Elements-file mode: hand back the combined NDJSON path instead of a body. Checked + # before the `elements is None` guard because `elements` is intentionally empty + # here -- nothing was parsed. + if self.ndjson_mode.get(operation_id, False): + ndjson_path = self.ndjson_output_path.get(operation_id) + if ndjson_path is None: + logger.warning( + "split_pdf event=ndjson_missing_output operation_id=%s falling_back=true", + operation_id, + ) + return response + return request_utils.create_elements_file_response(ndjson_path) + if elements is None: return response @@ -1586,6 +1637,8 @@ def _clear_operation(self, operation_id: str) -> None: self.allow_failed.pop(operation_id, None) self.cache_tmp_data_feature.pop(operation_id, None) self.cache_tmp_data_dir.pop(operation_id, None) + self.ndjson_mode.pop(operation_id, None) + self.ndjson_output_path.pop(operation_id, None) self.pending_operation_ids.pop(operation_id, None) future = self.operation_futures.pop(operation_id, None) loop_holder = self.operation_loops.pop(operation_id, None) diff --git a/src/unstructured_client/general.py b/src/unstructured_client/general.py index a6ec3099..d8cc4b1a 100644 --- a/src/unstructured_client/general.py +++ b/src/unstructured_client/general.py @@ -14,6 +14,53 @@ class PartitionAcceptEnum(str, Enum): APPLICATION_JSON = "application/json" TEXT_CSV = "text/csv" + APPLICATION_X_NDJSON = "application/x-ndjson" + r"""Elements as NDJSON, one per line. The response is written to a temp file and + returned as `PartitionResponse.elements_file` rather than parsed into `elements`, so + the caller can stream a large document instead of holding it in memory.""" + + +def _ndjson_elements_file(http_res) -> str: + """Resolve an NDJSON response to a path on disk, without parsing the elements. + + Two shapes arrive here. When the split-PDF hook ran, it has already combined the + per-chunk temp files into one NDJSON file and passes the path through + `ELEMENTS_FILE_HEADER` -- nothing to do but read the header, so the document is never + materialized. Otherwise this is a real body from the server, which is written to a + temp file in bounded chunks. + """ + import tempfile + + from unstructured_client._hooks.custom.request_utils import ELEMENTS_FILE_HEADER + + existing_path = http_res.headers.get(ELEMENTS_FILE_HEADER) + if existing_path: + return existing_path + + with tempfile.NamedTemporaryFile( + mode="wb", prefix="unst_elements_", suffix=".ndjson", delete=False + ) as out: + for byte_chunk in http_res.iter_bytes(): + out.write(byte_chunk) + return out.name + + +async def _ndjson_elements_file_async(http_res) -> str: + """Async counterpart of `_ndjson_elements_file`.""" + import tempfile + + from unstructured_client._hooks.custom.request_utils import ELEMENTS_FILE_HEADER + + existing_path = http_res.headers.get(ELEMENTS_FILE_HEADER) + if existing_path: + return existing_path + + with tempfile.NamedTemporaryFile( + mode="wb", prefix="unst_elements_", suffix=".ndjson", delete=False + ) as out: + async for byte_chunk in http_res.aiter_bytes(): + out.write(byte_chunk) + return out.name class General(BaseSDK): @@ -128,6 +175,13 @@ def partition( content_type=http_res.headers.get("Content-Type") or "", raw_response=http_res, ) + if utils.match_response(http_res, "200", "application/x-ndjson"): + return operations.PartitionResponse( + elements_file=_ndjson_elements_file(http_res), + status_code=http_res.status_code, + content_type=http_res.headers.get("Content-Type") or "", + raw_response=http_res, + ) if utils.match_response(http_res, "422", "application/json"): response_data = unmarshal_json_response( errors.HTTPValidationErrorData, http_res @@ -253,6 +307,13 @@ async def partition_async( content_type=http_res.headers.get("Content-Type") or "", raw_response=http_res, ) + if utils.match_response(http_res, "200", "application/x-ndjson"): + return operations.PartitionResponse( + elements_file=await _ndjson_elements_file_async(http_res), + status_code=http_res.status_code, + content_type=http_res.headers.get("Content-Type") or "", + raw_response=http_res, + ) if utils.match_response(http_res, "422", "application/json"): response_data = unmarshal_json_response( errors.HTTPValidationErrorData, http_res diff --git a/src/unstructured_client/models/operations/partition.py b/src/unstructured_client/models/operations/partition.py index de4664fb..e447f641 100644 --- a/src/unstructured_client/models/operations/partition.py +++ b/src/unstructured_client/models/operations/partition.py @@ -78,6 +78,10 @@ class PartitionResponseTypedDict(TypedDict): r"""Successful Response""" elements: NotRequired[List[Dict[str, Any]]] r"""Successful Response""" + elements_file: NotRequired[str] + r"""Path to an NDJSON file of elements, one per line. Set instead of `elements` when + `application/x-ndjson` was requested, so a large document never has to be held in + memory as a parsed list.""" class PartitionResponse(BaseModel): @@ -95,3 +99,8 @@ class PartitionResponse(BaseModel): elements: Optional[List[Dict[str, Any]]] = None r"""Successful Response""" + + elements_file: Optional[str] = None + r"""Path to an NDJSON file of elements, one per line. Set instead of `elements` when + `application/x-ndjson` was requested, so a large document never has to be held in + memory as a parsed list. The caller owns the file and should delete it when done.""" From 214e82f6a2806d84b4343adb6984ab34305fd475 Mon Sep 17 00:00:00 2001 From: Yao You Date: Wed, 29 Jul 2026 17:59:30 -0500 Subject: [PATCH 2/3] fix: decouple ndjson elements-file mode from cache_tmp_data ndjson_mode was gated on BOTH the Accept header AND cache_tmp_data. Those are set by different parties -- the caller picks the Accept header, cache_tmp_data is a separate split-PDF setting -- so they can disagree. When they did (NDJSON requested, caching off), the server returned NDJSON while the hook took the JSON path and res.json() raised on a body this client had itself asked for. Observed on an SnD whose effective split policy reports cache_mode=disabled: the partition node failed outright rather than degrading. ndjson_mode now depends only on the Accept header, and both caching modes are handled: cached chunks contribute their existing temp-file path, uncached chunks spill their body verbatim via write_chunk_body_to_temp (no parsing) so the combine step gets uniform input. Peak stays ~one chunk either way. Tests: 2 new regression tests for the uncached path; suite 226 passed, 1 xfailed. Co-Authored-By: Claude Opus 5 (1M context) --- .../unit/test_ndjson_combine.py | 32 +++++++++++++++++++ .../_hooks/custom/request_utils.py | 15 +++++++++ .../_hooks/custom/split_pdf_hook.py | 26 +++++++++++---- 3 files changed, 67 insertions(+), 6 deletions(-) diff --git a/_test_unstructured_client/unit/test_ndjson_combine.py b/_test_unstructured_client/unit/test_ndjson_combine.py index 64d40caa..7b4ce9ed 100644 --- a/_test_unstructured_client/unit/test_ndjson_combine.py +++ b/_test_unstructured_client/unit/test_ndjson_combine.py @@ -12,10 +12,13 @@ import pytest +import httpx + from unstructured_client._hooks.custom.request_utils import ( ELEMENTS_FILE_HEADER, combine_chunk_files_to_ndjson, create_elements_file_response, + write_chunk_body_to_temp, ) @@ -148,6 +151,35 @@ def test_non_ascii_is_preserved(tmp_path): assert _read_ndjson(out)[0]["text"] == "日本語 café ✓" +def test_write_chunk_body_to_temp_roundtrips(tmp_path): + """The cache_tmp_data=OFF path: an in-memory NDJSON body must spill verbatim. + + Regression guard. `ndjson_mode` used to also require cache_tmp_data, so with caching off + the server returned NDJSON while the hook took the JSON path and `res.json()` raised on a + body this client had itself requested. + """ + elements = _elements("x", 3) + body = "".join(json.dumps(e) + "\n" for e in elements).encode() + response = httpx.Response(status_code=200, content=body) + + path = write_chunk_body_to_temp(response, str(tmp_path)) + assert _read_ndjson(path) == elements + + +def test_combine_accepts_bodies_spilled_without_caching(tmp_path): + """End-to-end of the uncached path: spill two bodies, then combine them.""" + a, b = _elements("a", 2), _elements("b", 3) + ra = httpx.Response(200, content="".join(json.dumps(e) + "\n" for e in a).encode()) + rb = httpx.Response(200, content="".join(json.dumps(e) + "\n" for e in b).encode()) + paths = [write_chunk_body_to_temp(r, str(tmp_path)) for r in (ra, rb)] + + out = tmp_path / "combined.ndjson" + written = combine_chunk_files_to_ndjson(paths, str(out)) + + assert written == 5 + assert _read_ndjson(out) == a + b + + def test_elements_file_response_carries_path_in_header_and_body(tmp_path): path = str(tmp_path / "combined.ndjson") response = create_elements_file_response(path) diff --git a/src/unstructured_client/_hooks/custom/request_utils.py b/src/unstructured_client/_hooks/custom/request_utils.py index 993b2bcf..1fb24e79 100644 --- a/src/unstructured_client/_hooks/custom/request_utils.py +++ b/src/unstructured_client/_hooks/custom/request_utils.py @@ -4,6 +4,8 @@ import io import json import logging +import os +import tempfile from typing import Tuple, Any, BinaryIO, Optional from urllib.parse import urlparse @@ -333,6 +335,19 @@ def combine_chunk_files_to_ndjson(chunk_paths: list[str], out_path: str) -> int: return total +def write_chunk_body_to_temp(response: httpx.Response, dir_: Optional[str] = None) -> str: + """Spill an in-memory chunk response body to a temp file, returning its path. + + Needed for elements-file mode when `cache_tmp_data` is OFF: there is no cached file to + reference, so the body is written out verbatim (no parsing) to give + `combine_chunk_files_to_ndjson` the same uniform input it gets in the cached case. + """ + fd, path = tempfile.mkstemp(suffix=".ndjson", dir=dir_ or tempfile.gettempdir()) + with os.fdopen(fd, "wb") as f: + f.write(response.content) + return path + + def create_elements_file_response(elements_file: str) -> httpx.Response: """Synthetic 200 whose payload is a path to an NDJSON file of elements. diff --git a/src/unstructured_client/_hooks/custom/split_pdf_hook.py b/src/unstructured_client/_hooks/custom/split_pdf_hook.py index c14e9760..f77b7fdc 100644 --- a/src/unstructured_client/_hooks/custom/split_pdf_hook.py +++ b/src/unstructured_client/_hooks/custom/split_pdf_hook.py @@ -796,11 +796,13 @@ def _before_request_unlocked( self.cache_tmp_data_feature[operation_id] = cache_tmp_data_feature self.cache_tmp_data_dir[operation_id] = cache_tmp_data_dir self.concurrency_level[operation_id] = concurrency_level - # Elements-file mode requires the per-chunk temp files, so it is only available - # when chunk caching is on; otherwise fall back to the in-memory recombination. + # Depends ONLY on what the caller asked for. It must NOT also depend on + # cache_tmp_data: the Accept header is chosen by the caller while cache_tmp_data is + # a separate setting, so gating on both lets them disagree -- and when they do, the + # server returns NDJSON while the hook takes the JSON path and `res.json()` blows up + # on a body this very client requested. Both caching modes are handled below. self.ndjson_mode[operation_id] = ( request_utils.NDJSON_MEDIA_TYPE in request.headers.get("Accept", "") - and cache_tmp_data_feature ) timeout_seconds = _get_request_timeout_seconds(request) @@ -1404,9 +1406,21 @@ def _elements_from_task_responses( ) successful_responses.append(res) if ndjson_mode: - # Record the cached temp-file path only. Nothing is parsed here -- - # that is what keeps peak memory at ~one chunk during recombination. - chunk_paths.append(res.text) + # Nothing is parsed here either way -- that is what keeps peak memory at + # ~one chunk during recombination. + if self.cache_tmp_data_feature.get(operation_id, DEFAULT_CACHE_TMP_DATA): + # Cached: the body was already streamed to a temp file and `res.text` + # holds that path. + chunk_paths.append(res.text) + else: + # Not cached: the body is in memory. Spill it verbatim so the combine + # step is uniform; peak stays at ~one chunk body, the same order as + # the JSON path's `res.json()` would have cost. + chunk_paths.append( + request_utils.write_chunk_body_to_temp( + res, self.cache_tmp_data_dir.get(operation_id) + ) + ) elif self.cache_tmp_data_feature.get(operation_id, DEFAULT_CACHE_TMP_DATA): elements.append(load_elements_from_response(res)) else: From 74fc2306e0ab074650346d75d3cbed208daec8fa Mon Sep 17 00:00:00 2001 From: Yao You Date: Thu, 30 Jul 2026 14:52:36 -0500 Subject: [PATCH 3/3] fix: release spilled chunk bodies in ndjson elements-file mode Spilling each chunk body to disk did not reduce memory, because every response object is retained in api_successful_responses for failure bookkeeping -- so the in-memory body stayed alive alongside the file. On a 2500-page document split into 125 chunks of 32 MB, the whole 4 GB accumulated anyway and the partitioner plugin peaked at ~15 GiB against a 16 GiB limit. Overwrite _content with the spill path after writing, mirroring what the cached branch already does at the point of caching, so res.text means the same thing in both branches and the payload is reclaimed immediately. Measured on an SnD with cache_mode=disabled, 2500 pages / 125 chunks / 4 GB: buffered -> 15252 MiB, job FAILED streaming -> 15814 MiB, job succeeded (before this fix) Co-Authored-By: Claude Opus 5 (1M context) --- .../unit/test_ndjson_combine.py | 22 +++++++++++++++++++ .../_hooks/custom/split_pdf_hook.py | 19 ++++++++++------ 2 files changed, 34 insertions(+), 7 deletions(-) diff --git a/_test_unstructured_client/unit/test_ndjson_combine.py b/_test_unstructured_client/unit/test_ndjson_combine.py index 7b4ce9ed..3f2b62f0 100644 --- a/_test_unstructured_client/unit/test_ndjson_combine.py +++ b/_test_unstructured_client/unit/test_ndjson_combine.py @@ -180,6 +180,28 @@ def test_combine_accepts_bodies_spilled_without_caching(tmp_path): assert _read_ndjson(out) == a + b +def test_spilled_body_is_released_from_the_response(tmp_path): + """After spilling, the response must no longer hold the body. + + Regression guard for the real failure mode: every chunk response is retained in + `api_successful_responses` for failure bookkeeping, so spilling to disk without + releasing `_content` still accumulates the whole document in memory (125 chunks x + 32 MB = 4 GB observed on a 2500-page split). + """ + elements = _elements("a", 3) + body = "".join(json.dumps(e) + "\n" for e in elements).encode() + response = httpx.Response(status_code=200, content=body) + assert len(response.content) == len(body) + + path = write_chunk_body_to_temp(response, str(tmp_path)) + response._content = path.encode() + + # The body is on disk, and the response now costs a path rather than a payload. + assert _read_ndjson(path) == elements + assert response.text == path + assert len(response.content) < 512 + + def test_elements_file_response_carries_path_in_header_and_body(tmp_path): path = str(tmp_path / "combined.ndjson") response = create_elements_file_response(path) diff --git a/src/unstructured_client/_hooks/custom/split_pdf_hook.py b/src/unstructured_client/_hooks/custom/split_pdf_hook.py index f77b7fdc..df600b63 100644 --- a/src/unstructured_client/_hooks/custom/split_pdf_hook.py +++ b/src/unstructured_client/_hooks/custom/split_pdf_hook.py @@ -1413,14 +1413,19 @@ def _elements_from_task_responses( # holds that path. chunk_paths.append(res.text) else: - # Not cached: the body is in memory. Spill it verbatim so the combine - # step is uniform; peak stays at ~one chunk body, the same order as - # the JSON path's `res.json()` would have cost. - chunk_paths.append( - request_utils.write_chunk_body_to_temp( - res, self.cache_tmp_data_dir.get(operation_id) - ) + # Not cached: the body is in memory. Spill it verbatim (no parsing) + # and then RELEASE it -- every response object is retained in + # `successful_responses` for failure bookkeeping, so without this the + # spilled bodies stay resident and the whole document accumulates + # anyway (125 chunks x 32 MB = 4 GB observed). Overwriting `_content` + # with the path mirrors what the cached branch already does at the + # point of caching, so downstream `res.text` means the same thing in + # both branches. + spilled = request_utils.write_chunk_body_to_temp( + res, self.cache_tmp_data_dir.get(operation_id) ) + res._content = spilled.encode() # pylint: disable=protected-access + chunk_paths.append(spilled) elif self.cache_tmp_data_feature.get(operation_id, DEFAULT_CACHE_TMP_DATA): elements.append(load_elements_from_response(res)) else: