diff --git a/CHANGELOG.md b/CHANGELOG.md index e45c1498..d6e9292b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,9 @@ +## 0.1.11 + +### Fixes + +- **Bound gzip decompression memory**: gzip uploads are decompressed incrementally into a spooled temporary file instead of being materialized as a complete in-memory `bytes` value. Expanded outputs larger than 1 MiB remain disk-backed through MIME detection and partitioning, avoiding downstream document-sized copies while preserving the partition API. + ## 0.1.10 ### Fixes diff --git a/prepline_general/api/__version__.py b/prepline_general/api/__version__.py index f15b0cd2..346129b4 100644 --- a/prepline_general/api/__version__.py +++ b/prepline_general/api/__version__.py @@ -1 +1 @@ -__version__ = "0.1.10" # pragma: no cover +__version__ = "0.1.11" # pragma: no cover diff --git a/prepline_general/api/general.py b/prepline_general/api/general.py index df1b67ec..6ec94f65 100644 --- a/prepline_general/api/general.py +++ b/prepline_general/api/general.py @@ -7,10 +7,12 @@ import mimetypes import os import secrets +import shutil +import tempfile from base64 import b64encode from concurrent.futures import ThreadPoolExecutor from functools import partial -from typing import IO, Any, Dict, List, Mapping, Optional, Sequence, Tuple, Union, cast +from typing import IO, Any, BinaryIO, Dict, List, Mapping, Optional, Sequence, Tuple, Union, cast import backoff import pandas as pd @@ -31,6 +33,7 @@ from starlette.datastructures import Headers from starlette.types import Send +from prepline_general.api import __version__ as api_version from prepline_general.api.filetypes import get_validated_mimetype from prepline_general.api.models.form_params import GeneralFormParams from unstructured.documents.elements import Element @@ -41,11 +44,37 @@ elements_from_json, ) from unstructured_inference.models.base import UnknownModelException -from prepline_general.api import __version__ as api_version app = FastAPI() router = APIRouter() +_GZIP_COPY_CHUNK_SIZE = 1024 * 1024 +_GZIP_SPOOL_MAX_MEMORY_BYTES = 1024 * 1024 + + +class _SpooledFileProxy: + """Expose a spooled file without its concrete type triggering downstream copies.""" + + def __init__(self, file: IO[bytes], name: str | None): + self._file = file + self.name = name + + def __getattr__(self, name: str) -> Any: + return getattr(self._file, name) + + def __enter__(self) -> _SpooledFileProxy: + self._file.__enter__() + return self + + def __exit__(self, *args: Any) -> Any: + return self._file.__exit__(*args) + + def __iter__(self): + return iter(self._file) + + def __next__(self) -> bytes: + return next(self._file) + def is_compatible_response_type(media_type: str, response_type: type) -> bool: """True when `response_type` can be converted to `media_type` for HTTP Response.""" @@ -613,13 +642,30 @@ def return_content_type(filename: str): if filename.endswith(".gz"): filename = filename[:-3] - gzip_file = gzip.open(file.file).read() - return UploadFile( - file=io.BytesIO(gzip_file), - size=len(gzip_file), - filename=filename, - headers=Headers({"content-type": return_content_type(filename)}), + output_file = tempfile.SpooledTemporaryFile(max_size=_GZIP_SPOOL_MAX_MEMORY_BYTES) + try: + with gzip.open(file.file) as gzip_file: + shutil.copyfileobj(gzip_file, output_file, length=_GZIP_COPY_CHUNK_SIZE) + uncompressed_size = output_file.tell() + output_file.seek(0) + except Exception: + output_file.close() + raise + + # Reuse the request-owned UploadFile so FastAPI closes the decompressed spool when the + # request finishes. A newly-created UploadFile would not belong to the parsed FormData and + # would therefore remain open after the response. + file.file.close() + decompressed_file = ( + _SpooledFileProxy(output_file, filename) + if uncompressed_size > _GZIP_SPOOL_MAX_MEMORY_BYTES + else output_file ) + file.file = cast(BinaryIO, decompressed_file) + file.size = uncompressed_size + file.filename = filename + file.headers = Headers({"content-type": return_content_type(filename)}) + return file @router.get("/general/v0/general", include_in_schema=False) @@ -747,7 +793,7 @@ def join_responses( frames = [ pd.read_csv(io.BytesIO(response.body)) # pyright: ignore[reportUnknownMemberType] for response in responses - if response.body.strip() + if cast(bytes, response.body).strip() ] if not frames: return PlainTextResponse(responses[0].body) diff --git a/test_general/api/test_gzip.py b/test_general/api/test_gzip.py index 24940a51..7653cee8 100644 --- a/test_general/api/test_gzip.py +++ b/test_general/api/test_gzip.py @@ -9,13 +9,58 @@ import pandas as pd import pytest from deepdiff import DeepDiff +from fastapi import UploadFile from fastapi.testclient import TestClient from prepline_general.api.app import app +from prepline_general.api.general import _GZIP_SPOOL_MAX_MEMORY_BYTES, ungz_file MAIN_API_ROUTE = "general/v0/general" +def _gzip_upload(content: bytes, filename: str = "sample.txt.gz") -> UploadFile: + compressed = io.BytesIO(gzip.compress(content)) + return UploadFile(file=compressed, filename=filename) + + +@pytest.mark.parametrize( + ("content", "spills_to_disk"), + [ + pytest.param(b"small gzip payload", False, id="small-stays-in-memory"), + pytest.param(b"x" * (_GZIP_SPOOL_MAX_MEMORY_BYTES + 1), True, id="large-spills-to-disk"), + ], +) +def test_ungz_file_bounds_decompression_memory(content: bytes, spills_to_disk: bool): + upload = _gzip_upload(content) + compressed_file = upload.file + + result = ungz_file(upload) + + try: + assert result is upload + assert compressed_file.closed + assert result.filename == "sample.txt" + assert result.content_type == "text/plain" + assert result.size == len(content) + assert result.file.read() == content + assert result.file._rolled is spills_to_disk + assert isinstance(result.file, tempfile.SpooledTemporaryFile) is not spills_to_disk + finally: + result.file.close() + assert result.file.closed + + +def test_ungz_file_proxy_preserves_special_file_methods(): + result = ungz_file(_gzip_upload(b"first\nsecond\n" * _GZIP_SPOOL_MAX_MEMORY_BYTES)) + + try: + assert iter(result.file) is not result.file + assert next(result.file) == b"first\n" + assert result.file.__enter__() is result.file + finally: + result.file.close() + + @pytest.mark.xfail(reason="The outputs are different as of unstructured==0.13.5") @pytest.mark.parametrize("output_format", ["application/json", "text/csv"]) @pytest.mark.parametrize(