https handling in host field

This commit is contained in:
robertvoy
2025-12-30 12:11:12 +11:00
parent 9805049268
commit b768e7fc8a
3 changed files with 71 additions and 19 deletions
+22 -9
View File
@@ -4,6 +4,7 @@ from PIL import Image
import folder_paths
import os
import json
import re
import asyncio
import aiohttp
from aiohttp import web
@@ -60,6 +61,17 @@ def cleanup():
atexit.register(cleanup)
def normalize_host(value):
if value is None:
return None
if not isinstance(value, str):
return value
host = value.strip()
if not host:
return host
host = re.sub(r"^https?://", "", host, flags=re.IGNORECASE)
return host.split("/")[0]
# --- API Endpoints ---
@server.PromptServer.instance.routes.get("/distributed/config")
async def get_config_endpoint(request):
@@ -367,7 +379,7 @@ async def update_worker_endpoint(request):
if data["host"] is None:
worker.pop("host", None)
else:
worker["host"] = data["host"]
worker["host"] = normalize_host(data["host"])
# Handle cuda_device field - remove it if None
if "cuda_device" in data:
@@ -396,7 +408,7 @@ async def update_worker_endpoint(request):
new_worker = {
"id": worker_id,
"name": data["name"],
"host": data.get("host", "localhost"),
"host": normalize_host(data.get("host", "localhost")),
"port": data["port"],
"cuda_device": data["cuda_device"],
"enabled": data.get("enabled", False),
@@ -673,7 +685,8 @@ async def get_local_worker_status_endpoint(request):
for worker in config.get("workers", []):
# Only check local workers
if not worker.get("host") or worker.get("host") in ["localhost", "127.0.0.1"]:
host = normalize_host(worker.get("host")) or ""
if not host or host in ["localhost", "127.0.0.1"]:
worker_id = worker["id"]
port = worker["port"]
@@ -1324,7 +1337,7 @@ async def get_client_session():
# Local worker detection functions
async def is_local_worker(worker_config):
"""Check if a worker is running on the same machine as the master."""
host = worker_config.get('host', 'localhost')
host = normalize_host(worker_config.get('host', 'localhost')) or 'localhost'
if host in ['localhost', '127.0.0.1', '0.0.0.0', ''] or worker_config.get('type') == 'local':
return True
@@ -1341,7 +1354,7 @@ async def is_same_physical_host(worker_config):
master_machine_id = get_machine_id()
# Fetch worker's machine ID via API
host = worker_config.get('host', 'localhost')
host = normalize_host(worker_config.get('host', 'localhost')) or 'localhost'
port = worker_config.get('port', 8188)
session = await get_client_session()
@@ -1404,11 +1417,11 @@ async def get_comms_channel(worker_id, worker_config):
return f"http://127.0.0.1:{worker_config['port']}"
elif worker_config.get('type') == 'cloud' and is_runpod_environment():
# Runpod same-host optimization (if detected)
host = worker_config.get('host', 'localhost')
host = normalize_host(worker_config.get('host', 'localhost')) or 'localhost'
return f"http://{host}:{worker_config['port']}"
else:
# Remote worker: use configured endpoint
host = worker_config.get('host', 'localhost')
host = normalize_host(worker_config.get('host', 'localhost')) or 'localhost'
return f"http://{host}:{worker_config['port']}"
# Auto-launch workers if enabled
@@ -1435,7 +1448,7 @@ def auto_launch_workers():
worker_name = worker.get('name', f'Worker {worker_id}')
# Skip remote workers
host = worker.get('host', 'localhost').lower()
host = (normalize_host(worker.get('host', 'localhost')) or 'localhost').lower()
if host not in ['localhost', '127.0.0.1', '', None]:
debug_log(f"Skipping remote worker {worker_name} (host: {host})")
continue
@@ -2121,7 +2134,7 @@ class DistributedCollectorNode:
if not wrec:
debug_log(f"Collector probe: worker {wid} not found in config")
continue
host = wrec.get('host') or 'localhost'
host = normalize_host(wrec.get('host') or 'localhost') or 'localhost'
port = int(wrec.get('port', 8188))
url = f"http://{host}:{port}/prompt"
try:
+15 -3
View File
@@ -30,6 +30,14 @@ PUBLIC_URL_PATTERN = re.compile(r"(https?://[\w.-]+\.(?:trycloudflare\.com|cloud
TUNNEL_START_TIMEOUT = 25
def _normalize_host(value):
if not value or not isinstance(value, str):
return ""
value = value.strip()
value = re.sub(r"^https?://", "", value, flags=re.IGNORECASE)
return value.split("/")[0]
class CloudflareTunnelManager:
def __init__(self):
self.process = None
@@ -300,13 +308,14 @@ class CloudflareTunnelManager:
raise RuntimeError(error_msg)
debug_log(f"Cloudflare tunnel ready at {self.public_url}")
master_host = _normalize_host(self.public_url)
self._persist_state(
status="running",
public_url=self.public_url,
pid=self.pid,
log_file=self.log_file,
previous_host=self.previous_master_host or "",
master_host=self.public_url
master_host=master_host
)
return {
"status": self.status,
@@ -339,8 +348,11 @@ class CloudflareTunnelManager:
active_url = tunnel_cfg.get("public_url")
current_master_host = (config.get("master") or {}).get("host")
restore_host = None
if active_url and current_master_host == active_url:
restore_host = self.previous_master_host or ""
if active_url:
active_host = _normalize_host(active_url)
current_host = _normalize_host(current_master_host)
if current_host == active_host:
restore_host = self.previous_master_host or ""
self.status = "stopped"
self.public_url = None
+34 -7
View File
@@ -158,6 +158,24 @@ class DistributedExtension {
}
}
_parseHostInput(value) {
if (!value) {
return { host: "", port: null };
}
let cleaned = value.trim().replace(/^https?:\/\//i, "");
cleaned = cleaned.split("/")[0];
try {
const url = new URL(`http://${cleaned}`);
const port = url.port ? parseInt(url.port, 10) : null;
return {
host: url.hostname || cleaned,
port: Number.isFinite(port) ? port : null,
};
} catch (error) {
return { host: cleaned, port: null };
}
}
updateTunnelUIElements() {
const elements = this.tunnelElements || {};
const status = (this.tunnelStatus?.status || "stopped").toLowerCase();
@@ -501,7 +519,9 @@ class DistributedExtension {
// Helper to build worker URL
getWorkerUrl(worker, endpoint = '') {
const host = worker.host || window.location.hostname;
const parsed = this._parseHostInput(worker.host || window.location.hostname);
const host = parsed.host || window.location.hostname;
const resolvedPort = parsed.port || worker.port;
// Cloud workers always use HTTPS
const isCloud = worker.type === 'cloud';
@@ -524,13 +544,13 @@ class DistributedExtension {
}
// Determine protocol: HTTPS for cloud, Runpod proxies, or port 443
const useHttps = isCloud || isRunpodProxy || worker.port === 443;
const useHttps = isCloud || isRunpodProxy || resolvedPort === 443;
const protocol = useHttps ? 'https' : 'http';
// Only add port if non-standard
const defaultPort = useHttps ? 443 : 80;
const needsPort = !isRunpodProxy && worker.port !== defaultPort;
const portStr = needsPort ? `:${worker.port}` : '';
const needsPort = !isRunpodProxy && resolvedPort !== defaultPort;
const portStr = needsPort ? `:${resolvedPort}` : '';
return `${protocol}://${finalHost}${portStr}${endpoint}`;
}
@@ -974,7 +994,8 @@ class DistributedExtension {
return true;
}
// Otherwise check by host (backward compatibility)
const host = worker.host || window.location.hostname;
const parsed = this._parseHostInput(worker.host || window.location.hostname);
const host = parsed.host || window.location.hostname;
return host !== "localhost" && host !== "127.0.0.1" && host !== window.location.hostname;
}
@@ -1194,10 +1215,16 @@ class DistributedExtension {
const workerType = document.getElementById(`worker-type-${workerId}`).value;
const isRemote = workerType === 'remote' || workerType === 'cloud';
const isCloud = workerType === 'cloud';
const host = isRemote ? document.getElementById(`host-${workerId}`).value : window.location.hostname;
const port = parseInt(document.getElementById(`port-${workerId}`).value);
const rawHost = isRemote ? document.getElementById(`host-${workerId}`).value : window.location.hostname;
const parsedHost = isRemote ? this._parseHostInput(rawHost) : { host: window.location.hostname, port: null };
const host = isRemote ? parsedHost.host : window.location.hostname;
let port = parseInt(document.getElementById(`port-${workerId}`).value);
const cudaDevice = isRemote ? undefined : parseInt(document.getElementById(`cuda-${workerId}`).value);
const extraArgs = isRemote ? undefined : document.getElementById(`args-${workerId}`).value;
if (isRemote && Number.isFinite(parsedHost.port)) {
port = parsedHost.port;
}
// Validate
if (!name.trim()) {