diff --git a/buildkite/pipeline_generator/buildkite_step.py b/buildkite/pipeline_generator/buildkite_step.py index 7e05dc9f..553b393c 100644 --- a/buildkite/pipeline_generator/buildkite_step.py +++ b/buildkite/pipeline_generator/buildkite_step.py @@ -677,7 +677,7 @@ def convert_group_step_to_buildkite_step( no_gpu=amd_no_gpu, num_nodes=amd.get("num_nodes", step.num_nodes), soft_fail=amd.get("soft_fail", step.soft_fail or False), - parallelism=step.parallelism, + parallelism=amd.get("parallelism", step.parallelism), concurrency=amd.get("concurrency", step.concurrency), concurrency_group=amd.get( "concurrency_group", step.concurrency_group diff --git a/buildkite/tests/pipeline_generator/test_amd.py b/buildkite/tests/pipeline_generator/test_amd.py index 976b7864..fed9cf9c 100644 --- a/buildkite/tests/pipeline_generator/test_amd.py +++ b/buildkite/tests/pipeline_generator/test_amd.py @@ -541,3 +541,63 @@ def test_amd_mirror_without_timeout_uses_standard_timeout(fake_global_config): # An AMD mirror without its own timeout does not inherit the shorter main # timeout (AMD runs slower); it receives the standard AMD timeout. assert amd_command_step.timeout_in_minutes == 180 + + +@pytest.mark.parametrize( + ("amd_parallelism", "expected"), + [ + (None, 4), # unset -> inherit parent parallelism (existing behavior) + (1, 1), # explicit override -> AMD stays a single copy + (2, 2), # explicit override -> independent AMD shard count + ], +) +def test_amd_mirror_parallelism_override( + fake_global_config, amd_parallelism, expected +): + """NVIDIA sharding must not multiply AMD mirrors: mirror.amd.parallelism + overrides the inherited parent parallelism when set.""" + amd_config = { + "device": "mi325_1", + "depends_on": ["image-build-amd"], + # pooling-shaped mirror: custom unsharded AMD command preserved + "commands": ["pytest tests/pooling.py"], + } + if amd_parallelism is not None: + amd_config["parallelism"] = amd_parallelism + step = Step( + label="Pooling test", + group="Models", + key="pooling-test", + depends_on=["image-build"], + working_dir="/vllm-workspace/tests", + commands=[ + "pytest tests/pooling.py --shard-id=$$BUILDKITE_PARALLEL_JOB" + " --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT" + ], + parallelism=4, + source_file_dependencies=["vllm/"], + mirror={"amd": amd_config}, + ) + + group_steps = buildkite_step.convert_group_step_to_buildkite_step( + {step.group: [step]} + ) + nvidia_group = next(g for g in group_steps if g.group == "Models") + nvidia_step = next( + s + for s in nvidia_group.steps + if isinstance(s, buildkite_step.BuildkiteCommandStep) + ) + amd_group = next( + g for g in group_steps if g.group == "Hardware-AMD Tests" + ) + amd_step = next( + s for s in amd_group.steps if isinstance(s, buildkite_step.BuildkiteCommandStep) + ) + + # parent sharding always preserved + assert nvidia_step.parallelism == 4 + assert amd_step.parallelism == expected + # custom AMD command remains the unsharded original + assert "pytest tests/pooling.py" in amd_step.env["VLLM_TEST_COMMANDS"] + assert "--num-shards" not in amd_step.env["VLLM_TEST_COMMANDS"]