docker_init.sh

#!/bin/bash
set -x

apt-get update
apt-get install -y ca-certificates curl
install -m 0755 -d /etc/apt/keyrings
curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc
chmod a+r /etc/apt/keyrings/docker.asc

# Add the repository to Apt sources:
echo \
  "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] https://download.docker.com/linux/ubuntu \
  $(. /etc/os-release && echo "$VERSION_CODENAME") stable" | \
  tee /etc/apt/sources.list.d/docker.list > /dev/null
apt-get update
apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
service docker start
apt-get install -y nvidia-container-toolkit-base
apt-get install -y nvidia-container-runtime
nvidia-ctk runtime configure --runtime=docker
service docker restart

apt install -y tini

# Check if OpenSSH server is already installed
if ! command -v sshd &> /dev/null; then
    echo "OpenSSH server is not installed. Installing..."

    apt update
    apt install -y openssh-server

    echo "OpenSSH server installation complete."
else
    echo "OpenSSH server is already installed."
fi

# Check if POD_PASSWORD_ACCESS variable is set to "true"
if [[ "$POD_PASSWORD_ACCESS" == "true" ]]; then
    # Enable password authentication in SSH configuration
    sed -i '/^#PasswordAuthentication/c\PasswordAuthentication yes' /etc/ssh/sshd_config
    sed -i '/^PasswordAuthentication/c\PasswordAuthentication yes' /etc/ssh/sshd_config

    # Display a message indicating that user/password SSH access is enabled
    echo "User/password SSH access is enabled."
fi

# Set root password if POD_ROOT_PASSWORD is provided
if [[ -n "$POD_ROOT_PASSWORD" ]]; then
    echo "root:${POD_ROOT_PASSWORD}" | chpasswd
    echo "Root password has been set."
fi

# Check if POD_UBLIC_KEY variable is set and not empty
if [[ -n "$POD_PUBLIC_KEY" ]]; then
    # Create the .ssh directory and authorized_keys file if they don't exist
    if [ ! -d "$HOME/.ssh" ]; then
        mkdir -p "$HOME/.ssh"
        chmod 0700 "$HOME/.ssh"
        echo "Directory $HOME/.ssh created."
    fi
    if [ ! -f "$HOME/.ssh/authorized_keys" ]; then
        touch "$HOME/.ssh/authorized_keys"
        chmod 0600 "$HOME/.ssh/authorized_keys"
        echo "File $HOME/.ssh/authorized_keys created."
    fi
    # Check if the public key is not already present in authorized_keys
    if ! grep -q "${POD_PUBLIC_KEY}" "$HOME/.ssh/authorized_keys"; then
        # Append the public key to authorized_keys
        echo "$POD_PUBLIC_KEY" >> "$HOME/.ssh/authorized_keys"
        echo "Public key from env variable added."
    fi
fi

# Check if POD_PERMIT_ROOT_LOGIN is enabled
if [[ "$POD_PERMIT_ROOT_LOGIN" == "true" ]]; then
    # Enable root login in SSH configuration
    sed -i '/^#PermitRootLogin/c\PermitRootLogin yes' /etc/ssh/sshd_config
    sed -i '/^PermitRootLogin/c\PermitRootLogin yes' /etc/ssh/sshd_config
    echo "Root login is enabled."
fi

if [[ -n "$POD_JUPYTER_PASSWORD" ]]; then
    if pgrep jupyter-lab > /dev/null 2>&1; then
      echo "jupyter already started"
      return
    fi

    # Check if jupyter lab is already installed
    if ! command -v jupyter-lab &> /dev/null; then
        echo "jupyter lab is not installed. Installing..."

        apt update
        apt install python3 python3-pip -y
        pip install -U virtualenv --break-system-packages
        pip3 install jupyterlab --break-system-packages

        echo "jupyter lab installation complete."
    else
        echo "jupyter lab is already installed."
    fi

    jupyter lab --generate-config

    {
     echo "c.ServerApp.ip = '0.0.0.0'"
     echo "c.ServerApp.open_browser = False"
     echo "c.ServerApp.port = 18889"
    } >> ~/.jupyter/jupyter_lab_config.py

    # Set root password if LEPTON_POD_ROOT_PASSWORD is provided
    if [[ -n "$POD_JUPYTER_PASSWORD" ]]; then
        echo "c.ServerApp.token = '${POD_JUPYTER_PASSWORD}'" >> ~/.jupyter/jupyter_lab_config.py
        echo "Root password has been set."
    fi

    jupyter lab --allow-root > /var/log/jupyter.log 2>&1 &
fi

# turn off PAM to fix sshd login issue
sed -i 's/UsePAM yes/UsePAM no/' /etc/ssh/sshd_config

echo "Exposing ENV variables"
env | sed 's/=/="/' | sed 's/$/"/' > /etc/environment
echo "set -a; source /etc/environment; set +a;" >> /root/.bashrc

echo "Removing the init script"
rm ./sshd_init.sh

mkdir /run/sshd
chmod 0755 /run/sshd

echo "Starting sshD"
tini -s -- /usr/sbin/sshd -D -p 2222