From 486a0ac23f57e9684e9850c4aae50f0706e5563c Mon Sep 17 00:00:00 2001
From: James Meakin <12661555+jmsmkn@users.noreply.github.com>
Date: Wed, 3 Jul 2024 17:29:26 +0200
Subject: [PATCH] Reserve a constant 1GB for the system

---
 pyproject.toml        | 2 +-
 sagemaker_shim/cli.py | 5 ++++-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index cc93d49..1f088c0 100644
--- 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 <12661555+jmsmkn@users.noreply.github.com>"]
 license = "Apache-2.0"
diff --git a/sagemaker_shim/cli.py b/sagemaker_shim/cli.py
index a4fff04..b5d8217 100644
--- 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),
     )