Fix tests, improve distributed worker health check, add torch compile…

… options

comfy/cmd/extra_model_paths.py

@@ -1,25 +1,4 @@
-import os
-import yaml
-import logging
-
 def load_extra_path_config(yaml_path):
-    from . import folder_paths
+    from ..extra_config import load_extra_path_config
 
-    with open(yaml_path, 'r') as stream:
-        config = yaml.safe_load(stream)
-    for c in config:
-        conf = config[c]
-        if conf is None:
-            continue
-        base_path = None
-        if "base_path" in conf:
-            base_path = conf.pop("base_path")
-        for x in conf:
-            for y in conf[x].split("\n"):
-                if len(y) == 0:
-                    continue
-                full_path = y
-                if base_path is not None:
-                    full_path = os.path.join(base_path, full_path)
-                logging.info(f"Adding extra search path {x} ({full_path})")
-                folder_paths.add_model_folder_path(x, full_path)
+    return load_extra_path_config(yaml_path)

comfy/cmd/folder_paths.py

@@ -349,7 +349,7 @@ def invalidate_cache(folder_name):
     _filename_list_cache.pop(folder_name, None)
 
 
-def filter_files_content_types(files: list[str], content_types: Literal["image", "video", "audio"]) -> list[str]:
+def filter_files_content_types(files: list[str], content_types: list[Literal["image", "video", "audio"]]) -> list[str]:
     """
     Example:
         files = os.listdir(folder_paths.get_input_directory())

comfy/distributed/distributed_prompt_worker.py

@@ -41,7 +41,14 @@ def __init__(self, embedded_comfy_client: Optional[EmbeddedComfyClient] = None,
         self._health_check_site: Optional[web.TCPSite] = None
 
     async def _health_check(self, request):
-        return web.Response(text="OK", content_type="text/plain")
+        if self._connection is None:
+            return web.Response(text="UNHEALTHY: RabbitMQ connection is not established", status=503)
+
+        is_healthy = await self._is_connection_healthy()
+        if is_healthy:
+            return web.Response(text="HEALTHY", status=200)
+        else:
+            return web.Response(text="UNHEALTHY: RabbitMQ connection is not healthy", status=503)
 
     async def _start_health_check_server(self):
         app = web.Application()
@@ -85,9 +92,27 @@ async def _do_work_item(self, request: dict) -> dict:
         await self.on_did_complete_work_item(request_obj, reply)
         return asdict(reply)
 
+    async def _is_connection_healthy(self):
+        if self._connection is None:
+            return False
+
+        return (
+                not self._connection.is_closed
+                and self._connection.connected.is_set()
+                and await self._check_connection_ready()
+        )
+
+    async def _check_connection_ready(self):
+        try:
+            await asyncio.wait_for(self._connection.ready(), timeout=1.0)
+            return True
+        except asyncio.TimeoutError:
+            return False
+
     @tracer.start_as_current_span("Initialize Prompt Worker")
     async def init(self):
         await self._exit_stack.__aenter__()
+        await self._start_health_check_server()
         try:
             self._connection = await connect_robust(self._connection_uri, loop=self._loop)
         except AMQPConnectionError as connection_error:
@@ -102,7 +127,6 @@ async def init(self):
             await self._exit_stack.enter_async_context(self._embedded_comfy_client)
 
         await self._rpc.register(self._queue_name, self._do_work_item)
-        await self._start_health_check_server()
 
     async def __aenter__(self) -> "DistributedPromptWorker":
         await self.init()

utils/extra_config.py → comfy/extra_config.py

@@ -1,9 +1,12 @@
+import logging
 import os
+
 import yaml
-import folder_paths
-import logging
+
 
 def load_extra_path_config(yaml_path):
+    from .cmd import folder_paths
+
     with open(yaml_path, 'r') as stream:
         config = yaml.safe_load(stream)
     for c in config: