feat: Enable per dataset max-osl

[arekay-nv]

What does this PR do?

Enables the max-osl to be specified per-dataset. This enables running perf/accuracy runs with different OSL in the same benchmark.

Type of change

• Bug fix
• New feature
• Documentation update
• Refactor/cleanup

Related issues

Testing

• Tests added/updated
• All tests pass locally
• Manual testing completed

Checklist

• Code follows project style
• Pre-commit hooks pass
• Documentation updated (if needed)

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[gemini-code-assist]

Code Review

This pull request introduces a model_params_override field to the dataset configuration, allowing per-dataset overrides of the top-level model parameters. This is particularly useful for scenarios like MLPerf-style runs where accuracy and performance datasets require different configurations. The changes include validation of override keys, merging logic in effective_model_params, template updates, and comprehensive unit tests. Feedback from the reviewer highlights that the current implementation of effective_model_params performs a shallow merge, which could inadvertently overwrite entire nested dictionaries (such as osl_distribution). The reviewer suggests implementing a deep merge helper to preserve other nested configuration fields and adding corresponding unit tests to verify this behavior.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[arekay-nv]

Review Council — Claude-only review (Codex failed: cloud policy auth error)

[arekay-nv]

Review Council — Claude-only Code Review

Reviewed by: Claude (Codex failed — cloud workspace policy auth error) | Depth: standard (240 lines)

Found 2 issues across 2 files:

• 1 medium
• 1 low

#	File	Line	Severity	Category	Summary
1	src/inference_endpoint/openai/completions_adapter.py	65	medium	bug	skip_special_tokens: False in metadata dict never wired into TextCompletionRequest — silently dropped
2	src/inference_endpoint/commands/benchmark/execute.py	320	low	bug	FileNotFoundError handler uses performance_cfgs[0].path instead of the perf_cfg variable extracted on line 305 (outside diff hunk — not posted inline)

---

Issue 1 (inline): completions_adapter.py:65 — "skip_special_tokens": False is added to metadata and flows into query.data, but TextCompletionRequest has no such field and encode_query() never reads it. Either wire it through to the request or remove the entry.

Issue 2 (summary-only, line outside diff): execute.py:320 — the FileNotFoundError handler body still uses performance_cfgs[0].path instead of perf_cfg.path. This PR introduced perf_cfg = performance_cfgs[0] at line 305 to avoid repeated indexing, but the error message at line 320 was not updated. Minor consistency nit.

[nvzhihanj]

Review Council — first-principles WAR design review

Reviewed by: Claude (Codex review timed out — xhigh reasoning latency exceeds the 600s cap even on this small diff) · Depth: thorough

Scope: is the per-dataset max-OSL WAR correctly designed for a clean later refactor, given the long-term goal of per-phase generation config? 8 net-new findings (none duplicate the existing maintainer / gemini comments); verdict + long-term plan in the summary comment. The HIGH streaming data-integrity bug was verified against the source.

[nvzhihanj]

Review Council — first-principles WAR design review

Reviewed by: Claude · Depth: thorough
Codex review timed out (its xhigh reasoning latency exceeds the 600s cap even on this 381-line diff); Claude-led. The HIGH streaming bug was verified against the source.

This is a focused assessment of whether the per-dataset max-OSL WAR is correctly designed for a clean later refactor, given the long-term goal of per-phase generation config. 8 net-new findings posted inline — none duplicate the existing maintainer / gemini comments.

Verdict on the WAR

It achieves the immediate need (accuracy vs performance differ, since they're separate dataset entries) but is not yet "correct" as a WAR — three structural issues:

1. Wrong axis. It's keyed to dataset identity, which cannot express the stated long-term goal ("different generation config for the same dataset across phases") — phases derive from datasets and duplicate (name,type) entries are rejected, so you'd have to duplicate the dataset entry.
2. Over-broad surface. It accepts all ModelParams keys, including three that silently misbehave per-dataset: name (→ ISL/OSL tokenized with the wrong tokenizer), streaming (→ aggregator ignores it; TTFT/TPOT unrecorded), osl_distribution (no consumer + shallow-merge degrades). The field's own documented example advertises streaming: 'on', which doesn't work.
3. Naming (the maintainer's point): model_params_override mislabels generation config as model parameters.

Minimal clean WAR: rename → generation_config_override; restrict accepted keys to the per-dataset-meaningful set (max_new_tokens + sampling params), rejecting name/streaming/osl_distribution; add # TODO(post-mortem): generation config is per-phase, not per-dataset; trim the template/docstring example accordingly.

Long-term: split ModelParams → ModelIdentity {name, tokenizer} (per-run invariant) + GenerationConfig {max_new_tokens, temperature, top_p, top_k, penalties, streaming, osl_distribution} (per-phase). Carry GenerationConfig on PhaseConfig; adapters' dataset_transforms read the phase's GenerationConfig, and the dataset becomes generation-agnostic. "Same dataset, different gen-config per phase" then falls out naturally (per-phase streaming additionally needs per-phase aggregator handling).

🔴 Must-fix (for the WAR to be correct)

#	File	Line	Cat	Summary
A	config/schema.py	317	data-integrity	Per-dataset streaming override reaches request rows but the single aggregator reads the global streaming (execute.py:185) → TTFT/TPOT silently unrecorded. The documented example advertises this broken key.
B	config/schema.py	309	design	Override bound to dataset identity is a dead-end for "same dataset, different gen-config per phase." WAR: rename + TODO. Long-term: per-phase GenerationConfig.
C	config/schema.py	329	design	Accepts all ModelParams keys incl. name → per-dataset name mismatches the global tokenizer/aggregator. Restrict to a per-dataset allowlist.

🟡 Should-fix

#	File	Line	Cat	Summary
D	config/schema.py	353	bug	(extends gemini) Shallow-merge of osl_distribution silently resets siblings and passes validation (all defaults), so the docstring's "catches bad overrides" is false; also osl_distribution has no consumer → overriding it is a no-op.
E	tests/unit/commands/test_benchmark.py	1144	testing	Gaps: nested/deep-merge, streaming end-to-end, multi-turn loader, and the completions adapter max_tokens key (tests only assert the chat max_completion_tokens, but the PR targets the completions path).
F	config/templates/offline_template_full.yaml (×3)	29	design	WAR baked into all three generated templates + the template-parse test → load-bearing / hard to undo; also advertises broken streaming:'on'. Rename now; trim the example.
G	config/schema.py	344	error-handling	Override value validation deferred to load-time (after setup writes report dir / config.yaml / probes tokenizer); move into BenchmarkConfig._resolve_and_validate to fail atomically at construction.

🔵 Consider

#	File	Line	Cat	Summary
H	commands/benchmark/execute.py	293	error-handling	Narrow except Exception → (ValidationError, ValueError) so a future bug isn't mislabeled "invalid model_params_override."

Through-line: A, C, D, F are all symptoms of the same root cause — generation config is modeled as part of ModelParams and attached to the dataset, so the override surface is both too wide (accepts identity/per-run keys) and on the wrong axis. Renaming + key-allowlisting makes the WAR honest now; the ModelIdentity / GenerationConfig split on PhaseConfig is the long-term fix.

Dedup: builds on — does not duplicate — the maintainer's naming/per-run comment and gemini's shallow-merge note (D explicitly extends it). The skip_special_tokens dead-code was already flagged by a prior review and is omitted here.

[arekay-nv]

Review Council — Multi-AI Code Review

Reviewed by: Claude (Codex did not complete — see note) | Depth: standard | Head: 7fe333a

No blocking code-level issues found in the per-dataset generation_config_override change.

Verified during review:

• Layered validation is correct — unknown top-level keys are rejected in _validate_generation_config_override; nested / cross-field invariants (min_new_tokens > max_new_tokens, OSLDistribution extra="forbid") are rejected at merge time via ModelParams.model_validate, caught in execute.py and surfaced as InputValidationError.
• _deep_merge preserves nested siblings, handles base.osl_distribution is None, and enum values round-trip losslessly through model_dump() / model_validate.
• Templates regenerated and in sync (0 drift); the multi_turn → agentic_inference rename is consistent across all PR artifacts.
• Tests — 27 new tests pass, covering both chat (max_completion_tokens) and completions (max_tokens) adapter branches, no-override, override propagation, and invalid-value paths.
• Documented caveat (per-dataset name / streaming not honored by the single-aggregator path) matches the code — execute.py launches the aggregator/tokenizer from the global config.model_params.

🔵 One optional follow-up (non-blocking): the perf-path invalid-override branch (execute.py:329-334) is structurally identical to the tested accuracy branch but has no direct test — test_invalid_override_value_raises_input_validation_error only exercises the accuracy override. Low risk; a one-line parametrization would close it.

---

Notes:

• Codex review did not complete: under this host's enterprise-managed Codex profile, --yolo was overridden to a read-only sandbox that broke its git / /tmp access, and it hit the 10-min timeout with no findings emitted. This was a Claude-only pass.
• The council reviews the diff for new issues only — it does not resolve the existing CHANGES_REQUESTED (field-naming thread) or the PR's draft status.

[arekay-nv]

Review Council — Multi-AI Code Review

Reviewed by: Codex + Claude | Depth: standard

1 net-new issue (most council findings were already covered by existing review comments — see summary).

[arekay-nv]

Review Council — Multi-AI Code Review

Reviewed by: Codex + Claude | Depth: standard | Head: da55aff

Ran two independent reviews (Codex in an isolated worktree + Claude on the diff), verified every cited line against source at HEAD, and deduped against the PR's existing review comments.

Posted: 1 net-new issue. Most council findings were already raised in earlier review threads and were suppressed to avoid duplicate noise.

#	File	Line	Severity	Category	Reviewer(s)	Status
1	src/inference_endpoint/config/schema.py	463	medium	api-contract	Codex	✅ posted inline

Suppressed as duplicates of existing comments

Finding	Reviewer	Already covered by
Override surface too wide — accepts metric-desyncing keys (name/streaming/tokenizer_name)	Claude	existing "override surface is too wide" + "documented example advertises a key that silently half-works"
Override value validation deferred to load-time → inference-endpoint validate gives false confidence	Claude	existing "value validation of the override is deferred to dataset-load time"
Perf-dataset override error branch untested	Claude	existing test-coverage + except Exception (perf branch) comments

The net-new finding

Codex spotted that effective_generation_config re-validates the merged per-dataset override only as a ModelParams sub-model, so BenchmarkConfig's cross-field guard (schema.py:1015-1031) — which restricts min_new_tokens/skip_special_tokens to api_type: openai_completions — is bypassed. A per-dataset override of those fields is accepted under openai/sglang and silently ignored, while the same setting at top level is correctly rejected. Verified against HEAD.

✅ Commit hygiene: single squashed commit, no fixups.