Add oarsh/oarsh_shell cgroup-v2 support (POC)

oar-team · Dec 6, 2023 · d8da4c8 · d8da4c8
1 parent e65c6ac
commit d8da4c8
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 11 deletions.
diff --git a/oar/tools/oarsh/oarsh.in b/oar/tools/oarsh/oarsh.in
@@ -163,7 +163,7 @@ if [ "$OARSH_BYPASS_WHOLE_SECURITY" != "0" ]; then
 fi
 
 # -1- try connection using a user provided job key file for a job using the job key mechanism
-if [ -n "$OAR_JOB_KEY_FILE" ] 
+if [ -n "$OAR_JOB_KEY_FILE" ]
 then
     # first, get rid of remaining unused jobkey files if any...
     for f in $OAR_RUNTIME_DIRECTORY/oarsh.jobkey.*; do
@@ -187,8 +187,8 @@ then
 fi
 
 
-# -2- try connection using a job key pushed by OAR for a job using the job key mechanism. 
-# (oarsh is run from one of the node of the job) 
+# -2- try connection using a job key pushed by OAR for a job using the job key mechanism.
+# (oarsh is run from one of the node of the job)
 TMP_JOB_KEY_FILE="$OAR_RUNTIME_DIRECTORY/$OARDO_USER.jobkey"
 if [ -r $TMP_JOB_KEY_FILE ]; then
     umask $OLDUMASK
@@ -197,13 +197,15 @@ if [ -r $TMP_JOB_KEY_FILE ]; then
     exit 4
 fi
 
-if [ "$CPUSET_PATH" != "" ]; then
+if [ "$CPUSET_PATH" != "" ]; then #TODO adapt for cgroupv2
     if [ -r /proc/self/cpuset ]; then
         if [ -n "$GET_CURRENT_CPUSET_CMD" ]; then
             OAR_CPUSET=$(bash -c "$GET_CURRENT_CPUSET_CMD")
         else
-            OAR_CPUSET=$(< /proc/self/cpuset)
-        fi
+            OAR_CPUSET=$(< /proc/self/cpuset) #Support for cgroup-v1 and cgroup-v2
+            USER_JOB_ID=$(basename $OAR_CPUSET)
+            OAR_CPUSET="/oar/${USER_JOB_ID%.*}" # Remove .scope in cgroup-v2
+         fi
         if [ "${OAR_CPUSET%/*}" == "$CPUSET_PATH" ] || [ "${OAR_CPUSET%/*}" == "$CPUSET_PATH/" ]; then
             JOB_KEY_FILE="$OAR_RUNTIME_DIRECTORY/${OAR_CPUSET##*/}.jobkey"
             if [ -r $JOB_KEY_FILE ]; then

diff --git a/oar/tools/oarsh/oarsh_shell.in b/oar/tools/oarsh/oarsh_shell.in
@@ -20,6 +20,28 @@ JOBENVFILE="$OAR_RUNTIME_DIRECTORY/${OAR_CPUSET##*/}.env"
 OLDUMASK=$(umask)
 umask 0022
 
+# $1 = Name for cgroup name (systemd transient scope)
+# $2,$3,... = PIDs to add
+# Add PIDS into the cgroup $1
+function add_process_to_cgroup_v2() {
+    [ "$1" = "undef" ] && return
+    ######################
+    # Cgroup assignation #
+    ###########################################################################
+    user_jobid=$(basename $1) # /oar/user1_1 -> user1_1
+    CGROUPNAME="oar.slice/$user_jobid.scope"
+    if [ ! -f /sys/fs/cgroup/$CGROUPNAME/cgroup.procs ]; then
+        echo "oarsh: Cannot find cgroup.procs file: /sys/fs/cgroup/$CGROUPNAME/cgroup.procs" 1>&2
+        exit 61
+    fi
+    shift
+
+    # Add each processes to the right OAR cgroups
+    for p in $@; do
+      $OARDIR/oardodo/oardodo sh -c "echo $p > /sys/fs/cgroup/$CGROUPNAME/cgroup.procs" || return 1
+    done
+}
+
 # $1 = Name of the cpuset
 # $2,$3,... = PIDs to add
 # Add PIDS into the cpuset $1
@@ -63,12 +85,12 @@ then
     # It must be oar
     if [ "$OAR_CPUSET" != "" ]
     then
-        add_process_to_cpuset $OAR_CPUSET $$ $PPID || exit 62
+        add_process_to_cgroup_v2 $OAR_CPUSET $$ $PPID || exit 62
         [ -r "$JOBENVFILE" ] && source $JOBENVFILE
     fi
     $OARDIR/oardodo/oardodo renice 0 $$ $PPID > /dev/null 2>&1
     export SHELL=$DEFAULT_SHELL
-    
+
     umask $OLDUMASK
     exec $DEFAULT_SHELL "$@"
     echo "oarsh: exec failed" 1>&2
@@ -79,8 +101,8 @@ else
         echo "oarsh: OAR_CPUSET variable is empty; Is your sshd right configured with 'AcceptEnv OAR_CPUSET OAR_JOB_USER' on all computing nodes?" 1>&2
         exit 63
     fi
-    add_process_to_cpuset $OAR_CPUSET $$ $PPID || exit 62
-    
+    add_process_to_cgroup_v2 $OAR_CPUSET $$ $PPID || exit 62
+
     #Manage display
     if [ -n "$DISPLAY" ]
     then
@@ -115,4 +137,3 @@ fi
 
 echo "oarsh: Really bad error" 1>&2
 exit 67
-