Handle dynamic IP address for Vast.ai

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

build/COPY_ROOT_0/opt/ai-dock/bin/external-ip-address

@@ -0,0 +1 @@
+external-ip-address.sh

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"

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

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/*"
 

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

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
  }

build/COPY_ROOT_0/opt/ai-dock/fastapi/serviceportal/templates/partials/index/page.html

@@ -16,7 +16,7 @@ <h1 id="status">Service List
  <div>{{ service["service_name"] }}</div>
  <div>
  <ul>
- {% if context.direct_address %}
+ {% if context.direct_address == True %}
  <li id='cfqt-{{ service["proxy_port"] }}'
  hx-post="/direct"
  hx-vals='{"port": "{{ service["proxy_port"] }}", "path": "/ai-dock/pre-auth?token={{ context.auth_token }}" }'

build/COPY_ROOT_0/opt/caddy/public/500.html

@@ -0,0 +1,60 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+ <meta charset="UTF-8">
+ <meta name="viewport" content="width=device-width, initial-scale=1.0">
+ <title>Loading...</title>
+ <style>
+ body {
+ background-color: white;
+ display: flex;
+ flex-direction: column;
+ align-items: center;
+ justify-content: center;
+ height: 100vh;
+ margin: 0;
+ font-family: Arial, sans-serif;
+ }
+ .spinner {
+ border: 8px solid #F2F2F2; /* Light grey */
+ border-top: 8px solid grey; /* Dark grey */
+ border-radius: 50%;
+ width: 50px;
+ height: 50px;
+ animation: spin 1s linear infinite;
+ }
+ @keyframes spin {
+ 0% { transform: rotate(0deg); }
+ 100% { transform: rotate(360deg); }
+ }
+ .text {
+ margin-top: 20px;
+ color: #333;
+ }
+ </style>
+</head>
+<body>
+ <div class="spinner"></div>
+ <div class="text">Waiting for the service to respond</div>
+
+ <script>
+ // Function to check if the current page is up
+ function checkService() {
+ fetch(window.location.href, { method: 'HEAD' })
+ .then(response => {
+ // If the response status is not a 5xx error
+ if (response.status < 500 || response.status >= 600) {
+ // Reload the page
+ window.location.reload();
+ }
+ })
+ .catch(error => {
+ console.error('Error checking service:', error);
+ });
+ }
+
+ // Check the current page every 5 seconds
+ setInterval(checkService, 5000);
+ </script>
+</body>
+</html>

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 /*