From 926082d7ca04ae51fce59d5e6a3c479de25fc891 Mon Sep 17 00:00:00 2001
From: Omer Tuchfeld <omer@tuchfeld.dev>
Date: Tue, 6 Aug 2024 12:42:07 +0200
Subject: [PATCH] ilab-wrapper: Improve non-single subuid ranges error handling

# Background

df8885777d2257a6271ca1a6461673d1d96e99d6

# Issue

The current error handling for multiple subuid ranges is broken due to
surprising behavior of `wc -l` which always returns `1` even when the
input is empty.

# Solution

More carefully count the number of lines in the
`CURRENT_USER_SUBUID_RANGE` variable

# Additional changes

50fb00f26f921bc5e420bed4e43d275a600e9b4c forgot to modify the duplicated
wrapper script and had a small merge error, this commit fixes that.

Signed-off-by: Omer Tuchfeld <omer@tuchfeld.dev>
---
 training/ilab-wrapper/ilab                     | 18 +++++++++---------
 .../nvidia-bootc/duplicated/ilab-wrapper/ilab  | 18 +++++++++---------
 2 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/training/ilab-wrapper/ilab b/training/ilab-wrapper/ilab
index 406a939e..2018b167 100755
--- a/training/ilab-wrapper/ilab
+++ b/training/ilab-wrapper/ilab
@@ -72,21 +72,21 @@ else
         '$1 == current_user || $1 == current_uid {print $2 ":" $3}' \
         /etc/subuid)
 
-    # TODO: Handle multiple subuid ranges, for now, hard fail
-    if [[ $(wc -l <<<"$CURRENT_USER_SUBUID_RANGE") != 1 ]]; then
-        if [[ -z "$CURRENT_USER_SUBUID_RANGE" ]]; then
-            echo-err "No subuid range found for user $CURRENT_USER_NAME ($UID)"
-        else
-            echo-err "Multiple subuid ranges found for user $CURRENT_USER_NAME ($UID), this is currently unsupported"
-            echo-err "$CURRENT_USER_SUBUID_RANGE"
-        fi
+    NUMBER_OF_MATCHING_SUBUID_RANGES=$(if [[ -z "$CURRENT_USER_SUBUID_RANGE" ]]; then echo 0; else wc -l <<< "$CURRENT_USER_SUBUID_RANGE"; fi)
+
+    if [[ "$NUMBER_OF_MATCHING_SUBUID_RANGES" == 0 ]]; then
+        echo-err "No /etc/subuid range found for user $CURRENT_USER_NAME ($UID)"
+        exit 1
+    elif [[ "$NUMBER_OF_MATCHING_SUBUID_RANGES" != 1 ]]; then
+        # TODO: Handle multiple subuid ranges. But for now, hard fail
+        echo-err "Multiple /etc/subuid ranges found for user $CURRENT_USER_NAME ($UID), this is currently unsupported:"
+        echo-err "$CURRENT_USER_SUBUID_RANGE"
         exit 1
     fi
 
     IMPERSONATE_CURRENT_USER_PODMAN_FLAGS=("--uidmap" "0:$UID" "--uidmap" "1:$CURRENT_USER_SUBUID_RANGE")
 fi
 
-IMPERSONATE_CURRENT_USER_PODMAN_FLAGS=("--uidmap" "0:$UID" "--uidmap" "1:$CURRENT_USER_SUBUID_RANGE")
 PRESERVE_ENV="VLLM_LOGGING_LEVEL,NCCL_DEBUG,HOME,HF_TOKEN"
 PODMAN_COMMAND=("sudo" "--preserve-env=$PRESERVE_ENV" "podman" "run" "--rm" "-it"
     "${IMPERSONATE_CURRENT_USER_PODMAN_FLAGS[@]}"
diff --git a/training/nvidia-bootc/duplicated/ilab-wrapper/ilab b/training/nvidia-bootc/duplicated/ilab-wrapper/ilab
index 406a939e..2018b167 100755
--- a/training/nvidia-bootc/duplicated/ilab-wrapper/ilab
+++ b/training/nvidia-bootc/duplicated/ilab-wrapper/ilab
@@ -72,21 +72,21 @@ else
         '$1 == current_user || $1 == current_uid {print $2 ":" $3}' \
         /etc/subuid)
 
-    # TODO: Handle multiple subuid ranges, for now, hard fail
-    if [[ $(wc -l <<<"$CURRENT_USER_SUBUID_RANGE") != 1 ]]; then
-        if [[ -z "$CURRENT_USER_SUBUID_RANGE" ]]; then
-            echo-err "No subuid range found for user $CURRENT_USER_NAME ($UID)"
-        else
-            echo-err "Multiple subuid ranges found for user $CURRENT_USER_NAME ($UID), this is currently unsupported"
-            echo-err "$CURRENT_USER_SUBUID_RANGE"
-        fi
+    NUMBER_OF_MATCHING_SUBUID_RANGES=$(if [[ -z "$CURRENT_USER_SUBUID_RANGE" ]]; then echo 0; else wc -l <<< "$CURRENT_USER_SUBUID_RANGE"; fi)
+
+    if [[ "$NUMBER_OF_MATCHING_SUBUID_RANGES" == 0 ]]; then
+        echo-err "No /etc/subuid range found for user $CURRENT_USER_NAME ($UID)"
+        exit 1
+    elif [[ "$NUMBER_OF_MATCHING_SUBUID_RANGES" != 1 ]]; then
+        # TODO: Handle multiple subuid ranges. But for now, hard fail
+        echo-err "Multiple /etc/subuid ranges found for user $CURRENT_USER_NAME ($UID), this is currently unsupported:"
+        echo-err "$CURRENT_USER_SUBUID_RANGE"
         exit 1
     fi
 
     IMPERSONATE_CURRENT_USER_PODMAN_FLAGS=("--uidmap" "0:$UID" "--uidmap" "1:$CURRENT_USER_SUBUID_RANGE")
 fi
 
-IMPERSONATE_CURRENT_USER_PODMAN_FLAGS=("--uidmap" "0:$UID" "--uidmap" "1:$CURRENT_USER_SUBUID_RANGE")
 PRESERVE_ENV="VLLM_LOGGING_LEVEL,NCCL_DEBUG,HOME,HF_TOKEN"
 PODMAN_COMMAND=("sudo" "--preserve-env=$PRESERVE_ENV" "podman" "run" "--rm" "-it"
     "${IMPERSONATE_CURRENT_USER_PODMAN_FLAGS[@]}"