fix: Fix not reading limits when cgroup v2 used (#107)

Docker image was not picking up memory limits when cgroup v2 was used. This lead to incorrect java heap memory limits that lead to crashes due to OOMKilled as described in #106. Add extra check for limits in cgroup v2 file for the image to correctly read memory limits with both cgroup v1 and cgroup v2. Resolves #106 --- Signed-off-by: funbiscuit <[email protected]>
kserve · Sep 29, 2023 · 757b91e · 757b91e
1 parent de7836a
commit 757b91e
Showing 1 changed file with 9 additions and 5 deletions.
diff --git a/src/main/scripts/start.sh b/src/main/scripts/start.sh
@@ -309,11 +309,15 @@ fi
 echo $$ > ${ANCHOR_FILE}
 
 if [ "$MEM_LIMIT_MB" = "" ]; then
-	DOCKER_LIM_FILE="/sys/fs/cgroup/memory/memory.limit_in_bytes"
-
-	if [ -e "${DOCKER_LIM_FILE}" ]; then
-		MEM_LIMIT_MB=$(($(cat ${DOCKER_LIM_FILE})/1024/1024))
-		echo "Using process mem limit of ${MEM_LIMIT_MB}MiB from ${DOCKER_LIM_FILE}"
+	CGROUP_V1_LIM_FILE="/sys/fs/cgroup/memory/memory.limit_in_bytes"
+	CGROUP_V2_LIM_FILE="/sys/fs/cgroup/memory.max"
+
+	if [ -e "${CGROUP_V1_LIM_FILE}" ]; then
+		MEM_LIMIT_MB=$(($(cat ${CGROUP_V1_LIM_FILE})/1024/1024))
+		echo "Using process mem limit of ${MEM_LIMIT_MB}MiB from ${CGROUP_V1_LIM_FILE}"
+	elif [ -e "${CGROUP_V2_LIM_FILE}" ]; then
+		MEM_LIMIT_MB=$(($(cat ${CGROUP_V2_LIM_FILE})/1024/1024))
+		echo "Using process mem limit of ${MEM_LIMIT_MB}MiB from ${CGROUP_V2_LIM_FILE}"
 	else
     	MEM_LIMIT_MB="1536"
     	echo "No process mem limit provided or found, defaulting to ${MEM_LIMIT_MB}MiB"