File size: 5,125 Bytes
ff2893e 0998bb4 ff2893e 0998bb4 ff2893e 0998bb4 ff2893e 0998bb4 ff2893e 0998bb4 ff2893e 0998bb4 ff2893e 0998bb4 ff2893e 0998bb4 ff2893e 0998bb4 ff2893e 0998bb4 ff2893e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | import time
from datetime import datetime
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
from typing import Optional, Dict, Any
from config import MASTER_SECRET, HEARTBEAT_TIMEOUT
from core.state import workers, add_worker, update_worker, get_worker_by_id
from core.task_manager import get_next_task_for_worker
router = APIRouter(prefix="/api/workers", tags=["workers"])
class WorkerRegisterRequest(BaseModel):
worker_id: str
worker_type: str
capabilities: Dict[str, Any]
cpu_cores: int
memory_gb: float
class WorkerHeartbeatRequest(BaseModel):
worker_id: str
token: str
current_load: float
memory_used_gb: float
gpu_used_gb: Optional[float] = None
current_task: Optional[str] = None
backup_task: Optional[str] = None
progress: Optional[float] = 0.0
status_detail: Optional[str] = None
def generate_token(worker_id: str) -> str:
import hashlib
salt = str(time.time())
return hashlib.sha256(f"{worker_id}:{MASTER_SECRET}:{salt}".encode()).hexdigest()[:32]
def verify_token(worker_id: str, token: str) -> bool:
if worker_id not in workers:
return False
return workers[worker_id].get("token") == token
@router.post("/register")
async def register_worker(data: WorkerRegisterRequest):
try:
worker_id = data.worker_id
if worker_id in workers:
workers[worker_id]["token"] = generate_token(worker_id)
workers[worker_id]["last_heartbeat"] = time.time()
workers[worker_id]["status"] = "idle"
workers[worker_id]["status_detail"] = "idle"
workers[worker_id]["info"] = data.dict()
print(f"✅ Worker {worker_id} re-registered")
return {"status": "updated", "worker_id": worker_id, "token": workers[worker_id]["token"]}
token = generate_token(worker_id)
workers[worker_id] = {
"info": data.dict(),
"token": token,
"last_heartbeat": time.time(),
"status": "idle",
"status_detail": "idle",
"current_task": None,
"backup_task": None,
"progress": 0.0,
"completed_tasks": 0,
"registered_at": datetime.now().isoformat(),
}
print(f"✅ Worker {worker_id} registered")
return {"status": "registered", "worker_id": worker_id, "token": token}
except Exception as e:
print(f"❌ Register error: {e}")
raise HTTPException(500, str(e))
@router.post("/heartbeat")
async def worker_heartbeat(data: WorkerHeartbeatRequest):
try:
if data.worker_id not in workers:
raise HTTPException(404, "Worker not found")
if not verify_token(data.worker_id, data.token):
raise HTTPException(401, "Invalid token")
w = workers[data.worker_id]
w["last_heartbeat"] = time.time()
w["status"] = "busy" if data.current_task else "idle"
w["current_task"] = data.current_task
w["backup_task"] = data.backup_task
w["progress"] = data.progress or 0.0
w["status_detail"] = data.status_detail or "idle"
return {"status": "ok"}
except HTTPException:
raise
except Exception as e:
print(f"❌ Heartbeat error: {e}")
raise HTTPException(500, str(e))
@router.get("/list")
async def list_workers():
try:
now = time.time()
result = []
for wid, w in workers.items():
is_online = (now - w.get("last_heartbeat", 0)) < HEARTBEAT_TIMEOUT
result.append({
"worker_id": wid,
"worker_type": w.get("info", {}).get("worker_type", "unknown"),
"status": "online" if is_online else "offline",
"status_detail": w.get("status_detail", "idle"),
"current_task": w.get("current_task"),
"backup_task": w.get("backup_task"),
"progress": w.get("progress", 0),
"completed_tasks": w.get("completed_tasks", 0),
"last_heartbeat": datetime.fromtimestamp(w.get("last_heartbeat", 0)).isoformat() if w.get("last_heartbeat") else None,
})
return {"workers": result, "total": len(result)}
except Exception as e:
print(f"❌ List error: {e}")
return {"workers": [], "total": 0}
@router.get("/next-task")
async def get_next_task(worker_id: str, token: str):
try:
if worker_id not in workers:
raise HTTPException(404, "Worker not found")
if not verify_token(worker_id, token):
raise HTTPException(401, "Invalid token")
task = get_next_task_for_worker(worker_id)
if task:
workers[worker_id]["current_task"] = task["task_id"]
workers[worker_id]["status"] = "busy"
return {"current_task": task, "backup_task": None}
else:
return {"current_task": None, "backup_task": None}
except HTTPException:
raise
except Exception as e:
print(f"❌ Get next task error: {e}")
raise HTTPException(500, str(e)) |