From edee56820e3507234a743738f42a9942506ed8d1 Mon Sep 17 00:00:00 2001 From: Rob Ballantyne Date: Mon, 12 Aug 2024 17:57:30 +0100 Subject: [PATCH] Handle dynamic IP address for Vast.ai --- .../COPY_ROOT_0/opt/ai-dock/bin/direct-url.sh | 14 +++-- .../opt/ai-dock/bin/external-ip-address | 1 + .../opt/ai-dock/bin/external-ip-address.sh | 24 ++++++++ build/COPY_ROOT_0/opt/ai-dock/bin/init.sh | 34 ----------- .../opt/ai-dock/bin/supervisor-caddy.sh | 6 +- .../ai-dock/fastapi/serviceportal/helpers.py | 4 +- .../opt/ai-dock/fastapi/serviceportal/main.py | 2 +- .../templates/partials/index/page.html | 2 +- build/COPY_ROOT_0/opt/caddy/public/500.html | 60 +++++++++++++++++++ build/COPY_ROOT_0/opt/caddy/share/base_config | 7 +++ 10 files changed, 109 insertions(+), 45 deletions(-) create mode 120000 build/COPY_ROOT_0/opt/ai-dock/bin/external-ip-address create mode 100755 build/COPY_ROOT_0/opt/ai-dock/bin/external-ip-address.sh create mode 100644 build/COPY_ROOT_0/opt/caddy/public/500.html diff --git a/build/COPY_ROOT_0/opt/ai-dock/bin/direct-url.sh b/build/COPY_ROOT_0/opt/ai-dock/bin/direct-url.sh index 771e78d..7d2f79c 100755 --- a/build/COPY_ROOT_0/opt/ai-dock/bin/direct-url.sh +++ b/build/COPY_ROOT_0/opt/ai-dock/bin/direct-url.sh @@ -49,14 +49,18 @@ function get_url() { preset_url=$(jq -r ".service_url" "/run/http_ports/${port}") if [[ -n $preset_url ]]; then url="$preset_url" + elif [[ ${DIRECT_ADDRESS_GET_WAN,,} == "true" ]]; then + url="$(get_scheme)$(/opt/ai-dock/bin/external-ip-address):${port}" # Vast.ai - elif [[ $DIRECT_ADDRESS == "auto#vast-ai" ]]; then + elif env | grep 'VAST_TCP_PORT' > /dev/null 2>&1; then declare -n vast_mapped_port=VAST_TCP_PORT_${port} - if [[ -n $vast_mapped_port && -n $PUBLIC_IPADDR ]]; then - url="$(get_scheme)${PUBLIC_IPADDR}:${vast_mapped_port}" + if [[ -n $vast_mapped_port ]]; then + url="$(get_scheme)$(/opt/ai-dock/bin/external-ip-address):${vast_mapped_port}" + else + url="$(/opt/ai-dock/bin/cfqt-url -p $port)" fi # Runpod.io - elif [[ $DIRECT_ADDRESS == "auto#runpod-io" ]]; then + elif env | grep 'RUNPOD' > /dev/null 2>&1; then declare -n runpod_mapped_port=RUNPOD_TCP_PORT_${port} if [[ -n $runpod_mapped_port && -n $RUNPOD_PUBLIC_IP ]]; then url="$(get_scheme)${RUNPOD_PUBLIC_IP}:${runpod_mapped_port}" @@ -65,7 +69,7 @@ function get_url() { fi # Other cloud / local else - url="$(get_scheme)${DIRECT_ADDRESS}:${port}" + url="$(get_scheme)${DIRECT_ADDRESS:-localhost}:${port}" fi if [[ -n $url ]]; then diff --git a/build/COPY_ROOT_0/opt/ai-dock/bin/external-ip-address b/build/COPY_ROOT_0/opt/ai-dock/bin/external-ip-address new file mode 120000 index 0000000..a457ddc --- /dev/null +++ b/build/COPY_ROOT_0/opt/ai-dock/bin/external-ip-address @@ -0,0 +1 @@ +external-ip-address.sh \ No newline at end of file diff --git a/build/COPY_ROOT_0/opt/ai-dock/bin/external-ip-address.sh b/build/COPY_ROOT_0/opt/ai-dock/bin/external-ip-address.sh new file mode 100755 index 0000000..7bbf51f --- /dev/null +++ b/build/COPY_ROOT_0/opt/ai-dock/bin/external-ip-address.sh @@ -0,0 +1,24 @@ +#!/bin/bash + +current_time=$(date +%s) +cache_file="/tmp/external_ip_address" +cache_max_age=${EXTERNAL_IP_CACHE_SECS:-900} + +if [[ -f $cache_file ]]; then + file_mod_time=$(stat -c %Y "$cache_file") + if [[ $((current_time - file_mod_time)) -lt $cache_max_age ]]; then + ip=$(cat "$cache_file") + if [[ $ip =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}$ ]]; then + echo "$ip" + exit 0 + fi + fi +fi + +# Fetch new IP address +ip="$(dig whoami.cloudflare ch txt @1.1.1.1 +short | tr -d '"')" +[[ $ip =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}$ ]] || ip="$(dig myip.opendns.com @resolver1.opendns.com +short)" +[[ $ip =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}$ ]] || ip=$(curl -s ifconfig.me) + +# Cache the IP address +echo "$ip" | tee "$cache_file" 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 e1c3441..52c1e80 100755 --- a/build/COPY_ROOT_0/opt/ai-dock/bin/init.sh +++ b/build/COPY_ROOT_0/opt/ai-dock/bin/init.sh @@ -16,7 +16,6 @@ function init_main() { init_create_directories init_create_logfiles init_set_ssh_keys - init_direct_address init_set_web_config init_set_workspace init_count_gpus @@ -370,39 +369,6 @@ function init_toggle_supervisor_autostart() { done } -function init_direct_address() { - export EXTERNAL_IP_ADDRESS="$(dig +short myip.opendns.com @resolver1.opendns.com)" - if [[ -z $EXTERNAL_IP_ADDRESS ]];then - export EXTERNAL_IP_ADDRESS=$(curl -s ifconfig.me) - fi - if [[ ! -v DIRECT_ADDRESS ]]; then - DIRECT_ADDRESS="" - fi - - if [[ ${DIRECT_ADDRESS,,} == "false" ]]; then - export DIRECT_ADDRESS="" - elif [[ -z $DIRECT_ADDRESS || ${DIRECT_ADDRESS_GET_WAN,,} == 'true' ]]; then - if [[ ${DIRECT_ADDRESS_GET_WAN,,} == 'true' ]]; then - export DIRECT_ADDRESS="$EXTERNAL_IP_ADDRESS" - # Detected provider has direct connection method - elif env | grep 'VAST' > /dev/null 2>&1; then - export DIRECT_ADDRESS="auto#vast-ai" - export CLOUD_PROVIDER="vast.ai" - export EXTERNAL_IP_ADDRESS=${PUBLIC_IPADDR} - elif env | grep 'RUNPOD' > /dev/null 2>&1; then - export DIRECT_ADDRESS="auto#runpod-io" - export CLOUD_PROVIDER="runpod.io" - export EXTERNAL_IP_ADDRESS=${RUNPOD_PUBLIC_IP} - # Detected provider does not support direct connections - elif env | grep 'PAPERSPACE' > /dev/null 2>&1; then - export DIRECT_ADDRESS="" - export CLOUD_PROVIDER="paperspace.com" - else - export DIRECT_ADDRESS="localhost" - fi - fi -} - function init_create_directories() { mkdir -m 2770 -p /run/http_ports chown root.ai-dock /run/http_ports diff --git a/build/COPY_ROOT_0/opt/ai-dock/bin/supervisor-caddy.sh b/build/COPY_ROOT_0/opt/ai-dock/bin/supervisor-caddy.sh index 9a1d9a4..6d98265 100755 --- a/build/COPY_ROOT_0/opt/ai-dock/bin/supervisor-caddy.sh +++ b/build/COPY_ROOT_0/opt/ai-dock/bin/supervisor-caddy.sh @@ -7,14 +7,16 @@ function cleanup() { } function start() { + rm -f /tmp/external_ip_address source /opt/ai-dock/etc/environment.sh # Give processes time to register their ports sleep 2 - export SERVICEPORTAL_LOGIN=$(direct-url.sh -p "${SERVICEPORTAL_PORT_HOST:-1111}" -l "/login") - env-store SERVICEPORTAL_LOGIN export SERVICEPORTAL_HOME=$(direct-url.sh -p "${SERVICEPORTAL_PORT_HOST:-1111}") env-store SERVICEPORTAL_HOME + export SERVICEPORTAL_LOGIN="${SERVICEPORTAL_HOME}/login" + env-store SERVICEPORTAL_LOGIN + port_files="/run/http_ports/*" diff --git a/build/COPY_ROOT_0/opt/ai-dock/fastapi/serviceportal/helpers.py b/build/COPY_ROOT_0/opt/ai-dock/fastapi/serviceportal/helpers.py index 71c8a96..addaf23 100644 --- a/build/COPY_ROOT_0/opt/ai-dock/fastapi/serviceportal/helpers.py +++ b/build/COPY_ROOT_0/opt/ai-dock/fastapi/serviceportal/helpers.py @@ -50,8 +50,8 @@ def get_cfqt_url(port, path=""): def get_direct_url(port, path=""): try: - process = subprocess.run(['direct-url.sh', '-p', port, '-l', path], - stdout=subprocess.PIPE, + process = subprocess.run(['direct-url.sh', '-p', port, '-l', path], + stdout=subprocess.PIPE, universal_newlines=True) output = process.stdout.strip() scheme = urlparse(output).scheme diff --git a/build/COPY_ROOT_0/opt/ai-dock/fastapi/serviceportal/main.py b/build/COPY_ROOT_0/opt/ai-dock/fastapi/serviceportal/main.py index 48f598a..b1ffe61 100644 --- a/build/COPY_ROOT_0/opt/ai-dock/fastapi/serviceportal/main.py +++ b/build/COPY_ROOT_0/opt/ai-dock/fastapi/serviceportal/main.py @@ -84,7 +84,7 @@ def get_index_context(request, message=None): "auth_token": request.cookies.get(os.environ.get('CADDY_AUTH_COOKIE_NAME')), "services": services, "urlslug": os.environ.get('IMAGE_SLUG'), - "direct_address": os.environ.get('DIRECT_ADDRESS'), + "direct_address": False if os.environ.get('DIRECT_ADDRESS', 'true').lower() == "false" else True, 'quicktunnels': False if os.environ.get('CF_QUICK_TUNNELS', 'true').lower() == "false" else True, 'namedtunnels': False if not os.environ.get('CF_TUNNEL_TOKEN') else True } diff --git a/build/COPY_ROOT_0/opt/ai-dock/fastapi/serviceportal/templates/partials/index/page.html b/build/COPY_ROOT_0/opt/ai-dock/fastapi/serviceportal/templates/partials/index/page.html index e3cd1d0..be19757 100644 --- a/build/COPY_ROOT_0/opt/ai-dock/fastapi/serviceportal/templates/partials/index/page.html +++ b/build/COPY_ROOT_0/opt/ai-dock/fastapi/serviceportal/templates/partials/index/page.html @@ -16,7 +16,7 @@

Service List
{{ service["service_name"] }}
    - {% if context.direct_address %} + {% if context.direct_address == True %}
  • + + + + + Loading... + + + + +
    Waiting for the service to respond
    + + + + diff --git a/build/COPY_ROOT_0/opt/caddy/share/base_config b/build/COPY_ROOT_0/opt/caddy/share/base_config index 967d2e9..5712978 100644 --- a/build/COPY_ROOT_0/opt/caddy/share/base_config +++ b/build/COPY_ROOT_0/opt/caddy/share/base_config @@ -25,6 +25,7 @@ (universal-config) { {$CADDY_TLS_LISTEN_STRING} + root * /opt/caddy/public @authenticating_bearer { expression \ {http.request.header.authorization} == "Bearer {$WEB_TOKEN}" || \ @@ -52,6 +53,12 @@ @preauth { path /ai-dock/pre-auth } + + handle_errors 5xx { + rewrite * /500.html + file_server + } + @default { path /*