[Thangnguyenvn98]

What

Shard the spec-decode-draft-model job (Spec Decode Draft Model) with pytest-shard on its one whole directory, v1/e2e/spec_decode/draft_model/: 4 replicas on the primary lane and 2 on the AMD mirror. Recent main runs spend 34–43m in pytest on the primary lane and ~25m on the mirror, against the ~20m per-job target. No tests move and no filters are added.

Old

tests/v1/e2e/spec_decode/draft_model/     # 20 cases, one job, one replica per lane
├── test_async.py                         # 3 cases
├── test_draft_model.py                   # 16 cases (2 skip: need 2 GPUs)
└── test_lora.py                          # 1 case

pytest -v -s v1/e2e/spec_decode/draft_model/, key spec-decode-draft-model, mirror amd-spec-decode-draft-model.

New

Same tree and same key. The command gets --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB, and the labels get Shard %N.

key	parallelism	cases per replica (hash assignment)
spec-decode-draft-model	4	5 / 3 / 5 / 7
amd-spec-decode-draft-model	2 (explicit mirror.amd.parallelism)	10 / 10

Notes

• Pattern A (whole-directory hash shards) on the existing named job. A per-file split does not fit because test_draft_model.py alone takes ~29m of pytest, and splitting that monolith would add a larger diff with no benefit.
• Why 4 replicas: the per-case timings put the slowest replica at 22m of pytest with 2 replicas and 19m with 3 (10 of the 20 cases hash into one shard), which is too close to the target given the 34–43m spread between runs. With 4, the slowest replica is ~16m. Every replica gets at least 3 cases, so no shard is empty. The mirror's own timings put its slowest replica at ~14m with 2 replicas, so it overrides the inherited count instead of running 4.
• Pytest-shard keeps file order inside each shard, so no test runs earlier than it did before. The env changes in these tests all go through monkeypatch, so they are reverted after each test.
• Unchanged: the B200 nightly step and the legacy test-amd.yaml copies (optional) still run the directory unsharded.

Expected runtime and overhead (predictions)

These predictions use per-case pytest durations from two recent main builds (90437, 90426), taken from log timestamps. Per-replica overhead is 2–4m, the wall-minus-pytest range the unsharded job showed in those builds (checkout, image pull, container start). The cause breakdown is unattributed. Measured numbers will follow in a validation comment.

job / replica	pytest	est. running wall	est. non-pytest overhead	main overhead cause
primary shard 0	~6.7m	~10m	~3m	setup (unattributed)
primary shard 1	~6.2m	~9m	~3m	setup (unattributed)
primary shard 2	~13.8m	~17m	~3m	setup (unattributed)
primary shard 3	~15.8m	~19m	~3m	setup (unattributed)
mirror shard 0	~13.9m	~16m	~2.5m	setup (unattributed)
mirror shard 1	~11.8m	~14m	~2.5m	setup (unattributed)

• Primary lane: 44.4m wall before (build 90437) → 55m total worker-minutes after, **+10 extra** (range +5 to +12, from 3 extra setups). This is inside the ≤20 budget. Slowest replica ~16m pytest.
• Mirror lane: 26.8m wall before → 30m after, **+3 extra**. Slowest replica ~14m pytest.

AI assistance & verification

This PR was prepared with AI assistance (Claude Code); a human reviewed the diff. No existing PR covers this job.

• yaml.safe_load on spec_decode.yaml: passes.
• The coverage guard (check_area_test_coverage.py --tree tests/v1/e2e/spec_decode) accepts the shard flags on a whole-directory target for both base and mirror. It reports no new violations; the violations it does report were already there and belong to other jobs.
• The ci-infra generator reads amd.get("parallelism", step.parallelism), so the mirror's explicit 2 overrides the base 4.
• GPU validation: a scoped build will follow as a comment.

🤖 Generated with Claude Code

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[Thangnguyenvn98]

Validation

Build 90490 ran on abbcff92a5 with scoped keys spec-decode-draft-model and amd-spec-decode-draft-model, and passed. All 6 replicas passed. Each replica collected 3–10 cases, so none was empty, and the whole directory was covered on both lanes (20 cases each).

Baselines are the unsharded job on main build 90437, which ran the same commit range. Timings come from the Buildkite job start/finish times and the pytest summary line.

job / replica	cases	pytest (pred)	running wall	non-pytest overhead	main overhead cause
primary shard 1	5	7.6m (6.7)	8.9m	1.3m	image pull + checkout
primary shard 2	3	5.8m (6.2)	7.1m	1.3m	image pull + checkout
primary shard 3	5	17.0m (13.8)	18.3m	1.3m	image pull + checkout
primary shard 4	7	19.3m (15.8)	20.8m	1.5m	image pull + checkout
mirror shard 1	10	13.3m (13.9)	15.6m	2.3m	image pull + checkout
mirror shard 2	10	12.5m (11.8)	15.2m	2.7m	image pull + checkout

Primary lane. Total worker-minutes went from 44.4 to 55.1, +10.7 extra, inside the ≤20 budget. Non-pytest overhead came in under the prediction (1.3–1.5m against 2–4m). Most of the increase is inside pytest instead: 42.7m → 49.7m in total. It is concentrated in cases that used to run right after a sibling with the same model setup and now start cold in their own replica. For example, test_draft_model_quantization[True-draft_quantized] went from 1.5m to 3.2m, test_draft_model_realistic_example from 1.2m to 3.5m, and test_draft_model_correctness[True-args1] from 2.0m to 4.2m. Weight loading and dataset downloads are negligible in every replica, so this is engine start-up and warm-up. The exact phase is unattributed. The slowest replica ran 19.3m of pytest, which passes the ≤20m target but with only 0.7m of margin.

Mirror lane. Total worker-minutes went from 26.8 to 30.8, +4.0 extra, inside budget. The slowest replica ran 13.3m of pytest, inside the target and in line with the prediction.

Verdict. Both lanes pass the runtime and aggregate-cost budgets. The primary lane's slowest replica is close to the target, so a slow run could push it over.