[OMNIML-4788] specdec_bench: configuration.json provenance + upload_to_s3 (#1531)

> [!WARNING]
> **Breaking on-disk schema change (specdec_bench v1.0.0).** This PR
renames the acceptance-rate metric fields across `AcceptanceRate` /
`MTBench` / `SpecBench` writers:
>
> | Old (pre-1.0.0) | New (1.0.0) |
> |---|---|
> | `Request_AR` | `Request_AL` |
> | `Category_AR` | `Category_AL` |
> | `Average_AR` | `Average_AL` |
> | — | `Joint_Acceptance_Rate` (new) |
>
> The renamed values were always **acceptance length** (mean tokens
generated per inference step), not a rate, and the visualizer reads
`*_AL`. Pre-1.0.0 runs in S3 have `*_AR` and no `Joint_AR`; they must be
re-run or post-processed before comparing. The visualizer aggregates
runs by `specdec_bench` major version so accidental cross-methodology
comparison is blocked.

### What does this PR do?

Type of change: new feature

Adds reproducibility provenance to `specdec_bench/configuration.json`
and ports `upload_to_s3.py` from `iputterman/specdec_bench@main`
(personal-namespace fork) into upstream. This is the first PR in a
multi-stage migration off Izzy's fork now that he's left the team.
Tracked in [OMNIML-4788](https://jirasw.nvidia.com/browse/OMNIML-4788).

**Provenance fields added to configuration.json** (alongside existing
argv / engine_version / gpu / python_version):

- `specdec_bench_version` — methodology semver declared in
`specdec_bench/__init__.py`. Bump minor on additive metrics, major on
changed metric *definitions*. The visualizer (Phase 4 of the migration)
will aggregate runs by major version so plots don't accidentally compare
across methodology changes.
- `specdec_bench_sha`, `modelopt_sha`, `modelopt_version`,
`nmm_sandbox_sha`, `container_image` — code/runtime provenance. Each
prefers an env var set by the harness (`SPECDEC_BENCH_SHA`,
`MODELOPT_SHA`, `MODELOPT_VERSION`, `NMM_SANDBOX_SHA`,
`CONTAINER_IMAGE`) and falls back to `git rev-parse` /
`modelopt.__version__` when running standalone. The env-var preference
is necessary because the runtime container has no `.git/` (the launcher
packager tarballs source without git metadata) and may not have
`modelopt` installed.
- `checkpoint.{path, size_bytes, index_sha256, index_source}` — cheap
reproducibility fingerprint that hashes `model.safetensors.index.json`
(or `config.json` fallback). Changes whenever any tensor changes.
- `serving_config` — engine-level config dict captured after init via a
new `Model.get_serving_config()` method. VLLM dumps `AsyncEngineArgs` +
the live `vllm_config.to_dict()`; SGLANG dumps the `engine_kwargs`
passed to `sgl.Engine`; TRTLLM left at the base default `{}` for a later
iteration.
- `timestamp` — UTC ISO 8601.

**Other changes**
- `upload_to_s3.py` + `specdec_bench/s3_utils.py` ported from
iputterman/specdec_bench@main. Recognizes run dirs by sentinel files,
refuses to overwrite existing S3 prefixes.
- `_redact_config` allowlists `tokenizer`, `tokenizer_path`,
`tokenizer_mode`, `tokenizer_revision` so the model path stops being
redacted (latent bug from substring-matching `token` ⊂ `tokenizer`).
- `requirements_speed.txt`: `boto3`, `botocore` added (used by
`s3_utils`).

**Out of scope** (deferred to Phase 1b / Phase 2):
- `--sweep_config` driver that emits per-run-dir nesting
`<sweep>/<NNN_dataset_c<conc>>/`
- `--s3_upload` flag baked into `run.py` itself
- Launcher auto-injection of the provenance env vars (currently the
example YAML sets them statically)
- `container_digest` (enroot integration) and full GPU/driver inventory
- TRTLLM `get_serving_config()`

### Usage

```bash
# Run a smoke benchmark (Qwen3.5-4B + vLLM + MTP draft=3) — example YAML included
uv run launch.py --yaml examples/Qwen/Qwen3.5-4B/specdec_bench_mtp.yaml --yes

# After it lands, upload the run directory to S3:
S3_KEY_ID=team-specdec-workgroup \
S3_SECRET=... \
python upload_to_s3.py /path/to/sweep_dir s3://team-specdec-workgroup/results
```

### Testing

Cluster-tested end-to-end on cw-dfw (Slurm job 11978794, NeMo Run
experiment `cicd_1779403623`, ~19 min wall):

- Qwen3.5-4B + vLLM + MTP draft=3 + SPEED-Bench-Internal/qualitative (80
requests)
- `configuration.json` (22 KB) populated all eight new provenance fields
- `Request_AR` mean 3.327 (vs 3.330 on the pre-Phase-1a run — within
noise; methodology unchanged)
- `upload_to_s3.py` (real upload, not dry-run) landed
[s3://team-specdec-workgroup/results/qwen35_4_mtp_smoke_2026-05-21/specdec_bench_mtp/](https://app.s8k.io/buckets/team-specdec-workgroup/?prefix=results%2Fqwen35_4_mtp_smoke_2026-05-21%2F)
where the visualizer at http://10.131.132.205:8080 can pick it up.

### Before your PR is "Ready for review"

- Is this change backward compatible?: ✅
- `configuration.json` only gains fields. `upload_to_s3.py` /
`s3_utils.py` are new files. `Model.get_serving_config()` default = `{}`
so existing subclasses without an override behave as before.
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: ✅
- `boto3` / `botocore` are Apache 2.0 (permissive); `upload_to_s3.py` +
`s3_utils.py` are ported from a private NVIDIA repo with explicit
copyright headers retained.
- Did you write any new necessary tests?: ❌
- Validated by cluster smoke (see Testing). Will add unit-tests for
`dump_env` provenance fields and `upload_to_s3._discover_runs` in a
follow-up.
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A
  - Internal-facing tooling.
- Did you get Claude approval on this PR?: ❌ (triggering after open)

### Additional Information

Tracked in JIRA
[OMNIML-4788](https://jirasw.nvidia.com/browse/OMNIML-4788). The full
multi-phase plan is on that ticket's SPEC block — this PR is Phase 1a.

Cherry-picked alongside the harness change are two example YAMLs
(`examples/Qwen/Qwen3.5-4B/specdec_bench.yaml` for the NONE
autoregressive baseline, `..._mtp.yaml` for the MTP run) that gave us
cluster-test evidence. Can be split out if preferred.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added an S3 upload CLI for benchmark results with dry-run and
skip-existing options
* Automatic capture of run configuration, provenance and redacted
environment into saved config
  * Models now export serving configuration for reproducible runs
* New launcher entrypoint and example job configs for Qwen SPEED-Bench
runs

* **Documentation**
* README section describing S3 upload usage and supported local layouts

* **Bug Fixes / Changes**
  * Acceptance-rate metric keys renamed in output (AR -> AL)

* **Tests / CI**
* New tests for redaction and S3 utilities; CI now runs specdec_bench
examples

<!-- review_stack_entry_start -->

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](https://app.coderabbit.ai/change-stack/NVIDIA/Model-Optimizer/pull/1531?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack)

<!-- review_stack_entry_end -->
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: chenhany <chenhany@nvidia.com>

Author: ChenhanYu

.github/workflows/example_tests.yml

@@ -36,7 +36,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        example: [llm_distill, llm_qat, llm_sparsity, diffusers_sparsity]
+        example: [llm_distill, llm_qat, llm_sparsity, diffusers_sparsity, specdec_bench]
         include:
           - example: speculative_decoding
             docker_image: "26.01"

examples/specdec_bench/README.md

@@ -64,7 +64,7 @@ python3 run.py \
 
 ### Running [SPEED-Bench](https://huggingface.co/datasets/nvidia/SPEED-Bench) on Llama 3.3 70B + Eagle 3
 
-1. Install the requirements file using `pip install -r requirements_speed.txt`
+1. Install the requirements file using `pip install -r requirements.txt`
 
 2. Prepare the data using the provided script:
 
@@ -145,6 +145,33 @@ python3 run.py \
     --runtime_params runtime_args_long_context.yaml
 ```
 
+## Uploading results to S3
+
+Each `run.py` invocation writes a result directory containing `configuration.json`,
+`timing.json`, `acceptance_rate.json`, and (when applicable) `mtbench.json` / `specbench.json`.
+`upload_to_s3.py` is a single-file, drop-in tool that uploads one run — or an entire sweep —
+to any S3-compatible bucket:
+
+```bash
+python upload_to_s3.py /path/to/run_or_sweep_dir s3://your-bucket/some/prefix \
+    --endpoint https://your-s3-endpoint \
+    --key-id YOUR_KEY_ID \
+    --secret YOUR_SECRET
+```
+
+`--endpoint`, `--key-id`, and `--secret` default to the `S3_ENDPOINT`, `S3_KEY_ID`, and
+`S3_SECRET` environment variables. Omit `--endpoint` (or set `S3_ENDPOINT=""`) to use AWS S3's
+default endpoint. Use `--dry-run` to preview the upload plan, and `--skip-existing` to skip
+runs already present at the destination instead of failing.
+
+The tool handles two directory layouts and mirrors them into S3:
+
+- **Flat** — `LOCAL_DIR/run_name/{configuration,timing,...}.json`
+- **Sweep** — `LOCAL_DIR/sweep_name/run_name/{configuration,timing,...}.json`
+
+`LOCAL_DIR`'s basename is preserved in the destination prefix, so re-uploads from the same
+source land in the same place.
+
 ## Notes
 
 The goal of this benchmark is to provide an easy way to configure, run, and compare speculative implementations across frameworks in an apples-to-apples method.

…les/specdec_bench/requirements_speed.txt examples/specdec_bench/requirements.txtexamples/specdec_bench/requirements_speed.txt renamed to examples/specdec_bench/requirements.txt examples/specdec_bench/requirements_speed.txt renamed to examples/specdec_bench/requirements.txt

@@ -1,3 +1,5 @@
+boto3>=1.34.0
+botocore>=1.34.0
 datasets>=3.1.0
 rich>=14.2.0
 seaborn>=0.13.2

examples/specdec_bench/run.py

@@ -20,6 +20,7 @@
 from specdec_bench import datasets, metrics, models, runners
 from specdec_bench.utils import (
     decode_chat,
+    dump_env,
     encode_chat,
     get_tokenizer,
     postprocess_base,
@@ -174,6 +175,10 @@ def run_simple(args):
     if args.save_dir is not None:
         for metric in metrics_list:
             metric.update_directory(args.save_dir)
+        # Stamp configuration.json BEFORE the run loop so the file lands even
+        # when the run crashes mid-way. Engine init is already done, so the
+        # live serving_config from the model is available.
+        dump_env(args, args.save_dir, overrides={"serving_config": model.get_serving_config()})
 
     runner = runners.SimpleRunner(model, metrics=metrics_list)
 

examples/specdec_bench/specdec_bench/__init__.py

@@ -13,3 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Re-export modelopt's version so configuration.json's `specdec_bench_version`
+# tracks the parent package without a separate semver source of truth.
+# Breaking schema/methodology changes are recorded in commit messages and
+# fingerprinted by `specdec_bench_sha` in configuration.json.
+from modelopt import __version__

examples/specdec_bench/specdec_bench/datasets/speed.py

@@ -147,7 +147,7 @@ def __init__(
     ):
         if not_installed:
             raise ImportError(
-                "Additional packages are required to use SPEED-Bench. Please run `pip install -r requirements_speed.txt`"
+                "Additional packages are required to use SPEED-Bench. Please run `pip install -r requirements.txt`"
             )
         self.data: list[Request] = []
         self.num_samples = num_samples

examples/specdec_bench/specdec_bench/metrics/acceptance_rate.py

@@ -55,23 +55,31 @@ def _process_lengths(self, lengths):
             self.out["Conditional_Acceptance_Rate"][k] = running_len / sum_lengths / prev_ratio
             prev_ratio = running_len / sum_lengths
             running_len -= v
+        # Joint acceptance rate at step k = product of conditional acceptance
+        # rates at steps 1..k = probability that ≥k tokens are accepted in
+        # a row. The visualizer renders this as a separate panel.
+        self.out["Joint_Acceptance_Rate"] = {}
+        running_joint = 1.0
+        for k, cond_ar in self.out["Conditional_Acceptance_Rate"].items():
+            running_joint *= cond_ar
+            self.out["Joint_Acceptance_Rate"][k] = running_joint
 
     def process_final(self, text_outputs):
         all_ar = []
         lengths = {}
-        self.out["Request_AR"] = {}
+        self.out["Request_AL"] = {}
         self.prompt_ar = dict(sorted(self.prompt_ar.items(), key=lambda x: x[0]))
         for request_id, turns in self.prompt_ar.items():
-            self.out["Request_AR"][request_id] = {}
+            self.out["Request_AL"][request_id] = {}
             for turn_id, turn in turns.items():
                 ar = sum(turn) / len(turn)
-                self.out["Request_AR"][request_id][turn_id] = ar
+                self.out["Request_AL"][request_id][turn_id] = ar
                 all_ar.append(ar)
                 self._get_lengths(turn, lengths)
-                print(request_id, turn_id, self.out["Request_AR"][request_id][turn_id])
+                print(request_id, turn_id, self.out["Request_AL"][request_id][turn_id])
         average_ar = sum(all_ar) / len(all_ar)
-        print("Average AR:", average_ar)
-        self.out["Average_AR"] = average_ar
+        print("Average AL:", average_ar)
+        self.out["Average_AL"] = average_ar
         self._process_lengths(lengths)
         self.write()
         self._format_write_output(text_outputs)

examples/specdec_bench/specdec_bench/metrics/mtbench.py

@@ -34,29 +34,29 @@ class MTBench(AcceptanceRate):
     def process_final(self, text_outputs):
         i = 0
         lengths = {}
-        self.out["Request_AR"] = {}
+        self.out["Request_AL"] = {}
         self.prompt_ar = dict(sorted(self.prompt_ar.items(), key=lambda x: x[0]))
         for request_id, turns in self.prompt_ar.items():
             turn_1 = turns[0]
             turn_2 = turns[1]
             q_id = request_id
             mtbench_topic = MTBENCH_TOPICS[q_id // 10]
-            self.out["Request_AR"][request_id] = sum(turn_1 + turn_2) / len(turn_1 + turn_2)
+            self.out["Request_AL"][request_id] = sum(turn_1 + turn_2) / len(turn_1 + turn_2)
             self._get_lengths(turn_1, lengths)
             self._get_lengths(turn_2, lengths)
             print(mtbench_topic, sum(turn_1 + turn_2) / len(turn_1 + turn_2))
         per_category = [[] for _ in range(len(MTBENCH_TOPICS))]
-        for q_id, ar in self.out["Request_AR"].items():
+        for q_id, ar in self.out["Request_AL"].items():
             per_category[q_id // 10].append(ar)
-        self.out["Category_AR"] = {}
+        self.out["Category_AL"] = {}
         for i, category in enumerate(per_category):
             if len(category) > 0:
                 category_ar = sum(category) / len(category)
-                self.out["Category_AR"][MTBENCH_TOPICS[i]] = category_ar
-                print(f"{MTBENCH_TOPICS[i]} Average AR: {category_ar}")
-        average_ar = sum(self.out["Request_AR"].values()) / len(self.out["Request_AR"])
-        print("Average AR:", average_ar)
-        self.out["Average_AR"] = average_ar
+                self.out["Category_AL"][MTBENCH_TOPICS[i]] = category_ar
+                print(f"{MTBENCH_TOPICS[i]} Average AL: {category_ar}")
+        average_ar = sum(self.out["Request_AL"].values()) / len(self.out["Request_AL"])
+        print("Average AL:", average_ar)
+        self.out["Average_AL"] = average_ar
         self._process_lengths(lengths)
         self.write()
         self._format_write_output(text_outputs)

examples/specdec_bench/specdec_bench/metrics/specbench.py

@@ -44,26 +44,26 @@ def __init__(self, requests):
 
     def process_final(self, text_outputs):
         lengths = {}
-        self.out["Request_AR"] = {}
+        self.out["Request_AL"] = {}
         for request_id, request in enumerate(self.requests):
             turns = self.prompt_ar[request_id].values()
             assert len(turns) == len(request.turns), (
                 f"Number of turns {len(turns)} does not match number of turns in request {len(request.turns)}"
             )
-            self.out["Request_AR"][request.question_id] = mean(list(chain(*turns)))
+            self.out["Request_AL"][request.question_id] = mean(list(chain(*turns)))
             for turn in turns:
                 self._get_lengths(turn, lengths)
-            print(request.category, self.out["Request_AR"][request.question_id])
+            print(request.category, self.out["Request_AL"][request.question_id])
         per_category = defaultdict(list)
         for request in self.requests:
-            per_category[request.category].append(self.out["Request_AR"][request.question_id])
-        self.out["Category_AR"] = {}
+            per_category[request.category].append(self.out["Request_AL"][request.question_id])
+        self.out["Category_AL"] = {}
         for category_name, category_ar in per_category.items():
             if len(category_ar) > 0:
                 category_ar = mean(category_ar)
-                self.out["Category_AR"][category_name] = category_ar
-        average_ar = mean(self.out["Request_AR"].values())
-        self.out["Average_AR"] = average_ar
+                self.out["Category_AL"][category_name] = category_ar
+        average_ar = mean(self.out["Request_AL"].values())
+        self.out["Average_AL"] = average_ar
         self._process_lengths(lengths)
         self.write()
         self._format_write_output(text_outputs)
@@ -93,15 +93,15 @@ def _pretty_print_results(self):
             header_style="bold magenta",
         )
         table.add_column("Category", style="cyan", no_wrap=True)
-        table.add_column("Average AR", justify="right", style="green")
+        table.add_column("Average AL", justify="right", style="green")
 
         # Add category rows
-        for category_name, category_ar in sorted(self.out["Category_AR"].items()):
+        for category_name, category_ar in sorted(self.out["Category_AL"].items()):
             table.add_row(category_name, f"{category_ar:.4f}")
 
         # Add separator and summary row
         table.add_section()
-        table.add_row("[bold]Overall Average[/bold]", f"[bold]{self.out['Average_AR']:.4f}[/bold]")
+        table.add_row("[bold]Overall Average[/bold]", f"[bold]{self.out['Average_AL']:.4f}[/bold]")
 
         console.print(table)
 
@@ -124,8 +124,8 @@ def _create_visualizations(
 
         df_clean = pd.DataFrame.from_dict(
             {
-                "question_id": list(self.out["Request_AR"].keys()),
-                "acceptance_rate": list(self.out["Request_AR"].values()),
+                "question_id": list(self.out["Request_AL"].keys()),
+                "acceptance_rate": list(self.out["Request_AL"].values()),
                 "category": [request.category for request in self.requests],
                 "response_length": [
                     mean([len(c["content"]) for c in messages if c["role"] == "assistant"])

examples/specdec_bench/specdec_bench/models/base.py

@@ -27,5 +27,14 @@ async def run(self, prompt_ids, sampling_params, request_id, turn_id):
         """
         raise NotImplementedError
 
+    def get_serving_config(self):
+        """Return a JSON-serializable dict describing the engine's effective config.
+
+        Captured into configuration.json's `serving_config` for reproducibility.
+        Subclasses override to surface engine-specific defaults (max_model_len,
+        kv_cache_dtype, etc.) that don't appear in the CLI args. Default: empty.
+        """
+        return {}
+
     def stop(self):
         pass