Set the memory ulimit

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sagemaker-shim"
-version = "0.3.0"
+version = "0.3.1"
 description = "Adapts algorithms that implement the Grand Challenge inference API for running in SageMaker"
 authors = ["James Meakin <[email protected]>"]
 license = "Apache-2.0"
@@ -21,6 +21,7 @@ fastapi = "!=0.89.0"  # See https://github.com/DIAGNijmegen/rse-sagemaker-shim/i
 uvicorn = "*"
 click = "*"
 boto3 = "*"
+psutil = "*"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "!=8.0.0"  # pytest 8 is not yet supported by pytest-asyncio

sagemaker_shim/cli.py

@@ -1,12 +1,14 @@
 import asyncio
 import logging.config
+import resource
 import sys
 from collections.abc import Callable, Coroutine
 from functools import wraps
 from json import JSONDecodeError
 from typing import Any, TypeVar
 
 import click
+import psutil
 import uvicorn
 from botocore.exceptions import ClientError, NoCredentialsError
 from pydantic import ValidationError
@@ -101,7 +103,18 @@ async def invoke(tasks: str, file: str) -> None:
         raise click.UsageError("Empty task list provided")
 
 
+def set_memory_limits() -> None:
+    total_memory_bytes = psutil.virtual_memory().total
+
+    resource.setrlimit(
+        resource.RLIMIT_AS,
+        (int(total_memory_bytes * 0.85), int(total_memory_bytes * 0.95)),
+    )
+
+
 if __name__ == "__main__":
+    set_memory_limits()
+
     # https://pyinstaller.org/en/stable/runtime-information.html#run-time-information
     we_are_bundled = getattr(sys, "frozen", False) and hasattr(sys, "_MEIPASS")