diff --git a/api/worker_routes.py b/api/worker_routes.py index 162c081..aa1423c 100644 --- a/api/worker_routes.py +++ b/api/worker_routes.py @@ -4,6 +4,7 @@ import os import time import platform import subprocess +import socket import torch import aiohttp @@ -111,152 +112,147 @@ async def clear_launching_state(request): except Exception as e: return await handle_api_error(request, e, 500) + +def get_network_ips(): + """Get all network IPs, trying multiple methods.""" + ips = [] + hostname = socket.gethostname() + + # Method 1: Try socket.getaddrinfo + try: + addr_info = socket.getaddrinfo(hostname, None) + for info in addr_info: + ip = info[4][0] + if ip and ip not in ips and not ip.startswith('::'): # Skip IPv6 for now + ips.append(ip) + except (socket.gaierror, OSError): + pass + + # Method 2: Try to connect to external server and get local IP + try: + s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + s.connect(("8.8.8.8", 80)) # Google DNS + local_ip = s.getsockname()[0] + s.close() + if local_ip not in ips: + ips.append(local_ip) + except (OSError, socket.error): + pass + + # Method 3: Platform-specific commands + try: + if platform.system() == "Windows": + # Windows ipconfig + result = subprocess.run(["ipconfig"], capture_output=True, text=True) + lines = result.stdout.split('\n') + for i, line in enumerate(lines): + if 'IPv4' in line and i + 1 < len(lines): + ip = lines[i].split(':')[-1].strip() + if ip and ip not in ips: + ips.append(ip) + else: + # Unix/Linux/Mac ifconfig or ip addr + try: + result = subprocess.run(["ip", "addr"], capture_output=True, text=True) + except (FileNotFoundError, OSError): + try: + result = subprocess.run(["ifconfig"], capture_output=True, text=True) + except (FileNotFoundError, OSError): + result = None + + import re + ip_pattern = re.compile(r'inet\s+(\d+\.\d+\.\d+\.\d+)') + if result is not None: + for match in ip_pattern.finditer(result.stdout): + ip = match.group(1) + if ip and ip not in ips: + ips.append(ip) + except (OSError, subprocess.SubprocessError): + pass + + return ips + + +def get_recommended_ip(ips): + """Choose the best IP for master-worker communication.""" + # Priority order: + # 1. Private network ranges (192.168.x.x, 10.x.x.x, 172.16-31.x.x) + # 2. Other non-localhost IPs + # 3. Localhost as last resort + + private_ips = [] + public_ips = [] + + for ip in ips: + if ip.startswith('127.') or ip == 'localhost': + continue + elif (ip.startswith('192.168.') + or ip.startswith('10.') + or (ip.startswith('172.') and 16 <= int(ip.split('.')[1]) <= 31)): + private_ips.append(ip) + else: + public_ips.append(ip) + + # Prefer private IPs + if private_ips: + # Prefer 192.168 range as it's most common + for ip in private_ips: + if ip.startswith('192.168.'): + return ip + return private_ips[0] + elif public_ips: + return public_ips[0] + elif ips: + return ips[0] + else: + return None + + +def _get_cuda_info(): + """Detect CUDA device index and total physical GPU count. + + Returns (cuda_device, cuda_device_count, physical_device_count). + All three are 0/None if CUDA is unavailable. + """ + if not torch.cuda.is_available(): + return None, 0, 0 + try: + cuda_device_count = torch.cuda.device_count() + cuda_visible = os.environ.get('CUDA_VISIBLE_DEVICES', '') + if cuda_visible and cuda_visible.strip(): + visible_devices = [int(d.strip()) for d in cuda_visible.split(',') if d.strip().isdigit()] + if visible_devices: + cuda_device = visible_devices[0] + try: + result = subprocess.run( + ['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'], + capture_output=True, + text=True, + timeout=5, + ) + physical_device_count = ( + len(result.stdout.strip().split('\n')) + if result.returncode == 0 + else max(visible_devices) + 1 + ) + except (FileNotFoundError, OSError, subprocess.SubprocessError): + physical_device_count = max(visible_devices) + 1 + return cuda_device, cuda_device_count, physical_device_count + else: + return 0, cuda_device_count, cuda_device_count + else: + cuda_device = torch.cuda.current_device() + return cuda_device, cuda_device_count, cuda_device_count + except Exception as e: + debug_log(f"CUDA detection error: {e}") + return None, 0, 0 + + @server.PromptServer.instance.routes.get("/distributed/network_info") async def get_network_info_endpoint(request): """Get network interfaces and recommend best IP for master.""" - import socket - - # Get CUDA device if available - cuda_device = None - cuda_device_count = 0 - physical_device_count = 0 - - if torch.cuda.is_available(): - try: - import os - import subprocess - - # Get visible device count (what PyTorch sees) - cuda_device_count = torch.cuda.device_count() - - # Try to get actual physical device info - cuda_visible = os.environ.get('CUDA_VISIBLE_DEVICES', '') - - # Method 1: Parse CUDA_VISIBLE_DEVICES - if cuda_visible and cuda_visible.strip(): - visible_devices = [int(d.strip()) for d in cuda_visible.split(',') if d.strip().isdigit()] - if visible_devices: - # Get the first visible device as the actual physical device - cuda_device = visible_devices[0] - - # Try to get total physical device count using nvidia-smi - try: - result = subprocess.run(['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'], - capture_output=True, text=True, timeout=5) - if result.returncode == 0: - physical_device_count = len(result.stdout.strip().split('\n')) - else: - physical_device_count = max(visible_devices) + 1 # Best guess - except (FileNotFoundError, OSError, subprocess.SubprocessError): - physical_device_count = max(visible_devices) + 1 # Best guess - else: - cuda_device = 0 - physical_device_count = cuda_device_count - else: - # No CUDA_VISIBLE_DEVICES set, current device is actual device - cuda_device = torch.cuda.current_device() - physical_device_count = cuda_device_count - - except Exception as e: - debug_log(f"CUDA detection error: {e}") - cuda_device = None - cuda_device_count = 0 - physical_device_count = 0 - - def get_network_ips(): - """Get all network IPs, trying multiple methods.""" - ips = [] - hostname = socket.gethostname() - - # Method 1: Try socket.getaddrinfo - try: - addr_info = socket.getaddrinfo(hostname, None) - for info in addr_info: - ip = info[4][0] - if ip and ip not in ips and not ip.startswith('::'): # Skip IPv6 for now - ips.append(ip) - except (socket.gaierror, OSError): - pass - - # Method 2: Try to connect to external server and get local IP - try: - s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) - s.connect(("8.8.8.8", 80)) # Google DNS - local_ip = s.getsockname()[0] - s.close() - if local_ip not in ips: - ips.append(local_ip) - except (OSError, socket.error): - pass - - # Method 3: Platform-specific commands - - try: - if platform.system() == "Windows": - # Windows ipconfig - result = subprocess.run(["ipconfig"], capture_output=True, text=True) - lines = result.stdout.split('\n') - for i, line in enumerate(lines): - if 'IPv4' in line and i + 1 < len(lines): - ip = lines[i].split(':')[-1].strip() - if ip and ip not in ips: - ips.append(ip) - else: - # Unix/Linux/Mac ifconfig or ip addr - try: - result = subprocess.run(["ip", "addr"], capture_output=True, text=True) - except (FileNotFoundError, OSError): - try: - result = subprocess.run(["ifconfig"], capture_output=True, text=True) - except (FileNotFoundError, OSError): - result = None - - import re - ip_pattern = re.compile(r'inet\s+(\d+\.\d+\.\d+\.\d+)') - if result is not None: - for match in ip_pattern.finditer(result.stdout): - ip = match.group(1) - if ip and ip not in ips: - ips.append(ip) - except (OSError, subprocess.SubprocessError): - pass - - return ips - - def get_recommended_ip(ips): - """Choose the best IP for master-worker communication.""" - # Priority order: - # 1. Private network ranges (192.168.x.x, 10.x.x.x, 172.16-31.x.x) - # 2. Other non-localhost IPs - # 3. Localhost as last resort - - private_ips = [] - public_ips = [] - - for ip in ips: - if ip.startswith('127.') or ip == 'localhost': - continue - elif (ip.startswith('192.168.') or - ip.startswith('10.') or - (ip.startswith('172.') and 16 <= int(ip.split('.')[1]) <= 31)): - private_ips.append(ip) - else: - public_ips.append(ip) - - # Prefer private IPs - if private_ips: - # Prefer 192.168 range as it's most common - for ip in private_ips: - if ip.startswith('192.168.'): - return ip - return private_ips[0] - elif public_ips: - return public_ips[0] - elif ips: - return ips[0] - else: - return None - + cuda_device, cuda_device_count, physical_device_count = _get_cuda_info() + try: hostname = socket.gethostname() all_ips = get_network_ips() diff --git a/distributed_queue_api.py b/distributed_queue_api.py index af16ee5..c40aed5 100644 --- a/distributed_queue_api.py +++ b/distributed_queue_api.py @@ -439,6 +439,70 @@ async def _queue_master_prompt(prompt_obj, workflow_meta, client_id): return prompt_id +def _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map): + """Configure DistributedSeed nodes for master or worker role.""" + for node_id in prompt_index.nodes_for_class("DistributedSeed"): + inputs = prompt_copy.setdefault(node_id, {}).setdefault("inputs", {}) + inputs["is_worker"] = not is_master + if is_master: + inputs["worker_id"] = "" + else: + inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}" + + +def _override_collector_nodes( + prompt_copy, + prompt_index, + is_master, + participant_id, + job_id_map, + master_url, + enabled_json, + delegate_master, +): + """Configure DistributedCollector nodes for master or worker role.""" + for node_id in prompt_index.nodes_for_class("DistributedCollector"): + if prompt_index.has_upstream(node_id, "UltimateSDUpscaleDistributed"): + prompt_copy.setdefault(node_id, {}).setdefault("inputs", {})["pass_through"] = True + continue + + inputs = prompt_copy.setdefault(node_id, {}).setdefault("inputs", {}) + inputs["multi_job_id"] = job_id_map.get(node_id, node_id) + inputs["is_worker"] = not is_master + inputs["enabled_worker_ids"] = enabled_json + if is_master: + inputs["delegate_only"] = bool(delegate_master) + inputs.pop("master_url", None) + inputs.pop("worker_id", None) + else: + inputs["master_url"] = master_url + inputs["worker_id"] = participant_id + inputs["delegate_only"] = False + + +def _override_upscale_nodes( + prompt_copy, + prompt_index, + is_master, + participant_id, + job_id_map, + master_url, + enabled_json, +): + """Configure UltimateSDUpscaleDistributed nodes for master or worker role.""" + for node_id in prompt_index.nodes_for_class("UltimateSDUpscaleDistributed"): + inputs = prompt_copy.setdefault(node_id, {}).setdefault("inputs", {}) + inputs["multi_job_id"] = job_id_map.get(node_id, node_id) + inputs["is_worker"] = not is_master + inputs["enabled_worker_ids"] = enabled_json + if is_master: + inputs.pop("master_url", None) + inputs.pop("worker_id", None) + else: + inputs["master_url"] = master_url + inputs["worker_id"] = participant_id + + def _apply_participant_overrides( prompt_copy, participant_id, @@ -453,60 +517,54 @@ def _apply_participant_overrides( worker_index_map = {wid: idx for idx, wid in enumerate(enabled_worker_ids)} enabled_json = json.dumps(enabled_worker_ids) - # Distributed seed nodes - for node_id in prompt_index.nodes_for_class("DistributedSeed"): - node = prompt_copy.get(node_id, {}) - inputs = node.setdefault("inputs", {}) - inputs["is_worker"] = not is_master - if not is_master: - idx = worker_index_map.get(participant_id, 0) - inputs["worker_id"] = f"worker_{idx}" - else: - inputs["worker_id"] = "" - - # Distributed collectors - for node_id in prompt_index.nodes_for_class("DistributedCollector"): - node = prompt_copy.get(node_id, {}) - inputs = node.setdefault("inputs", {}) - - if prompt_index.has_upstream(node_id, "UltimateSDUpscaleDistributed"): - inputs["pass_through"] = True - continue - - unique_id = job_id_map.get(node_id, node_id) - inputs["multi_job_id"] = unique_id - inputs["is_worker"] = not is_master - - if is_master: - inputs["enabled_worker_ids"] = enabled_json - inputs["delegate_only"] = bool(delegate_master) - inputs.pop("master_url", None) - inputs.pop("worker_id", None) - else: - inputs["master_url"] = master_url - inputs["worker_id"] = participant_id - inputs["enabled_worker_ids"] = enabled_json - inputs["delegate_only"] = False - - # Distributed upscaler nodes - for node_id in prompt_index.nodes_for_class("UltimateSDUpscaleDistributed"): - node = prompt_copy.get(node_id, {}) - inputs = node.setdefault("inputs", {}) - unique_id = job_id_map.get(node_id, node_id) - inputs["multi_job_id"] = unique_id - inputs["is_worker"] = not is_master - if is_master: - inputs["enabled_worker_ids"] = enabled_json - inputs.pop("master_url", None) - inputs.pop("worker_id", None) - else: - inputs["master_url"] = master_url - inputs["worker_id"] = participant_id - inputs["enabled_worker_ids"] = enabled_json + _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map) + _override_collector_nodes( + prompt_copy, + prompt_index, + is_master, + participant_id, + job_id_map, + master_url, + enabled_json, + delegate_master, + ) + _override_upscale_nodes( + prompt_copy, + prompt_index, + is_master, + participant_id, + job_id_map, + master_url, + enabled_json, + ) return prompt_copy +async def _select_active_workers(workers, use_websocket, delegate_master): + """Probe workers and return (active_workers, updated_delegate_master). + + If all workers are offline and delegate_master was True, flips delegate_master to False. + """ + active_workers = [] + for worker in workers: + is_active = ( + await _worker_ws_is_active(worker) + if use_websocket + else await _worker_is_active(worker) + ) + if is_active: + active_workers.append(worker) + else: + log(f"[Distributed] Worker {worker['name']} ({worker['id']}) is offline, skipping.") + + if not active_workers and delegate_master: + debug_log("All workers offline while delegate-only requested; enabling master participation.") + delegate_master = False + + return active_workers, delegate_master + + async def orchestrate_distributed_execution( prompt_obj, workflow_meta, @@ -535,21 +593,7 @@ async def orchestrate_distributed_execution( debug_log("Delegate-only requested but no workers are enabled. Falling back to master execution.") delegate_master = False - # Filter to active workers - active_workers = [] - for worker in workers: - if use_websocket: - is_active = await _worker_ws_is_active(worker) - else: - is_active = await _worker_is_active(worker) - if is_active: - active_workers.append(worker) - else: - log(f"[Distributed] Worker {worker['name']} ({worker['id']}) is offline, skipping.") - - if not active_workers and delegate_master: - debug_log("All workers offline while delegate-only requested; enabling master participation.") - delegate_master = False + active_workers, delegate_master = await _select_active_workers(workers, use_websocket, delegate_master) enabled_ids = [worker["id"] for worker in active_workers] diff --git a/nodes/collector.py b/nodes/collector.py index dc558a3..b99300a 100644 --- a/nodes/collector.py +++ b/nodes/collector.py @@ -22,6 +22,8 @@ prompt_server = _server.PromptServer.instance class DistributedCollectorNode: + EMPTY_AUDIO = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100} + @classmethod def INPUT_TYPES(s): return { @@ -148,21 +150,76 @@ class DistributedCollectorNode: log(f"[Distributed] Master - Audio combination failed, returning silence: {e}") return empty_audio - async def execute(self, images, audio, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False): - empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100} + def _store_worker_result(self, worker_images: dict, item: dict) -> int: + """Store one queue item (batch or single image) in worker_images in-place. + Handles two formats: + - Batch mode: item has 'tensors' list and optional 'indices' list (index-aware) + - Single mode: item has 'image_index' int and 'tensor' + + Returns the number of images stored. + """ + worker_id = item['worker_id'] + tensors = item.get('tensors', []) + if tensors: + worker_images.setdefault(worker_id, {}) + indices = item.get('indices', []) + if indices: + for i, tensor in enumerate(tensors): + worker_images[worker_id][indices[i]] = tensor + else: + for idx, tensor in enumerate(tensors): + worker_images[worker_id][idx] = tensor + return len(tensors) + else: + worker_images.setdefault(worker_id, {}) + worker_images[worker_id][item['image_index']] = item['tensor'] + return 1 + + def _reorder_and_combine_tensors( + self, + worker_images: dict, + master_batch_size: int, + images_on_cpu, + delegate_mode: bool, + fallback_images, + ) -> torch.Tensor: + """Assemble final tensor: master images first, then workers sorted by ID and index.""" + ordered_tensors = [] + if not delegate_mode and images_on_cpu is not None: + for i in range(master_batch_size): + ordered_tensors.append(images_on_cpu[i:i+1]) + for worker_id_str in sorted(worker_images.keys()): + for idx in sorted(worker_images[worker_id_str].keys()): + ordered_tensors.append(worker_images[worker_id_str][idx]) + + cpu_tensors = [] + for t in ordered_tensors: + if t.is_cuda: + t = t.cpu() + t = ensure_contiguous(t) + cpu_tensors.append(t) + + if cpu_tensors: + return ensure_contiguous(torch.cat(cpu_tensors, dim=0)) + elif fallback_images is not None: + return ensure_contiguous(fallback_images) + else: + raise ValueError("No image data collected from master or workers") + + async def execute(self, images, audio, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False): if is_worker: # Worker mode: send images and audio to master in a single batch debug_log(f"Worker - Job {multi_job_id} complete. Sending {images.shape[0]} image(s) to master") await self.send_batch_to_master(images, audio, multi_job_id, master_url, worker_id) - return (images, audio if audio is not None else empty_audio) + return (images, audio if audio is not None else self.EMPTY_AUDIO) else: delegate_mode = delegate_only or is_master_delegate_only() # Master mode: collect images and audio from workers enabled_workers = json.loads(enabled_worker_ids) num_workers = len(enabled_workers) if num_workers == 0: - return (images, audio if audio is not None else empty_audio) + return (images, audio if audio is not None else self.EMPTY_AUDIO) if delegate_mode: master_batch_size = 0 @@ -226,37 +283,12 @@ class DistributedCollectorNode: # Check if batch mode tensors = result.get('tensors', []) - indices = result.get('indices', []) # Get the indices + count = self._store_worker_result(worker_images, result) + collected_count += count if tensors: - # Batch mode debug_log(f"Master - Got batch from worker {worker_id}, size={len(tensors)}, is_last={is_last}") - - if worker_id not in worker_images: - worker_images[worker_id] = {} - - # Use actual indices if available, otherwise fall back to sequential - if indices: - for i, tensor in enumerate(tensors): - actual_idx = indices[i] - worker_images[worker_id][actual_idx] = tensor - else: - # Fallback for backward compatibility - for idx, tensor in enumerate(tensors): - worker_images[worker_id][idx] = tensor - - collected_count += len(tensors) else: - # Single image mode (backward compat) - image_index = result['image_index'] - tensor = result['tensor'] - - debug_log(f"Master - Got single result from worker {worker_id}, image {image_index}, is_last={is_last}") - - if worker_id not in worker_images: - worker_images[worker_id] = {} - worker_images[worker_id][image_index] = tensor - - collected_count += 1 + debug_log(f"Master - Got single result from worker {worker_id}, image {result.get('image_index', 0)}, is_last={is_last}") # Collect audio data if present result_audio = result.get('audio') @@ -343,28 +375,8 @@ class DistributedCollectorNode: for item in remaining_items: worker_id = item['worker_id'] is_last = item.get('is_last', False) - - # Check if batch mode - tensors = item.get('tensors', []) - if tensors: - # Batch mode - if worker_id not in worker_images: - worker_images[worker_id] = {} - - for idx, tensor in enumerate(tensors): - worker_images[worker_id][idx] = tensor - - collected_count += len(tensors) - else: - # Single image mode - image_index = item['image_index'] - tensor = item['tensor'] - - if worker_id not in worker_images: - worker_images[worker_id] = {} - worker_images[worker_id][image_index] = tensor - - collected_count += 1 + + collected_count += self._store_worker_result(worker_images, item) if is_last: workers_done.add(worker_id) @@ -386,49 +398,18 @@ class DistributedCollectorNode: if multi_job_id in prompt_server.distributed_pending_jobs: del prompt_server.distributed_pending_jobs[multi_job_id] - # Reorder images according to seed distribution pattern - # Pattern: master img 1, master img 2, worker 1 img 1, worker 1 img 2, worker 2 img 1, worker 2 img 2, etc. - ordered_tensors = [] - - # Add master images first (if any) - if not delegate_mode and images_on_cpu is not None: - for i in range(master_batch_size): - ordered_tensors.append(images_on_cpu[i:i+1]) - - # Add worker images in order - # The worker IDs in worker_images are already strings (e.g., "1", "2") - # Just iterate through what we actually received - for worker_id_str in sorted(worker_images.keys()): - # Sort by image index for each worker - for idx in sorted(worker_images[worker_id_str].keys()): - ordered_tensors.append(worker_images[worker_id_str][idx]) - - # Ensure all tensors are on CPU and properly formatted before concatenation - cpu_tensors = [] - for t in ordered_tensors: - if t.is_cuda: - t = t.cpu() - # Ensure tensor is contiguous in memory - t = ensure_contiguous(t) - cpu_tensors.append(t) - try: - if cpu_tensors: - combined = torch.cat(cpu_tensors, dim=0) - else: - # No tensors collected (likely delegate mode with no worker output) - combined = ensure_contiguous(images) if images is not None else images - if combined is None: - raise ValueError("No image data collected from master or workers") - # Ensure the combined tensor is contiguous and properly formatted - combined = ensure_contiguous(combined) - debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total (master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})") + combined = self._reorder_and_combine_tensors( + worker_images, master_batch_size, images_on_cpu, delegate_mode, images + ) + debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total " + f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})") # Combine audio from master and workers - combined_audio = self._combine_audio(master_audio, worker_audio, empty_audio) + combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO) return (combined, combined_audio) except Exception as e: log(f"Master - Error combining images: {e}") # Return just the master images as fallback - return (images, audio if audio is not None else empty_audio) + return (images, audio if audio is not None else self.EMPTY_AUDIO)