From cc27f6304021ce9bd8a069a6273cb64bf526e5a8 Mon Sep 17 00:00:00 2001 From: Jason Hoku Date: Tue, 19 May 2026 15:23:17 -0700 Subject: [PATCH] refactor(distribution): delete distribution_worker.py (~815 lines) WorkerThread + _MasterFinishedError moved to ComfyUI-USCG-Distributed companion. Worker's Remote VAE call now imports comfyui_uscg_remote_vae directly, bypassing the Phase 1 facade alias. --- distribution_worker.py | 815 ----------------------------------------- 1 file changed, 815 deletions(-) delete mode 100644 distribution_worker.py diff --git a/distribution_worker.py b/distribution_worker.py deleted file mode 100644 index 7b52001..0000000 --- a/distribution_worker.py +++ /dev/null @@ -1,815 +0,0 @@ -""" -Distributed Worker Thread -Runs on remote ComfyUI instances. Polls the master for jobs, -processes them using the local GPU, and uploads results back. -""" - -import threading -import time -import uuid -import json -import io -import os -import random -from urllib.parse import quote as url_quote -from .network_utils import ( - distribution_get, - distribution_post_json, - distribution_post_multipart, - distribution_stream_download, -) - -import torch - - -class _MasterFinishedError(Exception): - """Raised when the master has deactivated (503 on submit). - Not a real error — means master already completed all jobs.""" - pass - - -class WorkerThread(threading.Thread): - """ - Background thread that polls master for jobs and processes them. - - Lifecycle: - 1. Register with master - 2. Claim job (GET /distribution/claim_job) - 3. Process job locally (model load → encode → sample → VAE decode) - 4. Submit result (POST /distribution/submit_result) - 5. Repeat from 2 until no more jobs - 6. Exit - """ - - def __init__(self, master_url): - super().__init__(daemon=True) - self.master_url = master_url.rstrip("/") - self.worker_id = f"w_{uuid.uuid4().hex[:8]}" - self.jobs_processed = 0 - self.current_model = None - self._stop_event = threading.Event() - self.poll_interval = 2 # seconds between polls when idle - self.heartbeat_interval = 30 # seconds between heartbeats - self.consecutive_empty = 0 # Track consecutive empty poll responses - self.max_empty_polls = 30 # Stop after this many consecutive empty polls - - # Cached state for model reuse between jobs - self._loaded_model = None - self._loaded_clip = None - self._loaded_vae = None - self._patched_model = None - self._patched_clip = None - self._cached_model_key = None - self._cached_lora_key = None - self._incompatible_loras = {} - - # Model sync: download missing models from master before processing - self._sync_models = False # Set by master when sync_models_to_workers enabled - self._prefetch_thread = None - self._prefetch_queue = [] # List of {category, filename} to pre-fetch - self._prefetch_lock = threading.Lock() - - # ========================================================================= - # MODEL SYNC HELPERS - # ========================================================================= - - def _download_model_from_master(self, category, filename): - """ - Download a single model file from the master instance. - Preserves full subdirectory structure. - - Args: - category: Model category (checkpoints, loras, vae, etc.) - filename: Relative path including subdirs (e.g. SDXL/Illustrious/model.safetensors) - - Returns: - str: Local path where file was saved, or None on failure - """ - import folder_paths - - # Determine target directory - # folder_paths stores base directories per category - base_dirs = folder_paths.get_folder_paths(category) - if not base_dirs: - print(f"[Worker {self.worker_id}] ❌ Unknown model category: {category}") - return None - target_dir = base_dirs[0] # Use first (primary) directory - target_path = os.path.join(target_dir, filename.replace("/", os.sep)) - - # Skip if already exists - if os.path.exists(target_path): - return target_path - - # Create subdirectories - os.makedirs(os.path.dirname(target_path), exist_ok=True) - - # Download from master via network gateway - try: - print(f"[Worker {self.worker_id}] ⬇️ Downloading {category}/{filename}...") - last_progress = 0 - - def on_progress(downloaded, total_size): - nonlocal last_progress - if total_size > 0: - pct = int(downloaded * 100 / total_size) - if pct >= last_progress + 10: - last_progress = pct - size_str = f"{total_size / (1024**3):.2f} GB" if total_size > 1024**3 else f"{total_size / (1024**2):.0f} MB" - print(f"[Worker {self.worker_id}] ⬇️ {category}/{filename}: {pct}% of {size_str}") - - distribution_stream_download( - url=f"{self.master_url}/distribution/download_model", - data_dict={ - "worker_id": self.worker_id, - "category": category, - "filename": filename - }, - target_path=target_path, - timeout=3600, - progress_callback=on_progress - ) - print(f"[Worker {self.worker_id}] ✅ Downloaded {category}/{filename}") - return target_path - - except Exception as e: - print(f"[Worker {self.worker_id}] ❌ Failed to download {category}/{filename}: {e}") - # Clean up partial download - temp_path = target_path + ".downloading" - if os.path.exists(temp_path): - os.remove(temp_path) - return None - - def _extract_required_models(self, config): - """ - Extract all model file references from a job config. - - Returns: - List of {category, filename} dicts - """ - import folder_paths - models = [] - - # 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": - category = "checkpoints" - elif model_type == "gguf": - category = "unet_gguf" - elif model_type == "diffusion_model": - category = "diffusion_models" - else: - category = "checkpoints" - models.append({"category": category, "filename": model_name}) - - # LoRA — combo string format: "name:model_str:clip_str + name2:model_str:clip_str" - # Parts can also be folder references like "FolderName/:0.90:0.90" which expand - # to all LoRA files in that folder, or random syntax "Folder/[3,0.5]" - lora = config.get("lora_expanded", config.get("lora", "None")) - if lora and lora != "None": - for lora_part in lora.split(" + "): - lora_part = lora_part.strip() - if not lora_part: - continue - # Extract just the filename (before first ":") - lora_filename = lora_part.split(":")[0].strip() - if not lora_filename or lora_filename == "None": - continue - - # Check if it's a folder reference (ends with / or /*) - clean_name = lora_filename.rstrip("/*").rstrip("/") - is_folder = lora_filename.endswith("/") or lora_filename.endswith("/*") - # Also check for random selection syntax: folder/[count,strength] - has_bracket = "[" in lora_filename and "]" in lora_filename - if has_bracket: - clean_name = lora_filename.split("[")[0].rstrip("/") - is_folder = True - - if is_folder: - # Expand folder to individual LoRA files - folder_prefix = clean_name.replace("\\", "/") + "/" - all_loras = folder_paths.get_filename_list("loras") - for lora_file in all_loras: - if lora_file.replace("\\", "/").startswith(folder_prefix): - models.append({"category": "loras", "filename": lora_file}) - else: - models.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:"): - models.append({"category": "vae", "filename": vae}) - - # Text encoders - text_encoders = config.get("text_encoders", []) - for te in text_encoders: - if te and te != "None": - models.append({"category": "text_encoders", "filename": te}) - - return models - - def _ensure_models_available(self, config): - """ - Check all required models exist locally, download missing ones from master. - Blocks until all models are available. - """ - import folder_paths - - required = self._extract_required_models(config) - if not required: - return - - missing = [] - for m in required: - path = folder_paths.get_full_path(m["category"], m["filename"]) - if path is None: - missing.append(m) - - if not missing: - return - - print(f"[Worker {self.worker_id}] 📦 Missing {len(missing)} model(s), downloading from master...") - for m in missing: - result = self._download_model_from_master(m["category"], m["filename"]) - if result is None: - raise RuntimeError( - f"Failed to download required model: {m['category']}/{m['filename']}. " - f"Check that the file exists on the master instance." - ) - - def _start_prefetch(self, upcoming_models): - """Start background thread to pre-fetch upcoming models.""" - if not upcoming_models: - return - - with self._prefetch_lock: - self._prefetch_queue = list(upcoming_models) - - if self._prefetch_thread and self._prefetch_thread.is_alive(): - return # Already running - - def _prefetch_worker(): - while True: - with self._prefetch_lock: - if not self._prefetch_queue: - break - model = self._prefetch_queue.pop(0) - try: - self._download_model_from_master(model["category"], model["filename"]) - except Exception as e: - print(f"[Worker {self.worker_id}] ⚠️ Pre-fetch failed: {e}") - - self._prefetch_thread = threading.Thread(target=_prefetch_worker, daemon=True) - self._prefetch_thread.start() - - def run(self): - """Main worker loop.""" - print(f"[Worker {self.worker_id}] 🚀 Starting worker thread, master: {self.master_url}") - - # Worker threads run outside ComfyUI's normal prompt execution context. - # The sampling callback (latent_preview.py → comfy/utils.py → main.py hook) - # expects PromptServer.instance.last_prompt_id to exist. Set a dummy value - # so the progress bar callback doesn't crash with AttributeError. - try: - from server import PromptServer - if PromptServer and hasattr(PromptServer, 'instance') and PromptServer.instance is not None: - if not hasattr(PromptServer.instance, 'last_prompt_id'): - PromptServer.instance.last_prompt_id = f"dist_worker_{self.worker_id}" - print(f"[Worker {self.worker_id}] 🔧 Set last_prompt_id for sampling callbacks") - except Exception: - pass - - # Register with master - self._register() - - # Start background heartbeat thread so heartbeats flow even during - # long-running job processing (model loading, generation, etc.). - # Without this, the main loop blocks on _process_job() and no heartbeats - # are sent, causing the master to wrongly reclaim jobs from slow workers. - self._heartbeat_thread = threading.Thread( - target=self._heartbeat_loop, daemon=True - ) - self._heartbeat_thread.start() - - while not self._stop_event.is_set(): - # Claim a job - job = self._claim_job() - - if job is None: - self.consecutive_empty += 1 - if self.consecutive_empty >= self.max_empty_polls: - print(f"[Worker {self.worker_id}] 🏁 No more jobs after " - f"{self.max_empty_polls} polls, stopping") - break - self._stop_event.wait(self.poll_interval) - continue - - self.consecutive_empty = 0 - - # Validate job has required fields - if not all(k in job for k in ("job_id", "config", "input_job")): - print(f"[Worker {self.worker_id}] ⚠️ Malformed job, skipping: {list(job.keys())}") - continue - - # Sync missing models from master if enabled - try: - if self._sync_models: - self._ensure_models_available(job["config"]) - - # Start background pre-fetch for upcoming models - if job.get("upcoming_models"): - self._start_prefetch(job["upcoming_models"]) - except Exception as e: - print(f"[Worker {self.worker_id}] ❌ Model sync failed for job {job['job_id']}: {e}") - self._report_failure(job["job_id"], f"Model sync failed: {e}") - continue - - # Process the job - try: - image_bytes, meta = self._process_job(job) - - # Submit result (with retry on network error) - self._submit_result(job["job_id"], image_bytes, meta) - self.jobs_processed += 1 - - print(f"[Worker {self.worker_id}] ✅ Job {job['job_id']} complete " - f"(total: {self.jobs_processed})") - - # Periodic VRAM cleanup - if self.jobs_processed % 5 == 0: - import gc - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - except _MasterFinishedError: - # Master has shut down (503 on submit). This is NOT a failure — - # the master already finished all jobs (possibly reclaimed ours). - # Stop gracefully instead of reporting failure. - print(f"[Worker {self.worker_id}] 🏁 Master finished before we could submit " - f"job {job['job_id']} — stopping gracefully") - break - - except Exception as e: - # Check if ComfyUI interrupted processing on this worker - try: - from comfy.model_management import InterruptProcessingException - if isinstance(e, InterruptProcessingException): - print(f"[Worker {self.worker_id}] 🛑 Processing interrupted, stopping worker") - self._report_failure(job["job_id"], "Worker interrupted") - break - except ImportError: - pass - - print(f"[Worker {self.worker_id}] ❌ Job {job['job_id']} failed: {e}") - import traceback - traceback.print_exc() - self._report_failure(job["job_id"], str(e)) - - # Cleanup - self._cleanup_models() - print(f"[Worker {self.worker_id}] 🛑 Worker stopped. " - f"Processed {self.jobs_processed} jobs.") - - def stop(self): - """Signal the worker to stop after its current job.""" - self._stop_event.set() - - def _register(self): - """Register with master. Retries on 503 (master still starting up).""" - max_retries = 5 - retry_delay = 3 # seconds between retries - - for attempt in range(1, max_retries + 1): - try: - status, resp_data = distribution_post_json( - f"{self.master_url}/distribution/register_worker", - {"worker_id": self.worker_id, "worker_url": ""}, - timeout=10 - ) - result = json.loads(resp_data.decode("utf-8")) - print(f"[Worker {self.worker_id}] 🤝 Registered with master " - f"(session: {result.get('session_name', 'unknown')})") - return # Success - except Exception as e: - if hasattr(e, 'code') and e.code == 503 and attempt < max_retries: - print(f"[Worker {self.worker_id}] ⏳ Master not ready yet " - f"(attempt {attempt}/{max_retries}), retrying in {retry_delay}s...") - self._stop_event.wait(retry_delay) - if self._stop_event.is_set(): - return - continue - print(f"[Worker {self.worker_id}] ⚠️ Registration failed: {e}") - return - - def _claim_job(self): - """Claim next job from master. Returns job dict or None.""" - url = (f"{self.master_url}/distribution/claim_job" - f"?worker_id={url_quote(self.worker_id)}") - - if self.consecutive_empty == 0 and self.jobs_processed == 0: - print(f"[Worker {self.worker_id}] 🔍 Claiming from: {url}") - - try: - status, resp_data = distribution_get(url, timeout=10) - if status == 204: - return None - return json.loads(resp_data.decode("utf-8")) - except Exception as e: - if hasattr(e, 'code'): - if e.code == 204: - return None - if e.code == 503: - # Distribution not active yet — master may still be starting up - # (e.g. pre-encoding prompts). Don't immediately stop; let the - # normal empty-poll counter handle eventual timeout. - if self.jobs_processed == 0 and self.consecutive_empty < 5: - print(f"[Worker {self.worker_id}] ⏳ Master not active yet, " - f"waiting... ({self.consecutive_empty + 1}/5)") - elif self.jobs_processed > 0: - # Already processed jobs — master has finished and deactivated. - # Stop immediately to avoid wasting time. - print(f"[Worker {self.worker_id}] ℹ️ Distribution ended on master") - self.consecutive_empty = self.max_empty_polls - else: - print(f"[Worker {self.worker_id}] ℹ️ Distribution not active on master") - self.consecutive_empty = self.max_empty_polls - return None - print(f"[Worker {self.worker_id}] ⚠️ Claim failed: HTTP {e.code}") - return None - print(f"[Worker {self.worker_id}] ⚠️ Claim failed: {e}") - return None - - @torch.inference_mode() - def _process_job(self, job): - """ - Process a single job using local ComfyUI infrastructure. - - Steps: - 1. Load model/LoRA (reuse if same as previous job) - 2. Encode prompts with CLIP - 3. Create latent noise tensor - 4. Generate image (KSampler) - 5. Decode with local VAE - 6. Return image bytes + metadata - - Args: - job: Job dict from master with config, input_job, gen_index - - Returns: - tuple: (image_bytes, metadata_dict) - - Note: @torch.inference_mode() is required because worker threads run - outside ComfyUI's execution engine which wraps everything in - inference_mode. Without it, Dynamic VRAM's lazy weight loading can - produce NaN values during VAE decode. - """ - from .model_loader import ( - load_checkpoint, load_loras, load_vae_by_name, - load_diffusion_model_and_clip, get_latent_channels - ) - from .image_generation import ( - generate_image, decode_latent_with_vae, create_image_metadata - ) - from .batch_encoding import encode_prompt_with_combinators - from .trigger_words import build_prompt_with_triggers - from .generation_orchestrator import get_model_cache_key - - config = job["config"] - input_job = job["input_job"] - w = input_job["width"] - h = input_job["height"] - batch_idx = input_job.get("batch_idx", 0) - - # Handle seed_behavior: "randomize" (same as main generation loop) - if config.get("seed_behavior") == "randomize": - config["seed"] = random.randint(0, 2**63 - 1) - - # Full run seed behavior is handled by the orchestrator/master before dispatching jobs. - # Workers just use whatever seed is in the config they receive. - - # ==== GPU COOLDOWN (if enabled in session settings) ==== - cooldown_config = config.get("_session_cooldown", {}) - if cooldown_config.get("enabled", False): - if not hasattr(self, '_gen_count'): - self._gen_count = 0 - self._gen_count += 1 - cooldown_every_n = int(cooldown_config.get("every_n", 1)) - if self._gen_count % cooldown_every_n == 0: - cooldown_seconds = int(cooldown_config.get("seconds", 5)) - clear_vram = cooldown_config.get("clear_vram", False) - print(f"[Worker {self.worker_id}] ❄️ GPU Cooldown: pausing {cooldown_seconds}s after {self._gen_count} generations") - if clear_vram: - import comfy.model_management as mm_cooldown - mm_cooldown.soft_empty_cache() - mm_cooldown.unload_all_models() - print(f"[Worker {self.worker_id}] ❄️ VRAM cleared") - self._cached_model_key = None # Force model reload on next iteration - import time as time_cooldown - time_cooldown.sleep(cooldown_seconds) - print(f"[Worker {self.worker_id}] ❄️ Cooldown complete, resuming") - - # --- Model Loading (with caching between jobs) --- - target_model_key = get_model_cache_key(config) - self.current_model = config.get("model", "unknown") - - if target_model_key != self._cached_model_key: - # Need to load a different model - self._cleanup_models() - - model_type = config.get("model_type", "checkpoint") - if model_type == "checkpoint": - self._loaded_model, self._loaded_clip, self._loaded_vae = load_checkpoint( - config["model"], config["model"], False, - None, None, None, None, None, None, None, - model_cache=None - ) - else: - self._loaded_model, self._loaded_clip, self._loaded_vae = load_diffusion_model_and_clip( - model_name=config["model"], - model_type=model_type, - text_encoder_paths=config.get("text_encoders", []), - clip_type_str=config.get("clip_type", "stable_diffusion"), - gguf_options=config.get("gguf_options"), - use_remote_vae=False, - optional_model=None, - optional_clip=None, - optional_vae=None, - model_cache=None - ) - - self._cached_model_key = target_model_key - self._cached_lora_key = None # Force LoRA reload on model switch - print(f"[Worker {self.worker_id}] 📦 Loaded model: {config['model']}") - - # --- Per-config VAE loading --- - # Supports three modes: - # 1. "Default" → use model's bundled VAE - # 2. "remote:https://..." → remote HuggingFace VAE endpoint (decoded at decode step) - # 3. "ae.safetensors" → load local VAE file - config_vae = config.get("vae", "Default") - _remote_vae_url = None - if config_vae != "Default": - if config_vae.startswith("remote:"): - _remote_vae_url = config_vae[len("remote:"):] - vae = None # No local VAE needed — will use remote endpoint - print(f"[Worker {self.worker_id}] 🌐 Using remote VAE: {_remote_vae_url}") - else: - vae = load_vae_by_name(config_vae) - else: - vae = self._loaded_vae - - # Validate VAE is available (non-checkpoint models like GGUF may not bundle a VAE) - # Skip this check when using remote VAE — no local VAE object needed. - if vae is None and _remote_vae_url is None: - raise RuntimeError( - f"No VAE available for model '{config.get('model', 'unknown')}'. " - f"Non-checkpoint models (GGUF/diffusion) require a VAE to be specified " - f"in the config builder, or connect one via optional_vae on the master." - ) - - # --- LoRA Loading (with caching) --- - lora_key = config.get("lora_expanded", config.get("lora", "None")) - if lora_key != self._cached_lora_key: - if lora_key != "None": - self._patched_model, self._patched_clip, should_skip = load_loras( - self._loaded_model, self._loaded_clip, lora_key, - config["model"], self._incompatible_loras, - model_cache=None - ) - if should_skip: - raise RuntimeError(f"LoRA incompatible: {lora_key}") - else: - self._patched_model = self._loaded_model - self._patched_clip = self._loaded_clip - - self._cached_lora_key = lora_key - print(f"[Worker {self.worker_id}] 🔗 Applied LoRA: {lora_key}") - - # --- Prompt Encoding --- - # Check for master pre-encoded conditionings (skip CLIP encoding if available) - encoded_pos = job.get("encoded_positive") - encoded_neg = job.get("encoded_negative") - - if encoded_pos is not None and encoded_neg is not None: - # Master already encoded these prompts — deserialize and use directly. - # Use the master's prompt text for metadata (avoids CivitAI trigger lookups - # on the worker since the master already applied trigger words during encoding). - from .generation_orchestrator import _serializable_to_conditioning - pos_cond = _serializable_to_conditioning(encoded_pos) - neg_cond = _serializable_to_conditioning(encoded_neg) - actual_positive = job.get("actual_positive_text", config.get("positive", "")) - actual_negative = job.get("actual_negative_text", config.get("negative", "")) - print(f"[Worker {self.worker_id}] 🧠 Using master pre-encoded conditionings") - else: - # Fallback: encode locally with CLIP (need to build prompt with trigger words) - 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", "") - clip_skip = config.get("clip_skip", 0) - pos_cond = encode_prompt_with_combinators(self._patched_clip, actual_positive, clip_skip) - neg_cond = encode_prompt_with_combinators(self._patched_clip, actual_negative, clip_skip) - - # --- Create Latent --- - latent_channels = get_latent_channels(self._loaded_model, None) - latent_in = {"samples": torch.zeros([1, latent_channels, h // 8, w // 8])} - - # --- Generate Image --- - attention_mode = config.get("attention_mode", "default") - result_latent, duration = generate_image( - self._patched_model, config["seed"], config["steps"], config["cfg"], - config["sampler"], config["scheduler"], pos_cond, neg_cond, - latent_in, config["denoise"], - attention_mode=attention_mode, - model_sampling_override=config.get("model_sampling_override", "none"), - model_sampling_shift=config.get("model_sampling_shift", 1.73), - model_sampling_flux_max_shift=config.get("model_sampling_flux_max_shift", 1.15), - model_sampling_flux_base_shift=config.get("model_sampling_flux_base_shift", 0.5), - use_advanced_sampling=config.get("use_advanced_sampling", False), - advanced_guider=config.get("advanced_guider", "cfg_guider"), - advanced_scheduler=config.get("advanced_scheduler", "basic"), - flux_guidance_value=config.get("flux_guidance_value", 0.0), - width=w, - height=h - ) - - # --- VAE Decode --- - # Clone latent samples and free UNet-related tensors BEFORE VAE decode. - # With Dynamic VRAM, the model manager needs to evict the UNet to load - # the VAE. Holding references to UNet output tensors on GPU can interfere - # with memory eviction, causing the VAE to run in a degraded state that - # produces NaN values. - latent_for_decode = result_latent["samples"].clone() - del result_latent, latent_in, pos_cond, neg_cond - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - if _remote_vae_url: - # Remote VAE: send latents to HuggingFace endpoint for decoding - from .remote_vae import _companion_decode as remote_decode_hf - from PIL import Image as PILImage - import numpy as np - if latent_for_decode.ndim == 3: - latent_for_decode = latent_for_decode.unsqueeze(0) - decoded = remote_decode_hf(_remote_vae_url, latent_for_decode, h, w) - del latent_for_decode - # Denormalize from [-1, 1] → [0, 1] (same as VaeImageProcessor.postprocess) - decoded = (decoded / 2 + 0.5).clamp(0, 1) - img_data = decoded[0].permute(1, 2, 0).cpu().numpy() - image = PILImage.fromarray((img_data * 255).round().astype(np.uint8)) - else: - image = decode_latent_with_vae(vae, latent_for_decode) - del latent_for_decode - - # --- Convert to bytes --- - # Honour image_format from session settings (injected by master into config) - _raw_fmt = config.get("_session_image_format", "webp") - _fmt = (_raw_fmt or "webp").lower().strip() - if _fmt not in ("webp", "png", "jpg", "jpeg"): - _fmt = "webp" - buf = io.BytesIO() - if _fmt == "png": - image.save(buf, format="PNG") - elif _fmt in ("jpg", "jpeg"): - image.save(buf, format="JPEG", quality=95) - else: - image.save(buf, format="WEBP", quality=80) - image_bytes = buf.getvalue() - del image, buf # Free PIL image and buffer - - # --- Create metadata --- - ts = int(time.time() * 100000) + random.randint(0, 1000) - meta = create_image_metadata( - config, w, h, duration, config["seed"], - batch_idx, actual_positive, actual_negative, - gen_index=job.get("gen_index") - ) - meta["id"] = ts - # Pass the chosen extension to the master so it can save with the right filename - _ext = "jpg" if _fmt in ("jpg", "jpeg") else _fmt - meta["_image_ext"] = _ext - - return image_bytes, meta - - def _submit_result(self, job_id, image_bytes, meta, max_retries=3): - """ - Upload image + metadata to master via multipart POST. - Uses network_utils.py gateway for all outbound requests. - Retries on transient network errors (connection reset, timeout). - """ - # JSON-safe meta (ensure no non-serializable types from config round-trip) - try: - safe_meta = json.loads(json.dumps(meta, default=str)) - except Exception: - safe_meta = meta - - boundary = f"----DistWorker{uuid.uuid4().hex[:16]}" - - # Build multipart body manually - body = b"" - - # Metadata part - body += f"--{boundary}\r\n".encode() - body += b'Content-Disposition: form-data; name="metadata"\r\n' - body += b'Content-Type: application/json\r\n\r\n' - body += json.dumps({ - "job_id": job_id, - "meta": safe_meta, - "worker_id": self.worker_id - }).encode("utf-8") - body += b"\r\n" - - # Image part - body += f"--{boundary}\r\n".encode() - body += b'Content-Disposition: form-data; name="image"; filename="image.webp"\r\n' - body += b'Content-Type: image/webp\r\n\r\n' - body += image_bytes - body += b"\r\n" - - body += f"--{boundary}--\r\n".encode() - - last_error = None - for attempt in range(max_retries): - try: - status, _ = distribution_post_multipart( - f"{self.master_url}/distribution/submit_result", - f"multipart/form-data; boundary={boundary}", - body, - timeout=60 - ) - if status != 200: - raise RuntimeError(f"Submit failed: HTTP {status}") - return # Success - - except Exception as e: - # 503 = master deactivated (all jobs done). Don't retry — it won't help. - # Raise a special exception so the main loop can stop gracefully. - if hasattr(e, 'code') and e.code == 503: - raise _MasterFinishedError( - f"Master distribution is no longer active (HTTP 503)" - ) - last_error = e - if attempt < max_retries - 1: - wait = 2 ** attempt # Exponential backoff: 1s, 2s, 4s - print(f"[Worker {self.worker_id}] ⚠️ Submit retry {attempt + 1}/{max_retries} " - f"for job {job_id}: {e} (waiting {wait}s)") - time.sleep(wait) - - raise RuntimeError(f"Submit failed after {max_retries} attempts: {last_error}") - - def _send_heartbeat(self): - """Send heartbeat to master.""" - try: - distribution_post_json( - f"{self.master_url}/distribution/heartbeat", - {"worker_id": self.worker_id}, - timeout=5 - ) - except Exception: - pass - - def _heartbeat_loop(self): - """ - Background heartbeat loop. Runs in a separate daemon thread so that - heartbeats are sent continuously even while _process_job() blocks the - main worker loop. This prevents the master from incorrectly reclaiming - jobs from slow-but-alive workers. - """ - while not self._stop_event.is_set(): - self._send_heartbeat() - self._stop_event.wait(self.heartbeat_interval) - - def _report_failure(self, job_id, error): - """Report job failure to master.""" - try: - distribution_post_json( - f"{self.master_url}/distribution/fail_job", - { - "job_id": job_id, - "error": error, - "worker_id": self.worker_id - }, - timeout=10 - ) - except Exception: - pass - - def _cleanup_models(self): - """Release cached model references.""" - self._loaded_model = None - self._loaded_clip = None - self._loaded_vae = None - self._patched_model = None - self._patched_clip = None - self._cached_model_key = None - self._cached_lora_key = None - self.current_model = None - - import gc - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache()