Integrate secure websocket cluster master into main vidai.py

- Modify ClusterMaster to accept host parameter
- Start cluster master in vidai.py when running as master
- Use --cluster-host and --cluster-port for websocket server binding
- Default to 0.0.0.0:5003 for cluster master
parent 772f6213
......@@ -44,6 +44,7 @@ from vidai.config import (
get_backend_host, set_backend_host, get_backend_web_port, set_backend_web_port,
get_backend_worker_port, set_backend_worker_port
)
from vidai.cluster_master import start_cluster_master
def main():
# Detect available GPU backends at startup
......@@ -432,6 +433,14 @@ Examples:
web_cmd = [sys.executable, '-m', 'vidai.web']
web_proc = subprocess.Popen(web_cmd)
# Start cluster master in a separate thread
import threading
cluster_host = args.cluster_host or '0.0.0.0'
cluster_thread = threading.Thread(target=start_cluster_master, args=(cluster_host, args.cluster_port))
cluster_thread.daemon = True
cluster_thread.start()
print(f"Started cluster master on {cluster_host}:{args.cluster_port}")
try:
# Wait for processes
backend_proc.wait()
......
......@@ -37,7 +37,8 @@ from collections import defaultdict
class ClusterMaster:
"""Master server for cluster coordination."""
def __init__(self, port: int = 5003, shared_dir: str = None, weight: int = 100, weight_explicit: bool = False):
def __init__(self, host: str = '0.0.0.0', port: int = 5003, shared_dir: str = None, weight: int = 100, weight_explicit: bool = False):
self.host = host
self.port = port
self.shared_dir = shared_dir
self.weight = weight # Master weight (default 100, changes to 0 when clients connect if not explicit)
......@@ -130,7 +131,7 @@ class ClusterMaster:
ssl_context = self._generate_ssl_cert()
# Start secure websocket server
start_server = websockets.serve(self._handle_client, '0.0.0.0', self.port, ssl=ssl_context)
start_server = websockets.serve(self._handle_client, self.host, self.port, ssl=ssl_context)
await start_server
# Start management loop in background
......@@ -669,7 +670,7 @@ class ClusterMaster:
cluster_master = ClusterMaster()
def start_cluster_master(port: int = 5003, shared_dir: str = None, weight=None, weight_explicit: bool = False) -> None:
def start_cluster_master(host: str = '0.0.0.0', port: int = 5003, shared_dir: str = None, weight=None, weight_explicit: bool = False) -> None:
"""Start the cluster master server."""
from .database import get_config
......@@ -696,6 +697,7 @@ def start_cluster_master(port: int = 5003, shared_dir: str = None, weight=None,
else:
weight_explicit = True
cluster_master.host = host
cluster_master.port = port
cluster_master.shared_dir = shared_dir
cluster_master.weight = weight
......@@ -707,6 +709,7 @@ if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description='VidAI Cluster Master')
parser.add_argument('--host', default='0.0.0.0', help='Host to bind to (default: 0.0.0.0)')
parser.add_argument('--port', type=int, default=5003, help='Port to listen on (default: 5003)')
parser.add_argument('--shared-dir', help='Shared directory for file transfers')
parser.add_argument('--weight', default='auto', help='Master weight for load balancing (default: auto, or specify number, or "auto" for automatic adjustment)')
......@@ -728,9 +731,9 @@ if __name__ == "__main__":
weight_explicit = False
if weight_explicit and weight_value != 'auto':
print(f"Starting VidAI Cluster Master on port {args.port} with explicit weight {weight_value}")
print(f"Starting VidAI Cluster Master on {args.host}:{args.port} with explicit weight {weight_value}")
else:
print(f"Starting VidAI Cluster Master on port {args.port} with {'default ' if not weight_explicit else ''}weight {weight_value} (will auto-adjust when clients connect)")
print(f"Starting VidAI Cluster Master on {args.host}:{args.port} with {'default ' if not weight_explicit else ''}weight {weight_value} (will auto-adjust when clients connect)")
if args.shared_dir:
print(f"Using shared directory: {args.shared_dir}")
......@@ -739,4 +742,4 @@ if __name__ == "__main__":
print(f"Warning: Shared directory {args.shared_dir} does not exist. Creating it.")
os.makedirs(args.shared_dir, exist_ok=True)
start_cluster_master(args.port, args.shared_dir, weight_value, weight_explicit)
\ No newline at end of file
start_cluster_master(args.host, args.port, args.shared_dir, weight_value, weight_explicit)
\ No newline at end of file
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment