pi_run: hard per-job memory cap via a systemd scope (MemoryMax) when the memory cgroup is enabled

The Pi now boots with cgroup_enable=memory cgroup_memory=1; the MemAvailable guard stays as a
second layer and as the fallback on systems without the memory controller.

scripts/pi_run.sh

@@ -1,23 +1,36 @@
 #!/usr/bin/env bash
-# Run a heavy command on the 8 GB Pi with a memory guard. Raspberry Pi OS boots with cgroup_disable=memory, so
-# systemd MemoryMax cannot be used; instead the command runs in its own process group and is stopped (TERM, then
-# KILL) as soon as the system's MemAvailable drops below MIN_FREE_MB -- before the kernel OOM killer would pick
-# victims such as the desktop session. Prints the peak resident memory of the command's process tree.
-#   MIN_FREE_MB=1200 bash scripts/pi_run.sh <command...>
+# Run a heavy command on the 8 GB Pi so that running out of memory kills only this command, never the desktop session.
+#   MEM=4G MIN_FREE_MB=1200 bash scripts/pi_run.sh <command...>
+# 1. When the memory cgroup is available (cmdline.txt: cgroup_enable=memory cgroup_memory=1), the command runs in a
+#    systemd user scope with MemoryMax=$MEM and no swap: the kernel kills it alone when it exceeds the cap.
+# 2. In every case a guard stops the command's process group when the system's MemAvailable drops below
+#    MIN_FREE_MB (covers other processes filling memory, and systems without the memory cgroup).
+# Prints the peak memory of the command when it ends.
 set -uo pipefail
+MEM=${MEM:-4G}
 MIN_FREE_MB=${MIN_FREE_MB:-1200}
-setsid "$@" &
+uid=$(id -u)
+unit="simurg-job-$$-$(date +%s)"
+cg="/sys/fs/cgroup/user.slice/user-$uid.slice/user@$uid.service/app.slice/$unit.scope"
+if grep -qw memory "/sys/fs/cgroup/user.slice/user-$uid.slice/user@$uid.service/cgroup.controllers" 2>/dev/null; then
+  setsid systemd-run --user --scope --unit="$unit" -p MemoryMax="$MEM" -p MemorySwapMax=0 --quiet -- "$@" &
+  mode="cgroup MemoryMax=$MEM"
+else
+  setsid "$@" &
+  mode="no memory cgroup"
+fi
 pid=$!
 peak=0
 tripped=0
-tree_rss() {  # kB, summed over the process group
-  ps -o rss= -g "$pid" 2>/dev/null | awk '{s+=$1} END {print s+0}'
-}
 while kill -0 "$pid" 2>/dev/null; do
-  rss=$(tree_rss); (( rss > peak )) && peak=$rss
+  if [[ -r $cg/memory.peak ]]; then
+    peak=$(( $(cat "$cg/memory.peak") / 1024 ))
+  else
+    rss=$(ps -o rss= -g "$pid" 2>/dev/null | awk '{s+=$1} END {print s+0}'); (( rss > peak )) && peak=$rss
+  fi
   avail=$(awk '/^MemAvailable:/ {print int($2/1024)}' /proc/meminfo)
   if (( avail < MIN_FREE_MB )); then
-    echo "pi_run: MemAvailable ${avail} MiB < ${MIN_FREE_MB} MiB -- stopping the job (process group $pid)" >&2
+    echo "pi_run: MemAvailable ${avail} MiB < ${MIN_FREE_MB} MiB -- stopping the job" >&2
     kill -TERM -- "-$pid" 2>/dev/null; sleep 5; kill -KILL -- "-$pid" 2>/dev/null
     tripped=1
     break
@@ -26,5 +39,7 @@ while kill -0 "$pid" 2>/dev/null; do
 done
 wait "$pid"; rc=$?
 (( tripped )) && rc=137
-echo "pi_run: exit $rc, peak RSS $((peak / 1024)) MiB, guard ${MIN_FREE_MB} MiB free" >&2
+(( rc == 137 && ! tripped )) && [[ $mode == cgroup* ]] \
+  && echo "pi_run: the job was killed (SIGKILL); most likely it exceeded MemoryMax=$MEM" >&2
+echo "pi_run: exit $rc, peak memory $((peak / 1024)) MiB ($mode, guard ${MIN_FREE_MB} MiB free)" >&2
 exit $rc