Robert Wojciechowski
5e724665a1
When master is in delegate-only mode, USDU nodes are swapped to a lightweight USDUDelegateCollector that skips model loading and only collects results from workers. Workers initialize the dynamic job queue on the master via a new /distributed/init_dynamic_job endpoint since the delegate master doesn't know the batch size.
345 lines
16 KiB
Python
345 lines
16 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import io
|
|
import time
|
|
|
|
from aiohttp import web
|
|
from PIL import Image
|
|
import server
|
|
|
|
from .request_guards import authorization_error_or_none
|
|
from .schemas import require_bool_literal
|
|
from ..upscale.job_models import BaseJobState, ImageJobState, TileJobState
|
|
from ..upscale.job_store import ensure_tile_jobs_initialized, init_dynamic_job
|
|
from ..upscale.payload_parsers import parse_tiles_from_form
|
|
from ..utils.logging import debug_log
|
|
from ..utils.network import handle_api_error
|
|
from ..utils.usdu_management import MAX_PAYLOAD_SIZE
|
|
|
|
|
|
def _parse_int_field(value, field_name: str, *, minimum: int | None = None) -> int:
|
|
try:
|
|
parsed = int(value)
|
|
except (TypeError, ValueError) as exc:
|
|
raise ValueError(f"{field_name} must be an integer") from exc
|
|
if minimum is not None and parsed < minimum:
|
|
raise ValueError(f"{field_name} must be >= {minimum}")
|
|
return parsed
|
|
|
|
|
|
def _parse_bool_field(
|
|
value,
|
|
field_name: str,
|
|
*,
|
|
default: bool = False,
|
|
) -> bool:
|
|
if value is None:
|
|
return default
|
|
return require_bool_literal(value, field_name=field_name)
|
|
|
|
|
|
def _is_worker_allowed(job_data: BaseJobState, worker_id: str) -> bool:
|
|
known_workers = {str(wid) for wid in getattr(job_data, "worker_status", {}).keys()}
|
|
known_workers.update(str(wid) for wid in getattr(job_data, "assigned_to_workers", {}).keys())
|
|
if not known_workers:
|
|
return True
|
|
return str(worker_id) in known_workers
|
|
|
|
|
|
@server.PromptServer.instance.routes.post("/distributed/init_dynamic_job")
|
|
async def init_dynamic_job_endpoint(request: web.Request) -> web.StreamResponse:
|
|
"""Allow workers to initialize a dynamic job queue on the master.
|
|
|
|
Called by the first worker to reach the USDU node in delegate-only mode.
|
|
Idempotent — subsequent calls are no-ops if the queue already exists.
|
|
"""
|
|
auth_error = await authorization_error_or_none(request)
|
|
if auth_error is not None:
|
|
return auth_error
|
|
try:
|
|
data = await request.json()
|
|
multi_job_id = data.get("multi_job_id")
|
|
if not multi_job_id:
|
|
return await handle_api_error(request, "Missing multi_job_id", 400)
|
|
|
|
batch_size = _parse_int_field(data.get("batch_size"), "batch_size", minimum=1)
|
|
enabled_workers = data.get("enabled_workers") or []
|
|
if isinstance(enabled_workers, str):
|
|
import json as _json
|
|
enabled_workers = _json.loads(enabled_workers)
|
|
|
|
await init_dynamic_job(multi_job_id, batch_size, enabled_workers)
|
|
return web.json_response({"status": "success"})
|
|
except Exception as e:
|
|
return await handle_api_error(request, e, 500)
|
|
|
|
|
|
@server.PromptServer.instance.routes.post("/distributed/heartbeat")
|
|
async def heartbeat_endpoint(request: web.Request) -> web.StreamResponse:
|
|
auth_error = await authorization_error_or_none(request)
|
|
if auth_error is not None:
|
|
return auth_error
|
|
try:
|
|
data = await request.json()
|
|
worker_id = data.get('worker_id')
|
|
multi_job_id = data.get('multi_job_id')
|
|
|
|
if not worker_id or not multi_job_id:
|
|
return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
|
|
|
|
prompt_server = ensure_tile_jobs_initialized()
|
|
async with prompt_server.distributed_tile_jobs_lock:
|
|
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
|
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
|
if isinstance(job_data, BaseJobState):
|
|
if not _is_worker_allowed(job_data, worker_id):
|
|
return await handle_api_error(request, "Unauthorized worker_id", 403)
|
|
job_data.worker_status[worker_id] = time.time()
|
|
debug_log(f"Heartbeat from worker {worker_id}")
|
|
return web.json_response({"status": "success"})
|
|
return await handle_api_error(request, "Worker status tracking not available", 400)
|
|
return await handle_api_error(request, "Job not found", 404)
|
|
except Exception as e:
|
|
return await handle_api_error(request, e, 500)
|
|
|
|
|
|
@server.PromptServer.instance.routes.post("/distributed/submit_tiles")
|
|
async def submit_tiles_endpoint(request: web.Request) -> web.StreamResponse:
|
|
"""Endpoint for workers to submit processed tiles in static mode."""
|
|
auth_error = await authorization_error_or_none(request)
|
|
if auth_error is not None:
|
|
return auth_error
|
|
try:
|
|
content_length = request.headers.get('content-length')
|
|
if content_length:
|
|
try:
|
|
payload_size = int(content_length)
|
|
except ValueError:
|
|
return await handle_api_error(request, "Invalid content-length header", 400)
|
|
if payload_size > MAX_PAYLOAD_SIZE:
|
|
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
|
|
|
|
data = await request.post()
|
|
multi_job_id = data.get('multi_job_id')
|
|
worker_id = data.get('worker_id')
|
|
try:
|
|
is_last = _parse_bool_field(data.get('is_last', False), "is_last", default=False)
|
|
except ValueError as e:
|
|
return await handle_api_error(request, str(e), 400)
|
|
|
|
if multi_job_id is None or worker_id is None:
|
|
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
|
|
|
|
prompt_server = ensure_tile_jobs_initialized()
|
|
|
|
try:
|
|
batch_size = _parse_int_field(data.get('batch_size', 0), "batch_size", minimum=0)
|
|
except ValueError as e:
|
|
return await handle_api_error(request, str(e), 400)
|
|
|
|
# Handle completion signal
|
|
if batch_size == 0 and is_last:
|
|
async with prompt_server.distributed_tile_jobs_lock:
|
|
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
|
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
|
if not isinstance(job_data, TileJobState):
|
|
return await handle_api_error(request, "Job not configured for tile submissions", 400)
|
|
if not _is_worker_allowed(job_data, worker_id):
|
|
return await handle_api_error(request, "Unauthorized worker_id", 403)
|
|
await job_data.queue.put({
|
|
'worker_id': worker_id,
|
|
'is_last': True,
|
|
'tiles': [],
|
|
})
|
|
debug_log(f"Received completion signal from worker {worker_id}")
|
|
return web.json_response({"status": "success"})
|
|
|
|
try:
|
|
tiles = parse_tiles_from_form(data)
|
|
except ValueError as e:
|
|
return await handle_api_error(request, str(e), 400)
|
|
|
|
# Submit tiles to queue
|
|
async with prompt_server.distributed_tile_jobs_lock:
|
|
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
|
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
|
if not isinstance(job_data, TileJobState):
|
|
return await handle_api_error(request, "Job not configured for tile submissions", 400)
|
|
if not _is_worker_allowed(job_data, worker_id):
|
|
return await handle_api_error(request, "Unauthorized worker_id", 403)
|
|
|
|
q = job_data.queue
|
|
if batch_size > 0 or len(tiles) > 0:
|
|
await q.put({
|
|
'worker_id': worker_id,
|
|
'tiles': tiles,
|
|
'is_last': is_last,
|
|
})
|
|
debug_log(f"Received {len(tiles)} tiles from worker {worker_id} (is_last={is_last})")
|
|
else:
|
|
await q.put({
|
|
'worker_id': worker_id,
|
|
'is_last': True,
|
|
'tiles': [],
|
|
})
|
|
|
|
return web.json_response({"status": "success"})
|
|
return await handle_api_error(request, "Job not found", 404)
|
|
except Exception as e:
|
|
return await handle_api_error(request, e, 500)
|
|
|
|
|
|
@server.PromptServer.instance.routes.post("/distributed/submit_image")
|
|
async def submit_image_endpoint(request: web.Request) -> web.StreamResponse:
|
|
"""Endpoint for workers to submit processed images in dynamic mode."""
|
|
auth_error = await authorization_error_or_none(request)
|
|
if auth_error is not None:
|
|
return auth_error
|
|
try:
|
|
content_length = request.headers.get('content-length')
|
|
if content_length:
|
|
try:
|
|
payload_size = int(content_length)
|
|
except ValueError:
|
|
return await handle_api_error(request, "Invalid content-length header", 400)
|
|
if payload_size > MAX_PAYLOAD_SIZE:
|
|
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
|
|
|
|
data = await request.post()
|
|
multi_job_id = data.get('multi_job_id')
|
|
worker_id = data.get('worker_id')
|
|
try:
|
|
is_last = _parse_bool_field(data.get('is_last', False), "is_last", default=False)
|
|
except ValueError as e:
|
|
return await handle_api_error(request, str(e), 400)
|
|
|
|
if multi_job_id is None or worker_id is None:
|
|
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
|
|
|
|
prompt_server = ensure_tile_jobs_initialized()
|
|
|
|
# Handle image submission
|
|
if 'full_image' in data and 'image_idx' in data:
|
|
try:
|
|
image_idx = _parse_int_field(data.get('image_idx'), "image_idx", minimum=0)
|
|
except ValueError as e:
|
|
return await handle_api_error(request, str(e), 400)
|
|
img_data = data['full_image'].file.read()
|
|
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
|
|
|
debug_log(f"Received full image {image_idx} from worker {worker_id}")
|
|
|
|
async with prompt_server.distributed_tile_jobs_lock:
|
|
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
|
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
|
if not isinstance(job_data, ImageJobState):
|
|
return await handle_api_error(request, "Job not configured for image submissions", 400)
|
|
if not _is_worker_allowed(job_data, worker_id):
|
|
return await handle_api_error(request, "Unauthorized worker_id", 403)
|
|
await job_data.queue.put({
|
|
'worker_id': worker_id,
|
|
'image_idx': image_idx,
|
|
'image': img,
|
|
'is_last': is_last,
|
|
})
|
|
return web.json_response({"status": "success"})
|
|
|
|
# Handle completion signal (no image data)
|
|
elif is_last:
|
|
async with prompt_server.distributed_tile_jobs_lock:
|
|
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
|
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
|
if not isinstance(job_data, ImageJobState):
|
|
return await handle_api_error(request, "Job not configured for image submissions", 400)
|
|
if not _is_worker_allowed(job_data, worker_id):
|
|
return await handle_api_error(request, "Unauthorized worker_id", 403)
|
|
await job_data.queue.put({
|
|
'worker_id': worker_id,
|
|
'is_last': True,
|
|
})
|
|
debug_log(f"Received completion signal from worker {worker_id}")
|
|
return web.json_response({"status": "success"})
|
|
else:
|
|
return await handle_api_error(request, "Missing image data or invalid request", 400)
|
|
|
|
return await handle_api_error(request, "Job not found", 404)
|
|
except Exception as e:
|
|
return await handle_api_error(request, e, 500)
|
|
|
|
|
|
@server.PromptServer.instance.routes.post("/distributed/request_image")
|
|
async def request_image_endpoint(request: web.Request) -> web.StreamResponse:
|
|
"""Endpoint for workers to request tasks (images in dynamic mode, tiles in static mode)."""
|
|
auth_error = await authorization_error_or_none(request)
|
|
if auth_error is not None:
|
|
return auth_error
|
|
try:
|
|
data = await request.json()
|
|
worker_id = data.get('worker_id')
|
|
multi_job_id = data.get('multi_job_id')
|
|
|
|
if not worker_id or not multi_job_id:
|
|
return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
|
|
|
|
prompt_server = ensure_tile_jobs_initialized()
|
|
async with prompt_server.distributed_tile_jobs_lock:
|
|
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
|
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
|
if not isinstance(job_data, BaseJobState):
|
|
return await handle_api_error(request, "Invalid job data structure", 500)
|
|
|
|
mode = job_data.mode
|
|
if isinstance(job_data, ImageJobState):
|
|
pending_queue = job_data.pending_images
|
|
elif isinstance(job_data, TileJobState):
|
|
pending_queue = job_data.pending_tasks
|
|
else:
|
|
return await handle_api_error(request, "Invalid job configuration", 400)
|
|
if not _is_worker_allowed(job_data, worker_id):
|
|
return await handle_api_error(request, "Unauthorized worker_id", 403)
|
|
|
|
try:
|
|
task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1)
|
|
job_data.assigned_to_workers.setdefault(worker_id, []).append(task_idx)
|
|
job_data.worker_status[worker_id] = time.time()
|
|
remaining = pending_queue.qsize()
|
|
|
|
if mode == 'dynamic':
|
|
debug_log(f"UltimateSDUpscale API - Assigned image {task_idx} to worker {worker_id}")
|
|
return web.json_response(
|
|
{
|
|
"kind": "image",
|
|
"task_idx": task_idx,
|
|
"estimated_remaining": remaining,
|
|
}
|
|
)
|
|
debug_log(f"UltimateSDUpscale API - Assigned tile {task_idx} to worker {worker_id}")
|
|
return web.json_response({
|
|
"kind": "tile",
|
|
"task_idx": task_idx,
|
|
"estimated_remaining": remaining,
|
|
"batched_static": job_data.batched_static,
|
|
})
|
|
except asyncio.TimeoutError:
|
|
return web.json_response({"kind": "none", "task_idx": None, "estimated_remaining": 0})
|
|
return await handle_api_error(request, "Job not found", 404)
|
|
except Exception as e:
|
|
return await handle_api_error(request, e, 500)
|
|
|
|
|
|
@server.PromptServer.instance.routes.get("/distributed/job_status")
|
|
async def job_status_endpoint(request: web.Request) -> web.StreamResponse:
|
|
"""Endpoint to check if a job is ready."""
|
|
auth_error = await authorization_error_or_none(request)
|
|
if auth_error is not None:
|
|
return auth_error
|
|
multi_job_id = request.query.get('multi_job_id')
|
|
if not multi_job_id:
|
|
return web.json_response({"ready": False})
|
|
prompt_server = ensure_tile_jobs_initialized()
|
|
async with prompt_server.distributed_tile_jobs_lock:
|
|
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
|
ready = bool(isinstance(job_data, BaseJobState) and job_data.queue is not None)
|
|
return web.json_response({"ready": ready})
|