Set RLIMIT_DATA with GRAND_CHALLENGE_COMPONENT_RESERVED_BYTES

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sagemaker-shim"
-version = "0.3.3"
+version = "0.3.4"
 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,15 @@
 import asyncio
 import logging.config
+import os
+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 +104,24 @@ async def invoke(tasks: str, file: str) -> None:
         raise click.UsageError("Empty task list provided")
 
 
+def set_memory_limits() -> None:
+    reserved_bytes = int(
+        os.environ.get(
+            "GRAND_CHALLENGE_COMPONENT_RESERVED_BYTES", 1_073_741_824
+        )
+    )
+
+    if reserved_bytes:
+        total_memory_bytes = psutil.virtual_memory().total
+
+        limit = total_memory_bytes - reserved_bytes
+
+        resource.setrlimit(resource.RLIMIT_DATA, (limit, limit))
+
+
 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")