Robert Wojciechowski
87c0cda9f6
Retain durable reservations until provider termination or explicit recovery, protect unresolved pools and credentials, and add scoped reconciliation. Consolidate lifecycle and secret policies and split prompt/media stages with regression coverage.
99 lines
4.4 KiB
Python
99 lines
4.4 KiB
Python
"""Master-only Runpod pool monitoring. Health checks never submit inference."""
|
|
from aiohttp import web
|
|
import server
|
|
|
|
from ..utils.config import load_config
|
|
from .orchestration.serverless_runtime import get_manager
|
|
from .orchestration.runpod_client import RunpodError
|
|
|
|
|
|
def _pool(worker_id):
|
|
return next((w for w in load_config().get("workers", [])
|
|
if w.get("id") == worker_id and w.get("type") == "serverless"), None)
|
|
|
|
|
|
async def _status(worker_id, fresh=False):
|
|
worker = _pool(worker_id)
|
|
if worker is None:
|
|
return web.json_response({"ok": False, "message": "Serverless worker not found"}, status=404)
|
|
jobs = []
|
|
try:
|
|
manager = get_manager()
|
|
jobs = manager.public_jobs(worker_id)
|
|
result = await manager.health(worker, fresh=fresh)
|
|
return web.json_response(dict(result, jobs=jobs))
|
|
except Exception as exc:
|
|
message = str(exc) if isinstance(exc, RunpodError) else "Check the saved endpoint ID and API key"
|
|
return web.json_response({"ok": False, "queueRemaining": 0, "message": message, "jobs": jobs})
|
|
|
|
|
|
@server.PromptServer.instance.routes.post("/distributed/serverless/test")
|
|
async def test_serverless_endpoint(request):
|
|
try:
|
|
data = await request.json()
|
|
worker_id = data.get("worker_id")
|
|
except Exception:
|
|
return web.json_response({"ok": False, "message": "Invalid request"}, status=400)
|
|
return await _status(worker_id, fresh=True)
|
|
|
|
|
|
@server.PromptServer.instance.routes.get("/distributed/serverless/status/{worker_id}")
|
|
async def serverless_status_endpoint(request):
|
|
return await _status(request.match_info["worker_id"])
|
|
|
|
|
|
@server.PromptServer.instance.routes.post("/distributed/serverless/cancel")
|
|
async def cancel_serverless_endpoint(request):
|
|
try:
|
|
data = await request.json()
|
|
worker_id = data.get("worker_id")
|
|
except Exception:
|
|
return web.json_response({"ok": False, "message": "Invalid request"}, status=400)
|
|
if _pool(worker_id) is None:
|
|
return web.json_response({"ok": False, "message": "Serverless worker not found"}, status=404)
|
|
try:
|
|
manager = get_manager()
|
|
count = await manager.cancel(worker_id=worker_id)
|
|
return web.json_response({"ok": True, "cancelled": count, "jobs": manager.public_jobs(worker_id),
|
|
"message": "Cancellation requested for this master's jobs only"})
|
|
except Exception:
|
|
return web.json_response({"ok": False, "message": "Cancellation could not be confirmed; check provider jobs"}, status=502)
|
|
|
|
|
|
async def _recover(request, *, acknowledge=False):
|
|
try:
|
|
data = await request.json()
|
|
worker_id = data.get("worker_id")
|
|
if not isinstance(worker_id, str) or not worker_id.strip():
|
|
raise ValueError()
|
|
record_id = data.get("record_id")
|
|
if acknowledge:
|
|
if not isinstance(record_id, str) or not record_id.strip() or data.get("provider_checked") is not True:
|
|
raise ValueError()
|
|
except Exception:
|
|
return web.json_response({"ok": False, "message": "Provide a worker ID; manual resolution also requires a record ID and provider_checked: true"}, status=400)
|
|
if _pool(worker_id) is None:
|
|
return web.json_response({"ok": False, "message": "Serverless worker not found"}, status=404)
|
|
try:
|
|
manager = get_manager()
|
|
if acknowledge:
|
|
await manager.acknowledge_recovery(worker_id, record_id, provider_checked=True)
|
|
result = {"message": "Recovery acknowledged by operator; not provider-confirmed cancellation"}
|
|
else:
|
|
result = {"resolved": await manager.reconcile(worker_id=worker_id)}
|
|
return web.json_response(dict(result, ok=True, jobs=manager.public_jobs(worker_id)))
|
|
except ValueError:
|
|
return web.json_response({"ok": False, "message": "Unresolved inactive job not found for this pool"}, status=409)
|
|
except Exception:
|
|
return web.json_response({"ok": False, "message": "Recovery could not be confirmed; retained jobs remain protected"}, status=502)
|
|
|
|
|
|
@server.PromptServer.instance.routes.post("/distributed/serverless/reconcile")
|
|
async def reconcile_serverless_endpoint(request):
|
|
return await _recover(request)
|
|
|
|
|
|
@server.PromptServer.instance.routes.post("/distributed/serverless/resolve")
|
|
async def resolve_serverless_endpoint(request):
|
|
return await _recover(request, acknowledge=True)
|