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 ( ...@@ -44,6 +44,7 @@ from vidai.config import (
get_backend_host, set_backend_host, get_backend_web_port, set_backend_web_port, get_backend_host, set_backend_host, get_backend_web_port, set_backend_web_port,
get_backend_worker_port, set_backend_worker_port get_backend_worker_port, set_backend_worker_port
) )
from vidai.cluster_master import start_cluster_master
def main(): def main():
# Detect available GPU backends at startup # Detect available GPU backends at startup
...@@ -432,6 +433,14 @@ Examples: ...@@ -432,6 +433,14 @@ Examples:
web_cmd = [sys.executable, '-m', 'vidai.web'] web_cmd = [sys.executable, '-m', 'vidai.web']
web_proc = subprocess.Popen(web_cmd) 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: try:
# Wait for processes # Wait for processes
backend_proc.wait() backend_proc.wait()
......
...@@ -37,7 +37,8 @@ from collections import defaultdict ...@@ -37,7 +37,8 @@ from collections import defaultdict
class ClusterMaster: class ClusterMaster:
"""Master server for cluster coordination.""" """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.port = port
self.shared_dir = shared_dir self.shared_dir = shared_dir
self.weight = weight # Master weight (default 100, changes to 0 when clients connect if not explicit) self.weight = weight # Master weight (default 100, changes to 0 when clients connect if not explicit)
...@@ -130,7 +131,7 @@ class ClusterMaster: ...@@ -130,7 +131,7 @@ class ClusterMaster:
ssl_context = self._generate_ssl_cert() ssl_context = self._generate_ssl_cert()
# Start secure websocket server # 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 await start_server
# Start management loop in background # Start management loop in background
...@@ -669,7 +670,7 @@ class ClusterMaster: ...@@ -669,7 +670,7 @@ class ClusterMaster:
cluster_master = 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.""" """Start the cluster master server."""
from .database import get_config from .database import get_config
...@@ -696,6 +697,7 @@ def start_cluster_master(port: int = 5003, shared_dir: str = None, weight=None, ...@@ -696,6 +697,7 @@ def start_cluster_master(port: int = 5003, shared_dir: str = None, weight=None,
else: else:
weight_explicit = True weight_explicit = True
cluster_master.host = host
cluster_master.port = port cluster_master.port = port
cluster_master.shared_dir = shared_dir cluster_master.shared_dir = shared_dir
cluster_master.weight = weight cluster_master.weight = weight
...@@ -707,6 +709,7 @@ if __name__ == "__main__": ...@@ -707,6 +709,7 @@ if __name__ == "__main__":
import argparse import argparse
parser = argparse.ArgumentParser(description='VidAI Cluster Master') 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('--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('--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)') 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__": ...@@ -728,9 +731,9 @@ if __name__ == "__main__":
weight_explicit = False weight_explicit = False
if weight_explicit and weight_value != 'auto': 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: 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: if args.shared_dir:
print(f"Using shared directory: {args.shared_dir}") print(f"Using shared directory: {args.shared_dir}")
...@@ -739,4 +742,4 @@ if __name__ == "__main__": ...@@ -739,4 +742,4 @@ if __name__ == "__main__":
print(f"Warning: Shared directory {args.shared_dir} does not exist. Creating it.") print(f"Warning: Shared directory {args.shared_dir} does not exist. Creating it.")
os.makedirs(args.shared_dir, exist_ok=True) os.makedirs(args.shared_dir, exist_ok=True)
start_cluster_master(args.port, args.shared_dir, weight_value, weight_explicit) start_cluster_master(args.host, args.port, args.shared_dir, weight_value, weight_explicit)
\ No newline at end of file \ 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