Robert Wojciechowski
7060e31edb
The checkbox read "Clear saved API key (disable pool first)", but the panel never defines "pool" and saving with it ticked on an enabled worker was rejected by the backend. - Disable the checkbox while the worker is enabled, with a tooltip saying to turn the worker off first. Also disable it when no key is saved. - Use "worker" instead of "pool" in user-facing serverless errors, the serverless guide and the README entry.
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 worker"}, 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)
|