From 12ff7bf39be438a57191c38774bbf7dce1fd24e7 Mon Sep 17 00:00:00 2001 From: Jason Hoku Date: Tue, 19 May 2026 15:26:06 -0700 Subject: [PATCH] refactor(distribution): delete distribution_manager.py (~650 lines) DistributionManager + DistributedJob + JobState moved to ComfyUI-USCG-Distributed companion. Manager now accepts job_completed_check/build_prompt_with_triggers/get_model_cache_key as callbacks to avoid bidirectional imports. --- distribution_manager.py | 650 ---------------------------------------- 1 file changed, 650 deletions(-) delete mode 100644 distribution_manager.py diff --git a/distribution_manager.py b/distribution_manager.py deleted file mode 100644 index a07a6c1..0000000 --- a/distribution_manager.py +++ /dev/null @@ -1,650 +0,0 @@ -""" -Distributed Task Manager -Thread-safe job queue for distributing generation work across multiple ComfyUI instances. -Uses a pull-based model: workers claim jobs from the master via API. -""" - -import threading -import time -import json -import os -import uuid - - -class JobState: - """Job lifecycle states.""" - PENDING = "pending" - CLAIMED = "claimed" - COMPLETED = "completed" - FAILED = "failed" - - -class DistributedJob: - """Represents a single job in the distributed queue.""" - __slots__ = [ - 'job_id', 'config', 'input_job', 'state', 'worker_id', - 'claimed_at', 'completed_at', 'error', 'gen_index', 'fail_count' - ] - - def __init__(self, job_id, config, input_job, gen_index): - self.job_id = job_id - self.config = config # Expanded config dict (sampler, model, lora, etc.) - self.input_job = input_job # {label, width, height, batch_idx} - self.state = JobState.PENDING - self.worker_id = None - self.claimed_at = None - self.completed_at = None - self.error = None - self.gen_index = gen_index # Sequential index for deterministic sort ordering - self.fail_count = 0 - - -class DistributionManager: - """ - Thread-safe job distribution coordinator. - - Manages an in-memory job list with atomic claim/complete/fail operations. - The master populates jobs at startup, then workers claim them one at a time. - Jobs that time out (worker crash) are automatically re-queued. - """ - - # Max times a job can fail before being permanently abandoned - MAX_FAIL_COUNT = 3 - - def __init__(self, session_name, paths, unique_id, existing_data, - claim_timeout_seconds=600): - self._lock = threading.Lock() - self._jobs = {} # job_id -> DistributedJob - self._pending_queue = [] # List of job_ids in FIFO order - self._workers = {} # worker_id -> {url, last_heartbeat, jobs_completed, ...} - self.session_name = session_name - self.paths = paths - self.unique_id = unique_id - self.existing_data = existing_data # Shared manifest reference - self.claim_timeout = claim_timeout_seconds - self._persistence_path = os.path.join(paths["base"], "distribution_state.json") - self._active = True - self._total_completed = 0 - # Pre-encoded conditionings from master, keyed by "(model_key|lora_key)" - # Structure: {"model|lora": {"positive": {prompt: serialized}, "negative": {...}}} - # Populated by master pre-encoding phase; attached to job claim responses - self._encoded_conditionings = {} - # Model sync: when enabled, job claim responses include upcoming_models for pre-fetching - self.sync_models_to_workers = False - # Session-level settings (upscaling, cooldown) passed through to workers - self._session_settings = None - - def populate_jobs(self, expanded_configs, input_jobs, existing_data, - overwrite_existing, has_optional_inputs, lora_triggerwords_mode): - """ - Build the full job list from expanded configs × input jobs. - Filters out already-completed jobs using check_if_job_completed(). - - Args: - expanded_configs: List of expanded config dicts (from expand_configs()) - input_jobs: List of input job dicts (from prepare_input_jobs()) - existing_data: Loaded manifest data for skip detection - overwrite_existing: If True, don't skip completed jobs - has_optional_inputs: If True, relaxed skip matching - lora_triggerwords_mode: Trigger word mode string - """ - from .generation_orchestrator import check_if_job_completed - from .trigger_words import build_prompt_with_triggers - - existing_items = existing_data.get("items", []) - gen_index_offset = len(existing_items) - skipped = 0 - total = 0 - - with self._lock: - job_index = 0 - for job_idx, job in enumerate(input_jobs): - w = job["width"] - h = job["height"] - batch_idx = job.get("batch_idx", 0) - - for conf_idx, conf in enumerate(expanded_configs): - total += 1 - - # Per-config resolution override - if conf.get("resolution") is not None: - if job_idx > 0: - continue - w = conf["resolution"][0] - h = conf["resolution"][1] - batch_idx = 0 - - # Compute actual prompts (with trigger words applied) - # build_prompt_with_triggers returns (final_positive, trigger_string) - try: - actual_positive, _ = build_prompt_with_triggers( - conf, lora_triggerwords_mode - ) - except Exception: - actual_positive = conf.get("positive", "") - actual_negative = conf.get("negative", "") - - # Skip detection (same logic as main generation loop) - if not overwrite_existing: - match_index = check_if_job_completed( - existing_items, conf, - conf.get("seed", 0), w, h, batch_idx, - actual_positive, actual_negative, - has_optional_inputs=has_optional_inputs - ) - if match_index != -1: - skipped += 1 - continue - - # Create job with serializable input_job (no latent tensors) - serializable_input_job = { - "label": job.get("label", f"{w}x{h}"), - "width": w, - "height": h, - "batch_idx": batch_idx - } - - job_id = str(uuid.uuid4())[:12] - gen_index = gen_index_offset + job_index - - distributed_job = DistributedJob( - job_id=job_id, - config=conf, - input_job=serializable_input_job, - gen_index=gen_index - ) - - self._jobs[job_id] = distributed_job - self._pending_queue.append(job_id) - job_index += 1 - - self._persist_state() - - print(f"[Distribution] 📋 Populated {len(self._jobs)} jobs " - f"({skipped} skipped, {total} total before filtering)") - - def claim_job(self, worker_id): - """ - Atomically claim the next pending job for a worker. - - Returns: - dict with job details, or None if no pending jobs. - """ - with self._lock: - # Release timed-out claimed jobs first - self._release_timed_out_jobs() - - if not self._pending_queue: - return None - - job_id = self._pending_queue.pop(0) - job = self._jobs[job_id] - job.state = JobState.CLAIMED - job.worker_id = worker_id - job.claimed_at = time.time() - - self._persist_state() - - return self._job_to_dict(job) - - def complete_job(self, job_id, metadata=None): - """ - Mark a job as completed. - - Guards against double-completion (can happen if a job times out and is - reclaimed by another worker, then the original worker finishes late). - - Args: - job_id: The job ID string - metadata: Optional metadata dict (not stored in job, just for logging) - - Returns: - True if job was found and completed, False otherwise. - """ - with self._lock: - job = self._jobs.get(job_id) - if not job: - print(f"[Distribution] ⚠️ complete_job: unknown job_id {job_id}") - return False - - # Prevent double-completion (job may have been reclaimed after timeout) - if job.state == JobState.COMPLETED: - print(f"[Distribution] ⚠️ complete_job: job {job_id} already completed, ignoring duplicate") - return False - - job.state = JobState.COMPLETED - job.completed_at = time.time() - self._total_completed += 1 - - # Update worker stats - if job.worker_id and job.worker_id in self._workers: - self._workers[job.worker_id]["jobs_completed"] += 1 - - self._persist_state() - return True - - def fail_job(self, job_id, error_msg=""): - """ - Mark a job as failed and re-queue to pending (up to MAX_FAIL_COUNT times). - - Args: - job_id: The job ID string - error_msg: Error description - - Returns: - True if job was found and re-queued, False otherwise. - """ - with self._lock: - job = self._jobs.get(job_id) - if not job: - print(f"[Distribution] ⚠️ fail_job: unknown job_id {job_id}") - return False - - job.fail_count += 1 - - # Update worker stats - if job.worker_id and job.worker_id in self._workers: - self._workers[job.worker_id]["jobs_failed"] += 1 - - if job.fail_count >= self.MAX_FAIL_COUNT: - job.state = JobState.FAILED - job.error = error_msg - print(f"[Distribution] ❌ Job {job_id} permanently failed after " - f"{job.fail_count} attempts: {error_msg}") - else: - job.state = JobState.PENDING - job.worker_id = None - job.claimed_at = None - job.error = error_msg - self._pending_queue.append(job_id) - print(f"[Distribution] 🔄 Job {job_id} re-queued " - f"(attempt {job.fail_count}/{self.MAX_FAIL_COUNT}): {error_msg}") - - self._persist_state() - return True - - def claim_batch_for_local(self, count): - """ - Claim a batch of jobs for the master's local processing. - Uses worker_id "master" to distinguish from remote workers. - - Args: - count: Maximum number of jobs to claim - - Returns: - List of job dicts (may be empty). - """ - with self._lock: - self._release_timed_out_jobs() - - batch = [] - for _ in range(min(count, len(self._pending_queue))): - job_id = self._pending_queue.pop(0) - job = self._jobs[job_id] - job.state = JobState.CLAIMED - job.worker_id = "master" - job.claimed_at = time.time() - batch.append(self._job_to_dict(job)) - - if batch: - self._persist_state() - - return batch - - def heartbeat(self, worker_id): - """Update worker heartbeat timestamp.""" - with self._lock: - if worker_id in self._workers: - self._workers[worker_id]["last_heartbeat"] = time.time() - - def register_worker(self, worker_id, worker_url=""): - """ - Register a new worker. - - Args: - worker_id: Unique worker identifier - worker_url: Worker's base URL (for display/monitoring only) - """ - with self._lock: - self._workers[worker_id] = { - "url": worker_url, - "last_heartbeat": time.time(), - "jobs_completed": 0, - "jobs_failed": 0, - "jobs_duplicated": 0, - "status": "idle" - } - print(f"[Distribution] 🤝 Worker registered: {worker_id} ({worker_url})") - - def record_late_submission(self, worker_id): - """ - Record a duplicate/late submission from a worker. - - Called when a worker submits a result for a job that was already completed - (e.g. because the job timed out, was reclaimed by master, and completed - before the worker finished). The worker did real work — this tracks it. - - Args: - worker_id: The worker that submitted late - """ - with self._lock: - if worker_id in self._workers: - self._workers[worker_id]["jobs_duplicated"] += 1 - - def get_status(self): - """Return overall progress status.""" - with self._lock: - total = len(self._jobs) - pending = len(self._pending_queue) - claimed = sum(1 for j in self._jobs.values() - if j.state == JobState.CLAIMED) - completed = sum(1 for j in self._jobs.values() - if j.state == JobState.COMPLETED) - failed = sum(1 for j in self._jobs.values() - if j.state == JobState.FAILED) - - return { - "total": total, - "pending": pending, - "claimed": claimed, - "completed": completed, - "failed": failed, - "workers": { - wid: info.copy() - for wid, info in self._workers.items() - } - } - - def deactivate(self): - """Mark the distribution manager as inactive (generation complete).""" - with self._lock: - self._active = False - - @property - def is_active(self): - """Check if distribution is still active.""" - with self._lock: - return self._active - - @property - def has_pending_or_claimed(self): - """True if any jobs are still pending or claimed (excluding failed).""" - with self._lock: - return any( - j.state in (JobState.PENDING, JobState.CLAIMED) - for j in self._jobs.values() - ) - - @property - def total_completed(self): - """Number of completed jobs.""" - with self._lock: - return self._total_completed - - def set_encoded_conditionings(self, encoded_dict): - """ - Store pre-encoded conditionings from the master's pre-encoding phase. - - Args: - encoded_dict: Dict keyed by "(model_cache_key|lora_key)" with structure: - {"positive": {prompt_text: serialized_list}, "negative": {prompt_text: serialized_list}} - """ - with self._lock: - self._encoded_conditionings = encoded_dict - total_groups = len(encoded_dict) - total_prompts = sum( - len(v.get("positive", {})) + len(v.get("negative", {})) - for v in encoded_dict.values() - ) - print(f"[Distribution] 🧠 Stored pre-encoded conditionings: " - f"{total_groups} model/LoRA groups, {total_prompts} total prompt encodings") - - def get_encoded_conditionings_for_job(self, job): - """ - Look up pre-encoded conditionings for a specific job. - - Uses the job's config to compute the group key and prompt text, - then returns the matching serialized positive/negative conditionings. - - Args: - job: DistributedJob instance - - Returns: - dict with "encoded_positive" and "encoded_negative" keys, or None if not available. - """ - if not self._encoded_conditionings: - return None - - try: - from .generation_orchestrator import get_model_cache_key - from .trigger_words import build_prompt_with_triggers - - config = job.config - - # Build the group key: "model_cache_key|lora_key" - model_key = get_model_cache_key(config) - lora_key = config.get("lora_expanded", config.get("lora", "None")) - group_key = f"{model_key}|{lora_key}" - - group = self._encoded_conditionings.get(group_key) - if not group: - return None - - # Build the actual prompt text (with trigger words applied) - lora_triggerwords_mode = config.get("_lora_triggerwords_mode", "None") - try: - actual_positive, _ = build_prompt_with_triggers(config, lora_triggerwords_mode) - except Exception: - actual_positive = config.get("positive", "") - actual_negative = config.get("negative", "") - - encoded_positive = group.get("positive", {}).get(actual_positive) - encoded_negative = group.get("negative", {}).get(actual_negative) - - if encoded_positive is not None and encoded_negative is not None: - return { - "encoded_positive": encoded_positive, - "encoded_negative": encoded_negative, - # Include final prompt text so workers can use it for metadata - # without needing to re-run trigger word lookups - "actual_positive_text": actual_positive, - "actual_negative_text": actual_negative, - } - - return None - - except Exception as e: - print(f"[Distribution] ⚠️ Error looking up conditionings for job {job.job_id}: {e}") - return None - - def _release_timed_out_jobs(self): - """ - Release jobs that have been claimed but not completed within timeout. - Must be called with lock held. - - Uses the worker's last heartbeat timestamp (not just job claimed_at) to - determine if the worker is still alive. A slow-but-alive worker that - sends heartbeats will NOT have its jobs reclaimed. Only workers that - have stopped heartbeating (crash/disconnect) will have jobs reclaimed. - """ - now = time.time() - released = 0 - for job in self._jobs.values(): - if (job.state == JobState.CLAIMED - and job.worker_id != "master" - and job.claimed_at - and now - job.claimed_at > self.claim_timeout): - - # Check if the worker is still heartbeating — if so, it's alive - # and just slow (e.g. slow GPU, first-time model loading). - # Only reclaim if the worker has gone silent. - worker = self._workers.get(job.worker_id) - if worker: - last_hb = worker.get("last_heartbeat", 0) - if now - last_hb < self.claim_timeout: - # Worker is still alive — don't reclaim - continue - - print(f"[Distribution] ⏰ Job {job.job_id} timed out " - f"from worker {job.worker_id} (no heartbeat), re-queuing") - job.state = JobState.PENDING - job.worker_id = None - job.claimed_at = None - self._pending_queue.append(job.job_id) - released += 1 - - if released > 0: - print(f"[Distribution] Released {released} timed-out jobs") - - def release_timed_out_jobs(self): - """Public wrapper — call from the wait loop to reclaim dead worker jobs.""" - with self._lock: - self._release_timed_out_jobs() - - def _persist_state(self): - """ - Persist job states to disk for crash recovery. - Must be called with lock held. - """ - state = { - "session_name": self.session_name, - "timestamp": time.time(), - "jobs": { - jid: { - "state": j.state, - "worker_id": j.worker_id, - "fail_count": j.fail_count - } - for jid, j in self._jobs.items() - } - } - try: - with open(self._persistence_path, "w") as f: - json.dump(state, f) - except Exception as e: - print(f"[Distribution] ⚠️ Persist error: {e}") - - def _get_upcoming_models(self, exclude_job_id=None, limit=10): - """ - Scan pending jobs and collect unique model references for pre-fetching. - Must be called with lock held. - - Args: - exclude_job_id: Job to skip (the current job being claimed) - limit: Maximum number of unique model references to return - - Returns: - List of {category, filename} dicts - """ - seen = set() - upcoming = [] - - for job_id in self._pending_queue: - if job_id == exclude_job_id: - continue - job = self._jobs.get(job_id) - if not job: - continue - - config = job.config - - # Main model/checkpoint - model_name = config.get("model", "None") - if model_name and model_name != "None": - model_type = config.get("model_type", "checkpoint") - if model_type == "checkpoint": - cat = "checkpoints" - elif model_type == "gguf": - cat = "unet_gguf" - elif model_type == "diffusion_model": - cat = "diffusion_models" - else: - cat = "checkpoints" - key = f"{cat}:{model_name}" - if key not in seen: - seen.add(key) - upcoming.append({"category": cat, "filename": model_name}) - - # LoRA — combo string: "name:str:str + name2:str:str", may include folders - lora = config.get("lora_expanded", config.get("lora", "None")) - if lora and lora != "None": - for lora_part in lora.split(" + "): - lora_filename = lora_part.strip().split(":")[0].strip() - if not lora_filename or lora_filename == "None": - continue - # Skip folder references — workers expand those locally - if lora_filename.endswith("/") or lora_filename.endswith("/*") or ("[" in lora_filename and "]" in lora_filename): - continue - key = f"loras:{lora_filename}" - if key not in seen: - seen.add(key) - upcoming.append({"category": "loras", "filename": lora_filename}) - - # VAE (skip remote URLs and "Default") - vae = config.get("vae", "Default") - if vae and vae != "Default" and not vae.startswith("remote:"): - key = f"vae:{vae}" - if key not in seen: - seen.add(key) - upcoming.append({"category": "vae", "filename": vae}) - - # Text encoders - text_encoders = config.get("text_encoders", []) - for te in text_encoders: - if te and te != "None": - key = f"text_encoders:{te}" - if key not in seen: - seen.add(key) - upcoming.append({"category": "text_encoders", "filename": te}) - - if len(upcoming) >= limit: - break - - return upcoming - - def _job_to_dict(self, job): - """ - Convert a DistributedJob to a serializable dict for API responses. - Note: input_job latent tensors are NOT included (workers create their own). - Config is round-tripped through JSON to ensure all values are serializable - (e.g. numpy int64 → Python int, sets → lists, etc.). - If pre-encoded conditionings are available, they are attached to the response - so workers can skip CLIP encoding entirely. - """ - try: - safe_config = json.loads(json.dumps(job.config, default=str)) - except Exception: - safe_config = job.config - - result = { - "job_id": job.job_id, - "config": safe_config, - "input_job": job.input_job, - "gen_index": job.gen_index - } - - # Attach session-level cooldown settings to worker config - if hasattr(self, '_session_settings') and self._session_settings and 'cooldown' in self._session_settings: - safe_config["_session_cooldown"] = self._session_settings["cooldown"] - - # Attach image_format to worker config so distributed workers honour the user's choice - if hasattr(self, '_session_settings') and self._session_settings: - safe_config["_session_image_format"] = self._session_settings.get("image_format", "webp") - - # Attach pre-encoded conditionings if available (master encoding feature) - encoded = self.get_encoded_conditionings_for_job(job) - if encoded: - result["encoded_positive"] = encoded["encoded_positive"] - result["encoded_negative"] = encoded["encoded_negative"] - # Include final prompt text (with triggers applied) so workers can - # use it for metadata without re-running CivitAI trigger lookups - result["actual_positive_text"] = encoded.get("actual_positive_text") - result["actual_negative_text"] = encoded.get("actual_negative_text") - - # Attach upcoming models list if sync is enabled (for pre-fetching) - if self.sync_models_to_workers: - result["upcoming_models"] = self._get_upcoming_models( - exclude_job_id=job.job_id, limit=10 - ) - - return result