nodes/robertvoy-ComfyUI-Distributed
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update multigpu.py

Browse Source
This commit is contained in:
Robert Wojciechowski
2025-06-30 07:37:43 +10:00
committed by GitHub
parent 573e266b1e
commit 691e1d028b
1 changed files with 161 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
multigpu.py
+161 -31
View File
@@ -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"
}