Skip to content

fix: preserve auth for CLI concurrency sweeps#972

Merged
ajcasagrande merged 1 commit into
ai-dynamo:mainfrom
abatilo:abatilo/fix-concurrency-sweep-auth
May 21, 2026
Merged

fix: preserve auth for CLI concurrency sweeps#972
ajcasagrande merged 1 commit into
ai-dynamo:mainfrom
abatilo:abatilo/fix-concurrency-sweep-auth

Conversation

@abatilo

@abatilo abatilo commented May 21, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Comma-list concurrency runs expand into separate benchmark configs before execution. Keep runtime-only expansion faithful so each generated run can authenticate while exported artifacts continue to redact secrets.

Signed-off-by: Aaron Batilo abatilo@coreweave.com
@abatilo
fix: preserve auth for CLI concurrency sweeps
be14801 
Comma-list concurrency runs expand into separate benchmark configs before execution. Keep runtime-only expansion faithful so each generated run can authenticate while exported artifacts continue to redact secrets.

Signed-off-by: Aaron Batilo <abatilo@coreweave.com>
@copy-pr-bot

copy-pr-bot Bot commented May 21, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.
@github-actions github-actions Bot added the fix label May 21, 2026
@github-actions

github-actions Bot commented May 21, 2026

Copy link
Copy Markdown

Try out this PR

Quick install:

pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@be14801e04bb7e36ff6676c962c3d834411c0a29

Recommended with virtual environment (using uv):

uv venv --python 3.12 && source .venv/bin/activate
uv pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@be14801e04bb7e36ff6676c962c3d834411c0a29

Last updated for commit: be14801 • Browse code
@abatilo abatilo marked this pull request as ready for review May 21, 2026 18:49
@coderabbitai

coderabbitai Bot commented May 21, 2026

Copy link
Copy Markdown
Contributor

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 44ea563d-6df3-400e-87b1-5125a708a479

📥 Commits

Reviewing files that changed from the base of the PR and between 83bdcd9 and be14801.

📒 Files selected for processing (2)
  • src/aiperf/config/loader/plan.py
  • tests/unit/config/test_cli_magic_list_sugar.py

Walkthrough

The pull request changes how benchmark plans preserve secrets during sweep expansion. build_benchmark_plan's fallback path now calls model_dump with mode="python" and context={"include_secrets": True} instead of mode="json", ensuring API keys and other secrets survive the prefill_concurrency sweep expansion. A new test validates this behavior.

Changes

Secret Preservation in Sweep Expansion

Layer / File(s) Summary
Envelope model_dump with secret context
src/aiperf/config/loader/plan.py		 build_benchmark_plan switches from json mode to python mode with include_secrets context when constructing the envelope dict used for sweep expansion and per-variation rendering.
API key preservation test for prefill_concurrency
tests/unit/config/test_cli_magic_list_sugar.py		 New test verifies that when prefill_concurrency comma-list expands into sweep variations, each resulting config retains the original api_key and profiling phase concurrency values match the expanded list.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Poem

🐰 A secret's safe when swept along,
Python mode keeps API strong,
No json here, just truth so plain,
Keys preserved through prefill's reign!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title 'fix: preserve auth for CLI concurrency sweeps' directly describes the main change: ensuring authentication/API keys are preserved when comma-list concurrency configurations are expanded into separate benchmark plans during sweep expansion.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.
@abatilo abatilo marked this pull request as draft May 21, 2026 18:51
@abatilo abatilo marked this pull request as ready for review May 21, 2026 18:52
@ajcasagrande ajcasagrande self-requested a review May 21, 2026 18:55
@codecov

codecov Bot commented May 21, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!
ajcasagrande
ajcasagrande approved these changes May 21, 2026
View reviewed changes

@ajcasagrande ajcasagrande left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the fix, @abatilo. This looks good to me.

I reviewed the build_benchmark_plan() path and the existing JSON-only api_key redaction serializer, then ran the focused tests/unit/config/test_cli_magic_list_sugar.py regression test. I also reproduced the flow with the real aiperf CLI against the in-repo mock server using --api-key sk-real-prod-key --concurrency 1,2.

To double-check the runtime behavior, I temporarily instrumented the mock server in a temp checkout to log the incoming Authorization header. The expanded sweep sent Authorization: Bearer sk-real-prod-key for requests from both concurrency variations, while generated run_config.json files kept the key redacted and a recursive artifact grep found no raw API key.

Nice targeted change, and the regression test covers the important bit.
@ajcasagrande ajcasagrande merged commit ea1ce37 into ai-dynamo:main May 21, 2026
20 of 21 checks passed
This was referenced May 22, 2026
Closed
Merged
@ajcasagrande ajcasagrande mentioned this pull request May 22, 2026
Merged
JimmyWhitaker added a commit to JimmyWhitaker/aiperf that referenced this pull request May 22, 2026
@JimmyWhitaker
fix(orchestrator): clear stale IPC env vars + preserve URL userinfo
656e8d9 
Addresses PR ai-dynamo#982 review feedback from ajcasagrande, coderabbitai,
and dynamo-ops.

Three issues raised:

1. Stale env-var leakage (CRITICAL — reproduced by ajcasagrande).
   `_run_benchmark_subprocess` copies the parent env wholesale, then
   only conditionally sets AIPERF_INJECTED_* when the current run has
   credentials. A stale value left in the parent shell (from a prior
   aiperf run, CI wrapper, etc.) would pass through and the subprocess
   would overlay it onto a run that configured no sensitive headers,
   sending an unintended credential to the upstream. Reproduced by
   ajcasagrande with the in-repo mock server:
   parent env AIPERF_INJECTED_HEADERS={"Authorization":"Bearer stale"} +
   child run with only X-Trace-Id => upstream saw Authorization. Fix:
   env.pop() all three internal injection vars after os.environ.copy(),
   before the conditional re-set.

2. URL userinfo redaction (dynamo-ops). EndpointConfig.urls has an
   unconditional _redact_urls field_serializer (no when_used="json"
   guard), so mode="python" dumps still strip user:pass@. Two parts:
   (a) planners now pass context={"include_secrets": True} alongside
   mode="python", matching the pattern PR ai-dynamo#972 used in plan.py;
   (b) local_executor adds AIPERF_INJECTED_ENDPOINT_URLS env-var IPC,
   set only when at least one URL would be redacted, mirroring the
   api_key + headers channels.

3. Malformed env-var hardening (coderabbitai). JSON parsing of the
   injected env vars moved into the try block so malformed payloads
   surface via the structured error envelope; payload shape now
   validated (dict for headers, list[str] for urls) before use.
   Extracted _parse_injected_dict and _parse_injected_str_list helpers
   to keep main() under the cyclomatic complexity threshold.

Also: type hints added to all new test functions per CLAUDE.md.

Test coverage: 9 new regression tests across executor + 3 planner
suites (stale-env isolation, URL userinfo end-to-end, malformed-payload
rejection, planner URL userinfo preservation). Full unit suite at
12,695 passing (was 12,686 on the original commit). Pre-commit fully
green including check-ergonomics and check-ruff-baselined.

Signed-off-by: Jimmy Whitaker <jimmy.whitaker@baseten.co>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

fix

3 participants

@abatilo @ajcasagrande @dynamo-ops