Harden exception handling and add distributed exception types

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)
-

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(

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):

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

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,
 )
 

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."""

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

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()
 

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

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}")