Reserve a constant 1GB for the system

DIAGNijmegen · Jul 3, 2024 · 486a0ac · 486a0ac
1 parent 606936f
commit 486a0ac
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sagemaker-shim"
-version = "0.3.1"
+version = "0.3.2"
 description = "Adapts algorithms that implement the Grand Challenge inference API for running in SageMaker"
 authors = ["James Meakin <[email protected]>"]
 license = "Apache-2.0"

diff --git a/sagemaker_shim/cli.py b/sagemaker_shim/cli.py
@@ -106,9 +106,12 @@ async def invoke(tasks: str, file: str) -> None:
 def set_memory_limits() -> None:
     total_memory_bytes = psutil.virtual_memory().total
 
+    # Reserve 1GB for the system
+    limit = total_memory_bytes - 1_073_741_824
+
     resource.setrlimit(
         resource.RLIMIT_AS,
-        (int(total_memory_bytes * 0.85), int(total_memory_bytes * 0.95)),
+        (limit, limit),
     )