diff --git a/api/worker_routes.py b/api/worker_routes.py index a121150..162c081 100644 --- a/api/worker_routes.py +++ b/api/worker_routes.py @@ -147,7 +147,7 @@ async def get_network_info_endpoint(request): physical_device_count = len(result.stdout.strip().split('\n')) else: physical_device_count = max(visible_devices) + 1 # Best guess - except: + except (FileNotFoundError, OSError, subprocess.SubprocessError): physical_device_count = max(visible_devices) + 1 # Best guess else: cuda_device = 0 @@ -175,7 +175,7 @@ async def get_network_info_endpoint(request): ip = info[4][0] if ip and ip not in ips and not ip.startswith('::'): # Skip IPv6 for now ips.append(ip) - except: + except (socket.gaierror, OSError): pass # Method 2: Try to connect to external server and get local IP @@ -186,7 +186,7 @@ async def get_network_info_endpoint(request): s.close() if local_ip not in ips: ips.append(local_ip) - except: + except (OSError, socket.error): pass # Method 3: Platform-specific commands @@ -205,16 +205,20 @@ async def get_network_info_endpoint(request): # 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) + 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+)') - for match in ip_pattern.finditer(result.stdout): - ip = match.group(1) - if ip and ip not in ips: - ips.append(ip) - except: + 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 @@ -551,4 +555,3 @@ async def get_worker_log_endpoint(request): except Exception as e: return await handle_api_error(request, e, 500) - diff --git a/distributed_queue_api.py b/distributed_queue_api.py index ccf7e45..af16ee5 100644 --- a/distributed_queue_api.py +++ b/distributed_queue_api.py @@ -310,7 +310,14 @@ async def _worker_is_active(worker): try: async with session.get(url, timeout=aiohttp.ClientTimeout(total=3)) as resp: return resp.status == 200 - except Exception: + except asyncio.TimeoutError: + debug_log(f"[Distributed] Worker probe timed out: {url}") + return False + except aiohttp.ClientConnectorError: + debug_log(f"[Distributed] Worker unreachable (connection refused): {url}") + return False + except Exception as e: + debug_log(f"[Distributed] Worker probe unexpected error: {e}") return False @@ -322,7 +329,14 @@ async def _worker_ws_is_active(worker): ws = await session.ws_connect(url, heartbeat=20, timeout=3) await ws.close() return True - except Exception: + except asyncio.TimeoutError: + debug_log(f"[Distributed] Worker WS probe timed out: {url}") + return False + except aiohttp.ClientConnectorError: + debug_log(f"[Distributed] Worker WS unreachable: {url}") + return False + except Exception as e: + debug_log(f"[Distributed] Worker WS probe unexpected error: {e}") return False @@ -385,6 +399,7 @@ async def _dispatch_worker_prompt(worker, prompt_obj, workflow_meta, client_id=N worker_id = worker.get("id") await _close_worker_ws(worker_id) log(f"[Distributed] Websocket dispatch failed for worker {worker_id}: {exc}") + raise session = await _get_client_session() async with session.post( diff --git a/nodes/collector.py b/nodes/collector.py index d29f3c8..dc558a3 100644 --- a/nodes/collector.py +++ b/nodes/collector.py @@ -145,7 +145,7 @@ class DistributedCollectorNode: debug_log(f"Master - Combined audio: {len(audio_pieces)} pieces, final shape={combined_waveform.shape}") return {"waveform": combined_waveform, "sample_rate": sample_rate} except Exception as e: - log(f"Master - Error combining audio: {e}") + 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): diff --git a/upscale/modes/dynamic.py b/upscale/modes/dynamic.py index 2a1e0ef..e5fe54f 100644 --- a/upscale/modes/dynamic.py +++ b/upscale/modes/dynamic.py @@ -1,4 +1,4 @@ -import asyncio, torch +import asyncio, time, torch from PIL import Image import comfy.model_management from ...utils.logging import debug_log, log diff --git a/upscale/modes/static.py b/upscale/modes/static.py index ae7be53..47867e4 100644 --- a/upscale/modes/static.py +++ b/upscale/modes/static.py @@ -9,7 +9,7 @@ from ...utils.constants import TILE_WAIT_TIMEOUT, TILE_SEND_TIMEOUT, MAX_BATCH from ...utils.usdu_managment import ( ensure_tile_jobs_initialized, init_static_job_batched, _mark_task_completed, _cleanup_job, _drain_results_queue, - _send_heartbeat_to_master, + _send_heartbeat_to_master, _get_completed_count, JOB_PENDING_TASKS, JOB_WORKER_STATUS, ) diff --git a/utils/exceptions.py b/utils/exceptions.py new file mode 100644 index 0000000..79dc96a --- /dev/null +++ b/utils/exceptions.py @@ -0,0 +1,42 @@ +"""Custom exceptions for ComfyUI-Distributed.""" + + +class DistributedError(Exception): + """Base exception for all ComfyUI-Distributed errors.""" + + +class WorkerError(DistributedError): + """Error related to a specific distributed worker.""" + + def __init__(self, message, worker_id=None): + super().__init__(message) + self.worker_id = worker_id + + +class WorkerTimeoutError(WorkerError): + """Worker did not respond within the expected timeout.""" + + +class WorkerNotAvailableError(WorkerError): + """Worker is unreachable or not running.""" + + +class JobQueueError(DistributedError): + """Error in distributed job queue management.""" + + +class TileCollectionError(DistributedError): + """Error collecting processed tiles from workers.""" + + +class ProcessError(DistributedError): + """Error managing a worker subprocess.""" + + def __init__(self, message, pid=None, worker_id=None): + super().__init__(message) + self.pid = pid + self.worker_id = worker_id + + +class TunnelError(DistributedError): + """Error managing the Cloudflare tunnel.""" diff --git a/utils/logging.py b/utils/logging.py index 3901a9d..f4d272c 100644 --- a/utils/logging.py +++ b/utils/logging.py @@ -26,7 +26,7 @@ def is_debug_enabled(): with open(CONFIG_FILE, 'r') as f: config = json.load(f) enabled = config.get("settings", {}).get("debug", False) - except: + except (OSError, json.JSONDecodeError, ValueError): pass _debug_cache = enabled diff --git a/workers/detection.py b/workers/detection.py index 91f138b..77aba9c 100644 --- a/workers/detection.py +++ b/workers/detection.py @@ -57,12 +57,12 @@ def get_machine_id(): try: # Method 1: MAC address-based UUID return str(uuid.getnode()) - except: + except Exception: try: # Method 2: Platform + hostname import socket return f"{platform.machine()}_{socket.gethostname()}" - except: + except Exception: # Fallback return platform.machine() diff --git a/workers/process_manager.py b/workers/process_manager.py index 18cb8cf..d707460 100644 --- a/workers/process_manager.py +++ b/workers/process_manager.py @@ -74,7 +74,7 @@ class WorkerProcessManager: if hasattr(folder_paths, 'base_path') and os.path.exists(os.path.join(folder_paths.base_path, "main.py")): debug_log(f"Found ComfyUI root via folder_paths: {folder_paths.base_path}") return folder_paths.base_path - except: + except Exception: pass # If all methods fail, log detailed information and return the best guess @@ -497,8 +497,8 @@ class WorkerProcessManager: try: subprocess.run(['taskkill', '/F', '/PID', child_pid], capture_output=True, check=False) - except: - pass + except (OSError, FileNotFoundError) as e: + debug_log(f"[Distributed] Warning: taskkill failed for PID {child_pid}: {e}") # Kill the parent with tree flag result = subprocess.run(['taskkill', '/F', '/PID', str(pid), '/T'], @@ -516,5 +516,6 @@ class WorkerProcessManager: subprocess.run(['pkill', '-KILL', '-P', str(pid)], check=False) os.kill(pid, signal.SIGKILL) return True - except: + except Exception as e: + log(f"[Distributed] Error killing process tree for PID {pid}: {e}") return False diff --git a/workers/startup.py b/workers/startup.py index 0c5d146..d42d226 100644 --- a/workers/startup.py +++ b/workers/startup.py @@ -98,7 +98,7 @@ async def async_cleanup_and_exit(signum=None): try: await cloudflare_tunnel_manager.stop_tunnel() except Exception as tunnel_error: - debug_log(f"Error stopping Cloudflare tunnel during shutdown: {tunnel_error}") + log(f"[Distributed] Warning: Cloudflare tunnel did not stop cleanly during shutdown: {tunnel_error}") except Exception as e: print(f"[Distributed] Error during cleanup: {e}") @@ -169,6 +169,6 @@ def sync_cleanup(): # No running loop; create a temporary one asyncio.run(cloudflare_tunnel_manager.stop_tunnel()) except Exception as tunnel_error: - debug_log(f"Error stopping Cloudflare tunnel during sync cleanup: {tunnel_error}") + log(f"[Distributed] Warning: Cloudflare tunnel did not stop cleanly during sync cleanup: {tunnel_error}") except Exception as e: print(f"[Distributed] Error during cleanup: {e}")