become the ubuntu user when root logs in (#3494)

docker/Dockerfile-cloud-uv

@@ -11,7 +11,7 @@ ENV HF_HUB_ENABLE_HF_TRANSFER="1"
 EXPOSE 8888
 EXPOSE 22
 
-COPY scripts/cloud-entrypoint.sh /home/ubuntu/cloud-entrypoint.sh
+COPY scripts/cloud-entrypoint.sh /etc/cloud-entrypoint.sh
 COPY scripts/motd /etc/motd
 
 RUN uv pip install jupyterlab notebook ipywidgets && \
@@ -22,14 +22,14 @@ RUN apt update && \
     rm -rf /var/lib/apt/lists/* && \
     mkdir -p /home/ubuntu/.ssh && \
     chmod 700 /home/ubuntu/.ssh && \
-    printf "\n[[ -z \"\$TMUX\"  ]] && { tmux attach-session -t ssh_tmux || tmux new-session -s ssh_tmux; exit; }\n" >> /home/ubuntu/.bashrc && \
+    printf "\n[[ -z \"\$TMUX\" ]] && tty -s && { tmux attach-session -t ssh_tmux || tmux new-session -s ssh_tmux; exit; }\n" >> /home/ubuntu/.bashrc && \
     printf "[ ! -z \"\$TERM\" -a -r /etc/motd ] && cat /etc/motd\n" >> /home/ubuntu/.bashrc && \
-    chmod +x /workspace/axolotl/scripts/cloud-entrypoint.sh && \
+    printf "\n[[ -z \"\$AXOLOTL_SKIP_SWITCH\" ]] && exec sudo -u ubuntu AXOLOTL_SKIP_SWITCH=1 -i\n" >> /root/.bashrc && \
-    chmod +x /home/ubuntu/cloud-entrypoint.sh && \
+    chmod +x /etc/cloud-entrypoint.sh && \
     echo 'set-option -g history-limit 5000' >> /home/ubuntu/.tmux.conf && \
     chown -R ubuntu:ubuntu /home/ubuntu /workspace
 
-USER ubuntu
+# USER ubuntu
 
-ENTRYPOINT ["/home/ubuntu/cloud-entrypoint.sh"]
+ENTRYPOINT ["/etc/cloud-entrypoint.sh"]
 CMD ["sleep", "infinity"]

scripts/cloud-entrypoint.sh

@@ -3,24 +3,35 @@
 # Detect if running as non-root and set sudo prefix accordingly
 if [ "$(id -u)" -ne 0 ]; then
     SUDO="sudo"
+    RUN_AS_USER=""
 else
     SUDO=""
+    RUN_AS_USER="sudo -u ubuntu"
 fi
 
 # Export specific ENV variables to /etc/rp_environment
 echo "Exporting environment variables..."
 printenv | grep -E '^HF_|^BNB_|^CUDA_|^NCCL_|^NV|^RUNPOD_|^PATH=|^_=' | sed 's/^\([^=]*\)=\(.*\)$/export \1="\2"/' | grep -v 'printenv' | $SUDO tee /etc/rp_environment > /dev/null
-echo 'source /etc/rp_environment' >> ~/.bashrc
+
+# Add rp_environment sourcing to ubuntu's bashrc (if ubuntu user exists and line not already present)
+if id ubuntu &>/dev/null; then
+    grep -q 'source /etc/rp_environment' /home/ubuntu/.bashrc 2>/dev/null || \
+        echo 'source /etc/rp_environment' >> /home/ubuntu/.bashrc
+fi
+# Also add to current user's bashrc if different from ubuntu
+grep -q 'source /etc/rp_environment' ~/.bashrc 2>/dev/null || \
+    echo 'source /etc/rp_environment' >> ~/.bashrc
 
 add_keys_to_authorized() {
     local key_value=$1
+    local target_home=$2
 
-    # Create the ~/.ssh directory and set permissions
+    # Create the .ssh directory and set permissions
-    mkdir -p ~/.ssh
+    mkdir -p "$target_home/.ssh"
-    chmod 700 ~/.ssh
+    chmod 700 "$target_home/.ssh"
 
     # Create the authorized_keys file if it doesn't exist
-    touch ~/.ssh/authorized_keys
+    touch "$target_home/.ssh/authorized_keys"
 
     # Initialize an empty key variable
     local key=""
@@ -31,7 +42,7 @@ add_keys_to_authorized() {
         if [[ $word == ssh-* ]]; then
             # If there's a key being built, add it to the authorized_keys file
             if [[ -n $key ]]; then
-                echo $key >> ~/.ssh/authorized_keys
+                echo $key >> "$target_home/.ssh/authorized_keys"
             fi
             # Start a new key
             key=$word
@@ -43,12 +54,25 @@ add_keys_to_authorized() {
 
     # Add the last key to the authorized_keys file
     if [[ -n $key ]]; then
-        echo $key >> ~/.ssh/authorized_keys
+        echo $key >> "$target_home/.ssh/authorized_keys"
     fi
 
     # Set the correct permissions
-    chmod 600 ~/.ssh/authorized_keys
+    chmod 600 "$target_home/.ssh/authorized_keys"
-    chmod 700 -R ~/.ssh
+    chmod 700 -R "$target_home/.ssh"
+}
+
+setup_ssh_keys() {
+    local key_value=$1
+
+    # Set up keys for the current user
+    add_keys_to_authorized "$key_value" "$HOME"
+
+    # Also set up keys for ubuntu user if we're root and ubuntu exists
+    if [ "$(id -u)" -eq 0 ] && id ubuntu &>/dev/null; then
+        add_keys_to_authorized "$key_value" "/home/ubuntu"
+        chown -R ubuntu:ubuntu /home/ubuntu/.ssh
+    fi
 }
 
 # Set SSH port
@@ -58,12 +82,12 @@ fi
 
 if [[ $PUBLIC_KEY ]]; then
     # runpod, prime intellect
-    add_keys_to_authorized "$PUBLIC_KEY"
+    setup_ssh_keys "$PUBLIC_KEY"
     # Start the SSH service in the background
     $SUDO service ssh start
 elif [[ $SSH_KEY ]]; then
     # latitude.sh
-    add_keys_to_authorized "$SSH_KEY"
+    setup_ssh_keys "$SSH_KEY"
     # Start the SSH service in the background
     $SUDO service ssh start
 else
@@ -77,12 +101,16 @@ if [ -n "$JUPYTER_PASSWORD" ]; then
 fi
 
 if [ "$JUPYTER_DISABLE" != "1" ]; then
-    # Run Jupyter Lab in the background
+    # Run Jupyter Lab as ubuntu user when possible
     JUPYTER_ARGS="--port=8888 --ip=* --ServerApp.allow_origin=*"
-    if [ "$(id -u)" -eq 0 ]; then
+    if [ "$(id -u)" -eq 0 ] && id ubuntu &>/dev/null; then
-        JUPYTER_ARGS="$JUPYTER_ARGS --allow-root"
+        sudo -u ubuntu bash -c "JUPYTER_TOKEN='$JUPYTER_TOKEN' jupyter lab $JUPYTER_ARGS" &
+    else
+        if [ "$(id -u)" -eq 0 ]; then
+            JUPYTER_ARGS="$JUPYTER_ARGS --allow-root"
+        fi
+        jupyter lab $JUPYTER_ARGS &
     fi
-    jupyter lab $JUPYTER_ARGS &
 fi
 
 if [ ! -d "/workspace/data/axolotl-artifacts" ]; then
@@ -91,6 +119,7 @@ fi
 if [ ! -L "/workspace/axolotl/outputs" ]; then
     ln -sf /workspace/data/axolotl-artifacts /workspace/axolotl/outputs
 fi
+chown -R ubuntu:ubuntu /workspace 2>/dev/null || true
 
 # start the runpod slurm init
 SLURM_INIT="${SLURM_INIT:-/slurm-init.sh}"
@@ -100,5 +129,9 @@ if [[ -f "$SLURM_INIT" ]]; then
   $SUDO bash "$SLURM_INIT"
 fi
 
-# Execute the passed arguments (CMD)
+# Execute the passed arguments (CMD) as ubuntu when possible
-exec "$@"
+if [ "$(id -u)" -eq 0 ] && id ubuntu &>/dev/null; then
+    exec sudo -u ubuntu "$@"
+else
+    exec "$@"
+fi