Update script 90-limits.sh

- Use command prlimit instead of ulimit
- Add option to not limit address space

NOTES.md

@@ -39,14 +39,16 @@ The following startup hooks are put in place:
 * [/usr/local/bin/before-notebook.d/71-tensorboard.sh](base/scripts/usr/local/bin/before-notebook.d/71-tensorboard.sh)
   to use Jupyter Server Proxy for TensorBoard.
 * [/usr/local/bin/before-notebook.d/90-limits.sh](base/scripts/usr/local/bin/before-notebook.d/90-limits.sh)
-  * *soft* limit the *address space* based on the amount of *physical* and
-    *virtual memory* of the host. (default: command `ulimit -v`)
+  * *soft* limit the *address space* based on the amount of *physical memory*
+    (`MEM_LIMIT`) and *virtual memory* (`SWAP_ENABLE`, `SWAP_FACTOR`). (default:
+    command `prlimit -v`)
+    * Do not limit if `NO_AS_LIMIT` or `NO_MEM_LIMIT` is set to `1` or `yes`.
   * limit the number of *file descriptors* according to environment variable
-    `NOFILE_LIMIT`. (default: command `ulimit -n`)
+    `NOFILE_LIMIT`. (default: command `prlimit -n`)
   * limit the number of *processes* according to environment variable
-    `NPROC_LIMIT`. (default: command `ulimit -u`)
+    `NPROC_LIMIT`. (default: command `prlimit -u`)
   * limit the number of *pending signals* according to environment variable
-    `SIGPEN_LIMIT`. (default: command `ulimit -i`)
+    `SIGPEN_LIMIT`. (default: command `prlimit -i`)
 * [/usr/local/bin/before-notebook.d/95-misc.sh](base/scripts/usr/local/bin/before-notebook.d/95-misc.sh)
   to export environment variables to `/tmp/environment`.
 

base/scripts/usr/local/bin/before-notebook.d/90-limits.sh

@@ -4,32 +4,38 @@
 
 set -e
 
-DIVISOR=1024
+if [ "$(id -u)" != 0 ]; then
+  soft_limit=1
+fi
 
 if [[ "$SWAP_ENABLE" == "1" || "$SWAP_ENABLE" == "yes" ]]; then
-  FACTOR=$(echo 1 "${SWAP_FACTOR:-1}" | awk '{ printf "%.1f", $1 + $2 }')
+  factor=$(echo 1 "${SWAP_FACTOR:-1}" | awk '{ printf "%.1f", $1 + $2 }')
 else
-  FACTOR=1
+  factor=1
 fi
 
 # Limit address space: Soft when run as root and as other user
 if [ -n "$MEM_LIMIT" ]; then
-  ulimit -Sv "$(echo "$MEM_LIMIT" "$DIVISOR" "$FACTOR" |
-    awk '{ printf "%.0f", $1 / $2 * $3 }')"
+  NO_AS_LIMIT=${NO_AS_LIMIT:-$NO_MEM_LIMIT}
+  # Do not limit if NO_AS_LIMIT or NO_MEM_LIMIT is set to 1 or yes
+  if [[ "$NO_AS_LIMIT" != "1" && "$NO_AS_LIMIT" != "yes" ]]; then
+    prlimit --pid $$ --as="$(echo "$MEM_LIMIT" "$factor" |
+      awk '{ printf "%.0f", $1 * $2 }')":
+  fi
 fi
 
 # Other limits: Hard when run as root user; Soft when run as other user
 # pending signals
 if [ -n "$SIGPEN_LIMIT" ]; then
-  ulimit -i "$(printf %.0f "$SIGPEN_LIMIT")"
+  prlimit --pid $$ --sigpending="$(printf %.0f "$SIGPEN_LIMIT")"${soft_limit:+:}
 fi
 
 # file descriptors
 if [ -n "$NOFILE_LIMIT" ]; then
-  ulimit -n "$(printf %.0f "$NOFILE_LIMIT")"
+  prlimit --pid $$ --nofile="$(printf %.0f "$NOFILE_LIMIT")"${soft_limit:+:}
 fi
 
 # processes
 if [ -n "$NPROC_LIMIT" ]; then
-  ulimit -u "$(printf %.0f "$NPROC_LIMIT")"
+  prlimit --pid $$ --nproc="$(printf %.0f "$NPROC_LIMIT")"${soft_limit:+:}
 fi