From a634ba07a6a592aca216ba48d9b4618b7838f3d3 Mon Sep 17 00:00:00 2001 From: Rob Ballantyne Date: Thu, 7 May 2026 11:24:14 +0100 Subject: [PATCH] Support BENCHMARK_JSON_PATH for provisioning-supplied benchmarks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- workers/comfyui-json/README.md | 17 ++++++++------ workers/comfyui-json/worker.py | 43 +++++++++++++++++++++++++++------- 2 files changed, 44 insertions(+), 16 deletions(-) diff --git a/workers/comfyui-json/README.md b/workers/comfyui-json/README.md index c91b09d..b63d5ea 100644 --- a/workers/comfyui-json/README.md +++ b/workers/comfyui-json/README.md @@ -104,13 +104,15 @@ 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) +- **Fork this repository** and commit your workflow to `workers/comfyui-json/misc/benchmark.json`. +- **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) @@ -120,9 +122,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 only if `misc/benchmark.json` is absent. | +| 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. diff --git a/workers/comfyui-json/worker.py b/workers/comfyui-json/worker.py index 2e78a5c..944fe0f 100644 --- a/workers/comfyui-json/worker.py +++ b/workers/comfyui-json/worker.py @@ -2,10 +2,15 @@ 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 - workflow (recommended: match the workflow your endpoint will actually - serve, so the autoscaler's performance estimate is meaningful). - 2. Otherwise an SD1.5 Text2Image fallback runs, parameterised by the + 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. 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``. @@ -53,17 +58,37 @@ TEST_PROMPTS = MISC_DIR / "test_prompts.txt" 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: - """Build a payload from ``misc/benchmark.json``, or None if unavailable.""" - if not BENCHMARK_FILE.exists(): + """Build a payload from a custom benchmark workflow JSON, or None if unavailable.""" + path = _resolve_benchmark_path() + if path is None: return None try: - with open(BENCHMARK_FILE) as f: + with open(path) as f: workflow = json.load(f) 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 - log.info("Using custom benchmark workflow from %s", BENCHMARK_FILE) + log.info("Using custom benchmark workflow from %s", path) return { "input": { "request_id": f"test-{random.randint(1000, 99999)}",