From e65faca8174d4e31697f3a97c45bdda59e538cd3 Mon Sep 17 00:00:00 2001
From: doctorpangloss <@hiddenswitch.com>
Date: Tue, 18 Feb 2025 15:07:02 -0800
Subject: [PATCH] Distributed setup now defaults to panicking when out of
 memory now, to facilitate graceful recovery

---
 comfy/cli_args.py       | 2 +-
 comfy/cli_args_types.py | 2 +-
 docker-compose.yml      | 2 ++
 3 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/comfy/cli_args.py b/comfy/cli_args.py
index eb5dbc545..1ba607dc4 100644
--- a/comfy/cli_args.py
+++ b/comfy/cli_args.py
@@ -194,7 +194,7 @@ def _create_parser() -> EnhancedConfigArgParser:
         '--panic-when',
         action='append',
         help="""
-        List of fully qualified exception class names to panic (os.exit(1)) when a workflow raises it.
+        List of fully qualified exception class names to panic (sys.exit(1)) when a workflow raises it.
         Example: --panic-when=torch.cuda.OutOfMemoryError. Can be specified multiple times or as a 
         comma-separated list.""",
         type=str,
diff --git a/comfy/cli_args_types.py b/comfy/cli_args_types.py
index d6e0e6859..dce80a22d 100644
--- a/comfy/cli_args_types.py
+++ b/comfy/cli_args_types.py
@@ -122,7 +122,7 @@ class Configuration(dict):
         anthropic_api_key (str): Configures the Anthropic API key for its nodes related to Claude functionality. Visit https://console.anthropic.com/settings/keys to create this key.
         user_directory (Optional[str]): Set the ComfyUI user directory with an absolute path.
         log_stdout (bool): Send normal process output to stdout instead of stderr (default)
-        panic_when (list[str]): List of fully qualified exception class names to panic (os.exit(1)) when a workflow raises it.
+        panic_when (list[str]): List of fully qualified exception class names to panic (sys.exit(1)) when a workflow raises it.
     """
 
     def __init__(self, **kwargs):
diff --git a/docker-compose.yml b/docker-compose.yml
index d968a3156..37fd9038f 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -16,6 +16,8 @@ services:
               capabilities: [ gpu ]
     environment:
       - COMFYUI_DISTRIBUTED_QUEUE_CONNECTION_URI=amqp://guest:guest@rabbitmq:5672
+      - COMFYUI_EXECUTOR_FACTORY=ProcessPoolExecutor
+      - COMFYUI_PANIC_WHEN=torch.cuda.OutOfMemoryError
       - COMFYUI_LOGGING_LEVEL=ERROR
     command:
       - comfyui-worker