From 691e1d028b222787476d6a1bbfc85a460c8aabe2 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 30 Jun 2025 07:37:43 +1000 Subject: [PATCH] Update multigpu.py --- multigpu.py | 192 +++++++++++++++++++++++++++++++++++++++++++--------- 1 file changed, 161 insertions(+), 31 deletions(-) diff --git a/multigpu.py b/multigpu.py index f5896fa..6abbd92 100644 --- a/multigpu.py +++ b/multigpu.py @@ -190,6 +190,9 @@ async def job_complete_endpoint(request): data = await request.post() multi_job_id = data.get('multi_job_id') image_file = data.get('image') + worker_id = data.get('worker_id') + image_index = data.get('image_index') + is_last = data.get('is_last', 'False').lower() == 'true' if not all([multi_job_id, image_file]): return await handle_api_error(request, "Missing job_id or image data", 400) @@ -202,8 +205,13 @@ async def job_complete_endpoint(request): async with PENDING_JOBS_LOCK: if multi_job_id in PENDING_JOBS: - await PENDING_JOBS[multi_job_id].put(tensor) - print(f"[MultiGPU] Received result for job {multi_job_id}") + await PENDING_JOBS[multi_job_id].put({ + 'tensor': tensor, + 'worker_id': worker_id, + 'image_index': int(image_index) if image_index else 0, + 'is_last': is_last + }) + print(f"[MultiGPU] Received result for job {multi_job_id} from worker {worker_id} (last: {is_last})") return web.json_response({"status": "success"}) else: return await handle_api_error(request, "Job not found or already complete", 404) @@ -221,9 +229,8 @@ class MultiGPUCollectorNode: "is_worker": ("BOOLEAN", {"default": False}), "master_url": ("STRING", {"default": ""}), "enabled_worker_ids": ("STRING", {"default": "[]"}), - # --- FIX [7 of 7]: Type Mismatch --- - # Changed from FLOAT to INT for type safety and clarity. "worker_batch_size": ("INT", {"default": 1, "min": 1, "max": 1024}), + "worker_id": ("STRING", {"default": ""}), }, } @@ -231,7 +238,7 @@ class MultiGPUCollectorNode: FUNCTION = "run" CATEGORY = "image" - def run(self, images, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1): + def run(self, images, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id=""): if not multi_job_id: return (images,) @@ -240,7 +247,7 @@ class MultiGPUCollectorNode: async def run_and_set_future(): try: - result = await self.execute(images, multi_job_id, is_worker, master_url, enabled_worker_ids, worker_batch_size) + result = await self.execute(images, multi_job_id, is_worker, master_url, enabled_worker_ids, worker_batch_size, worker_id) future.set_result(result) except Exception as e: future.set_exception(e) @@ -248,7 +255,7 @@ class MultiGPUCollectorNode: asyncio.run_coroutine_threadsafe(run_and_set_future(), main_loop) return future.result() - async def send_image_to_master(self, image_tensor, multi_job_id, master_url, image_index): + async def send_image_to_master(self, image_tensor, multi_job_id, master_url, image_index, worker_id, is_last=False): """Helper method to send a single image to the master""" img_np = (image_tensor.cpu().numpy() * 255.).astype(np.uint8) img = Image.fromarray(img_np) @@ -258,6 +265,9 @@ class MultiGPUCollectorNode: data = aiohttp.FormData() data.add_field('multi_job_id', multi_job_id) + data.add_field('worker_id', str(worker_id)) + data.add_field('image_index', str(image_index)) + data.add_field('is_last', str(is_last)) data.add_field('image', byte_io, filename=f'image_{image_index}.png', content_type='image/png') try: @@ -267,48 +277,168 @@ class MultiGPUCollectorNode: except Exception as e: print(f"[MultiGPU Worker] Failed to send image {image_index+1} to master: {e}") - async def execute(self, images, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1): + async def execute(self, images, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id=""): if is_worker: # Worker mode: send images to master - print(f"[MultiGPU Worker] Job {multi_job_id} complete. Sending {images.shape[0]} image(s) to master at {master_url}") + print(f"[MultiGPU Worker] Job {multi_job_id} complete. Sending {images.shape[0]} image(s) to master") + + # Send all images, marking the last one + for i in range(images.shape[0]): + is_last = (i == images.shape[0] - 1) + await self.send_image_to_master(images[i], multi_job_id, master_url, i, worker_id, is_last) - tasks = [self.send_image_to_master(images[i], multi_job_id, master_url, i) for i in range(images.shape[0])] - await asyncio.gather(*tasks) return (images,) else: # Master mode: collect images from workers - num_workers = len(json.loads(enabled_worker_ids)) + enabled_workers = json.loads(enabled_worker_ids) + num_workers = len(enabled_workers) if num_workers == 0: return (images,) images_on_cpu = images.cpu() - total_images_to_collect = int(num_workers * worker_batch_size) - print(f"[MultiGPU Master] Job {multi_job_id}: Waiting for {total_images_to_collect} image(s)...") + print(f"[MultiGPU Master] Job {multi_job_id}: Collecting images from {num_workers} workers...") - all_tensors = [images_on_cpu] + # Store master images first + master_batch_size = images.shape[0] + + # Initialize storage for collected images + worker_images = {} # Dict to store images by worker_id and index async with PENDING_JOBS_LOCK: if multi_job_id not in PENDING_JOBS: PENDING_JOBS[multi_job_id] = asyncio.Queue() q = PENDING_JOBS[multi_job_id] - try: - # Collect all images with a potentially dynamic timeout - for i in range(total_images_to_collect): - worker_tensor = await asyncio.wait_for(q.get(), timeout=300.0) - all_tensors.append(worker_tensor) - print(f"[MultiGPU Master] Collected image {i+1}/{total_images_to_collect}") - except asyncio.TimeoutError: - print(f"[MultiGPU Master] Timeout waiting for worker images (got {len(all_tensors)-1}/{total_images_to_collect})") - finally: - # Clean up job queue - async with PENDING_JOBS_LOCK: - if multi_job_id in PENDING_JOBS: - del PENDING_JOBS[multi_job_id] + # Collect images until all workers report they're done + collected_count = 0 + workers_done = set() + + # Use a reasonable timeout for the first image + timeout = 60.0 + + while len(workers_done) < num_workers: + try: + result = await asyncio.wait_for(q.get(), timeout=timeout) + worker_id = result['worker_id'] + image_index = result['image_index'] + tensor = result['tensor'] + is_last = result.get('is_last', False) + + if worker_id not in worker_images: + worker_images[worker_id] = {} + worker_images[worker_id][image_index] = tensor + + collected_count += 1 + + # Once we start receiving images, use shorter timeout + timeout = 10.0 + + if is_last: + workers_done.add(worker_id) + print(f"[MultiGPU Master] Worker {worker_id} done. Collected {len(worker_images[worker_id])} images") + else: + print(f"[MultiGPU Master] Collected image {image_index + 1} from worker {worker_id}") + + except asyncio.TimeoutError: + missing_workers = set(enabled_workers) - workers_done + print(f"[MultiGPU Master] Timeout. Still waiting for workers: {list(missing_workers)}") + break + + total_collected = sum(len(imgs) for imgs in worker_images.values()) + print(f"[MultiGPU Master] Collection complete. Received {total_collected} images from {len(workers_done)} workers") + + # Clean up job queue + async with PENDING_JOBS_LOCK: + if multi_job_id in PENDING_JOBS: + del PENDING_JOBS[multi_job_id] - combined = torch.cat(all_tensors, dim=0) + # 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 + for i in range(master_batch_size): + ordered_tensors.append(images_on_cpu[i:i+1]) + + # Add worker images in order + for worker_id in enabled_workers: + # Convert worker_id to string since it comes as string from form data + worker_id_str = str(worker_id) + if worker_id_str in worker_images: + # 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]) + + combined = torch.cat(ordered_tensors, dim=0) print(f"[MultiGPU Master] Job {multi_job_id} complete. Combined {combined.shape[0]} images total.") return (combined,) -NODE_CLASS_MAPPINGS = { "MultiGPUCollector": MultiGPUCollectorNode } -NODE_DISPLAY_NAME_MAPPINGS = { "MultiGPUCollector": "Multi-GPU Collector" } +# --- Distributor Node --- +class MultiGPUDistributor: + """ + Distributes seed values across multiple GPUs. + On master: passes through the original seed. + On workers: adds offset based on worker ID. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "seed": ("INT", { + "default": 0, + "min": -2147483648, + "max": 2147483647, + "forceInput": False # Widget by default, can be converted to input + }), + "batch_size": ("INT", { + "default": 1, + "min": 1, + "max": 4096, + "forceInput": False # Widget by default, can be converted to input + }), + }, + "hidden": { + "is_worker": ("BOOLEAN", {"default": False}), + "worker_id": ("STRING", {"default": ""}), + }, + } + + RETURN_TYPES = ("INT", "INT") + RETURN_NAMES = ("seed", "batch_size") + FUNCTION = "distribute" + CATEGORY = "utils" + + def distribute(self, seed, batch_size, is_worker=False, worker_id=""): + if not is_worker: + # Master node: pass through original values + print(f"[MultiGPU Distributor] Master: seed={seed}, batch={batch_size}") + return (seed, batch_size) + else: + # Worker node: apply offset based on worker index + # Find worker index from enabled_worker_ids + try: + # Worker IDs are passed as "worker_0", "worker_1", etc. + if worker_id.startswith("worker_"): + worker_index = int(worker_id.split("_")[1]) + else: + # Fallback: try to parse as direct index + worker_index = int(worker_id) + + offset = worker_index + 1 + new_seed = seed + offset + print(f"[MultiGPU Distributor] Worker {worker_index}: seed={seed} → {new_seed}, batch={batch_size}") + return (new_seed, batch_size) + except (ValueError, IndexError) as e: + print(f"[MultiGPU Distributor] Error parsing worker_id '{worker_id}': {e}") + # Fallback: return original seed + return (seed, batch_size) + +NODE_CLASS_MAPPINGS = { + "MultiGPUCollector": MultiGPUCollectorNode, + "MultiGPUDistributor": MultiGPUDistributor +} +NODE_DISPLAY_NAME_MAPPINGS = { + "MultiGPUCollector": "Multi-GPU Collector", + "MultiGPUDistributor": "Multi-GPU Distributor" +}