diff --git a/.github/workflows/build-custom.yml b/.github/workflows/build-custom.yml new file mode 100644 index 000000000..d84217e39 --- /dev/null +++ b/.github/workflows/build-custom.yml @@ -0,0 +1,48 @@ +name: Build Custom Docker Image + +on: + # Trigger on push to main branch when Dockerfile.custom changes + push: + branches: [main] + paths: + - 'Dockerfile.custom' + # Allow manual trigger + workflow_dispatch: + +jobs: + build-and-push: + runs-on: ubuntu-latest + steps: + - name: Free up disk space + run: | + sudo rm -rf /usr/share/dotnet + sudo rm -rf /usr/local/lib/android + sudo rm -rf /opt/ghc + sudo rm -rf /opt/hostedtoolcache/CodeQL + sudo docker image prune --all --force + df -h + + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Login to Docker Hub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + - name: Build and push + uses: docker/build-push-action@v5 + with: + context: . + file: Dockerfile.custom + platforms: linux/amd64 + push: true + tags: | + ${{ secrets.DOCKERHUB_USERNAME }}/runpod-comfy-worker:latest + ${{ secrets.DOCKERHUB_USERNAME }}/runpod-comfy-worker:${{ github.sha }} + cache-from: type=gha + cache-to: type=gha,mode=max diff --git a/Dockerfile.custom b/Dockerfile.custom new file mode 100644 index 000000000..f3e754fef --- /dev/null +++ b/Dockerfile.custom @@ -0,0 +1,12 @@ +# Custom ComfyUI worker with video/animation nodes +# No models included - uses network volume +FROM runpod/worker-comfyui:5.5.0-base + +# Install custom nodes for video/frame processing workflows +RUN comfy-node-install \ + comfyui-kjnodes \ + comfyui-videohelpersuite \ + comfyui-frame-interpolation \ + comfyui_essentials \ + comfyui-gimm-vfi \ + loadloramodelonlywithurl diff --git a/docs/configuration.md b/docs/configuration.md index 25bed9e70..8fc06348e 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -21,6 +21,7 @@ This document outlines the environment variables available for configuring the ` | Environment Variable | Description | Default | | ------------------------------ | ---------------------------------------------------------------------------------------------------------------------- | ------- | +| `WEBSOCKET_MESSAGE_TIMEOUT` | Maximum seconds to wait for a single WebSocket message before checking ComfyUI health. Heavy workloads (e.g. video frame interpolation) can produce long silences on the socket. Set to `0` to disable (recv blocks indefinitely). | `600` | | `WEBSOCKET_RECONNECT_ATTEMPTS` | Number of websocket reconnection attempts when connection drops during job execution. | `5` | | `WEBSOCKET_RECONNECT_DELAY_S` | Delay in seconds between websocket reconnection attempts. | `3` | | `WEBSOCKET_TRACE` | Enable low-level websocket frame tracing for protocol debugging. Set to `true` only when diagnosing connection issues. | `false` | diff --git a/handler.py b/handler.py index 65c8390b0..7aba622b2 100644 --- a/handler.py +++ b/handler.py @@ -38,6 +38,11 @@ # If the respective env-vars are not supplied we fall back to sensible defaults ("5" and "3"). WEBSOCKET_RECONNECT_ATTEMPTS = int(os.environ.get("WEBSOCKET_RECONNECT_ATTEMPTS", 5)) WEBSOCKET_RECONNECT_DELAY_S = int(os.environ.get("WEBSOCKET_RECONNECT_DELAY_S", 3)) +# Maximum seconds to wait for a single WebSocket message before checking +# whether ComfyUI is still alive. Heavy workloads (e.g. video frame +# interpolation) can produce long silences on the socket. Default: 600 s +# (10 minutes). Set to 0 to disable (recv blocks indefinitely). +WEBSOCKET_MESSAGE_TIMEOUT = int(os.environ.get("WEBSOCKET_MESSAGE_TIMEOUT", 600)) # Extra verbose websocket trace logs (set WEBSOCKET_TRACE=true to enable) if os.environ.get("WEBSOCKET_TRACE", "false").lower() == "true": @@ -68,6 +73,31 @@ def _comfy_server_status(): return {"reachable": False, "error": str(exc)} +def _is_job_in_queue(prompt_id): + """Check if a prompt is still running or pending in ComfyUI's queue. + + Polls the ``/queue`` endpoint and looks for *prompt_id* in both + ``queue_running`` and ``queue_pending``. Returns ``True`` if found, + ``False`` if not found, and ``None`` if the check itself failed (e.g. + network error). + """ + try: + resp = requests.get(f"http://{COMFY_HOST}/queue", timeout=5) + resp.raise_for_status() + queue_data = resp.json() + for item in queue_data.get("queue_running", []): + # Each item is a list; the prompt_id is at index 1 + if len(item) >= 2 and item[1] == prompt_id: + return True + for item in queue_data.get("queue_pending", []): + if len(item) >= 2 and item[1] == prompt_id: + return True + return False + except Exception as exc: + logger.warning("Error polling /queue: %s", exc) + return None + + def _attempt_websocket_reconnect(ws_url, max_attempts, delay_s, initial_error): """ Attempts to reconnect to the WebSocket server after a disconnect. @@ -564,7 +594,9 @@ def handler(job): print(f"worker-comfyui - Connecting to websocket: {ws_url}") ws = websocket.WebSocket() ws.connect(ws_url, timeout=10) - print(f"worker-comfyui - Websocket connected") + if WEBSOCKET_MESSAGE_TIMEOUT > 0: + ws.settimeout(WEBSOCKET_MESSAGE_TIMEOUT) + print(f"worker-comfyui - Websocket connected (message timeout: {WEBSOCKET_MESSAGE_TIMEOUT}s)") # Queue the workflow try: @@ -627,8 +659,37 @@ def handler(job): else: continue except websocket.WebSocketTimeoutException: - print(f"worker-comfyui - Websocket receive timed out. Still waiting...") - continue + # No message received within WEBSOCKET_MESSAGE_TIMEOUT seconds. + # 1) Check whether ComfyUI HTTP is alive at all. + srv_status = _comfy_server_status() + if not srv_status["reachable"]: + raise ValueError( + f"ComfyUI became unreachable while waiting for WebSocket messages " + f"(no message for {WEBSOCKET_MESSAGE_TIMEOUT}s). " + f"Server status: {srv_status}" + ) + + # 2) Check whether our job is still in the queue (running/pending). + in_queue = _is_job_in_queue(prompt_id) + if in_queue: + print( + f"worker-comfyui - No WebSocket message for {WEBSOCKET_MESSAGE_TIMEOUT}s, " + f"but job {prompt_id} is still in ComfyUI queue. Continuing to wait..." + ) + continue + elif in_queue is None: + # Queue check failed (network blip) — give benefit of the doubt. + print( + f"worker-comfyui - No WebSocket message for {WEBSOCKET_MESSAGE_TIMEOUT}s. " + f"Queue check failed; ComfyUI HTTP is reachable so continuing to wait..." + ) + continue + else: + # Job is NOT in queue and NOT completed via WS — truly stuck. + raise ValueError( + f"WebSocket message timeout: no message for {WEBSOCKET_MESSAGE_TIMEOUT}s " + f"and job {prompt_id} is no longer in ComfyUI queue." + ) except websocket.WebSocketConnectionClosedException as closed_err: try: # Attempt to reconnect @@ -639,6 +700,8 @@ def handler(job): closed_err, ) + if WEBSOCKET_MESSAGE_TIMEOUT > 0: + ws.settimeout(WEBSOCKET_MESSAGE_TIMEOUT) print( "worker-comfyui - Resuming message listening after successful reconnect." )