Support BENCHMARK_JSON_PATH for provisioning-supplied benchmarks

start_server.sh clones pyworker into /workspace/vast-pyworker after the
provisioning phase has run, so a provisioning script that wants to ship
a custom benchmark workflow cannot write to misc/benchmark.json — that
path doesn't exist yet at provisioning time, and pre-creating it would
make the subsequent clone fail.

Allow provisioning to drop the workflow anywhere (e.g. /workspace) and
point the worker at it via the BENCHMARK_JSON_PATH env var. The in-tree
file still takes precedence (so forks with a baked-in benchmark keep
working unchanged); the env var is consulted only as a second choice,
and a misconfigured path logs a warning rather than silently degrading
to the SD1.5 fallback.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Rob Ballantyne
2026-05-07 11:24:14 +01:00
parent 2dd4f7fc38
commit a634ba07a6
2 changed files with 44 additions and 16 deletions
+10 -7
View File
@@ -104,13 +104,15 @@ Images will be saved locally AND uploaded to `s3://{bucket}/comfyui/{filename}`.
### Custom Benchmark Workflows ### 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:** **Ways to provide the benchmark file:**
- Fork this repository and add your `benchmark.json` file - **Fork this repository** and commit your workflow to `workers/comfyui-json/misc/benchmark.json`.
- Write the file during worker provisioning (onstart script or setup phase) - **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.
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. If both are present, the in-tree `misc/benchmark.json` wins; `BENCHMARK_JSON_PATH` is consulted only when no in-tree file exists. If the env var is set but points at a missing or unreadable file, the worker logs a warning and falls back to the default benchmark.
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) ### Default Benchmark (Fallback)
@@ -120,9 +122,10 @@ The default benchmark uses Stable Diffusion v1.5 with ComfyUI's standard text-to
| Environment Variable | Default Value | Description | | Environment Variable | Default Value | Description |
| -------------------- | ------------- | ----------- | | -------------------- | ------------- | ----------- |
| BENCHMARK_TEST_WIDTH | 512 | Image width (pixels) | | BENCHMARK_JSON_PATH | (unset) | Path to a custom workflow file outside the pyworker tree. Used only if `misc/benchmark.json` is absent. |
| BENCHMARK_TEST_HEIGHT | 512 | Image height (pixels) | | BENCHMARK_TEST_WIDTH | 512 | Fallback benchmark: image width (pixels) |
| BENCHMARK_TEST_STEPS | 20 | Number of denoising steps | | 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. Each benchmark run uses a random prompt from `misc/test_prompts.txt` and a random seed to ensure consistent GPU load patterns.
+34 -9
View File
@@ -2,10 +2,15 @@
Each worker runs a benchmark on warm-up. The payload is selected as follows: Each worker runs a benchmark on warm-up. The payload is selected as follows:
1. If ``misc/benchmark.json`` exists, it is used as a custom ComfyUI 1. If ``misc/benchmark.json`` exists in the cloned worker tree, it is
workflow (recommended: match the workflow your endpoint will actually used as a custom ComfyUI workflow. Use this if you fork the repo and
serve, so the autoscaler's performance estimate is meaningful). bake in your workflow.
2. Otherwise an SD1.5 Text2Image fallback runs, parameterised by the 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. Otherwise an SD1.5 Text2Image fallback runs, parameterised by the
``BENCHMARK_TEST_{WIDTH,HEIGHT,STEPS}`` env vars and a random prompt ``BENCHMARK_TEST_{WIDTH,HEIGHT,STEPS}`` env vars and a random prompt
from ``misc/test_prompts.txt``. from ``misc/test_prompts.txt``.
@@ -53,17 +58,37 @@ TEST_PROMPTS = MISC_DIR / "test_prompts.txt"
log = logging.getLogger(__name__) log = logging.getLogger(__name__)
def _resolve_benchmark_path() -> Path | None:
"""Return the path to the custom benchmark workflow, or None if absent.
See module docstring for the precedence rule. ``$BENCHMARK_JSON_PATH``
is logged as a warning when set but missing, so a misconfigured
provisioning script doesn't silently degrade to the fallback benchmark.
"""
if BENCHMARK_FILE.exists():
return BENCHMARK_FILE
env_path = os.getenv("BENCHMARK_JSON_PATH")
if not env_path:
return None
path = Path(env_path)
if not path.exists():
log.warning("BENCHMARK_JSON_PATH=%s does not exist; falling back to default benchmark", path)
return None
return path
def _custom_workflow_payload() -> dict | None: def _custom_workflow_payload() -> dict | None:
"""Build a payload from ``misc/benchmark.json``, or None if unavailable.""" """Build a payload from a custom benchmark workflow JSON, or None if unavailable."""
if not BENCHMARK_FILE.exists(): path = _resolve_benchmark_path()
if path is None:
return None return None
try: try:
with open(BENCHMARK_FILE) as f: with open(path) as f:
workflow = json.load(f) workflow = json.load(f)
except (json.JSONDecodeError, OSError) as e: except (json.JSONDecodeError, OSError) as e:
log.error("Failed to load %s: %s; falling back to default benchmark", BENCHMARK_FILE, e) log.error("Failed to load %s: %s; falling back to default benchmark", path, e)
return None return None
log.info("Using custom benchmark workflow from %s", BENCHMARK_FILE) log.info("Using custom benchmark workflow from %s", path)
return { return {
"input": { "input": {
"request_id": f"test-{random.randint(1000, 99999)}", "request_id": f"test-{random.randint(1000, 99999)}",