Handle errors from model for comfyui-json

Fix healthcheck endpoint URL

Tested and merged by Colter

lib/backend.py

@@ -11,7 +11,7 @@ from functools import cached_property
 from distutils.util import strtobool
 
 from anyio import open_file
-from aiohttp import web, ClientResponse, ClientSession, ClientConnectorError
+from aiohttp import web, ClientResponse, ClientSession, ClientConnectorError, ClientTimeout, TCPConnector
 
 import requests
 from Crypto.Signature import pkcs1_15
@@ -75,7 +75,13 @@ class Backend:
     @cached_property
     def session(self):
         log.debug(f"starting session with {self.model_server_url}")
-        return ClientSession(self.model_server_url)
+        connector = TCPConnector(
+            force_close=True, # Required for long running jobs
+            enable_cleanup_closed=True,
+        )
+        
+        timeout = ClientTimeout(total=None)
+        return ClientSession(self.model_server_url, timeout=timeout, connector=connector)
 
     def create_handler(
         self,
@@ -184,18 +190,30 @@ class Backend:
             log.debug(f"Exception in main handler loop {e}")
             return web.Response(status=500)
 
+    @cached_property  
+    def healthcheck_session(self):
+        """Dedicated session for healthchecks to avoid conflicts with API session"""
+        log.debug("creating dedicated healthcheck session")
+        connector = TCPConnector(
+            force_close=True,  # Keep this for isolation
+            enable_cleanup_closed=True,
+        )
+        timeout = ClientTimeout(total=10)  # Reasonable timeout for healthchecks
+        return ClientSession(timeout=timeout, connector=connector)
+
     async def __healthcheck(self):
         health_check_url = self.benchmark_handler.healthcheck_endpoint
         if health_check_url is None:
             log.debug("No healthcheck endpoint defined, skipping healthcheck")
             return
+
         while True:
             await sleep(10)
             if self.__start_healthcheck is False:
                 continue
             try:
                 log.debug(f"Performing healthcheck on {health_check_url}")
-                async with self.session.get(health_check_url) as response:
+                async with self.healthcheck_session.get(health_check_url) as response:
                     if response.status == 200:
                         log.debug("Healthcheck successful")
                     elif response.status == 503:
@@ -204,7 +222,6 @@ class Backend:
                             f"Healthcheck failed with status: {response.status}"
                         )
                     else:
-                        # endpoint not ready yet so bail
                         log.debug(f"Healthcheck Endpoint not ready: {response.status}")
             except Exception as e:
                 log.debug(f"Healthcheck failed with exception: {e}")

start_server.sh

@@ -59,12 +59,12 @@ then
     fi
 
     # Fork testing
-    git clone "${PYWORKER_REPO:-https://github.com/vast-ai/pyworker}" "$SERVER_DIR"
+    [[ ! -d $SERVER_DIR ]] && git clone "${PYWORKER_REPO:-https://github.com/vast-ai/pyworker}" "$SERVER_DIR"
     if [[ -n ${PYWORKER_REF:-} ]]; then
         (cd "$SERVER_DIR" && git checkout "$PYWORKER_REF")
     fi
 
-    uv venv --managed-python "$ENV_PATH" -p 3.10
+    uv venv --python-preference only-managed "$ENV_PATH" -p 3.10
     source "$ENV_PATH/bin/activate"
 
     uv pip install -r "${SERVER_DIR}/requirements.txt"

workers/comfyui-json/server.py

@@ -33,33 +33,39 @@ log = logging.getLogger(__file__)
 async def generate_client_response(
         client_request: web.Request, model_response: ClientResponse
     ) -> Union[web.Response, web.StreamResponse]:
-        # Check if the response is actually streaming based on response headers/content-type
-        is_streaming_response = (
-            model_response.content_type == "text/event-stream"
-            or model_response.content_type == "application/x-ndjson"
-            or model_response.headers.get("Transfer-Encoding") == "chunked"
-            or "stream" in model_response.content_type.lower()
-        )
+        match model_response.status:
+            case 200:
+                log.debug("SUCCESS")
+                # Check if the response is actually streaming based on response headers/content-type
+                is_streaming_response = (
+                    model_response.content_type == "text/event-stream"
+                    or model_response.content_type == "application/x-ndjson"
+                    or model_response.headers.get("Transfer-Encoding") == "chunked"
+                    or "stream" in model_response.content_type.lower()
+                )
 
-        if is_streaming_response:
-            log.debug("Detected streaming response...")
-            res = web.StreamResponse()
-            res.content_type = model_response.content_type
-            await res.prepare(client_request)
-            async for chunk in model_response.content:
-                await res.write(chunk)
-            await res.write_eof()
-            log.debug("Done streaming response")
-            return res
-        else:
-            log.debug("Detected non-streaming response...")
-            content = await model_response.read()
-            return web.Response(
-                body=content,
-                status=model_response.status,
-                content_type=model_response.content_type
-            )
-            
+                if is_streaming_response:
+                    log.debug("Detected streaming response...")
+                    res = web.StreamResponse()
+                    res.content_type = model_response.content_type
+                    await res.prepare(client_request)
+                    async for chunk in model_response.content:
+                        await res.write(chunk)
+                    await res.write_eof()
+                    log.debug("Done streaming response")
+                    return res
+                else:
+                    log.debug("Detected non-streaming response...")
+                    content = await model_response.read()
+                    return web.Response(
+                        body=content,
+                        status=model_response.status,
+                        content_type=model_response.content_type
+                    )
+            case code:
+                log.debug(f"Model responded with error {code}")
+                return web.Response(status=code)
+                
 
 @dataclasses.dataclass
 class ComfyWorkflowHandler(EndpointHandler[ComfyWorkflowData]):
@@ -70,7 +76,7 @@ class ComfyWorkflowHandler(EndpointHandler[ComfyWorkflowData]):
 
     @property
     def healthcheck_endpoint(self) -> Optional[str]:
-        return "/health"
+        return f"{MODEL_SERVER_URL}/health"
 
     @classmethod
     def payload_cls(cls) -> Type[ComfyWorkflowData]: