From fddffd7a8c67a45eab52391baea03d0dcd35373d Mon Sep 17 00:00:00 2001 From: JasonHoku Date: Mon, 2 Mar 2026 15:15:54 -0800 Subject: [PATCH] Add worker model sync and upcoming models pre-fetching Worker downloads missing models from master via HTTP before processing. Background pre-fetch thread downloads upcoming models during job execution. Distribution manager scans pending queue for unique model references. Job claim responses include upcoming_models list when sync enabled. Master passes sync_models_to_workers flag to workers via start_worker. Co-Authored-By: Claude Opus 4.6 --- distribution_manager.py | 79 +++++++++++++++ distribution_routes.py | 8 +- distribution_worker.py | 199 +++++++++++++++++++++++++++++++++++++ generation_orchestrator.py | 7 +- 4 files changed, 290 insertions(+), 3 deletions(-) diff --git a/distribution_manager.py b/distribution_manager.py index d220ebd..a1d7877 100644 --- a/distribution_manager.py +++ b/distribution_manager.py @@ -69,6 +69,8 @@ class DistributionManager: # 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 def populate_jobs(self, expanded_configs, input_jobs, existing_data, overwrite_existing, has_optional_inputs, lora_triggerwords_mode): @@ -515,6 +517,77 @@ class DistributionManager: 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 + lora = config.get("lora_expanded", config.get("lora", "None")) + if lora and lora != "None": + key = f"loras:{lora}" + if key not in seen: + seen.add(key) + upcoming.append({"category": "loras", "filename": lora}) + + # 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. @@ -546,4 +619,10 @@ class DistributionManager: 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 diff --git a/distribution_routes.py b/distribution_routes.py index 9833870..8aa4747 100644 --- a/distribution_routes.py +++ b/distribution_routes.py @@ -336,6 +336,8 @@ async def start_worker(request): return web.Response(status=400, text="Missing master_url") thread = WorkerThread(master_url) + # Pass sync models flag to worker + thread._sync_models = data.get("sync_models_to_workers", False) set_worker_thread(thread) thread.start() @@ -549,7 +551,7 @@ def _send_distribution_status(manager): pass -def notify_workers_to_start(worker_urls, master_url, session_name): +def notify_workers_to_start(worker_urls, master_url, session_name, sync_models_to_workers=False): """ Notify remote worker ComfyUI instances to start processing. Sends POST /distribution/start_worker to each worker URL. @@ -560,6 +562,7 @@ def notify_workers_to_start(worker_urls, master_url, session_name): worker_urls: List of worker base URLs master_url: This master's base URL session_name: Current session name + sync_models_to_workers: Whether workers should download missing models from master Returns: List of (url, success, message) tuples @@ -577,7 +580,8 @@ def notify_workers_to_start(worker_urls, master_url, session_name): try: data = json.dumps({ "master_url": master_url, - "session_name": session_name + "session_name": session_name, + "sync_models_to_workers": sync_models_to_workers }).encode("utf-8") req = urllib.request.Request( diff --git a/distribution_worker.py b/distribution_worker.py index 9c8b6a8..51f3f97 100644 --- a/distribution_worker.py +++ b/distribution_worker.py @@ -59,6 +59,192 @@ class WorkerThread(threading.Thread): 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 + url = f"{self.master_url}/distribution/download_model" + payload = json.dumps({ + "worker_id": self.worker_id, + "category": category, + "filename": filename + }).encode("utf-8") + + req = urllib.request.Request( + url, data=payload, + headers={"Content-Type": "application/json"}, + method="POST" + ) + + try: + print(f"[Worker {self.worker_id}] ⬇️ Downloading {category}/{filename}...") + with urllib.request.urlopen(req, timeout=3600) as resp: + file_size = int(resp.headers.get("Content-Length", 0)) + if file_size > 0: + size_str = f"{file_size / (1024**3):.2f} GB" if file_size > 1024**3 else f"{file_size / (1024**2):.0f} MB" + print(f"[Worker {self.worker_id}] ⬇️ Size: {size_str}") + + # Stream to temp file, then rename (atomic-ish) + temp_path = target_path + ".downloading" + downloaded = 0 + last_progress = 0 + with open(temp_path, "wb") as f: + while True: + chunk = resp.read(8 * 1024 * 1024) # 8 MB chunks + if not chunk: + break + f.write(chunk) + downloaded += len(chunk) + if file_size > 0: + progress = int(downloaded / file_size * 100) + if progress >= last_progress + 10: + print(f"[Worker {self.worker_id}] ⬇️ {category}/{filename}: {progress}%") + last_progress = progress + + os.replace(temp_path, target_path) + 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 + """ + 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 + lora = config.get("lora_expanded", config.get("lora", "None")) + if lora and lora != "None": + models.append({"category": "loras", "filename": lora}) + + # 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}") @@ -108,6 +294,19 @@ class WorkerThread(threading.Thread): 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) diff --git a/generation_orchestrator.py b/generation_orchestrator.py index e98aee1..5dfc4e2 100644 --- a/generation_orchestrator.py +++ b/generation_orchestrator.py @@ -1570,7 +1570,12 @@ def _run_distributed_generation( print(f"[Distribution] 🌐 Master URL: {master_url}") print(f"[Distribution] 🌐 Notifying {len(worker_urls)} worker(s)...") - worker_results = notify_workers_to_start(worker_urls, master_url, session_name) + sync_models = distribution_config.get("sync_models_to_workers", False) + if sync_models: + manager.sync_models_to_workers = True + print(f"[Distribution] ☁️ Model sync enabled — workers will download missing models from master") + + worker_results = notify_workers_to_start(worker_urls, master_url, session_name, sync_models_to_workers=sync_models) successful_workers = sum(1 for _, ok, _ in worker_results if ok) print(f"[Distribution] ✅ {successful_workers}/{len(worker_urls)} workers started")