feat: Return a 503 if available memory is low (#203)

Add `UNSTRUCTURED_MEMORY_FREE_MINIMUM_MB`. When this is set, and the
host available memory is below this value, we'll start returning 503.
The exception is if the request originates from `10.4.x.x` or
`10.5.x.x`. These are parallel mode requests for a doc already in
progress (hosted api).

Also, move the `safetensors` pin to a new `constraints.in` file.

To test, set the variable super high and confirm that you get a 503:
```
export UNSTRUCTURED_MEMORY_FREE_MINIMUM_MB=30000
make run-web-app

curl 'http://localhost:8000/general/v0/general' --form files=@sample-docs/layout-parser-paper-fast.pdf

{"detail":"Server is under heavy load. Please try again later."}%      
```

Author: awalker4

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.0.40-dev0
+
+* Reject traffic when we're overloaded via `UNSTRUCTURED_MEMORY_FREE_MINIMUM_MB`
+
 ## 0.0.39
 
 * Fix wrong handleing on param skip_infer_table_types

pipeline-notebooks/pipeline-general.ipynb

@@ -553,6 +553,7 @@
     "from pypdf import PdfReader, PdfWriter\n",
     "from unstructured.partition.auto import partition\n",
     "from unstructured.staging.base import convert_to_isd, convert_to_dataframe, elements_from_json\n",
+    "import psutil\n",
     "import requests\n",
     "import time"
    ]
@@ -767,6 +768,22 @@
     "        \"m_xml_keep_tags\": m_xml_keep_tags,\n",
     "        }, default=str)\n",
     "    ))\n",
+    "\n",
+    "    # If this var is set, reject traffic when free memory is below minimum\n",
+    "    # Allow internal requests - these are parallel calls already in progress\n",
+    "    mem = psutil.virtual_memory()\n",
+    "    memory_free_minimum = int(os.environ.get(\"UNSTRUCTURED_MEMORY_FREE_MINIMUM_MB\", 0))\n",
+    "\n",
+    "    if memory_free_minimum > 0 and mem.available <= memory_free_minimum * 1024 * 1024:\n",
+    "        # Note(yuming): Use X-Forwarded-For header to find the orginal IP for external API\n",
+    "        # requests,since LB forwards requests in AWS\n",
+    "        origin_ip = request.headers.get(\"X-Forwarded-For\") or request.client.host\n",
+    "\n",
+    "        if not origin_ip.startswith(\"10.\"):\n",
+    "            raise HTTPException(\n",
+    "                status_code=503, detail=\"Server is under heavy load. Please try again later.\"\n",
+    "            )\n",
+    "\n",
     "    if filename.endswith(\".msg\"):\n",
     "        # Note(yuming): convert file type for msg files\n",
     "        # since fast api might sent the wrong one.\n",

prepline_general/api/general.py

@@ -24,6 +24,7 @@
 from pypdf import PdfReader, PdfWriter
 from unstructured.partition.auto import partition
 from unstructured.staging.base import convert_to_isd, convert_to_dataframe, elements_from_json
+import psutil
 import requests
 import time
 from unstructured_inference.models.chipper import MODEL_TYPES as CHIPPER_MODEL_TYPES
@@ -233,6 +234,22 @@ def pipeline_api(
             )
         )
     )
+
+    # If this var is set, reject traffic when free memory is below minimum
+    # Allow internal requests - these are parallel calls already in progress
+    mem = psutil.virtual_memory()
+    memory_free_minimum = int(os.environ.get("UNSTRUCTURED_MEMORY_FREE_MINIMUM_MB", 0))
+
+    if memory_free_minimum > 0 and mem.available <= memory_free_minimum * 1024 * 1024:
+        # Note(yuming): Use X-Forwarded-For header to find the orginal IP for external API
+        # requests,since LB forwards requests in AWS
+        origin_ip = request.headers.get("X-Forwarded-For") or request.client.host
+
+        if not origin_ip.startswith("10."):
+            raise HTTPException(
+                status_code=503, detail="Server is under heavy load. Please try again later."
+            )
+
     if filename.endswith(".msg"):
         # Note(yuming): convert file type for msg files
         # since fast api might sent the wrong one.
@@ -485,7 +502,7 @@ def return_content_type(filename):
 
 
 @router.post("/general/v0/general")
-@router.post("/general/v0.0.39/general")
+@router.post("/general/v0.0.40/general")
 def pipeline_1(
     request: Request,
     gz_uncompressed_content_type: Optional[str] = Form(default=None),

preprocessing-pipeline-family.yaml

@@ -1,2 +1,2 @@
 name: general
-version: 0.0.39
+version: 0.0.40

requirements/base.in

@@ -1,3 +1,4 @@
+-c constraints.in
 unstructured[local-inference]>=0.8.1
 unstructured-api-tools>=0.10.10
 pydantic<2.0.2
@@ -9,6 +10,5 @@ ratelimit
 requests
 pypdf
 pycryptodome
-# Note(austin) higher versions need rust installed? pinning this until we address in the dockerfile
-safetensors<0.3.3
+psutil
 

requirements/base.txt

@@ -123,7 +123,7 @@ jupyter-core==5.3.1
     #   nbformat
 jupyterlab-pygments==0.2.2
     # via nbconvert
-kiwisolver==1.4.4
+kiwisolver==1.4.5
     # via matplotlib
 layoutparser[layoutmodels,tesseract]==0.3.4
     # via unstructured-inference
@@ -230,6 +230,8 @@ portalocker==2.7.0
     # via iopath
 protobuf==4.24.1
     # via onnxruntime
+psutil==5.9.5
+    # via -r requirements/base.in
 pycocotools==2.0.7
     # via effdet
 pycparser==2.21
@@ -306,7 +308,7 @@ rpds-py==0.9.2
     #   referencing
 safetensors==0.3.2
     # via
-    #   -r requirements/base.in
+    #   -c requirements/constraints.in
     #   timm
     #   transformers
 scipy==1.10.1
@@ -397,7 +399,7 @@ uvicorn[standard]==0.23.2
     # via unstructured-api-tools
 uvloop==0.17.0
     # via uvicorn
-watchfiles==0.19.0
+watchfiles==0.20.0
     # via uvicorn
 webencodings==0.5.1
     # via

requirements/constraints.in

@@ -0,0 +1,7 @@
+####################################################################################################
+# This file can house global constraints that aren't *direct* requirements of the package or any
+# extras. Putting a dependency here will only affect dependency sets that contain them -- in other
+# words, if something does not require a constraint, it will not be installed.
+####################################################################################################
+# Note(austin) - preventing a build error with 0.3.3
+safetensors<0.3.3

requirements/test.in

@@ -1,3 +1,4 @@
+-c constraints.in
 black
 # NOTE(mrobinson) - Pinning click due to a unicode issue in black
 # can remove after black drops support for Python 3.6
@@ -9,4 +10,4 @@ pytest-cov
 pytest-mock
 nbdev
 jupyter
-httpx
+httpx

requirements/test.txt

@@ -330,7 +330,7 @@ jupyterlab-server==2.24.0
     #   notebook
 jupyterlab-widgets==3.0.8
     # via ipywidgets
-kiwisolver==1.4.4
+kiwisolver==1.4.5
     # via
     #   -r requirements/base.txt
     #   matplotlib
@@ -544,7 +544,9 @@ protobuf==4.24.1
     #   -r requirements/base.txt
     #   onnxruntime
 psutil==5.9.5
-    # via ipykernel
+    # via
+    #   -r requirements/base.txt
+    #   ipykernel
 ptyprocess==0.7.0
     # via
     #   pexpect
@@ -698,6 +700,7 @@ rpds-py==0.9.2
     #   referencing
 safetensors==0.3.2
     # via
+    #   -c requirements/constraints.in
     #   -r requirements/base.txt
     #   timm
     #   transformers
@@ -872,7 +875,7 @@ uvloop==0.17.0
     #   uvicorn
 watchdog==3.0.0
     # via nbdev
-watchfiles==0.19.0
+watchfiles==0.20.0
     # via
     #   -r requirements/base.txt
     #   uvicorn

test_general/api/test_app.py

@@ -387,6 +387,39 @@ def test_general_api_returns_400_bad_pdf():
     tmp.close()
 
 
+def test_general_api_returns_503(monkeypatch, mocker):
+    """
+    When available memory is below the minimum. return a 503, unless our origin ip is 10.{4,5}.x.x
+    """
+    monkeypatch.setenv("UNSTRUCTURED_MEMORY_FREE_MINIMUM_MB", "30000")
+
+    client = TestClient(app)
+    test_file = Path("sample-docs") / "fake-xml.xml"
+    response = client.post(
+        MAIN_API_ROUTE,
+        files=[("files", (str(test_file), open(test_file, "rb")))],
+    )
+
+    assert response.status_code == 503
+
+    mock_client = mocker.patch("fastapi.Request.client")
+    mock_client.host = "10.5.0.0"
+    response = client.post(
+        MAIN_API_ROUTE,
+        files=[("files", (str(test_file), open(test_file, "rb")))],
+    )
+
+    assert response.status_code == 200
+
+    mock_client.host = "10.4.0.0"
+    response = client.post(
+        MAIN_API_ROUTE,
+        files=[("files", (str(test_file), open(test_file, "rb")))],
+    )
+
+    assert response.status_code == 200
+
+
 class MockResponse:
     def __init__(self, status_code):
         self.status_code = status_code