From e1ff756245c16a91fd25e04c0774dd715cc3ec3b Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Fri, 13 Mar 2026 09:06:54 -0400 Subject: [PATCH] become the ubuntu user when root logs in (#3494) --- docker/Dockerfile-cloud-uv | 12 +++---- scripts/cloud-entrypoint.sh | 67 +++++++++++++++++++++++++++---------- 2 files changed, 56 insertions(+), 23 deletions(-) diff --git a/docker/Dockerfile-cloud-uv b/docker/Dockerfile-cloud-uv index c38b21e6f..d2af55485 100644 --- a/docker/Dockerfile-cloud-uv +++ b/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 && \ - chmod +x /home/ubuntu/cloud-entrypoint.sh && \ + printf "\n[[ -z \"\$AXOLOTL_SKIP_SWITCH\" ]] && exec sudo -u ubuntu AXOLOTL_SKIP_SWITCH=1 -i\n" >> /root/.bashrc && \ + 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"] diff --git a/scripts/cloud-entrypoint.sh b/scripts/cloud-entrypoint.sh index e88762e13..8eb056acb 100755 --- a/scripts/cloud-entrypoint.sh +++ b/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 - mkdir -p ~/.ssh - chmod 700 ~/.ssh + # Create the .ssh directory and set permissions + mkdir -p "$target_home/.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 700 -R ~/.ssh + chmod 600 "$target_home/.ssh/authorized_keys" + 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 - JUPYTER_ARGS="$JUPYTER_ARGS --allow-root" + if [ "$(id -u)" -eq 0 ] && id ubuntu &>/dev/null; then + 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) -exec "$@" +# Execute the passed arguments (CMD) as ubuntu when possible +if [ "$(id -u)" -eq 0 ] && id ubuntu &>/dev/null; then + exec sudo -u ubuntu "$@" +else + exec "$@" +fi