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 <noreply@anthropic.com>

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

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(

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)

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")