robertvoy-ComfyUI-Distributed/upscale/job_store.py

import asyncio
import os
import time
from typing import List, Optional

import server

from ..utils.logging import debug_log
from .job_models import BaseJobState, ImageJobState, TileJobState

# Configure maximum payload size (50MB default, configurable via environment variable)
MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024)))


def ensure_tile_jobs_initialized():
    """Ensure tile job storage is initialized on the server instance."""
    prompt_server = server.PromptServer.instance
    if not hasattr(prompt_server, 'distributed_pending_tile_jobs'):
        debug_log("Initializing persistent tile job queue on server instance.")
        prompt_server.distributed_pending_tile_jobs = {}
        prompt_server.distributed_tile_jobs_lock = asyncio.Lock()
    else:
        invalid_job_ids = [
            job_id
            for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items()
            if not isinstance(job_data, BaseJobState)
        ]
        for job_id in invalid_job_ids:
            debug_log(f"Removing invalid job state for {job_id}")
            del prompt_server.distributed_pending_tile_jobs[job_id]
    return prompt_server


async def _init_job_queue(
    multi_job_id,
    mode,
    batch_size=None,
    num_tiles_per_image=None,
    all_indices=None,
    enabled_workers=None,
    batched_static: bool = False,
):
    """Unified initialization for job queues in static and dynamic modes."""
    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:
            debug_log(f"Queue already exists for {multi_job_id}")
            return

        if mode == 'dynamic':
            job_data = ImageJobState(multi_job_id=multi_job_id)
        elif mode == 'static':
            job_data = TileJobState(multi_job_id=multi_job_id)
        else:
            raise ValueError(f"Unknown mode: {mode}")

        job_data.worker_status = {w: time.time() for w in enabled_workers or []}
        job_data.assigned_to_workers = {w: [] for w in enabled_workers or []}

        if mode == 'dynamic':
            job_data.batch_size = int(batch_size or 0)
            pending_queue = job_data.pending_images
            for i in (all_indices or range(int(batch_size or 0))):
                await pending_queue.put(i)
            debug_log(f"Initialized image queue with {batch_size} pending items")
        elif mode == 'static':
            job_data.num_tiles_per_image = int(num_tiles_per_image or 0)
            job_data.batch_size = int(batch_size or 0)
            job_data.batched_static = bool(batched_static)
            # For batched static distribution, populate only tile ids [0..num_tiles_per_image-1]
            pending_queue = job_data.pending_tasks
            if batched_static and num_tiles_per_image is not None:
                for i in range(num_tiles_per_image):
                    await pending_queue.put(i)
            else:
                total_tiles = int(batch_size or 0) * int(num_tiles_per_image or 0)
                for i in range(total_tiles):
                    await pending_queue.put(i)

        prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data


async def init_dynamic_job(
    multi_job_id: str,
    batch_size: int,
    enabled_workers: List[str],
    all_indices: Optional[List[int]] = None,
):
    """Initialize queue for dynamic mode (per-image), with collector fields."""
    await _init_job_queue(
        multi_job_id,
        'dynamic',
        batch_size=batch_size,
        all_indices=all_indices or list(range(batch_size)),
        enabled_workers=enabled_workers,
    )
    debug_log(f"Job {multi_job_id} initialized with {batch_size} images")


async def init_static_job_batched(
    multi_job_id: str,
    batch_size: int,
    num_tiles_per_image: int,
    enabled_workers: List[str],
):
    """Initialize queue for static mode (batched-per-tile)."""
    await _init_job_queue(
        multi_job_id,
        'static',
        batch_size=batch_size,
        num_tiles_per_image=num_tiles_per_image,
        enabled_workers=enabled_workers,
        batched_static=True,
    )


async def _drain_results_queue(multi_job_id):
    """Drain pending results from queue and update completed_tasks. Returns count drained."""
    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 not isinstance(job_data, BaseJobState):
            return 0
        q = job_data.queue
        completed_tasks = job_data.completed_tasks

        collected = 0
        while True:
            try:
                result = q.get_nowait()
            except asyncio.QueueEmpty:
                break

            worker_id = result['worker_id']
            is_last = result.get('is_last', False)

            if 'image_idx' in result and 'image' in result:
                task_id = result['image_idx']
                if task_id not in completed_tasks:
                    completed_tasks[task_id] = result['image']
                    collected += 1
            elif 'tiles' in result:
                for tile_data in result['tiles']:
                    task_id = tile_data.get('global_idx', tile_data['tile_idx'])
                    if task_id not in completed_tasks:
                        completed_tasks[task_id] = tile_data
                        collected += 1
            if is_last:
                if worker_id in job_data.worker_status:
                    del job_data.worker_status[worker_id]

        return collected


async def _get_completed_count(multi_job_id):
    """Get count of completed tasks."""
    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, BaseJobState):
            return len(job_data.completed_tasks)
        return 0


async def _mark_task_completed(multi_job_id, task_id, result):
    """Mark a task as completed."""
    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, BaseJobState):
            job_data.completed_tasks[task_id] = result


async def _cleanup_job(multi_job_id):
    """Cleanup the job data."""
    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:
            del prompt_server.distributed_pending_tile_jobs[multi_job_id]
            debug_log(f"Cleaned up job {multi_job_id}")