Refactor distributed module into packages and optimize utilities
This commit is contained in:
Robert Wojciechowski
@@ -0,0 +1,4 @@
|
||||
from . import config_routes # noqa: F401
|
||||
from . import tunnel_routes # noqa: F401
|
||||
from . import worker_routes # noqa: F401
|
||||
from . import job_routes # noqa: F401
|
||||
@@ -0,0 +1,197 @@
|
||||
import json
|
||||
|
||||
from aiohttp import web
|
||||
import server
|
||||
|
||||
from ..utils.config import load_config, save_config
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import handle_api_error, normalize_host
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/config")
|
||||
async def get_config_endpoint(request):
|
||||
config = load_config()
|
||||
return web.json_response(config)
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/queue_status/{job_id}")
|
||||
async def queue_status_endpoint(request):
|
||||
"""Check if a job queue is initialized."""
|
||||
try:
|
||||
job_id = request.match_info['job_id']
|
||||
|
||||
# Import to ensure initialization
|
||||
from ..distributed_upscale import ensure_tile_jobs_initialized
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
exists = job_id in prompt_server.distributed_pending_tile_jobs
|
||||
|
||||
debug_log(f"Queue status check for job {job_id}: {'exists' if exists else 'not found'}")
|
||||
return web.json_response({"exists": exists, "job_id": job_id})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/config/update_worker")
|
||||
async def update_worker_endpoint(request):
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get("worker_id")
|
||||
|
||||
if worker_id is None:
|
||||
return await handle_api_error(request, "Missing worker_id", 400)
|
||||
|
||||
config = load_config()
|
||||
worker_found = False
|
||||
|
||||
for worker in config.get("workers", []):
|
||||
if worker["id"] == worker_id:
|
||||
# Update all provided fields
|
||||
if "enabled" in data:
|
||||
worker["enabled"] = data["enabled"]
|
||||
if "name" in data:
|
||||
worker["name"] = data["name"]
|
||||
if "port" in data:
|
||||
worker["port"] = data["port"]
|
||||
|
||||
# Handle host field - remove it if None
|
||||
if "host" in data:
|
||||
if data["host"] is None:
|
||||
worker.pop("host", None)
|
||||
else:
|
||||
worker["host"] = normalize_host(data["host"])
|
||||
|
||||
# Handle cuda_device field - remove it if None
|
||||
if "cuda_device" in data:
|
||||
if data["cuda_device"] is None:
|
||||
worker.pop("cuda_device", None)
|
||||
else:
|
||||
worker["cuda_device"] = data["cuda_device"]
|
||||
|
||||
# Handle extra_args field - remove it if None
|
||||
if "extra_args" in data:
|
||||
if data["extra_args"] is None:
|
||||
worker.pop("extra_args", None)
|
||||
else:
|
||||
worker["extra_args"] = data["extra_args"]
|
||||
|
||||
# Handle type field
|
||||
if "type" in data:
|
||||
worker["type"] = data["type"]
|
||||
|
||||
worker_found = True
|
||||
break
|
||||
|
||||
if not worker_found:
|
||||
# If worker not found and all required fields are provided, create new worker
|
||||
if all(key in data for key in ["name", "port", "cuda_device"]):
|
||||
new_worker = {
|
||||
"id": worker_id,
|
||||
"name": data["name"],
|
||||
"host": normalize_host(data.get("host", "localhost")),
|
||||
"port": data["port"],
|
||||
"cuda_device": data["cuda_device"],
|
||||
"enabled": data.get("enabled", False),
|
||||
"extra_args": data.get("extra_args", ""),
|
||||
"type": data.get("type", "local")
|
||||
}
|
||||
if "workers" not in config:
|
||||
config["workers"] = []
|
||||
config["workers"].append(new_worker)
|
||||
worker_found = True
|
||||
else:
|
||||
return await handle_api_error(request, f"Worker {worker_id} not found and missing required fields for creation", 404)
|
||||
|
||||
if save_config(config):
|
||||
return web.json_response({"status": "success"})
|
||||
else:
|
||||
return await handle_api_error(request, "Failed to save config")
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 400)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/config/delete_worker")
|
||||
async def delete_worker_endpoint(request):
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get("worker_id")
|
||||
|
||||
if worker_id is None:
|
||||
return await handle_api_error(request, "Missing worker_id", 400)
|
||||
|
||||
config = load_config()
|
||||
workers = config.get("workers", [])
|
||||
|
||||
# Find and remove the worker
|
||||
worker_index = -1
|
||||
for i, worker in enumerate(workers):
|
||||
if worker["id"] == worker_id:
|
||||
worker_index = i
|
||||
break
|
||||
|
||||
if worker_index == -1:
|
||||
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
|
||||
|
||||
# Remove the worker
|
||||
removed_worker = workers.pop(worker_index)
|
||||
|
||||
if save_config(config):
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"message": f"Worker {removed_worker.get('name', worker_id)} deleted"
|
||||
})
|
||||
else:
|
||||
return await handle_api_error(request, "Failed to save config")
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 400)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/config/update_setting")
|
||||
async def update_setting_endpoint(request):
|
||||
"""Updates a specific key in the settings object."""
|
||||
try:
|
||||
data = await request.json()
|
||||
key = data.get("key")
|
||||
value = data.get("value")
|
||||
|
||||
if not key or value is None:
|
||||
return await handle_api_error(request, "Missing 'key' or 'value' in request", 400)
|
||||
|
||||
config = load_config()
|
||||
if 'settings' not in config:
|
||||
config['settings'] = {}
|
||||
|
||||
config['settings'][key] = value
|
||||
|
||||
if save_config(config):
|
||||
return web.json_response({"status": "success", "message": f"Setting '{key}' updated."})
|
||||
else:
|
||||
return await handle_api_error(request, "Failed to save config")
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 400)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/config/update_master")
|
||||
async def update_master_endpoint(request):
|
||||
"""Updates master configuration."""
|
||||
try:
|
||||
data = await request.json()
|
||||
|
||||
config = load_config()
|
||||
if 'master' not in config:
|
||||
config['master'] = {}
|
||||
|
||||
# Update all provided fields
|
||||
if "name" in data:
|
||||
config['master']['name'] = data['name']
|
||||
if "host" in data:
|
||||
config['master']['host'] = data['host']
|
||||
if "port" in data:
|
||||
config['master']['port'] = data['port']
|
||||
if "cuda_device" in data:
|
||||
config['master']['cuda_device'] = data['cuda_device']
|
||||
if "extra_args" in data:
|
||||
config['master']['extra_args'] = data['extra_args']
|
||||
|
||||
if save_config(config):
|
||||
return web.json_response({"status": "success", "message": "Master configuration updated."})
|
||||
else:
|
||||
return await handle_api_error(request, "Failed to save config")
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 400)
|
||||
@@ -0,0 +1,404 @@
|
||||
import json
|
||||
import asyncio
|
||||
import io
|
||||
import uuid
|
||||
import time
|
||||
import os
|
||||
|
||||
from aiohttp import web
|
||||
import server
|
||||
import server as _server
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.image import pil_to_tensor, ensure_contiguous
|
||||
from ..utils.network import handle_api_error
|
||||
from ..utils.constants import MEMORY_CLEAR_DELAY
|
||||
from ..utils.async_helpers import queue_prompt_payload
|
||||
from ..distributed_queue_api import orchestrate_distributed_execution
|
||||
|
||||
prompt_server = _server.PromptServer.instance
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/prepare_job")
|
||||
async def prepare_job_endpoint(request):
|
||||
try:
|
||||
data = await request.json()
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
if not multi_job_id:
|
||||
return await handle_api_error(request, "Missing multi_job_id", 400)
|
||||
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
if multi_job_id not in prompt_server.distributed_pending_jobs:
|
||||
prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue()
|
||||
|
||||
debug_log(f"Prepared queue for job {multi_job_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/clear_memory")
|
||||
async def clear_memory_endpoint(request):
|
||||
debug_log("Received request to clear VRAM.")
|
||||
try:
|
||||
# Use ComfyUI's prompt server queue system like the /free endpoint does
|
||||
if hasattr(server.PromptServer.instance, 'prompt_queue'):
|
||||
server.PromptServer.instance.prompt_queue.set_flag("unload_models", True)
|
||||
server.PromptServer.instance.prompt_queue.set_flag("free_memory", True)
|
||||
debug_log("Set queue flags for memory clearing.")
|
||||
|
||||
# Wait a bit for the queue to process
|
||||
await asyncio.sleep(MEMORY_CLEAR_DELAY)
|
||||
|
||||
# Also do direct cleanup as backup, but with error handling
|
||||
import gc
|
||||
import comfy.model_management as mm
|
||||
|
||||
try:
|
||||
mm.unload_all_models()
|
||||
except AttributeError as e:
|
||||
debug_log(f"Warning during model unload: {e}")
|
||||
|
||||
try:
|
||||
mm.soft_empty_cache()
|
||||
except Exception as e:
|
||||
debug_log(f"Warning during cache clear: {e}")
|
||||
|
||||
for _ in range(3):
|
||||
gc.collect()
|
||||
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.ipc_collect()
|
||||
|
||||
debug_log("VRAM cleared successfully.")
|
||||
return web.json_response({"status": "success", "message": "GPU memory cleared."})
|
||||
except Exception as e:
|
||||
# Even if there's an error, try to do basic cleanup
|
||||
import gc
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
debug_log(f"Partial VRAM clear completed with warning: {e}")
|
||||
return web.json_response({"status": "success", "message": "GPU memory cleared (with warnings)"})
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/queue")
|
||||
async def distributed_queue_endpoint(request):
|
||||
"""Queue a distributed workflow, mirroring the UI orchestration pipeline."""
|
||||
try:
|
||||
data = await request.json()
|
||||
except Exception as exc:
|
||||
return await handle_api_error(request, f"Invalid JSON payload: {exc}", 400)
|
||||
|
||||
prompt = data.get("prompt")
|
||||
if not isinstance(prompt, dict):
|
||||
return await handle_api_error(request, "Field 'prompt' must be an object", 400)
|
||||
|
||||
workflow_meta = data.get("workflow")
|
||||
client_id = data.get("client_id")
|
||||
delegate_master = data.get("delegate_master")
|
||||
enabled_ids = data.get("enabled_worker_ids")
|
||||
|
||||
if enabled_ids is not None:
|
||||
if not isinstance(enabled_ids, list):
|
||||
return await handle_api_error(request, "enabled_worker_ids must be a list of worker IDs", 400)
|
||||
enabled_ids = [str(worker_id) for worker_id in enabled_ids]
|
||||
|
||||
try:
|
||||
prompt_id, worker_count = await orchestrate_distributed_execution(
|
||||
prompt,
|
||||
workflow_meta,
|
||||
client_id,
|
||||
enabled_worker_ids=enabled_ids,
|
||||
delegate_master=delegate_master,
|
||||
)
|
||||
return web.json_response({
|
||||
"prompt_id": prompt_id,
|
||||
"worker_count": worker_count,
|
||||
})
|
||||
except Exception as exc:
|
||||
return await handle_api_error(request, exc, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/load_image")
|
||||
async def load_image_endpoint(request):
|
||||
"""Load an image or video file and return it as base64 data with hash."""
|
||||
try:
|
||||
data = await request.json()
|
||||
image_path = data.get("image_path")
|
||||
|
||||
if not image_path:
|
||||
return await handle_api_error(request, "Missing image_path", 400)
|
||||
|
||||
import folder_paths
|
||||
import base64
|
||||
import io
|
||||
import hashlib
|
||||
|
||||
# Use ComfyUI's folder paths to find the file
|
||||
full_path = folder_paths.get_annotated_filepath(image_path)
|
||||
|
||||
if not os.path.exists(full_path):
|
||||
return await handle_api_error(request, f"File not found: {image_path}", 404)
|
||||
|
||||
# Calculate file hash
|
||||
hash_md5 = hashlib.md5()
|
||||
with open(full_path, 'rb') as f:
|
||||
for chunk in iter(lambda: f.read(4096), b""):
|
||||
hash_md5.update(chunk)
|
||||
file_hash = hash_md5.hexdigest()
|
||||
|
||||
# Check if it's a video file
|
||||
video_extensions = {'.mp4', '.avi', '.mov', '.mkv', '.webm'}
|
||||
file_ext = os.path.splitext(full_path)[1].lower()
|
||||
|
||||
if file_ext in video_extensions:
|
||||
# For video files, read the raw bytes
|
||||
with open(full_path, 'rb') as f:
|
||||
file_data = f.read()
|
||||
|
||||
# Determine MIME type
|
||||
mime_types = {
|
||||
'.mp4': 'video/mp4',
|
||||
'.avi': 'video/x-msvideo',
|
||||
'.mov': 'video/quicktime',
|
||||
'.mkv': 'video/x-matroska',
|
||||
'.webm': 'video/webm'
|
||||
}
|
||||
mime_type = mime_types.get(file_ext, 'video/mp4')
|
||||
|
||||
# Return base64 encoded video with data URL
|
||||
video_base64 = base64.b64encode(file_data).decode('utf-8')
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"image_data": f"data:{mime_type};base64,{video_base64}",
|
||||
"hash": file_hash
|
||||
})
|
||||
else:
|
||||
# For images, use PIL
|
||||
from PIL import Image
|
||||
|
||||
# Load and convert to base64
|
||||
with Image.open(full_path) as img:
|
||||
# Convert to RGB if needed
|
||||
if img.mode not in ('RGB', 'RGBA'):
|
||||
img = img.convert('RGB')
|
||||
|
||||
# Save to bytes
|
||||
buffer = io.BytesIO()
|
||||
img.save(buffer, format='PNG', compress_level=1) # Fast compression
|
||||
img_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
|
||||
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"image_data": f"data:image/png;base64,{img_base64}",
|
||||
"hash": file_hash
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/check_file")
|
||||
async def check_file_endpoint(request):
|
||||
"""Check if a file exists and matches the given hash."""
|
||||
try:
|
||||
data = await request.json()
|
||||
filename = data.get("filename")
|
||||
expected_hash = data.get("hash")
|
||||
|
||||
if not filename or not expected_hash:
|
||||
return await handle_api_error(request, "Missing filename or hash", 400)
|
||||
|
||||
import folder_paths
|
||||
import hashlib
|
||||
|
||||
# Use ComfyUI's folder paths to find the file
|
||||
full_path = folder_paths.get_annotated_filepath(filename)
|
||||
|
||||
if not os.path.exists(full_path):
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"exists": False
|
||||
})
|
||||
|
||||
# Calculate file hash
|
||||
hash_md5 = hashlib.md5()
|
||||
with open(full_path, 'rb') as f:
|
||||
for chunk in iter(lambda: f.read(4096), b""):
|
||||
hash_md5.update(chunk)
|
||||
file_hash = hash_md5.hexdigest()
|
||||
|
||||
# Check if hash matches
|
||||
hash_matches = file_hash == expected_hash
|
||||
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"exists": True,
|
||||
"hash_matches": hash_matches
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/job_complete")
|
||||
async def job_complete_endpoint(request):
|
||||
try:
|
||||
data = await request.post()
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
worker_id = data.get('worker_id')
|
||||
is_last = data.get('is_last', 'False').lower() == 'true'
|
||||
|
||||
if multi_job_id is None or worker_id is None:
|
||||
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
|
||||
|
||||
# Check for batch mode
|
||||
batch_size = int(data.get('batch_size', 0))
|
||||
tensors = []
|
||||
|
||||
if batch_size > 0:
|
||||
# Batch mode: Extract multiple images
|
||||
debug_log(f"job_complete received batch - job_id: {multi_job_id}, worker: {worker_id}, batch_size: {batch_size}")
|
||||
|
||||
# Check for JSON metadata
|
||||
metadata_field = data.get('images_metadata')
|
||||
metadata = None
|
||||
if metadata_field:
|
||||
# New JSON metadata format
|
||||
try:
|
||||
# Handle different types of metadata field
|
||||
if hasattr(metadata_field, 'file'):
|
||||
# File-like object
|
||||
metadata_str = metadata_field.file.read().decode('utf-8')
|
||||
elif isinstance(metadata_field, (bytes, bytearray)):
|
||||
# Direct bytes/bytearray
|
||||
metadata_str = metadata_field.decode('utf-8')
|
||||
else:
|
||||
# String
|
||||
metadata_str = str(metadata_field)
|
||||
|
||||
metadata = json.loads(metadata_str)
|
||||
if len(metadata) != batch_size:
|
||||
return await handle_api_error(request, "Metadata length mismatch", 400)
|
||||
debug_log(f"Using JSON metadata for batch from worker {worker_id}")
|
||||
except Exception as e:
|
||||
log(f"Error parsing JSON metadata from worker {worker_id}: {e}")
|
||||
return await handle_api_error(request, f"Metadata parsing error: {e}", 400)
|
||||
else:
|
||||
# Legacy format - log deprecation warning
|
||||
debug_log(f"WARNING: Worker {worker_id} using legacy field format. Please update to use JSON metadata.")
|
||||
|
||||
# Process images with per-item error handling
|
||||
image_data_list = [] # List of tuples: (index, tensor)
|
||||
|
||||
for i in range(batch_size):
|
||||
image_field = data.get(f'image_{i}')
|
||||
if image_field is None:
|
||||
log(f"Missing image_{i} from worker {worker_id}, skipping")
|
||||
continue
|
||||
|
||||
try:
|
||||
img_data = image_field.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
tensor = pil_to_tensor(img)
|
||||
tensor = ensure_contiguous(tensor)
|
||||
|
||||
# Get expected index from metadata or use position
|
||||
if metadata and i < len(metadata):
|
||||
expected_idx = metadata[i].get('index', i)
|
||||
if i != expected_idx:
|
||||
debug_log(f"Warning: Image order mismatch at position {i}, expected index {expected_idx}")
|
||||
else:
|
||||
expected_idx = i
|
||||
|
||||
image_data_list.append((expected_idx, tensor))
|
||||
except Exception as e:
|
||||
log(f"Error processing image {i} from worker {worker_id}: {e}, skipping this image")
|
||||
# Continue processing other images instead of failing entire batch
|
||||
continue
|
||||
|
||||
# Check if we got any valid images
|
||||
if not image_data_list:
|
||||
return await handle_api_error(request, "No valid images in batch", 400)
|
||||
|
||||
# Sort by index to ensure correct order
|
||||
image_data_list.sort(key=lambda x: x[0])
|
||||
tensors = [tensor for _, tensor in image_data_list]
|
||||
|
||||
# Keep the indices for later use
|
||||
indices = [idx for idx, _ in image_data_list]
|
||||
|
||||
# Validate final order
|
||||
if metadata:
|
||||
debug_log(f"Reordered {len(tensors)} images based on metadata indices")
|
||||
else:
|
||||
# Fallback for single-image mode (backward compatibility)
|
||||
image_file = data.get('image')
|
||||
image_index = data.get('image_index')
|
||||
|
||||
debug_log(f"job_complete received single - job_id: {multi_job_id}, worker: {worker_id}, index: {image_index}, is_last: {is_last}")
|
||||
|
||||
if not image_file:
|
||||
return await handle_api_error(request, "Missing image data", 400)
|
||||
|
||||
try:
|
||||
img_data = image_file.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
tensor = pil_to_tensor(img)
|
||||
tensor = ensure_contiguous(tensor)
|
||||
tensors = [tensor]
|
||||
except Exception as e:
|
||||
log(f"Error processing image from worker {worker_id}: {e}")
|
||||
return await handle_api_error(request, f"Image processing error: {e}", 400)
|
||||
|
||||
# Parse audio data if present
|
||||
audio_data = None
|
||||
audio_waveform_field = data.get('audio_waveform')
|
||||
audio_sample_rate = data.get('audio_sample_rate')
|
||||
if audio_waveform_field is not None:
|
||||
try:
|
||||
waveform_bytes = audio_waveform_field.file.read()
|
||||
waveform_buffer = io.BytesIO(waveform_bytes)
|
||||
try:
|
||||
waveform_tensor = torch.load(waveform_buffer, weights_only=True)
|
||||
except TypeError:
|
||||
waveform_tensor = torch.load(waveform_buffer)
|
||||
sample_rate = int(audio_sample_rate) if audio_sample_rate else 44100
|
||||
audio_data = {"waveform": waveform_tensor, "sample_rate": sample_rate}
|
||||
debug_log(f"Received audio from worker {worker_id}: shape={waveform_tensor.shape}, sample_rate={sample_rate}")
|
||||
except Exception as e:
|
||||
log(f"Error parsing audio from worker {worker_id}: {e}")
|
||||
|
||||
# Put batch into queue
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
debug_log(f"Current pending jobs: {list(prompt_server.distributed_pending_jobs.keys())}")
|
||||
if multi_job_id in prompt_server.distributed_pending_jobs:
|
||||
if batch_size > 0:
|
||||
# Put batch as single item with indices
|
||||
await prompt_server.distributed_pending_jobs[multi_job_id].put({
|
||||
'worker_id': worker_id,
|
||||
'tensors': tensors,
|
||||
'indices': indices,
|
||||
'is_last': is_last,
|
||||
'audio': audio_data
|
||||
})
|
||||
debug_log(f"Received batch result for job {multi_job_id} from worker {worker_id}, size={len(tensors)}")
|
||||
else:
|
||||
# Put single image (backward compat)
|
||||
await prompt_server.distributed_pending_jobs[multi_job_id].put({
|
||||
'tensor': tensors[0],
|
||||
'worker_id': worker_id,
|
||||
'image_index': int(image_index) if image_index else 0,
|
||||
'is_last': is_last,
|
||||
'audio': audio_data
|
||||
})
|
||||
debug_log(f"Received single result for job {multi_job_id} from worker {worker_id}")
|
||||
|
||||
debug_log(f"Queue size after put: {prompt_server.distributed_pending_jobs[multi_job_id].qsize()}")
|
||||
return web.json_response({"status": "success"})
|
||||
else:
|
||||
log(f"ERROR: Job {multi_job_id} not found in distributed_pending_jobs")
|
||||
return await handle_api_error(request, "Job not found or already complete", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e)
|
||||
@@ -0,0 +1,51 @@
|
||||
from aiohttp import web
|
||||
import server
|
||||
|
||||
from ..utils.cloudflare import cloudflare_tunnel_manager
|
||||
from ..utils.config import load_config
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import handle_api_error
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/tunnel/status")
|
||||
async def tunnel_status_endpoint(request):
|
||||
"""Return Cloudflare tunnel status and last known details."""
|
||||
try:
|
||||
status = cloudflare_tunnel_manager.get_status()
|
||||
config = load_config()
|
||||
master_host = (config.get("master") or {}).get("host")
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"tunnel": status,
|
||||
"master_host": master_host
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/tunnel/start")
|
||||
async def tunnel_start_endpoint(request):
|
||||
"""Start a Cloudflare tunnel pointing at the current ComfyUI server."""
|
||||
try:
|
||||
result = await cloudflare_tunnel_manager.start_tunnel()
|
||||
config = load_config()
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"tunnel": result,
|
||||
"master_host": (config.get("master") or {}).get("host")
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/tunnel/stop")
|
||||
async def tunnel_stop_endpoint(request):
|
||||
"""Stop the managed Cloudflare tunnel if running."""
|
||||
try:
|
||||
result = await cloudflare_tunnel_manager.stop_tunnel()
|
||||
config = load_config()
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"tunnel": result,
|
||||
"master_host": (config.get("master") or {}).get("host")
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
@@ -0,0 +1,554 @@
|
||||
import json
|
||||
import asyncio
|
||||
import os
|
||||
import time
|
||||
import platform
|
||||
import subprocess
|
||||
|
||||
import torch
|
||||
import aiohttp
|
||||
from aiohttp import web
|
||||
import server
|
||||
|
||||
from ..utils.config import load_config
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import handle_api_error, normalize_host, get_client_session
|
||||
from ..utils.constants import CHUNK_SIZE
|
||||
from ..workers import get_worker_manager
|
||||
from ..workers.detection import (
|
||||
is_local_worker,
|
||||
is_same_physical_host,
|
||||
get_machine_id,
|
||||
is_docker_environment,
|
||||
is_runpod_environment,
|
||||
get_comms_channel,
|
||||
)
|
||||
from ..utils.async_helpers import queue_prompt_payload
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/worker_ws")
|
||||
async def worker_ws_endpoint(request):
|
||||
"""WebSocket endpoint for worker prompt dispatch."""
|
||||
ws = web.WebSocketResponse(heartbeat=30)
|
||||
await ws.prepare(request)
|
||||
|
||||
async for msg in ws:
|
||||
if msg.type == aiohttp.WSMsgType.TEXT:
|
||||
try:
|
||||
data = json.loads(msg.data or "{}")
|
||||
except json.JSONDecodeError:
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
"request_id": None,
|
||||
"ok": False,
|
||||
"error": "Invalid JSON payload.",
|
||||
})
|
||||
continue
|
||||
|
||||
if data.get("type") != "dispatch_prompt":
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
"request_id": data.get("request_id"),
|
||||
"ok": False,
|
||||
"error": "Unsupported websocket message type.",
|
||||
})
|
||||
continue
|
||||
|
||||
prompt = data.get("prompt")
|
||||
if not isinstance(prompt, dict):
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
"request_id": data.get("request_id"),
|
||||
"ok": False,
|
||||
"error": "Field 'prompt' must be an object.",
|
||||
})
|
||||
continue
|
||||
|
||||
try:
|
||||
prompt_id = await queue_prompt_payload(
|
||||
prompt,
|
||||
workflow_meta=data.get("workflow"),
|
||||
client_id=data.get("client_id"),
|
||||
)
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
"request_id": data.get("request_id"),
|
||||
"ok": True,
|
||||
"prompt_id": prompt_id,
|
||||
})
|
||||
except Exception as exc:
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
"request_id": data.get("request_id"),
|
||||
"ok": False,
|
||||
"error": str(exc),
|
||||
})
|
||||
elif msg.type == aiohttp.WSMsgType.ERROR:
|
||||
log(f"[Distributed] Worker websocket error: {ws.exception()}")
|
||||
|
||||
return ws
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/worker/clear_launching")
|
||||
async def clear_launching_state(request):
|
||||
"""Clear the launching flag when worker is confirmed running."""
|
||||
try:
|
||||
wm = get_worker_manager()
|
||||
data = await request.json()
|
||||
worker_id = str(data.get('worker_id'))
|
||||
|
||||
if not worker_id:
|
||||
return await handle_api_error(request, "worker_id is required", 400)
|
||||
|
||||
# Clear launching flag in managed processes
|
||||
if worker_id in wm.processes:
|
||||
if 'launching' in wm.processes[worker_id]:
|
||||
del wm.processes[worker_id]['launching']
|
||||
wm.save_processes()
|
||||
debug_log(f"Cleared launching state for worker {worker_id}")
|
||||
|
||||
return web.json_response({"status": "success"})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@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:
|
||||
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:
|
||||
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:
|
||||
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:
|
||||
result = subprocess.run(["ifconfig"], capture_output=True, text=True)
|
||||
|
||||
import re
|
||||
ip_pattern = re.compile(r'inet\s+(\d+\.\d+\.\d+\.\d+)')
|
||||
for match in ip_pattern.finditer(result.stdout):
|
||||
ip = match.group(1)
|
||||
if ip and ip not in ips:
|
||||
ips.append(ip)
|
||||
except:
|
||||
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
|
||||
|
||||
try:
|
||||
hostname = socket.gethostname()
|
||||
all_ips = get_network_ips()
|
||||
recommended_ip = get_recommended_ip(all_ips)
|
||||
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"hostname": hostname,
|
||||
"all_ips": all_ips,
|
||||
"recommended_ip": recommended_ip,
|
||||
"cuda_device": cuda_device,
|
||||
"cuda_device_count": physical_device_count if physical_device_count > 0 else cuda_device_count,
|
||||
"message": "Auto-detected network configuration"
|
||||
})
|
||||
except Exception as e:
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"message": str(e),
|
||||
"hostname": "unknown",
|
||||
"all_ips": [],
|
||||
"recommended_ip": None
|
||||
})
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/system_info")
|
||||
async def get_system_info_endpoint(request):
|
||||
"""Get system information including machine ID for local worker detection."""
|
||||
try:
|
||||
import socket
|
||||
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"hostname": socket.gethostname(),
|
||||
"machine_id": get_machine_id(),
|
||||
"platform": {
|
||||
"system": platform.system(),
|
||||
"machine": platform.machine(),
|
||||
"node": platform.node(),
|
||||
"path_separator": os.sep, # Add path separator
|
||||
"os_name": os.name # Add OS name (posix, nt, etc.)
|
||||
},
|
||||
"is_docker": is_docker_environment(),
|
||||
"is_runpod": is_runpod_environment(),
|
||||
"runpod_pod_id": os.environ.get('RUNPOD_POD_ID')
|
||||
})
|
||||
except Exception as e:
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"message": str(e)
|
||||
}, status=500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/launch_worker")
|
||||
async def launch_worker_endpoint(request):
|
||||
"""Launch a worker process from the UI."""
|
||||
try:
|
||||
wm = get_worker_manager()
|
||||
data = await request.json()
|
||||
worker_id = data.get("worker_id")
|
||||
|
||||
if not worker_id:
|
||||
return await handle_api_error(request, "Missing worker_id", 400)
|
||||
|
||||
# Find worker config
|
||||
config = load_config()
|
||||
worker = next((w for w in config.get("workers", []) if w["id"] == worker_id), None)
|
||||
if not worker:
|
||||
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
|
||||
|
||||
# Ensure consistent string ID
|
||||
worker_id_str = str(worker_id)
|
||||
|
||||
# Check if already running (managed by this instance)
|
||||
if worker_id_str in wm.processes:
|
||||
proc_info = wm.processes[worker_id_str]
|
||||
process = proc_info.get('process')
|
||||
|
||||
# Check if still running
|
||||
is_running = False
|
||||
if process:
|
||||
is_running = process.poll() is None
|
||||
else:
|
||||
# Restored process without subprocess object
|
||||
is_running = wm._is_process_running(proc_info['pid'])
|
||||
|
||||
if is_running:
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"message": "Worker already running (managed by UI)",
|
||||
"pid": proc_info['pid'],
|
||||
"log_file": proc_info.get('log_file')
|
||||
}, status=409)
|
||||
else:
|
||||
# Process is dead, remove it
|
||||
del wm.processes[worker_id_str]
|
||||
wm.save_processes()
|
||||
|
||||
# Launch the worker
|
||||
try:
|
||||
pid = wm.launch_worker(worker)
|
||||
log_file = wm.processes[worker_id_str].get('log_file')
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"pid": pid,
|
||||
"message": f"Worker {worker['name']} launched",
|
||||
"log_file": log_file
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, f"Failed to launch worker: {str(e)}", 500)
|
||||
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 400)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/stop_worker")
|
||||
async def stop_worker_endpoint(request):
|
||||
"""Stop a worker process that was launched from the UI."""
|
||||
try:
|
||||
wm = get_worker_manager()
|
||||
data = await request.json()
|
||||
worker_id = data.get("worker_id")
|
||||
|
||||
if not worker_id:
|
||||
return await handle_api_error(request, "Missing worker_id", 400)
|
||||
|
||||
success, message = wm.stop_worker(worker_id)
|
||||
|
||||
if success:
|
||||
return web.json_response({"status": "success", "message": message})
|
||||
else:
|
||||
return web.json_response({"status": "error", "message": message},
|
||||
status=404 if "not managed" in message else 409)
|
||||
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 400)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/managed_workers")
|
||||
async def get_managed_workers_endpoint(request):
|
||||
"""Get list of workers managed by this UI instance."""
|
||||
try:
|
||||
managed = get_worker_manager().get_managed_workers()
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"managed_workers": managed
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/local-worker-status")
|
||||
async def get_local_worker_status_endpoint(request):
|
||||
"""Check status of all local workers (localhost/no host specified)."""
|
||||
try:
|
||||
config = load_config()
|
||||
worker_statuses = {}
|
||||
|
||||
for worker in config.get("workers", []):
|
||||
# Only check local workers
|
||||
host = normalize_host(worker.get("host")) or ""
|
||||
if not host or host in ["localhost", "127.0.0.1"]:
|
||||
worker_id = worker["id"]
|
||||
port = worker["port"]
|
||||
|
||||
# Check if worker is enabled
|
||||
if not worker.get("enabled", False):
|
||||
worker_statuses[worker_id] = {
|
||||
"online": False,
|
||||
"enabled": False,
|
||||
"processing": False,
|
||||
"queue_count": 0
|
||||
}
|
||||
continue
|
||||
|
||||
# Try to connect to worker
|
||||
try:
|
||||
session = await get_client_session()
|
||||
async with session.get(
|
||||
f"http://localhost:{port}/prompt",
|
||||
timeout=aiohttp.ClientTimeout(total=2)
|
||||
) as resp:
|
||||
if resp.status == 200:
|
||||
data = await resp.json()
|
||||
queue_remaining = data.get("exec_info", {}).get("queue_remaining", 0)
|
||||
worker_statuses[worker_id] = {
|
||||
"online": True,
|
||||
"enabled": True,
|
||||
"processing": queue_remaining > 0,
|
||||
"queue_count": queue_remaining
|
||||
}
|
||||
else:
|
||||
worker_statuses[worker_id] = {
|
||||
"online": False,
|
||||
"enabled": True,
|
||||
"processing": False,
|
||||
"queue_count": 0,
|
||||
"error": f"HTTP {resp.status}"
|
||||
}
|
||||
except asyncio.TimeoutError:
|
||||
worker_statuses[worker_id] = {
|
||||
"online": False,
|
||||
"enabled": True,
|
||||
"processing": False,
|
||||
"queue_count": 0,
|
||||
"error": "Timeout"
|
||||
}
|
||||
except Exception as e:
|
||||
worker_statuses[worker_id] = {
|
||||
"online": False,
|
||||
"enabled": True,
|
||||
"processing": False,
|
||||
"queue_count": 0,
|
||||
"error": str(e)
|
||||
}
|
||||
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"worker_statuses": worker_statuses
|
||||
})
|
||||
except Exception as e:
|
||||
debug_log(f"Error checking local worker status: {e}")
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/worker_log/{worker_id}")
|
||||
async def get_worker_log_endpoint(request):
|
||||
"""Get log content for a specific worker."""
|
||||
try:
|
||||
wm = get_worker_manager()
|
||||
worker_id = request.match_info['worker_id']
|
||||
|
||||
# Ensure worker_id is string
|
||||
worker_id = str(worker_id)
|
||||
|
||||
# Check if we manage this worker
|
||||
if worker_id not in wm.processes:
|
||||
return await handle_api_error(request, f"Worker {worker_id} not managed by UI", 404)
|
||||
|
||||
proc_info = wm.processes[worker_id]
|
||||
log_file = proc_info.get('log_file')
|
||||
|
||||
if not log_file or not os.path.exists(log_file):
|
||||
return await handle_api_error(request, "Log file not found", 404)
|
||||
|
||||
# Read last N lines (or full file if small)
|
||||
lines_to_read = int(request.query.get('lines', 1000))
|
||||
|
||||
try:
|
||||
# Get file size
|
||||
file_size = os.path.getsize(log_file)
|
||||
|
||||
with open(log_file, 'r', encoding='utf-8', errors='replace') as f:
|
||||
if lines_to_read > 0 and file_size > 1024 * 1024: # If file > 1MB and limited lines requested
|
||||
# Read last N lines efficiently
|
||||
lines = []
|
||||
# Start from end and work backwards
|
||||
f.seek(0, 2) # Go to end
|
||||
file_length = f.tell()
|
||||
|
||||
# Read chunks from end
|
||||
chunk_size = min(CHUNK_SIZE, file_length)
|
||||
while len(lines) < lines_to_read and f.tell() > 0:
|
||||
# Move back and read chunk
|
||||
current_pos = max(0, f.tell() - chunk_size)
|
||||
f.seek(current_pos)
|
||||
chunk = f.read(chunk_size)
|
||||
|
||||
# Process chunk
|
||||
chunk_lines = chunk.splitlines()
|
||||
if current_pos > 0:
|
||||
# Partial line at beginning, combine with next chunk
|
||||
chunk_lines = chunk_lines[1:]
|
||||
|
||||
lines = chunk_lines + lines
|
||||
|
||||
# Move back for next chunk
|
||||
f.seek(current_pos)
|
||||
|
||||
# Take only last N lines
|
||||
content = '\n'.join(lines[-lines_to_read:])
|
||||
truncated = len(lines) > lines_to_read
|
||||
else:
|
||||
# Read entire file
|
||||
content = f.read()
|
||||
truncated = False
|
||||
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"content": content,
|
||||
"log_file": log_file,
|
||||
"file_size": file_size,
|
||||
"truncated": truncated,
|
||||
"lines_shown": lines_to_read if truncated else content.count('\n') + 1
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, f"Error reading log file: {str(e)}", 500)
|
||||
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
Reference in New Issue
Block a user