comfyui-json: key readiness off api-wrapper's BACKENDS_READY token

Rather than tailing for "Uvicorn running on", which only confirms the
api-wrapper's own HTTP listener is bound, watch for the api-wrapper's
new structured tokens that reflect actual end-to-end reachability:

  MODEL_LOAD_LOG_MSG  = ["BACKENDS_READY"]
  MODEL_ERROR_LOG_MSGS includes:
    - "BACKENDS_READY_TIMEOUT"   (backends never came up)
    - "BACKEND_UNRECOVERABLE"    (CUDA fault latched on a backend)
    - "Application startup failed" (kept; uvicorn's own ASGI failure)

Closes the race observed on a live test where the pyworker fired
benchmark the moment uvicorn bound, every request inside the
api-wrapper hit Cannot-connect-to-host on ComfyUI, and the SDK
counted the resulting fast 502s as a fast worker (perf=200).

Tokens are emitted by ai-dock/comfyui-api-wrapper#11 and onward;
earlier wrapper versions won't emit BACKENDS_READY so warm-up stalls
indefinitely — pin to a wrapper that includes that change.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

requirements.txt

@@ -1,13 +1,2 @@
-aiohttp==3.10.1
-aiodns~=3.6.0
-pycares~=4.11.0
-anyio~=4.4
-lib~=4.0
-nltk~=3.9
-psutil~=6.0
-pycryptodome~=3.20
-Requests~=2.32
-transformers~=4.52
-utils==1.0.*
-hf_transfer>=0.1.9
 vastai-sdk>=0.3.0
+nltk==3.9.4 

start_server.sh

@@ -2,10 +2,17 @@
 
 set -e -o pipefail
 
+# Check for force update flag
+FORCE_UPDATE=false
+if [ -f "/.force_update" ]; then
+    echo "Force update flag detected at /.force_update"
+    FORCE_UPDATE=true
+fi
+
 WORKSPACE_DIR="${WORKSPACE_DIR:-/workspace}"
 
 SERVER_DIR="$WORKSPACE_DIR/vast-pyworker"
-ENV_PATH="$WORKSPACE_DIR/worker-env"
+ENV_PATH="${ENV_PATH:-$WORKSPACE_DIR/worker-env}"
 DEBUG_LOG="$WORKSPACE_DIR/debug.log"
 PYWORKER_LOG="$WORKSPACE_DIR/pyworker.log"
 
@@ -47,29 +54,38 @@ JSON
 }
 
 function install_vastai_sdk() {
-    # If SDK_BRANCH is set, install vastai-sdk from the vast-sdk repo at that branch/tag/commit.
+    local uv_flags=()
+    if [ "${USE_SYSTEM_PYTHON:-}" = "true" ]; then
+        uv_flags+=(--system --break-system-packages)
+    fi
+    if [ "$FORCE_UPDATE" = true ]; then
+        uv_flags+=(--force-reinstall)
+        echo "Force reinstalling vastai"
+    fi
+
+    # If SDK_BRANCH is set, install vastai from the vast-cli repo at that branch/tag/commit.
     if [ -n "${SDK_BRANCH:-}" ]; then
         if [ -n "${SDK_VERSION:-}" ]; then
             echo "WARNING: Both SDK_BRANCH and SDK_VERSION are set; using SDK_BRANCH=${SDK_BRANCH}"
         fi
-        echo "Installing vastai-sdk from https://github.com/vast-ai/vast-sdk/ @ ${SDK_BRANCH}"
-        if ! uv pip install "vastai-sdk @ git+https://github.com/vast-ai/vast-sdk.git@${SDK_BRANCH}"; then
-            report_error_and_exit "Failed to install vastai-sdk from vast-ai/vast-sdk@${SDK_BRANCH}"
+        echo "Installing vastai from https://github.com/vast-ai/vast-cli/ @ ${SDK_BRANCH}"
+        if ! uv pip install "${uv_flags[@]}" "vastai @ git+https://github.com/vast-ai/vast-cli.git@${SDK_BRANCH}"; then
+            report_error_and_exit "Failed to install vastai from vast-ai/vast-cli@${SDK_BRANCH}"
         fi
         return 0
     fi
 
     if [ -n "${SDK_VERSION:-}" ]; then
-        echo "Installing vastai-sdk version ${SDK_VERSION}"
-        if ! uv pip install "vastai-sdk==${SDK_VERSION}"; then
-            report_error_and_exit "Failed to install vastai-sdk==${SDK_VERSION}"
+        echo "Installing vastai version ${SDK_VERSION}"
+        if ! uv pip install "${uv_flags[@]}" "vastai==${SDK_VERSION}"; then
+            report_error_and_exit "Failed to install vastai==${SDK_VERSION}"
         fi
         return 0
     fi
 
-    echo "Installing default vastai-sdk"
-    if ! uv pip install vastai-sdk; then
-        report_error_and_exit "Failed to install vastai-sdk"
+    echo "Installing default vastai"
+    if ! uv pip install "${uv_flags[@]}" vastai; then
+        report_error_and_exit "Failed to install vastai"
     fi
 }
 
@@ -90,7 +106,8 @@ echo_var DEBUG_LOG
 echo_var PYWORKER_LOG
 echo_var MODEL_LOG
 
-if [ -e "$MODEL_LOG" ]; then
+ROTATE_MODEL_LOG="${ROTATE_MODEL_LOG:-false}"
+if [ "$ROTATE_MODEL_LOG" = "true" ] && [ -e "$MODEL_LOG" ]; then
     echo "Rotating model log at $MODEL_LOG to $MODEL_LOG.old"
     if ! cat "$MODEL_LOG" >> "$MODEL_LOG.old"; then
         report_error_and_exit "Failed to rotate model log"
@@ -111,8 +128,21 @@ if ! grep -q "VAST" /etc/environment; then
     fi
 fi
 
-if [ ! -d "$ENV_PATH" ]
-then
+if [ "${USE_SYSTEM_PYTHON:-}" = "true" ]; then
+    echo "Using system Python: $(which python3)"
+    if ! which uv > /dev/null 2>&1; then
+        if ! curl -LsSf https://astral.sh/uv/install.sh | sh; then
+            report_error_and_exit "Failed to install uv package manager"
+        fi
+        if [[ -f ~/.local/bin/env ]]; then
+            if ! source ~/.local/bin/env; then
+                report_error_and_exit "Failed to source uv environment"
+            fi
+        fi
+    fi
+    install_vastai_sdk
+    touch ~/.no_auto_tmux
+elif [ ! -d "$ENV_PATH" ]; then
     echo "setting up venv"
     if ! which uv; then
         if ! curl -LsSf https://astral.sh/uv/install.sh | sh; then
@@ -131,12 +161,29 @@ then
         if ! git clone "${PYWORKER_REPO:-https://github.com/vast-ai/pyworker}" "$SERVER_DIR"; then
             report_error_and_exit "Failed to clone pyworker repository"
         fi
+    elif [ "$FORCE_UPDATE" = true ]; then
+        echo "Force updating pyworker repository"
+        if ! (cd "$SERVER_DIR" && git fetch --all); then
+            report_error_and_exit "Failed to fetch pyworker repository updates"
+        fi
     fi
     if [[ -n ${PYWORKER_REF:-} ]]; then
+        if [ "$FORCE_UPDATE" = true ]; then
+            echo "Force updating to pyworker reference: $PYWORKER_REF"
+            if ! (cd "$SERVER_DIR" && git checkout "$PYWORKER_REF" && git pull); then
+                report_error_and_exit "Failed to force update pyworker reference: $PYWORKER_REF"
+            fi
+        else
             if ! (cd "$SERVER_DIR" && git checkout "$PYWORKER_REF"); then
                 report_error_and_exit "Failed to checkout pyworker reference: $PYWORKER_REF"
             fi
         fi
+    elif [ "$FORCE_UPDATE" = true ]; then
+        echo "Force updating pyworker to latest"
+        if ! (cd "$SERVER_DIR" && git pull); then
+            report_error_and_exit "Failed to pull latest pyworker changes"
+        fi
+    fi
 
     if ! uv venv --python-preference only-managed "$ENV_PATH" -p 3.10; then
         report_error_and_exit "Failed to create virtual environment"
@@ -161,11 +208,44 @@ else
             report_error_and_exit "Failed to source uv environment"
         fi
     fi
-    if ! source "$WORKSPACE_DIR/worker-env/bin/activate"; then
+    if ! source "$ENV_PATH/bin/activate"; then
         report_error_and_exit "Failed to activate existing virtual environment"
     fi
     echo "environment activated"
     echo "venv: $VIRTUAL_ENV"
+
+    # Handle force update for existing environment
+    if [ "$FORCE_UPDATE" = true ]; then
+        echo "Performing force update on existing environment"
+
+        if [[ -d $SERVER_DIR ]]; then
+            echo "Force updating pyworker repository"
+            if ! (cd "$SERVER_DIR" && git fetch --all); then
+                report_error_and_exit "Failed to fetch pyworker repository updates"
+            fi
+
+            if [[ -n ${PYWORKER_REF:-} ]]; then
+                echo "Force updating to pyworker reference: $PYWORKER_REF"
+                if ! (cd "$SERVER_DIR" && git checkout "$PYWORKER_REF" && git pull); then
+                    report_error_and_exit "Failed to force update pyworker reference: $PYWORKER_REF"
+                fi
+            else
+                echo "Force updating pyworker to latest"
+                if ! (cd "$SERVER_DIR" && git pull); then
+                    report_error_and_exit "Failed to pull latest pyworker changes"
+                fi
+            fi
+        fi
+
+        install_vastai_sdk
+    fi
+fi
+
+# Remove force update flag after successful update
+if [ "$FORCE_UPDATE" = true ]; then
+    echo "Removing force update flag"
+    rm -f "/.force_update"
+    echo "Force update completed successfully"
 fi
 
 if [ "$USE_SSL" = true ]; then
@@ -203,16 +283,51 @@ EOF
         report_error_and_exit "Failed to generate SSL certificate request"
     fi
 
-    if ! curl --header 'Content-Type: application/octet-stream' \
+    max_retries=5
+    retry_delay=2
+    for attempt in $(seq 1 "$max_retries"); do
+        http_code=$(curl -sS -o /etc/instance.crt -w '%{http_code}' \
+            --header 'Content-Type: application/octet-stream' \
             --data-binary @/etc/instance.csr \
-        -X \
-        POST "https://console.vast.ai/api/v0/sign_cert/?instance_id=$CONTAINER_ID" > /etc/instance.crt; then
-        report_error_and_exit "Failed to sign SSL certificate"
+            -X POST "https://console.vast.ai/api/v0/sign_cert/?instance_id=$CONTAINER_ID")
+        if [ "$http_code" -ge 200 ] && [ "$http_code" -lt 300 ]; then
+            break
         fi
+        echo "SSL cert signing attempt $attempt/$max_retries failed (HTTP $http_code)"
+        if [ "$attempt" -eq "$max_retries" ]; then
+            report_error_and_exit "Failed to sign SSL certificate after $max_retries attempts (HTTP $http_code)"
+        fi
+        sleep "$retry_delay"
+        retry_delay=$((retry_delay * 2))
+    done
 fi
 
 export REPORT_ADDR WORKER_PORT USE_SSL UNSECURED
 
+# ─── SDK Deployment Mode ───────────────────────────────────────────────
+if [ "$IS_DEPLOYMENT" = "true" ]; then
+    echo "=== SDK Deployment Mode ==="
+    echo "DEPLOYMENT_ID: $DEPLOYMENT_ID"
+
+    DEPLOY_DIR="/workspace/deployment"
+    mkdir -p "$DEPLOY_DIR"
+
+    VAST_API_BASE="${VAST_API_BASE:-https://console.vast.ai}"
+
+    # Download deployment code, retrying until the blob is available on S3.
+    # The s3_key exists in the DB as soon as the deployment is created, but the
+    # actual upload may still be in flight from the client side.
+
+    # Install SDK (uses the install_vastai_sdk function which supports SDK_BRANCH/SDK_VERSION)
+    install_vastai_sdk
+    # Run deployment in serve mode
+    export VAST_DEPLOYMENT_MODE=serve
+    echo "Starting deployment: python3 $DEPLOY_DIR/deployment.py"
+    serve-vast-deployment
+    exit $?
+fi
+# ─── End SDK Deployment Mode ───────────────────────────────────────────
+
 if ! cd "$SERVER_DIR"; then
     report_error_and_exit "Failed to cd into SERVER_DIR: $SERVER_DIR"
 fi
@@ -224,19 +339,19 @@ set +e
 PY_STATUS=1
 
 if [ -f "$SERVER_DIR/worker.py" ]; then
-    echo "trying worker.py"
+    echo "Running worker.py"
     python3 -m "worker" |& tee -a "$PYWORKER_LOG"
     PY_STATUS=${PIPESTATUS[0]}
 fi
 
 if [ "${PY_STATUS}" -ne 0 ] && [ -f "$SERVER_DIR/workers/$BACKEND/worker.py" ]; then
-    echo "trying workers.${BACKEND}.worker"
+    echo "Running workers.${BACKEND}.worker"
     python3 -m "workers.${BACKEND}.worker" |& tee -a "$PYWORKER_LOG"
     PY_STATUS=${PIPESTATUS[0]}
 fi
 
 if [ "${PY_STATUS}" -ne 0 ] && [ -f "$SERVER_DIR/workers/$BACKEND/server.py" ]; then
-    echo "trying workers.${BACKEND}.server"
+    echo "Running workers.${BACKEND}.server"
     python3 -m "workers.${BACKEND}.server" |& tee -a "$PYWORKER_LOG"
     PY_STATUS=${PIPESTATUS[0]}
 fi
@@ -250,4 +365,4 @@ if [ "${PY_STATUS}" -ne 0 ]; then
     report_error_and_exit "PyWorker exited with status ${PY_STATUS}"
 fi
 
-echo "launching PyWorker server done"
+echo "PyWorker bootstrap complete"

workers/comfyui-json/README.md

@@ -104,13 +104,17 @@ Images will be saved locally AND uploaded to `s3://{bucket}/comfyui/{filename}`.
 
 ### Custom Benchmark Workflows
 
-You can provide a custom ComfyUI workflow for benchmarking by creating `workers/comfyui-json/misc/benchmark.json`. This allows you to test performance using your preferred models and workflow complexity.
+You can provide a custom ComfyUI workflow for benchmarking. This allows you to test performance using your preferred models and workflow complexity.
 
-**Ways to provide the benchmark file:**
-- Fork this repository and add your `benchmark.json` file
-- Write the file during worker provisioning (onstart script or setup phase)
+**Ways to provide the benchmark file** (in resolution order — first match wins):
 
-An example file is provided in the repository. To ensure varied generations, use the placeholder `__RANDOM_INT__` in place of static seed values - it will be replaced with a random integer for each benchmark run.
+1. **Fork this repository** and commit your workflow to `workers/comfyui-json/misc/benchmark.json`.
+2. **Write the file during provisioning** to a path *outside* the pyworker tree (e.g. `/workspace/benchmark.json`) and export `BENCHMARK_JSON_PATH` so the worker can find it. The pyworker repo is cloned by `start_server.sh` *after* provisioning runs, so provisioning cannot write into `misc/` directly — the destination would be clobbered, or the clone would fail.
+3. **Run on the vast.ai ComfyUI base image.** Its `convert-workflows.sh` maintains `/opt/comfyui-api-wrapper/workflows/pyworker_benchmark.json` as a symlink to the first provisioned workflow; the worker reads this automatically when neither of the above is set. No env var required.
+
+If `BENCHMARK_JSON_PATH` is set but points at a missing or unreadable file, the worker logs a warning and falls through to the next tier rather than going straight to the SD1.5 fallback.
+
+An example workflow is provided at `workers/comfyui-json/misc/benchmark.json.example`. To ensure varied generations, use the placeholder `__RANDOM_INT__` in place of static seed values — it will be replaced with a random integer for each benchmark run.
 
 ### Default Benchmark (Fallback)
 
@@ -120,9 +124,10 @@ The default benchmark uses Stable Diffusion v1.5 with ComfyUI's standard text-to
 
 | Environment Variable | Default Value | Description |
 | -------------------- | ------------- | ----------- |
-| BENCHMARK_TEST_WIDTH | 512 | Image width (pixels) |
-| BENCHMARK_TEST_HEIGHT | 512 | Image height (pixels) |
-| BENCHMARK_TEST_STEPS | 20 | Number of denoising steps |
+| BENCHMARK_JSON_PATH | (unset) | Path to a custom workflow file outside the pyworker tree. Used if `misc/benchmark.json` is absent. Falls through to `/opt/comfyui-api-wrapper/workflows/pyworker_benchmark.json` if set but missing. |
+| BENCHMARK_TEST_WIDTH | 512 | Fallback benchmark: image width (pixels) |
+| BENCHMARK_TEST_HEIGHT | 512 | Fallback benchmark: image height (pixels) |
+| BENCHMARK_TEST_STEPS | 20 | Fallback benchmark: number of denoising steps |
 
 Each benchmark run uses a random prompt from `misc/test_prompts.txt` and a random seed to ensure consistent GPU load patterns.
 

workers/comfyui-json/misc/benchmark.json.example

@@ -0,0 +1,107 @@
+{
+    "3": {
+        "inputs": {
+            "seed": "__RANDOM_INT__",
+            "steps": 20,
+            "cfg": 8,
+            "sampler_name": "euler",
+            "scheduler": "normal",
+            "denoise": 1,
+            "model": [
+            "4",
+            0
+            ],
+            "positive": [
+            "6",
+            0
+            ],
+            "negative": [
+            "7",
+            0
+            ],
+            "latent_image": [
+            "5",
+            0
+            ]
+        },
+        "class_type": "KSampler",
+        "_meta": {
+            "title": "KSampler"
+        }
+    },
+    "4": {
+        "inputs": {
+            "ckpt_name": "v1-5-pruned-emaonly-fp16.safetensors"
+        },
+        "class_type": "CheckpointLoaderSimple",
+        "_meta": {
+            "title": "Load Checkpoint"
+        }
+    },
+    "5": {
+        "inputs": {
+            "width": 512,
+            "height": 512,
+            "batch_size": 1
+        },
+        "class_type": "EmptyLatentImage",
+        "_meta": {
+            "title": "Empty Latent Image"
+        }
+    },
+    "6": {
+        "inputs": {
+            "text": "beautiful scenery nature glass bottle landscape, , purple galaxy bottle,",
+            "clip": [
+            "4",
+            1
+            ]
+        },
+        "class_type": "CLIPTextEncode",
+        "_meta": {
+            "title": "CLIP Text Encode (Prompt)"
+        }
+    },
+    "7": {
+        "inputs": {
+            "text": "text, watermark",
+            "clip": [
+            "4",
+            1
+            ]
+        },
+        "class_type": "CLIPTextEncode",
+        "_meta": {
+            "title": "CLIP Text Encode (Prompt)"
+        }
+    },
+    "8": {
+        "inputs": {
+            "samples": [
+            "3",
+            0
+            ],
+            "vae": [
+            "4",
+            2
+            ]
+        },
+        "class_type": "VAEDecode",
+        "_meta": {
+            "title": "VAE Decode"
+        }
+    },
+    "9": {
+        "inputs": {
+            "filename_prefix": "ComfyUI",
+            "images": [
+            "8",
+            0
+            ]
+        },
+        "class_type": "SaveImage",
+        "_meta": {
+            "title": "Save Image"
+        }
+    }
+}

workers/comfyui-json/misc/test_prompts.txt

@@ -0,0 +1,34 @@
+cartoon character of a person with a hoodie , in style of cytus and deemo, ork, gold chains, realistic anime cat, dripping black goo, lineage revolution style, thug life, cute anthropomorphic bunny, balrog, arknights, aliased, very buff, black and red and yellow paint, painting illustration collage style, character composition in vector with white background
+stardew valley, fine details
+2D Vector Illustration of a child with soccer ball Art for Sublimation, Design Art, Chrome Art, Painting and Stunning Artwork, Highly Detailed Digital Painting, Airbrush Art, Highly Detailed Digital Artwork, Dramatic Artwork, stained antique yellow copper paint, digital airbrush art, detailed by Mark Brooks, Chicano airbrush art, Swagger! snake Culture
+realistic futuristic city-downtown with short buildings, sunset
+seascape by Ray Collins and artgerm, front view of a perfect wave, sunny background, ultra detailed water
+inspired by realflow-cinema4d editor features, create image of a transparent luxury cup with ice fruits and mint, connected with white, yellow and pink cream, Slow - High Speed MO Photography, YouTube Video Screenshot, Abstract Clay, Transparent Cup , molecular gastronomy, wheel, 3D fluid,Simulation rendering, still video, 4k polymer clay futras photography, very surreal, Houdini Fluid Simulation, hyperrealistic CGI and FLUIDS & MULTIPHYSICS SIMULATION effect, with Somali Stain Lurex, Metallic Jacquard, Gold Thread, Mulberry Silk, Toub Saree, Warm background, a fantastic image worthy of an award.
+biker with backpack on his back riding a motorcycle, Style by Ade Santora, Oilpunk, Cover photo, craig mullins style, on the cover of a magazine, Outdoor Magazine, inspired by Alex Petruk APe, image of a male biker, Cover of an award-winning magazine, the man has a backpack, photo for magazine, with a backpack, magazine cover
+generate a collage-style illustration inspired by the Procreate raster graphic editor, photographic illustration with the theme, 2D vector, art for textile sublimation, containing surrealistic cartoon cat wearing a baseball cap and jeans standing in front of a poster, inspired by Sadao Watanabe, Doraemon, Japanese cartoon style, Eichiro Oda, Iconic high detail character, Director: Nakahara Nantenbō, Kastuhiro Otomo, image detailed, by Miyamoto, Hidetaka Miyazaki, Katsuhiro illustration, 8k, masterpiece, Minimize noise and grain in photo quality without lose quality and increase brightness and lighting,Symmetry and Alignment, Avoid asymmetrical shapes and out-of-focus points. Focus and Sharpness: Make sure the image is focused and sharp and encourages the viewer to see it as a work of art printed on fabric.
+fantasy medieval village world inside a glass sphere , high detail, fantasy, realistic, light effect, hyper detail, volumetric lighting, cinematic, macro, depth of field, blur, red light and clouds from the back, highly detailed epic cinematic concept art cg render made in maya, blender and photoshop, octane render, excellent composition, dynamic dramatic cinematic lighting, aesthetic, very inspirational, world inside a glass sphere by james gurney by artgerm with james jean, joe fenton and tristan eaton by ross tran, fine details
+armored hero with a glowing axe, mecha science_fiction, jungle background, dynamic lighting, detailed shading, digital texture painting, masterpiece, studio quality, 6k
+elderly figure in a leather jacket DJing in a smoky nightclub, mixing live on a giant console, dramatic stage lighting, a masterpiece
+elderly figure in a leather jacket on a motorcycle, magazine cover lighting, a masterpiece
+a young pilot ordering a burger and fries from a futuristic space cantina
+I want to generate a group avatar for a Feishu group chat. The role of this group is daily software technical communication. Now the subject technology stacks that members of this group discuss daily include: algorithms, data structures, optimization, functional programming, and the programming languages often discussed are: TypeScript, Java, python, etc. I hope this avatar has a simple aesthetic, this avatar is a single person avatar
+portrait Anime black girl cute-fine-face, pretty face, realistic shaded Perfect face, fine details. Anime. realistic shaded lighting by Ilya Kuvshinov Giuseppe Dangelico Pino and Michael Garmash and Rob Rey, IAMAG premiere, WLOP matte print, cute freckles, masterpiece
+young woman in modern fashion editorial, beige miniskirt and dark brown turtleneck sweater, soft studio lighting, brown hair, grey eyes, fine details, magazine cover style, a masterpiece
+Cute small cat sitting in a movie theater eating chicken wiggs watching a movie ,unreal engine, cozy indoor lighting, artstation, detailed, digital painting,cinematic,character design by mark ryden and pixar and hayao miyazaki, unreal 5, daz, hyperrealistic, octane render
+Cute small dog sitting in a movie theater eating popcorn watching a movie ,unreal engine, cozy indoor lighting, artstation, detailed, digital painting,cinematic,character design by mark ryden and pixar and hayao miyazaki, unreal 5, daz, hyperrealistic, octane render
+fox bracelet made of buckskin with fox features, rich details, fine carvings, studio lighting
+crane buckskin bracelet with crane features, rich details, fine carvings, studio lighting
+london luxurious interior living-room, light walls
+Parisian luxurious interior penthouse bedroom, dark walls, wooden panels
+cute girl, crop-top, blond hair, black glasses, stretching, with background by greg rutkowski makoto shinkai kyoto animation key art feminine mid shot
+houses in front, houses background, straight houses, digital art, smooth, sharp focus, gravity falls style, doraemon style, shinchan style, anime style
+Simplified technical drawing, Leonardo da Vinci, Mechanical Dinosaur Skeleton, Minimalistic annotations, Hand-drawn illustrations, Basic design and engineering, Wonder and curiosity
+High quality 8K painting impressionist style of a Japanese modern city street with a girl on the foreground wearing a traditional wedding dress with a fox mask, staring at the sky, daylight
+a landscape from the Moon with the Earth setting on the horizon, realistic, detailed
+Isometric Atlantis city,great architecture with columns, great details, ornaments,seaweed, blue ambiance, 3D cartoon style, soft light, 45° view
+A hyper realistic avatar of a guy riding on a black honda cbr 650r in leather suit,high detail, high quality,8K,photo realism
+the street of amedieval fantasy town, at dawn, dark, highly detailed
+overwhelmingly beautiful eagle framed with vector flowers, long shiny wavy flowing hair, polished, ultra detailed vector floral illustration mixed with hyper realism, muted pastel colors, vector floral details in background, muted colors, hyper detailed ultra intricate overwhelming realism in detailed complex scene with magical fantasy atmosphere, no signature, no watermark
+a highly detailed matte painting of a man on a hill watching a rocket launch in the distance by studio ghibli, makoto shinkai, by artgerm, by wlop, by greg rutkowski, volumetric lighting, octane render, 4 k resolution, trending on artstation, masterpiece | hyperrealism| highly detailed| insanely detailed| intricate| cinematic lighting| depth of field
+electronik robot and ofice ,unreal engine, cozy indoor lighting, artstation, detailed, digital painting,cinematic,character design by mark ryden and pixar and hayao miyazaki, unreal 5, daz, hyperrealistic, octane render
+exquisitely intricately detailed illustration, of a small world with a lake and a rainbow, inside a closed glass jar.

workers/comfyui-json/worker.py

@@ -1,60 +1,225 @@
+"""ComfyUI worker for the vast.ai PyWorker SDK.
+
+Each worker runs a benchmark on warm-up. The payload is selected as follows:
+
+  1. If ``misc/benchmark.json`` exists in the cloned worker tree, it is
+     used as a custom ComfyUI workflow. Use this if you fork the repo and
+     bake in your workflow.
+  2. Else, if ``$BENCHMARK_JSON_PATH`` is set and points at a readable
+     file, it is used. Use this from a provisioning script — provisioning
+     runs before pyworker is cloned, so it cannot write into ``misc/``,
+     but it can drop the workflow elsewhere (e.g. ``/workspace/``) and
+     export this env var.
+  3. Else, if the well-known path
+     ``/opt/comfyui-api-wrapper/workflows/pyworker_benchmark.json`` exists,
+     it is used. The vast.ai ComfyUI base image's ``convert-workflows.sh``
+     maintains this as a symlink to the first provisioned workflow, so on
+     that image no env var is needed.
+  4. Otherwise an SD1.5 Text2Image fallback runs, parameterised by the
+     ``BENCHMARK_TEST_{WIDTH,HEIGHT,STEPS}`` env vars and a random prompt
+     from ``misc/test_prompts.txt``.
+
+``__RANDOM_INT__`` placeholders in custom workflows are substituted
+server-side by ai-dock/comfyui-api-wrapper, so this worker does not handle
+them itself.
+"""
+
+import json
+import logging
+import os
 import random
 import sys
+from pathlib import Path
 
 from vastai import Worker, WorkerConfig, HandlerConfig, LogActionConfig, BenchmarkConfig
 
-# ComyUI model configuration
+# ComfyUI model configuration. The model server is ai-dock's
+# comfyui-api-wrapper sitting in front of ComfyUI itself, not ComfyUI's
+# own port (18188). We tail the api-wrapper's log rather than ComfyUI's
+# and key off the api-wrapper's own structured readiness/fault signals:
+#
+#   BACKENDS_READY            — api-wrapper has confirmed every ComfyUI
+#                               backend passes HTTP+WS probes. Until
+#                               this fires, posting to /generate/sync
+#                               can hit "Cannot connect to host" inside
+#                               the api-wrapper, which the SDK can't
+#                               recover from since __call_backend
+#                               doesn't retry connection-refused.
+#   BACKENDS_READY_TIMEOUT    — backends never reachable within
+#                               api-wrapper's deadline. Worker is
+#                               unrecoverable; mark errored.
+#   BACKEND_UNRECOVERABLE     — CUDA fault / illegal memory access on a
+#                               backend's GPU. Same fate.
+#   Application startup failed — uvicorn's own ASGI lifespan failed.
+#
+# These tokens are emitted by ai-dock/comfyui-api-wrapper >= the
+# "feat/backend-readiness-log-signals" change. Older wrappers won't
+# emit BACKENDS_READY, so warm-up will stall — pin the wrapper version
+# accordingly.
 MODEL_SERVER_URL           = 'http://127.0.0.1'
 MODEL_SERVER_PORT          = 18288
-MODEL_LOG_FILE             = '/var/log/portal/comfyui.log'
+MODEL_LOG_FILE             = '/var/log/portal/api-wrapper.log'
 MODEL_HEALTHCHECK_ENDPOINT = "/health"
 
-# ComyUI-specific log messages
+# Trigger benchmark only after the full stack (api-wrapper + ComfyUI
+# backends) is reachable. See BACKENDS_READY in the comment above.
 MODEL_LOAD_LOG_MSG = [
-    "To see the GUI go to: "
+    "BACKENDS_READY",
 ]
 
+# LogAction.ModelError is fatal: the SDK calls backend_errored() and
+# locks the worker into a permanent error state. Patterns must
+# therefore only match conditions where the api-wrapper genuinely
+# cannot serve any request — supervisord restarts on uvicorn exit, so
+# a real failure self-heals rather than dragging the worker down.
+#
+# Notably *not* matched here:
+#   - per-request errors (PreprocessWorker failures, ComfyUI workflow
+#     validation, "Value not in list:") — one malformed client payload
+#     would otherwise kill the worker
+#   - "CUDA out of memory" — surfaces both as a misconfigured GPU
+#     (which the benchmark-failure path already catches via
+#     backend_errored) and as a too-greedy client request, which is
+#     indistinguishable from a substring match
+#   - convert-workflows.sh warnings — that script is not load-bearing
+#     for serving
 MODEL_ERROR_LOG_MSGS = [
-    "MetadataIncompleteBuffer",
-    "Value not in list: ",
-    "[ERROR] Provisioning Script failed"
+    "BACKENDS_READY_TIMEOUT",       # backends never reachable
+    "BACKEND_UNRECOVERABLE",        # CUDA fault latched per backend
+    "Application startup failed",   # uvicorn ASGI lifespan startup failed
 ]
 
-MODEL_INFO_LOG_MSGS = [
-    '"message":"Downloading'
-]
+# LogAction.Info is purely informational (echoes log lines into the vast
+# console). Nothing in api-wrapper.log is currently worth surfacing —
+# model downloads are upstream in provisioning, per-request logs are
+# too noisy.
+MODEL_INFO_LOG_MSGS = []
 
-benchmark_prompts = [
-    "Cartoon hoodie hero; orc, anime cat, bunny; black goo; buff; vector on white.",
-    "Cozy farming-game scene with fine details.",
-    "2D vector child with soccer ball; airbrush chrome; swagger; antique copper.",
-    "Realistic futuristic downtown of low buildings at sunset.",
-    "Perfect wave front view; sunny seascape; ultra-detailed water; artful feel.",
-    "Clear cup with ice, fruit, mint; creamy swirls; fluid-sim CGI; warm glow.",
-    "Male biker with backpack on motorcycle; oilpunk; award-worthy magazine cover.",
-    "Collage for textile; surreal cartoon cat in cap/jeans before poster; crisp.",
-    "Medieval village inside glass sphere; volumetric light; macro focus.",
-    "Iron Man with glowing axe; mecha sci-fi; jungle scene; dynamic light.",
-    "Pope Francis DJ in leather jacket, mixing on giant console; dramatic.",
-]
+# Benchmark assets shipped alongside this worker. Resolved relative to this
+# file so the worker keeps working regardless of the launch cwd.
+MISC_DIR       = Path(__file__).parent / "misc"
+BENCHMARK_FILE = MISC_DIR / "benchmark.json"
+TEST_PROMPTS   = MISC_DIR / "test_prompts.txt"
+
+# Well-known location maintained by the vast.ai ComfyUI base image.
+# convert-workflows.sh symlinks this to the first provisioned workflow,
+# letting the base image work out-of-the-box without any env var.
+WELLKNOWN_BENCHMARK = Path("/opt/comfyui-api-wrapper/workflows/pyworker_benchmark.json")
+
+log = logging.getLogger(__name__)
+
+# Used when test_prompts.txt is unreadable or empty. Bare and generic
+# on purpose — this is a benchmark seed, not a creative output.
+_FALLBACK_PROMPT = "a still life on a wooden table, soft daylight"
 
 
+def _env_int(name: str, default: int) -> int:
+    """Read an integer env var, warning + falling back on bad values."""
+    raw = os.getenv(name)
+    if raw is None or raw == "":
+        return default
+    try:
+        return int(raw)
+    except ValueError:
+        log.warning("ignoring %s=%r (not an int); using default %d", name, raw, default)
+        return default
 
-benchmark_dataset = [
-    {
+
+def _try_load_workflow(path: Path) -> dict | None:
+    """Load and return a benchmark workflow from ``path``.
+
+    Returns None on any failure (path missing, not a regular file,
+    unreadable, invalid JSON) so the caller can fall through to the
+    next tier rather than dropping straight to the SD1.5 default.
+    """
+    if not path.is_file():
+        return None
+    try:
+        with open(path) as f:
+            return json.load(f)
+    except (json.JSONDecodeError, OSError) as e:
+        log.warning("Failed to load %s: %s; trying next tier", path, e)
+        return None
+
+
+def _custom_workflow_payload() -> dict | None:
+    """Try each benchmark workflow tier in order; return the first one
+    that loads cleanly as a payload, or None if every tier is absent /
+    unreadable. Tiers (in order): in-tree ``misc/benchmark.json``,
+    ``$BENCHMARK_JSON_PATH``, well-known base-image symlink.
+    """
+    env_path = os.getenv("BENCHMARK_JSON_PATH")
+    candidates = [("misc", BENCHMARK_FILE)]
+    if env_path:
+        candidates.append(("env", Path(env_path)))
+    candidates.append(("well-known", WELLKNOWN_BENCHMARK))
+
+    for label, path in candidates:
+        # Surface a warning specifically when the operator pointed
+        # BENCHMARK_JSON_PATH at something we can't use — silent
+        # fall-through there is a footgun (typo => SD1.5 fallback,
+        # operator wonders why custom benchmark didn't take).
+        if not path.is_file():
+            if label == "env":
+                log.warning(
+                    "BENCHMARK_JSON_PATH=%s is not a readable file; trying fallbacks", path
+                )
+            continue
+        workflow = _try_load_workflow(path)
+        if workflow is None:
+            continue
+        log.info("Using custom benchmark workflow from %s (%s)", path, label)
+        return {
+            "input": {
+                "request_id": f"test-{random.randint(1000, 99999)}",
+                "workflow_json": workflow,
+            }
+        }
+    return None
+
+
+def _load_prompts() -> list[str]:
+    """Read misc/test_prompts.txt; defensive against missing/empty file."""
+    try:
+        with open(TEST_PROMPTS) as f:
+            prompts = [line.strip() for line in f if line.strip()]
+    except OSError as e:
+        log.warning("could not read %s: %s; using built-in fallback prompt", TEST_PROMPTS, e)
+        return [_FALLBACK_PROMPT]
+    if not prompts:
+        log.warning("%s is empty; using built-in fallback prompt", TEST_PROMPTS)
+        return [_FALLBACK_PROMPT]
+    return prompts
+
+
+def _default_payload() -> dict:
+    """Build the SD1.5 Text2Image fallback payload."""
+    prompts = _load_prompts()
+    return {
         "input": {
             "request_id": f"test-{random.randint(1000, 99999)}",
             "modifier": "Text2Image",
             "modifications": {
-                "prompt": prompt,
-                "width": 512,
-                "height": 512,
-                "steps": 20,
-                "seed": random.randint(0, sys.maxsize)
+                "prompt": random.choice(prompts),
+                "width":  _env_int("BENCHMARK_TEST_WIDTH",  512),
+                "height": _env_int("BENCHMARK_TEST_HEIGHT", 512),
+                "steps":  _env_int("BENCHMARK_TEST_STEPS",  20),
+                "seed":   random.randint(0, sys.maxsize),
             }
         }
-    } for prompt in benchmark_prompts
-]
+    }
+
+
+def make_benchmark_payload() -> dict:
+    """Build one benchmark request payload.
+
+    Called once per benchmark run by the SDK; using a generator (rather
+    than a static ``dataset=``) lets each run re-pick a prompt and re-roll
+    the seed, and avoids holding multiple copies of a large workflow JSON
+    in memory.
+    """
+    return _custom_workflow_payload() or _default_payload()
+
 
 worker_config = WorkerConfig(
     model_server_url=MODEL_SERVER_URL,
@@ -67,7 +232,7 @@ worker_config = WorkerConfig(
             allow_parallel_requests=False,
             max_queue_time=10.0,
             benchmark_config=BenchmarkConfig(
-                dataset=benchmark_dataset,
+                generator=make_benchmark_payload,
             )
         )
     ],

workers/tgi/worker.py

@@ -35,7 +35,7 @@ def benchmark_generator() -> dict:
     benchmark_data = {
         "inputs": prompt,
         "parameters": {
-            "max_new_tokens": 128,
+            "max_new_tokens": 500,
             "temperature": 0.7,
             "return_full_text": False
         }