136 lines
5.7 KiB
Python
136 lines
5.7 KiB
Python
import asyncio
|
|
|
|
from ..utils.logging import debug_log
|
|
from .job_store import ensure_tile_jobs_initialized
|
|
from .job_timeout import _check_and_requeue_timed_out_workers as _requeue_usdu
|
|
from .job_models import ImageJobState, TileJobState
|
|
|
|
|
|
class JobStateMixin:
|
|
async def _get_job_data(self, multi_job_id):
|
|
"""Return current job data reference while holding lock briefly."""
|
|
prompt_server = ensure_tile_jobs_initialized()
|
|
async with prompt_server.distributed_tile_jobs_lock:
|
|
return prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
|
|
|
async def _get_all_completed_tasks(self, multi_job_id):
|
|
"""Helper to retrieve all completed tasks from the job data."""
|
|
job_data = await self._get_job_data(multi_job_id)
|
|
if isinstance(job_data, TileJobState):
|
|
return dict(job_data.completed_tasks)
|
|
if isinstance(job_data, ImageJobState):
|
|
return dict(job_data.completed_images)
|
|
return {}
|
|
|
|
async def _get_next_image_index(self, multi_job_id):
|
|
"""Get next image index from pending queue for master."""
|
|
prompt_server = ensure_tile_jobs_initialized()
|
|
pending_queue = None
|
|
async with prompt_server.distributed_tile_jobs_lock:
|
|
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
|
if isinstance(job_data, ImageJobState):
|
|
pending_queue = job_data.pending_images
|
|
|
|
if pending_queue is None:
|
|
return None
|
|
|
|
try:
|
|
return await asyncio.wait_for(pending_queue.get(), timeout=1.0)
|
|
except asyncio.TimeoutError:
|
|
return None
|
|
|
|
async def _get_next_tile_index(self, multi_job_id):
|
|
"""Get next tile index from pending queue for master in static mode."""
|
|
prompt_server = ensure_tile_jobs_initialized()
|
|
pending_queue = None
|
|
async with prompt_server.distributed_tile_jobs_lock:
|
|
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
|
if isinstance(job_data, TileJobState):
|
|
pending_queue = job_data.pending_tasks
|
|
|
|
if pending_queue is None:
|
|
return None
|
|
|
|
try:
|
|
return await asyncio.wait_for(pending_queue.get(), timeout=0.1)
|
|
except asyncio.TimeoutError:
|
|
return None
|
|
|
|
async def _get_total_completed_count(self, multi_job_id):
|
|
"""Get total count of all completed images (master + workers)."""
|
|
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)
|
|
if isinstance(job_data, ImageJobState):
|
|
return len(job_data.completed_images)
|
|
if isinstance(job_data, TileJobState):
|
|
return len(job_data.completed_tasks)
|
|
return 0
|
|
|
|
async def _get_all_completed_images(self, multi_job_id):
|
|
"""Get all completed images."""
|
|
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)
|
|
if isinstance(job_data, ImageJobState):
|
|
return job_data.completed_images.copy()
|
|
return {}
|
|
|
|
async def _get_pending_count(self, multi_job_id):
|
|
"""Get count of pending images in the queue."""
|
|
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)
|
|
if isinstance(job_data, ImageJobState):
|
|
return job_data.pending_images.qsize()
|
|
if isinstance(job_data, TileJobState):
|
|
return job_data.pending_tasks.qsize()
|
|
return 0
|
|
|
|
async def _drain_worker_results_queue(self, multi_job_id):
|
|
"""Drain pending worker results from queue and update completed images."""
|
|
prompt_server = ensure_tile_jobs_initialized()
|
|
worker_queue = None
|
|
async with prompt_server.distributed_tile_jobs_lock:
|
|
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
|
if isinstance(job_data, ImageJobState):
|
|
worker_queue = job_data.queue
|
|
|
|
if worker_queue is None:
|
|
return 0
|
|
|
|
drained_results = []
|
|
while True:
|
|
try:
|
|
drained_results.append(worker_queue.get_nowait())
|
|
except asyncio.QueueEmpty:
|
|
break
|
|
|
|
if not drained_results:
|
|
return 0
|
|
|
|
collected = 0
|
|
async with prompt_server.distributed_tile_jobs_lock:
|
|
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
|
if not isinstance(job_data, ImageJobState):
|
|
return 0
|
|
|
|
for result in drained_results:
|
|
worker_id = result.get("worker_id")
|
|
if "image_idx" in result and "image" in result:
|
|
image_idx = result["image_idx"]
|
|
image_pil = result["image"]
|
|
if image_idx not in job_data.completed_images:
|
|
job_data.completed_images[image_idx] = image_pil
|
|
collected += 1
|
|
debug_log(f"Drained image {image_idx} from worker {worker_id}")
|
|
|
|
if collected > 0:
|
|
debug_log(f"Drained {collected} worker images during retry")
|
|
|
|
return collected
|
|
|
|
async def _check_and_requeue_timed_out_workers(self, multi_job_id, batch_size):
|
|
"""Check for timed out workers and requeue their assigned images."""
|
|
return await _requeue_usdu(multi_job_id, batch_size)
|