108 lines
4.4 KiB
Python
108 lines
4.4 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
|
|
|
|
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
|
|
|
|
logger.info(f"[{worker_id}] Starting task assignment process")
|
|
|
|
try:
|
|
redis_client.set(f"worker_status_{worker_id}", "busy")
|
|
logger.debug(f"[{worker_id}] Marked as 'busy' in Redis")
|
|
|
|
for attempt in range(max_lock_retries):
|
|
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:
|
|
logger.info(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:
|
|
logger.info(f"[{worker_id}] Task selected: {selected.id}")
|
|
break
|
|
|
|
logger.warning(f"[{worker_id}] Lock contention, retrying...")
|
|
time.sleep(lock_retry_wait)
|
|
lock_retry_wait *= 2
|
|
else:
|
|
logger.error(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:
|
|
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}")
|