feat: Only allow one Chipper call at a time (#296)

Chipper V2 is very memory hungry. While we work to optimize this, we
need to restrict the server to one call at a time. While the model is in
use, we'll return a 503 "Please try again". Our hosted API should scale
up to meet demand, so the next call should route to an available server.

This includes a refactor to how partition_kwargs are passed to either
parallel mode, local partition, or local partition with the new Chipper
protection.

To verify, try calling Chipper twice:
```
curl -X POST 'http://localhost:8000/general/v0/general' --form files="@$file" --form strategy=hi_res --form hi_res_model_name=chipper &
curl -X POST 'http://localhost:8000/general/v0/general' --form files="@$file" --form strategy=hi_res --form hi_res_model_name=chipper
```

The second call will get a 503 response.

Other changes:
* Return a 400 error if Chipper isn't loaded. The model is private, make
sure we explain this for users who self host
* Pass the huggingface token to `make docker-start-api` for better dev
experience
* Add a `make docker-start-bash` while we're in here

Author: awalker4

CHANGELOG.md

@@ -1,7 +1,8 @@
-## 0.0.55-dev1
+## 0.0.55
 
-* Bump unstructured to 0.0.26
+* Bump unstructured to 0.10.26
 * Bring parent_id metadata field back after fixing a backwards compatibility bug
+* Restrict Chipper usage to one at a time. The model is very resource intense, and this will prevent issues while we improve it.
 
 ## 0.0.54
 

Makefile

@@ -66,7 +66,15 @@ docker-build:
 
 .PHONY: docker-start-api
 docker-start-api:
-	docker run -p 8000:8000 --mount type=bind,source=$(realpath .),target=/home/notebook-user/local -it --rm pipeline-family-${PIPELINE_FAMILY}-dev:latest scripts/app-start.sh
+	docker run -p 8000:8000 \
+	-it --rm  \
+	--mount type=bind,source=$(realpath .),target=/home/notebook-user/local \
+	-e UNSTRUCTURED_HF_TOKEN=${UNSTRUCTURED_HF_TOKEN} \
+	pipeline-family-${PIPELINE_FAMILY}-dev:latest scripts/app-start.sh
+
+.PHONY: docker-start-bash
+docker-start-bash:
+	docker run -p 8000:8000 -it --rm --mount type=bind,source=$(realpath .),target=/home/notebook-user/local --entrypoint /bin/bash pipeline-family-${PIPELINE_FAMILY}-dev:latest
 
 .PHONY: docker-test
 docker-test:

prepline_general/api/general.py

@@ -223,6 +223,32 @@ def partition_pdf_splits(
     return results
 
 
+IS_CHIPPER_PROCESSING = False
+
+
+class ChipperMemoryProtection:
+    """
+    Chipper calls are expensive, and right now we can only do one call at a time.
+    If the model is in use, return a 503 error. The API should scale up and the user can try again
+    on a different server.
+    """
+
+    def __enter__(self):
+        global IS_CHIPPER_PROCESSING
+        if IS_CHIPPER_PROCESSING:
+            # Log here so we can track how often it happens
+            logger.error("Chipper is already is use")
+            raise HTTPException(
+                status_code=503, detail="Server is under heavy load. Please try again later."
+            )
+
+        IS_CHIPPER_PROCESSING = True
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        global IS_CHIPPER_PROCESSING
+        IS_CHIPPER_PROCESSING = False
+
+
 def pipeline_api(
     file,
     request=None,
@@ -403,52 +429,52 @@ def pipeline_api(
             )
         )
 
-        # Be careful of naming differences in api params vs partition params!
-        # These kwargs are going back into the api, not into partition
-        # If there's a difference, remap the param in partition_pdf_splits
+        partition_kwargs = {
+            "file": file,
+            "metadata_filename": filename,
+            "content_type": file_content_type,
+            "encoding": encoding,
+            "include_page_breaks": include_page_breaks,
+            "model_name": hi_res_model_name,
+            "ocr_languages": ocr_languages,
+            "pdf_infer_table_structure": pdf_infer_table_structure,
+            "skip_infer_table_types": skip_infer_table_types,
+            "strategy": strategy,
+            "xml_keep_tags": xml_keep_tags,
+            "languages": languages,
+            "chunking_strategy": chunking_strategy,
+            "multipage_sections": multipage_sections,
+            "combine_under_n_chars": combine_under_n_chars,
+            "new_after_n_chars": new_after_n_chars,
+        }
+
         if file_content_type == "application/pdf" and pdf_parallel_mode_enabled:
+            # Be careful of naming differences in api params vs partition params!
+            # These kwargs are going back into the api, not into partition
+            # They need to be switched back in partition_pdf_splits
+            if partition_kwargs.get("model_name"):
+                partition_kwargs["hi_res_model_name"] = partition_kwargs.pop("model_name")
+
             elements = partition_pdf_splits(
-                request,
+                request=request,
                 pdf_pages=pdf.pages,
-                file=file,
-                metadata_filename=filename,
-                content_type=file_content_type,
                 coordinates=show_coordinates,
-                # partition_kwargs
-                encoding=encoding,
-                include_page_breaks=include_page_breaks,
-                hi_res_model_name=hi_res_model_name,
-                ocr_languages=ocr_languages,
-                pdf_infer_table_structure=pdf_infer_table_structure,
-                skip_infer_table_types=skip_infer_table_types,
-                strategy=strategy,
-                xml_keep_tags=xml_keep_tags,
-                languages=languages,
-                chunking_strategy=chunking_strategy,
-                multipage_sections=multipage_sections,
-                combine_under_n_chars=combine_under_n_chars,
-                new_after_n_chars=new_after_n_chars,
+                **partition_kwargs,
             )
+        elif hi_res_model_name and hi_res_model_name in CHIPPER_MODEL_TYPES:
+            with ChipperMemoryProtection():
+                elements = partition(**partition_kwargs)
         else:
-            elements = partition(
-                file=file,
-                metadata_filename=filename,
-                content_type=file_content_type,
-                # partition_kwargs
-                encoding=encoding,
-                include_page_breaks=include_page_breaks,
-                model_name=hi_res_model_name,
-                ocr_languages=ocr_languages,
-                pdf_infer_table_structure=pdf_infer_table_structure,
-                skip_infer_table_types=skip_infer_table_types,
-                strategy=strategy,
-                xml_keep_tags=xml_keep_tags,
-                languages=languages,
-                chunking_strategy=chunking_strategy,
-                multipage_sections=multipage_sections,
-                combine_under_n_chars=combine_under_n_chars,
-                new_after_n_chars=new_after_n_chars,
+            elements = partition(**partition_kwargs)
+
+    except OSError as e:
+        if "chipper-fast-fine-tuning is not a local folder" in e.args[0]:
+            raise HTTPException(
+                status_code=400,
+                detail="The Chipper model is not available for download. It can be accessed via the official hosted API.",
             )
+
+        raise e
     except ValueError as e:
         if "Invalid file" in e.args[0]:
             raise HTTPException(
@@ -459,6 +485,7 @@ def pipeline_api(
                 status_code=400,
                 detail="Json schema does not match the Unstructured schema",
             )
+
         raise e
     except zipfile.BadZipFile:
         raise HTTPException(

scripts/app-start.sh

@@ -1,14 +1,5 @@
 #!/usr/bin/env bash
 
-UNSTRUCTURED_DOWNLOAD_CHIPPER=${UNSTRUCTURED_DOWNLOAD_CHIPPER:-"false"}
-
-if [[ "$(echo "${UNSTRUCTURED_DOWNLOAD_CHIPPER}" | tr '[:upper:]' '[:lower:]')" == "true" ]]; then
-  echo "warming chipper model"
-  # NOTE(crag): in the cloud, this could add a minute to startup time
-  UNSTRUCTURED_HI_RES_SUPPORTED_MODEL=chipper python3.8 -c \
-    "from unstructured.ingest.doc_processor.generalized import initialize; initialize()"
-fi
-
 uvicorn prepline_general.api.app:app \
 	--log-config logger_config.yaml \
         --host 0.0.0.0

test_general/api/test_app.py

@@ -4,7 +4,9 @@
 import io
 import pytest
 import requests
+import time
 import pandas as pd
+from concurrent.futures import ThreadPoolExecutor
 from fastapi.testclient import TestClient
 from fastapi import HTTPException
 from pypdf import PdfWriter, PdfReader
@@ -784,3 +786,39 @@ def test_general_api_returns_400_bad_json(tmpdir):
     )
     assert "Unstructured schema" in response.json().get("detail")
     assert response.status_code == 400
+
+
+def test_chipper_memory_protection(monkeypatch, mocker):
+    """
+    For now, only 1 Chipper call is allowed at a time.
+    Assert that we return a 503 while it's in use.
+    """
+
+    def mock_partition(*args, **kwargs):
+        time.sleep(2)
+        return {}
+
+    monkeypatch.setattr(
+        general,
+        "partition",
+        mock_partition,
+    )
+
+    client = TestClient(app)
+    test_file = Path("sample-docs") / "layout-parser-paper-fast.pdf"
+
+    def make_request(*args):
+        return client.post(
+            MAIN_API_ROUTE,
+            files=[("files", (str(test_file), open(test_file, "rb"), "application/pdf"))],
+            data={"strategy": "hi_res", "hi_res_model_name": "chipper"},
+        )
+
+    with ThreadPoolExecutor() as executor:
+        responses = list(executor.map(make_request, range(3)))
+
+        status_codes = [response.status_code for response in responses]
+
+        # Assert only one call got through
+        assert status_codes.count(200) == 1
+        assert status_codes.count(503) == 2