coryHawkvelt
0237024d1c
Introduce `ENABLE_TASK_ASSIGNMENT_DEBUG` configuration to control verbose logging in task assignment, Redis subscription, and worker management flows. This reduces log noise by default while allowing detailed debugging when enabled.
117 lines
4.8 KiB
Python
117 lines
4.8 KiB
Python
# websocket_server/task_assigner.py
|
|
|
|
import time
|
|
from datetime import datetime
|
|
from sqlalchemy.exc import OperationalError
|
|
from sqlalchemy import or_
|
|
from websocket_server.config import get_db_session, get_redis_client
|
|
from websocket_server.models import Task
|
|
import logging
|
|
from websocket_server.shared_state import connected_workers, worker_lock
|
|
from websocket_server.events import base
|
|
from websocket_server.config import ENABLE_TASK_ASSIGNMENT_DEBUG
|
|
|
|
logger = logging.getLogger("websocket_server")
|
|
|
|
def assign_task_to_worker(worker_id):
|
|
"""
|
|
Attempt to assign a pending task to a connected worker.
|
|
Ensures locking, dependency satisfaction, and Redis status tracking.
|
|
"""
|
|
redis_client = get_redis_client()
|
|
max_lock_retries = 5
|
|
lock_retry_wait = 0.2
|
|
|
|
if ENABLE_TASK_ASSIGNMENT_DEBUG:
|
|
logger.debug(f"[{worker_id}] Starting task assignment process")
|
|
|
|
try:
|
|
redis_client.set(f"worker_status_{worker_id}", "busy")
|
|
if ENABLE_TASK_ASSIGNMENT_DEBUG:
|
|
logger.debug(f"[{worker_id}] Marked as 'busy' in Redis")
|
|
|
|
for attempt in range(max_lock_retries):
|
|
if ENABLE_TASK_ASSIGNMENT_DEBUG:
|
|
logger.debug(f"[{worker_id}] Lock attempt {attempt + 1}/{max_lock_retries}")
|
|
with get_db_session() as session:
|
|
now = datetime.utcnow()
|
|
candidates = session.query(Task).filter(
|
|
Task.worker_id == worker_id,
|
|
Task.status == 'pending',
|
|
or_(Task.not_before == None, Task.not_before <= now)
|
|
).all()
|
|
|
|
ready = []
|
|
for task in candidates:
|
|
if not task.depends_on:
|
|
ready.append(task)
|
|
else:
|
|
dep = session.query(Task).filter(
|
|
Task.id == task.depends_on,
|
|
Task.status == 'acknowledged'
|
|
).first()
|
|
if dep:
|
|
ready.append(task)
|
|
|
|
if not ready:
|
|
if ENABLE_TASK_ASSIGNMENT_DEBUG:
|
|
logger.debug(f"[{worker_id}] No ready tasks with satisfied dependencies")
|
|
redis_client.set(f"worker_queue_{worker_id}", "False")
|
|
redis_client.set(f"worker_status_{worker_id}", "idle")
|
|
return
|
|
|
|
# New isolated session for locking and updating
|
|
with get_db_session() as session:
|
|
selected = session.query(Task).filter(
|
|
Task.id.in_([t.id for t in ready]),
|
|
Task.status == 'pending'
|
|
).with_for_update(skip_locked=True).limit(1).one_or_none()
|
|
|
|
if selected:
|
|
if ENABLE_TASK_ASSIGNMENT_DEBUG:
|
|
logger.debug(f"[{worker_id}] Task selected: {selected.id}")
|
|
break
|
|
|
|
if ENABLE_TASK_ASSIGNMENT_DEBUG:
|
|
logger.debug(f"[{worker_id}] Lock contention, retrying...")
|
|
time.sleep(lock_retry_wait)
|
|
lock_retry_wait *= 2
|
|
else:
|
|
if ENABLE_TASK_ASSIGNMENT_DEBUG:
|
|
logger.debug(f"[{worker_id}] Failed to acquire task after {max_lock_retries} attempts")
|
|
redis_client.set(f"worker_status_{worker_id}", "idle")
|
|
redis_client.set(f"worker_queue_{worker_id}", "False")
|
|
return
|
|
|
|
with worker_lock:
|
|
if worker_id in connected_workers:
|
|
if ENABLE_TASK_ASSIGNMENT_DEBUG:
|
|
logger.debug(f"[{worker_id}] Dispatching task {selected.id}")
|
|
base.socketio.emit("task", {
|
|
"task_id": selected.id,
|
|
"worker_id": worker_id,
|
|
"job_details": selected.job_details,
|
|
"type": selected.task_type
|
|
}, to=connected_workers[worker_id])
|
|
|
|
with get_db_session() as session:
|
|
task = session.query(Task).filter(Task.id == selected.id).first()
|
|
task.start_time = datetime.utcnow()
|
|
task.wait_time = (task.start_time - task.creation_time).total_seconds()
|
|
task.status = "in-progress"
|
|
session.add(task)
|
|
session.commit()
|
|
|
|
remaining = session.query(Task).filter(
|
|
Task.worker_id == worker_id,
|
|
Task.status == 'pending'
|
|
).count()
|
|
|
|
redis_client.set(f"worker_queue_{worker_id}", "True" if remaining > 0 else "False")
|
|
logger.info(f"[{worker_id}] Task {selected.id} marked in-progress")
|
|
else:
|
|
logger.warning(f"[{worker_id}] Worker disconnected before dispatch")
|
|
|
|
except Exception as e:
|
|
logger.exception(f"[{worker_id}] Error during task assignment: {e}")
|