96 lines
3.2 KiB
Python
96 lines
3.2 KiB
Python
# websocket_server/task_api.py
|
|
|
|
from flask import request, jsonify
|
|
from datetime import datetime
|
|
import json
|
|
import logging
|
|
|
|
from websocket_server.models import Task
|
|
from websocket_server.config import get_db_session, get_redis_client
|
|
|
|
logger = logging.getLogger("websocket_server")
|
|
redis_client = get_redis_client()
|
|
|
|
# Import dispatch logic
|
|
from websocket_server.task_assigner import assign_task_to_worker
|
|
from websocket_server.shared_state import connected_workers, worker_lock
|
|
|
|
|
|
def register_routes(app):
|
|
"""
|
|
Register all Flask routes to the given app.
|
|
"""
|
|
|
|
@app.route("/api/create_task", methods=["POST"])
|
|
def create_task():
|
|
"""
|
|
Create a new task and insert it into the database.
|
|
"""
|
|
data = request.get_json()
|
|
worker_id = data.get("worker_id")
|
|
job_details = json.dumps(data.get("job_details"))
|
|
task_type = data.get("task_type", "default")
|
|
depends_on = data.get("depends_on")
|
|
not_before = data.get("not_before")
|
|
|
|
logger.info(f"Creating new task for worker {worker_id}")
|
|
|
|
if not worker_id or not job_details:
|
|
return jsonify({"error": "worker_id and job_details are required"}), 400
|
|
|
|
try:
|
|
not_before_dt = None
|
|
if not_before:
|
|
not_before_dt = datetime.fromisoformat(not_before)
|
|
|
|
with get_db_session() as session:
|
|
task = Task(
|
|
worker_id=worker_id,
|
|
job_details=job_details,
|
|
task_type=task_type,
|
|
status="pending",
|
|
success=None,
|
|
creation_time=datetime.utcnow(),
|
|
depends_on=depends_on,
|
|
not_before=not_before_dt
|
|
)
|
|
session.add(task)
|
|
session.commit()
|
|
|
|
logger.info(f"Created task {task.id} for worker {worker_id}")
|
|
redis_client.set(f"worker_queue_{worker_id}", "True")
|
|
|
|
return jsonify({"message": "Task created", "task_id": task.id}), 201
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error creating task: {e}")
|
|
return jsonify({"error": "Internal error creating task"}), 500
|
|
|
|
@app.route("/api/assign_task", methods=["POST"])
|
|
def manually_assign_task():
|
|
"""
|
|
Force task assignment for a specific worker via HTTP trigger.
|
|
"""
|
|
data = request.get_json()
|
|
worker_id = data.get("worker_id")
|
|
|
|
if not worker_id:
|
|
return jsonify({"error": "Missing worker_id"}), 400
|
|
|
|
logger.info(f"Manual trigger for assign_task to worker {worker_id}")
|
|
assign_task_to_worker(worker_id)
|
|
return jsonify({"message": f"Dispatched task assignment to worker {worker_id}."})
|
|
|
|
@app.route("/api/connected_clients", methods=["GET"])
|
|
def get_connected_clients():
|
|
"""
|
|
Returns currently connected worker IDs and socket session IDs.
|
|
"""
|
|
logger.debug("Listing connected clients")
|
|
with worker_lock:
|
|
result = [
|
|
{"worker_id": wid, "socket_id": sid, "ip": "unknown"}
|
|
for wid, sid in connected_workers.items()
|
|
]
|
|
return jsonify({"connected_clients": result})
|