import os import time import threading from fastapi import FastAPI, Request from fastapi.responses import HTMLResponse, JSONResponse from fastapi.middleware.cors import CORSMiddleware import uvicorn import traceback from config import REPO_ID, BUCKET_ID, MODEL_PRESETS, DATA_SOURCES, HEARTBEAT_TIMEOUT from core.state import workers, tasks from api.workers import router as workers_router from api.tasks import router as tasks_router from api.results import router as results_router from api.console import router as console_router app = FastAPI(title="Nightglow AI Master", version="4.0") app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) app.include_router(workers_router) app.include_router(tasks_router) app.include_router(results_router) app.include_router(console_router) def health_check_loop(): while True: try: now = time.time() for worker_id, w in list(workers.items()): if now - w.get("last_heartbeat", 0) > HEARTBEAT_TIMEOUT: if w.get("status") != "offline": print(f"⚠️ Worker {worker_id} heartbeat timeout, marking offline") w["status"] = "offline" w["status_detail"] = "offline" current_task = w.get("current_task") if current_task and current_task in tasks: if tasks[current_task].get("status") in ["assigned", "running"]: tasks[current_task]["status"] = "pending" tasks[current_task]["assigned_to"] = None w["current_task"] = None time.sleep(30) except Exception as e: print(f"❌ Health check error: {e}") time.sleep(30) threading.Thread(target=health_check_loop, daemon=True).start() print("✅ Health check thread started") @app.exception_handler(Exception) async def global_exception_handler(request: Request, exc: Exception): print(f"❌ Global error: {exc}") print(traceback.format_exc()) return JSONResponse( status_code=500, content={"error": str(exc), "detail": "Internal Server Error"} ) @app.get("/") async def index(): html = """ Nightglow AI Master

🌙 Nightglow AI Master

Distributed Training · Multi-Worker · Real Training

📊 Training Status

🖥️ Workers

📟 Console

Type help for available commands

""" return HTMLResponse(content=html) if __name__ == "__main__": port = int(os.getenv("PORT", 7860)) print("=" * 60) print("🌙 Nightglow AI Master starting") print("=" * 60) print(f" Port: {port}") print(f" Dataset repo: {REPO_ID}") print(f" Bucket: {BUCKET_ID}") print(f" Model presets: {len(MODEL_PRESETS)}") print("=" * 60) uvicorn.run(app, host="0.0.0.0", port=port)