Skip to content

Commit

Permalink
Handle dynamic IP address for Vast.ai
Browse files Browse the repository at this point in the history
  • Loading branch information
robballantyne committed Aug 12, 2024
1 parent 29d5c99 commit edee568
Show file tree
Hide file tree
Showing 10 changed files with 109 additions and 45 deletions.
14 changes: 9 additions & 5 deletions build/COPY_ROOT_0/opt/ai-dock/bin/direct-url.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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}"
Expand All @@ -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
Expand Down
1 change: 1 addition & 0 deletions build/COPY_ROOT_0/opt/ai-dock/bin/external-ip-address
24 changes: 24 additions & 0 deletions build/COPY_ROOT_0/opt/ai-dock/bin/external-ip-address.sh
Original file line number Diff line number Diff line change
@@ -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"
34 changes: 0 additions & 34 deletions build/COPY_ROOT_0/opt/ai-dock/bin/init.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
6 changes: 4 additions & 2 deletions build/COPY_ROOT_0/opt/ai-dock/bin/supervisor-caddy.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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/*"

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}" }'
Expand Down
60 changes: 60 additions & 0 deletions build/COPY_ROOT_0/opt/caddy/public/500.html
Original file line number Diff line number Diff line change
@@ -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>
7 changes: 7 additions & 0 deletions build/COPY_ROOT_0/opt/caddy/share/base_config
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@

(universal-config) {
{$CADDY_TLS_LISTEN_STRING}
root * /opt/caddy/public
@authenticating_bearer {
expression \
{http.request.header.authorization} == "Bearer {$WEB_TOKEN}" || \
Expand Down Expand Up @@ -52,6 +53,12 @@
@preauth {
path /ai-dock/pre-auth
}

handle_errors 5xx {
rewrite * /500.html
file_server
}


@default {
path /*
Expand Down

0 comments on commit edee568

Please sign in to comment.