From 4c4ebdb23645894444565fff5f21494f20473cf7 Mon Sep 17 00:00:00 2001
From: Rob Ballantyne <rob@dynamedia.uk>
Date: Thu, 1 Aug 2024 20:02:23 +0100
Subject: [PATCH] Improve startup logic

---
 .github/workflows/docker-build.yml            | 58 ++++++++++++-------
 README.md                                     |  2 +-
 .../supervisord/conf.d/cloudflared.conf       |  2 +-
 build/COPY_ROOT_0/opt/ai-dock/bin/init.sh     | 26 ++++++---
 .../opt/ai-dock/bin/supervisor-sshd.sh        |  5 --
 .../opt/ai-dock/bin/supervisor-syncthing.sh   |  5 --
 .../opt/ai-dock/etc/environment.sh            |  3 +
 build/Dockerfile                              |  1 +
 8 files changed, 61 insertions(+), 41 deletions(-)

diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml
index cfbfd15..c9f5fa9 100644
--- a/.github/workflows/docker-build.yml
+++ b/.github/workflows/docker-build.yml
@@ -34,19 +34,26 @@ jobs:
       -
         name: Env Setter
         run: |
-          echo "PACKAGE_NAME=${GITHUB_REPOSITORY,,}" >> ${GITHUB_ENV}
+          REPO=${GITHUB_REPOSITORY,,}
+          echo "REPO_NAMESPACE=${REPO%%/*}" >> ${GITHUB_ENV}
+          echo "REPO_NAME=${REPO#*/}" >> ${GITHUB_ENV}
       -
         name: Checkout
         uses: actions/checkout@v3
       -
         name: Permissions fixes
         run: |
-          reponame="$(basename ${GITHUB_REPOSITORY})"
-          target="${HOME}/work/${reponame}/${reponame}/build/COPY*"
+          target="${HOME}/work/${{ env.REPO_NAME }}/${{ env.REPO_NAME }}/build/COPY*"
           chmod -R ug+rwX ${target}
+      -
+        name: Login to DockerHub
+        uses: docker/login-action@v3
+        with:
+          username: ${{ secrets.DOCKERHUB_USER }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
       -
         name: Login to GitHub Container Registry
-        uses: docker/login-action@v2
+        uses: docker/login-action@v3
         with:
           registry: ghcr.io
           username: ${{ github.actor }}
@@ -54,15 +61,16 @@ jobs:
       -
         name: Set tags
         run: |
-          img_path="ghcr.io/${{ env.PACKAGE_NAME }}"
-          
+          img_path_ghcr="ghcr.io/${{ env.REPO_NAMESPACE }}/${{ env.REPO_NAME }}"
+          img_path_dhub="${{ secrets.DOCKERHUB_USER }}/${{ env.REPO_NAME }}"
           base_tag="v2-cpu-${{ env.UBUNTU_VERSION }}"
 
           if [[ ${{ matrix.build.latest }} == "true" ]]; then
               echo "Marking latest"
-              TAGS="${img_path}:${base_tag}, ${img_path}:latest-cpu"
-          else
-              TAGS="${img_path}:${base_tag}"
+              TAGS="${img_path_ghcr}:${base_tag}, ${img_path_ghcr}:latest-cpu" 
+              TAGS="${TAGS}, ${img_path_dhub}:${base_tag}, ${img_path_dhub}:latest-cpu"
+          else  
+              TAGS="${img_path_ghcr}:${base_tag}, ${img_path_dhub}:${base_tag}"
           fi
           echo "TAGS=${TAGS}" >> ${GITHUB_ENV}
       -
@@ -106,15 +114,16 @@ jobs:
       -
         name: Env Setter
         run: |
-          echo "PACKAGE_NAME=${GITHUB_REPOSITORY,,}" >> ${GITHUB_ENV}
+          REPO=${GITHUB_REPOSITORY,,}
+          echo "REPO_NAMESPACE=${REPO%%/*}" >> ${GITHUB_ENV}
+          echo "REPO_NAME=${REPO#*/}" >> ${GITHUB_ENV}
       -
         name: Checkout
         uses: actions/checkout@v3
       -
         name: Permissions fixes
         run: |
-          reponame="$(basename ${GITHUB_REPOSITORY})"
-          target="${HOME}/work/${reponame}/${reponame}/build/COPY*"
+          target="${HOME}/work/${{ env.REPO_NAME }}/${{ env.REPO_NAME }}/build/COPY*"
           chmod -R ug+rwX ${target}
       -
         name: Login to GitHub Container Registry
@@ -126,15 +135,17 @@ jobs:
       -
         name: Set tags
         run: |
-          img_path="ghcr.io/${{ env.PACKAGE_NAME }}"
+          img_path_ghcr="ghcr.io/${{ env.REPO_NAMESPACE }}/${{ env.REPO_NAME }}"
+          img_path_dhub="${{ secrets.DOCKERHUB_USER }}/${{ env.REPO_NAME }}"
           
           base_tag="v2-cuda-${{ matrix.build.cuda }}-${{ env.UBUNTU_VERSION }}"
 
           if [[ ${{ matrix.build.latest }} == "true" ]]; then
               echo "Marking latest"
-              TAGS="${img_path}:${base_tag}, ${img_path}:latest-cuda, ${img_path}:latest"
-          else
-              TAGS="${img_path}:${base_tag}"
+              TAGS="${img_path_ghcr}:${base_tag}, ${img_path_ghcr}:latest-cuda"
+              TAGS="${TAGS}, ${img_path_dhub}:${base_tag}, ${img_path_dhub}:latest-cuda"
+          else  
+              TAGS="${img_path_ghcr}:${base_tag}, ${img_path_dhub}:${base_tag}"
           fi
           echo "TAGS=${TAGS}" >> ${GITHUB_ENV}
       -
@@ -178,15 +189,16 @@ jobs:
       -
         name: Env Setter
         run: |
-          echo "PACKAGE_NAME=${GITHUB_REPOSITORY,,}" >> ${GITHUB_ENV}
+          REPO=${GITHUB_REPOSITORY,,}
+          echo "REPO_NAMESPACE=${REPO%%/*}" >> ${GITHUB_ENV}
+          echo "REPO_NAME=${REPO#*/}" >> ${GITHUB_ENV}
       -
         name: Checkout
         uses: actions/checkout@v3
       -
         name: Permissions fixes
         run: |
-          reponame="$(basename ${GITHUB_REPOSITORY})"
-          target="${HOME}/work/${reponame}/${reponame}/build/COPY*"
+          target="${HOME}/work/${{ env.REPO_NAME }}/${{ env.REPO_NAME }}/build/COPY*"
           chmod -R ug+rwX ${target}
       -
         name: Login to GitHub Container Registry
@@ -198,13 +210,15 @@ jobs:
       -
         name: Set tags
         run: |
-          img_path="ghcr.io/${{ env.PACKAGE_NAME }}"
-          
+          img_path_ghcr="ghcr.io/${{ env.REPO_NAMESPACE }}/${{ env.REPO_NAME }}"
+          img_path_dhub="${{ secrets.DOCKERHUB_USER }}/${{ env.REPO_NAME }}"
+
           base_tag="v2-rocm-${{ matrix.build.rocm }}-${{ env.UBUNTU_VERSION }}"
 
           if [[ ${{ matrix.build.latest }} == "true" ]]; then
               echo "Marking latest"
-              TAGS="${img_path}:${base_tag}, ${img_path}:latest-rocm"
+              TAGS="${img_path_ghcr}:${base_tag}, ${img_path_ghcr}:latest-rocm"
+              TAGS="${TAGS}, ${img_path_dhub}:${base_tag}, ${img_path_dhub}:latest-rocm"
           else
               TAGS="${img_path}:${base_tag}"
           fi
diff --git a/README.md b/README.md
index ca20e01..c6b8f46 100644
--- a/README.md
+++ b/README.md
@@ -8,7 +8,7 @@ This file should form the basis for the README.md for all extended images, with
 
 ## Documentation
 
-All AI-Dock containers share a common base which is designed to make running on cloud services such as [vast.ai](https://link.ai-dock.org/vast.ai) and [runpod.io](https://link.ai-dock.org/runpod.io) as straightforward and user friendly as possible.
+All AI-Dock containers share a common base which is designed to make running on container-first cloud services such as [vast.ai](https://link.ai-dock.org/vast.ai) as straightforward and user friendly as possible.
 
 Common features and options are documented in the [base wiki](https://github.com/ai-dock/base-image/wiki) but any additional features unique to this image will be detailed below.
 
diff --git a/build/COPY_ROOT_0/etc/supervisor/supervisord/conf.d/cloudflared.conf b/build/COPY_ROOT_0/etc/supervisor/supervisord/conf.d/cloudflared.conf
index 1306fb4..3842a1f 100644
--- a/build/COPY_ROOT_0/etc/supervisor/supervisord/conf.d/cloudflared.conf
+++ b/build/COPY_ROOT_0/etc/supervisor/supervisord/conf.d/cloudflared.conf
@@ -3,7 +3,7 @@ user=$USER_NAME
 environment=PROC_NAME="%(program_name)s",USER=$USER_NAME,HOME=/home/$USER_NAME
 command=supervisor-cloudflared.sh
 process_name=%(program_name)s
-numprocs=%(ENV_SUPERVISOR_START_CLOUDFLARED)s
+numprocs=1
 directory=/home/$USER_NAME
 priority=100
 autostart=true
diff --git a/build/COPY_ROOT_0/opt/ai-dock/bin/init.sh b/build/COPY_ROOT_0/opt/ai-dock/bin/init.sh
index 02d749a..80c3323 100755
--- a/build/COPY_ROOT_0/opt/ai-dock/bin/init.sh
+++ b/build/COPY_ROOT_0/opt/ai-dock/bin/init.sh
@@ -21,7 +21,7 @@ function init_main() {
     init_set_workspace
     init_count_gpus
     init_count_quicktunnels
-    init_set_cf_tunnel_wanted
+    init_toggle_supervisor_autostart
     touch /run/container_config
     touch /run/workspace_sync
     init_write_environment
@@ -44,6 +44,11 @@ function init_main() {
 }
 
 function init_set_envs() {
+    # Common services that we don't want in serverless mode
+    if [[ ${SERVERLESS,,} == "true" && -z $SUPERVISOR_NO_AUTOSTART ]]; then
+        export SUPERVISOR_NO_AUTOSTART="syncthing,jupyter,quicktunnel,cloudflared"
+    fi
+
     for i in "$@"; do
         IFS="=" read -r key val <<< "$i"
         if [[ -n $key && -n $val ]]; then
@@ -252,6 +257,7 @@ function init_create_user() {
     usermod -a -G sgx $USER_NAME
     # See the README (in)security notice
     printf "%s ALL=(ALL) NOPASSWD: ALL\n" ${USER_NAME} >> /etc/sudoers
+    sed -i 's/^Defaults[ \t]*secure_path/#Defaults secure_path/' /etc/sudoers
     if [[ ! -e ${home_dir}/.bashrc ]]; then
         cp -f /root/.bashrc ${home_dir}
         cp -f /root/.profile ${home_dir}
@@ -346,12 +352,18 @@ init_sync_opt() {
   fi
 }
 
-function init_set_cf_tunnel_wanted() {
-    if [[ -n $CF_TUNNEL_TOKEN ]]; then
-        export SUPERVISOR_START_CLOUDFLARED=1 
-    else
-        export SUPERVISOR_START_CLOUDFLARED=0
+function init_toggle_supervisor_autostart() {
+    if [[ -z $CF_TUNNEL_TOKEN ]]; then
+        SUPERVISOR_NO_AUTOSTART="${SUPERVISOR_NO_AUTOSTART:+$SUPERVISOR_NO_AUTOSTART,}cloudflared"
     fi
+
+    IFS="," read -r -a no_autostart <<< "$SUPERVISOR_NO_AUTOSTART"
+    for service in "${no_autostart[@]}"; do
+        file="/etc/supervisor/supervisord/conf.d/${service,,}.conf"
+        if [[ -f $file ]]; then
+            sed -i '/^autostart=/c\autostart=false' $file
+        fi
+    done
 }
 
 function init_direct_address() {
@@ -408,7 +420,7 @@ function init_source_preflight_scripts() {
 
 function init_write_environment() {
     # Ensure all variables available for interactive sessions
-    printf "# RUNTIME INSTRUCTIONS\n" >> /opt/ai-dock/etc/environment.sh
+    sed -i '7,$d' /opt/ai-dock/etc/environment.sh
     while IFS='=' read -r -d '' key val; do
         if [[  $key != "HOME" ]]; then
             env-store "$key"
diff --git a/build/COPY_ROOT_0/opt/ai-dock/bin/supervisor-sshd.sh b/build/COPY_ROOT_0/opt/ai-dock/bin/supervisor-sshd.sh
index 83c4a55..17ef0b0 100755
--- a/build/COPY_ROOT_0/opt/ai-dock/bin/supervisor-sshd.sh
+++ b/build/COPY_ROOT_0/opt/ai-dock/bin/supervisor-sshd.sh
@@ -10,11 +10,6 @@ function cleanup() {
 function start() {
     source /opt/ai-dock/etc/environment.sh
     
-    if [[ ${SERVERLESS,,} = "true" ]]; then
-        printf "Refusing to start SSH service in serverless mode\n"
-        exec sleep 6
-    fi
-    
     ak_file="/root/.ssh/authorized_keys"
     if [[ ! $(ssh-keygen -l -f $ak_file) ]]; then
         printf "Skipping SSH server: No public key\n" 1>&2
diff --git a/build/COPY_ROOT_0/opt/ai-dock/bin/supervisor-syncthing.sh b/build/COPY_ROOT_0/opt/ai-dock/bin/supervisor-syncthing.sh
index b61e8a1..93b466e 100755
--- a/build/COPY_ROOT_0/opt/ai-dock/bin/supervisor-syncthing.sh
+++ b/build/COPY_ROOT_0/opt/ai-dock/bin/supervisor-syncthing.sh
@@ -27,11 +27,6 @@ function run_with_retry() {
 function start() {
     source /opt/ai-dock/etc/environment.sh
     
-    if [[ ${SERVERLESS,,} = "true" ]]; then
-        printf "Refusing to start $SERVICE_NAME in serverless mode\n"
-        exec sleep 6
-    fi
-    
     file_content="$(
       jq --null-input \
         --arg listen_port "${LISTEN_PORT}" \
diff --git a/build/COPY_ROOT_0/opt/ai-dock/etc/environment.sh b/build/COPY_ROOT_0/opt/ai-dock/etc/environment.sh
index 31fa228..f082532 100644
--- a/build/COPY_ROOT_0/opt/ai-dock/etc/environment.sh
+++ b/build/COPY_ROOT_0/opt/ai-dock/etc/environment.sh
@@ -3,3 +3,6 @@ if [[ ! -f ~/.gitconfig ]]; then
     git config --global --add safe.directory "*"
 fi
 
+# RUNTIME INSTRUCTIONS
+
+
diff --git a/build/Dockerfile b/build/Dockerfile
index 5d5377b..bab5f11 100644
--- a/build/Dockerfile
+++ b/build/Dockerfile
@@ -59,6 +59,7 @@ ENV VENV_DIR=/opt/environments/python
 ENV SERVICEPORTAL_VENV=${VENV_DIR}/serviceportal
 ENV SERVICEPORTAL_VENV_PYTHON=${SERVICEPORTAL_VENV}/bin/python
 ENV SERVICEPORTAL_VENV_PIP=${SERVICEPORTAL_VENV}/bin/pip
+ENV SUPERVISOR_NO_AUTOSTART=
 
 ENV APT_INSTALL="apt-get install -y --no-install-recommends"
 ENV RCLONE_CONFIG="/etc/rclone/rclone.conf"