298 lines
11 KiB
Python
298 lines
11 KiB
Python
# websocket_server/worker_manager.py
|
|
|
|
import requests
|
|
import uuid
|
|
import time
|
|
import threading
|
|
import logging
|
|
import json
|
|
from websocket_server.shared_state import connected_workers, worker_lock, ping_tracker
|
|
from websocket_server.config import get_redis_client,PING_INTERVAL_SECONDS, ASSIGN_INTERVAL_SECONDS, LIVENESS_CHECK_INTERVAL_SECONDS, PING_EXPIRY_SECONDS, ENABLE_WEBSOCKET_PING_DEBUG, ENABLE_TASK_ASSIGNMENT_DEBUG
|
|
from websocket_server.task_assigner import assign_task_to_worker
|
|
from websocket_server.events import base
|
|
from websocket_server.models import Task
|
|
from sqlalchemy.orm import sessionmaker
|
|
from websocket_server.config import engine
|
|
|
|
|
|
logger = logging.getLogger("websocket_server")
|
|
|
|
# API server endpoint
|
|
api_server_url = "http://172.17.0.1:5000/api"
|
|
|
|
# Tracks threads for each worker
|
|
worker_threads = {}
|
|
|
|
redis = get_redis_client()
|
|
|
|
# Create session factory for database operations
|
|
Session = sessionmaker(bind=engine)
|
|
|
|
# Last timestamps
|
|
last_ping_time = 0
|
|
last_assign_time = 0
|
|
|
|
|
|
def notify_worker_online(worker_id):
|
|
"""
|
|
Inform the API server that a worker has connected and start reconciliation.
|
|
"""
|
|
logger.debug(f"[{worker_id}] Notifying API server: online")
|
|
payload = {"status": "online"}
|
|
headers = {"Content-Type": "application/json"}
|
|
|
|
try:
|
|
response = requests.put(f"{api_server_url}/workload_hosts/{worker_id}", json=payload, headers=headers)
|
|
logger.info(f"[{worker_id}] API server acknowledged online state")
|
|
|
|
# Start host reconciliation process
|
|
start_host_reconciliation(worker_id)
|
|
|
|
except Exception as e:
|
|
logger.error(f"[{worker_id}] Failed to notify API server of online status: {e}")
|
|
|
|
|
|
def notify_worker_disconnect(worker_id):
|
|
"""
|
|
Inform the API server that a worker has disconnected.
|
|
"""
|
|
logger.debug(f"[{worker_id}] Notifying API server: offline")
|
|
payload = {"status": "offline"}
|
|
headers = {"Content-Type": "application/json"}
|
|
|
|
try:
|
|
response = requests.put(f"{api_server_url}/workload_hosts/{worker_id}", json=payload, headers=headers)
|
|
logger.debug(response.text)
|
|
logger.info(f"[{worker_id}] API server acknowledged offline state")
|
|
except Exception as e:
|
|
logger.error(f"[{worker_id}] Failed to notify API server of disconnect: {e}")
|
|
|
|
|
|
def start_host_reconciliation(worker_id):
|
|
"""
|
|
Start the host reconciliation process when a worker connects.
|
|
This moves the host status to 'reconciling' and initiates the reconciliation workflow.
|
|
"""
|
|
logger.info(f"[{worker_id}] Starting host reconciliation process")
|
|
|
|
try:
|
|
# Step 1: Move host status to 'reconciling'
|
|
update_host_status(worker_id, "reconciling")
|
|
|
|
# Step 2: Get containers that should be on this host
|
|
expected_containers = get_expected_containers_for_host(worker_id)
|
|
|
|
# Step 3: Send pod-update tasks for each pod containing containers for this host
|
|
send_pod_update_tasks(worker_id, expected_containers)
|
|
|
|
# Step 4: Send reconcile_and_delete task
|
|
send_reconcile_and_delete_task(worker_id, expected_containers)
|
|
|
|
logger.info(f"[{worker_id}] Host reconciliation process initiated")
|
|
|
|
except Exception as e:
|
|
logger.error(f"[{worker_id}] Error during host reconciliation: {e}")
|
|
# If reconciliation fails, move host back to online status
|
|
update_host_status(worker_id, "online")
|
|
|
|
|
|
def update_host_status(worker_id, status):
|
|
"""
|
|
Update the host status in the API server.
|
|
"""
|
|
logger.debug(f"[{worker_id}] Updating host status to: {status}")
|
|
payload = {"status": status}
|
|
headers = {"Content-Type": "application/json"}
|
|
|
|
try:
|
|
response = requests.put(f"{api_server_url}/workload_hosts/{worker_id}", json=payload, headers=headers)
|
|
if response.status_code == 200:
|
|
logger.info(f"[{worker_id}] Host status updated to: {status}")
|
|
else:
|
|
logger.warning(f"[{worker_id}] Failed to update host status: {response.status_code} - {response.text}")
|
|
except Exception as e:
|
|
logger.error(f"[{worker_id}] Error updating host status: {e}")
|
|
|
|
|
|
def get_expected_containers_for_host(worker_id):
|
|
"""
|
|
Get all containers that should be running on this host from the API server.
|
|
"""
|
|
logger.debug(f"[{worker_id}] Fetching expected containers for host")
|
|
|
|
try:
|
|
response = requests.get(f"{api_server_url}/workload_hosts/{worker_id}/container_workloads")
|
|
if response.status_code == 200:
|
|
data = response.json()
|
|
if data.get('success') and 'data' in data:
|
|
containers_data = data['data']
|
|
logger.info(f"[{worker_id}] Found {len(containers_data.get('job_details', {}).get('containers', []))} expected containers")
|
|
return containers_data
|
|
else:
|
|
logger.warning(f"[{worker_id}] API response format unexpected: {data}")
|
|
return {"job_details": {"containers": []}}
|
|
else:
|
|
logger.error(f"[{worker_id}] Failed to fetch containers: {response.status_code}")
|
|
return {"job_details": {"containers": []}}
|
|
except Exception as e:
|
|
logger.error(f"[{worker_id}] Error fetching expected containers: {e}")
|
|
return {"job_details": {"containers": []}}
|
|
|
|
|
|
def send_pod_update_tasks(worker_id, expected_containers):
|
|
"""
|
|
Send pod-update tasks for each pod that contains containers for this host.
|
|
"""
|
|
containers = expected_containers.get('job_details', {}).get('containers', [])
|
|
|
|
# Group containers by pod_id
|
|
pods = {}
|
|
for container in containers:
|
|
pod_id = container.get('pod_id', 'no-pod')
|
|
if pod_id not in pods:
|
|
pods[pod_id] = []
|
|
pods[pod_id].append(container)
|
|
|
|
logger.info(f"[{worker_id}] Sending pod-update tasks for {len(pods)} pods")
|
|
|
|
# Create a task for each pod
|
|
for pod_id, pod_containers in pods.items():
|
|
try:
|
|
task = Task(
|
|
worker_id=worker_id,
|
|
task_type="pod-update",
|
|
job_details=json.dumps({
|
|
"pod_id": pod_id,
|
|
"containers": pod_containers
|
|
}),
|
|
status="pending"
|
|
)
|
|
|
|
session = Session()
|
|
session.add(task)
|
|
session.commit()
|
|
session.close()
|
|
|
|
logger.debug(f"[{worker_id}] Created pod-update task for pod {pod_id} with {len(pod_containers)} containers")
|
|
|
|
except Exception as e:
|
|
logger.error(f"[{worker_id}] Error creating pod-update task for pod {pod_id}: {e}")
|
|
|
|
|
|
def send_reconcile_and_delete_task(worker_id, expected_containers):
|
|
"""
|
|
Send the final reconcile_and_delete task with the list of expected container IDs.
|
|
"""
|
|
containers = expected_containers.get('job_details', {}).get('containers', [])
|
|
expected_container_ids = [container.get('container_id') for container in containers if container.get('container_id')]
|
|
|
|
logger.info(f"[{worker_id}] Creating reconcile_and_delete task for {len(expected_container_ids)} containers")
|
|
|
|
try:
|
|
task = Task(
|
|
worker_id=worker_id,
|
|
task_type="reconcile_and_delete",
|
|
job_details=json.dumps({
|
|
"expected_container_ids": expected_container_ids
|
|
}),
|
|
status="pending"
|
|
)
|
|
|
|
session = Session()
|
|
session.add(task)
|
|
session.commit()
|
|
session.close()
|
|
|
|
logger.debug(f"[{worker_id}] Created reconcile_and_delete task")
|
|
|
|
except Exception as e:
|
|
logger.error(f"[{worker_id}] Error creating reconcile_and_delete task: {e}")
|
|
|
|
|
|
def handle_reconcile_and_delete_completion(worker_id):
|
|
"""
|
|
Handle the completion of the reconcile_and_delete task by moving host to online status.
|
|
"""
|
|
logger.info(f"[{worker_id}] Reconcile and delete completed, moving host to online status")
|
|
update_host_status(worker_id, "online")
|
|
|
|
|
|
def worker_dispatch_flag_check():
|
|
from websocket_server.redis_utils import redis_subscribe, thread_stop_flags, worker_dispatch_flags
|
|
|
|
"""
|
|
Continuously checks for workers with dispatch flags and triggers task assignment.
|
|
"""
|
|
logger.info("Worker dispatch flag thread started")
|
|
while True:
|
|
for wid in list(worker_dispatch_flags.keys()):
|
|
if ENABLE_TASK_ASSIGNMENT_DEBUG:
|
|
logger.debug(f"[{wid}] Dispatch flag set. Triggering task assign.")
|
|
del worker_dispatch_flags[wid]
|
|
assign_task_to_worker(wid)
|
|
threading.Event().wait(0.01)
|
|
|
|
|
|
def all_worker_watchdog():
|
|
"""
|
|
Watchdog loop that:
|
|
- Sends a liveness ping every `PING_INTERVAL_SECONDS`
|
|
- Triggers task reassignment every `ASSIGN_INTERVAL_SECONDS`
|
|
"""
|
|
global last_ping_time, last_assign_time
|
|
logger.info("Watchdog thread started")
|
|
|
|
redis = get_redis_client()
|
|
last_liveness_check = 0
|
|
|
|
while True:
|
|
now = time.time()
|
|
worker_ids = list(connected_workers.keys())
|
|
|
|
# 1. Send custom ping to all workers
|
|
if now - last_ping_time >= PING_INTERVAL_SECONDS:
|
|
if worker_ids:
|
|
if ENABLE_WEBSOCKET_PING_DEBUG:
|
|
logger.debug("Sending liveness pings to all workers")
|
|
for wid in worker_ids:
|
|
ping_id = str(uuid.uuid4())
|
|
ping_tracker[wid] = {"ping_id": ping_id, "sent_time": now}
|
|
sid = connected_workers.get(wid)
|
|
if sid:
|
|
try:
|
|
base.socketio.emit("worker_ping", {"ping_id": ping_id}, to=sid)
|
|
if ENABLE_WEBSOCKET_PING_DEBUG:
|
|
logger.debug(f"[{wid}] Sent ping ID: {ping_id} to {wid}")
|
|
except Exception as e:
|
|
logger.warning(f"[{wid}] Ping send failed: {e}")
|
|
last_ping_time = now
|
|
|
|
# 2. Trigger task assignment if due
|
|
if now - last_assign_time >= ASSIGN_INTERVAL_SECONDS:
|
|
if worker_ids:
|
|
if ENABLE_TASK_ASSIGNMENT_DEBUG:
|
|
logger.debug("Checking for task assignment across all workers")
|
|
for wid in worker_ids:
|
|
if ENABLE_TASK_ASSIGNMENT_DEBUG:
|
|
logger.debug(f"[{wid}] Triggering assign check")
|
|
assign_task_to_worker(wid)
|
|
last_assign_time = now
|
|
|
|
|
|
|
|
time.sleep(1)
|
|
|
|
|
|
def start_worker_thread(worker_id):
|
|
from websocket_server.redis_utils import redis_subscribe, thread_stop_flags, worker_dispatch_flags
|
|
|
|
"""
|
|
Launch Redis subscription thread for the given worker.
|
|
"""
|
|
logger.info(f"[{worker_id}] Starting Redis listener thread")
|
|
|
|
thread_stop_flags[worker_id] = False
|
|
sub_thread = threading.Thread(target=redis_subscribe, args=(worker_id,), daemon=True)
|
|
worker_threads[worker_id] = sub_thread
|
|
sub_thread.start()
|