diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index 0100b99993ef..26eee8dcdbc7 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -115,21 +115,23 @@ steps: - label: ":nvidia: (L4) Hybrid SSM NixlConnector PD accuracy" key: hybrid-ssm-nixlconnector-pd-accuracy-tests-4-gpus - timeout_in_minutes: 60 + timeout_in_minutes: 30 working_dir: "/vllm-workspace/tests" device: l4 num_devices: 4 + parallelism: 4 source_file_dependencies: - vllm/distributed/kv_transfer/kv_connector/v1/nixl/ - tests/v1/kv_connector/nixl_integration/ commands: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh + - HYBRID_SSM=1 CONFIG_SHARD_ID=$$BUILDKITE_PARALLEL_JOB CONFIG_NUM_SHARDS=$$BUILDKITE_PARALLEL_JOB_COUNT bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh mirror: amd: label: ":amd: (MI300) Hybrid SSM NixlConnector PD accuracy" dind: false device: mi300_4 + parallelism: 1 timeout_in_minutes: 55 depends_on: - image-build-amd diff --git a/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh index 8acc96413327..96e6e45e194a 100755 --- a/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh @@ -55,6 +55,32 @@ else configs=("${tp_configs[@]}") fi +# Opt-in CI sharding; unset means run every config for existing callers. +if [[ -n "${CONFIG_SHARD_ID:-}" || -n "${CONFIG_NUM_SHARDS:-}" ]]; then + if [[ ! "${CONFIG_SHARD_ID:-}" =~ ^(0|[1-9][0-9]*)$ ]] || + [[ ! "${CONFIG_NUM_SHARDS:-}" =~ ^[1-9][0-9]*$ ]]; then + echo "CONFIG_SHARD_ID and CONFIG_NUM_SHARDS must be non-negative and positive integers, respectively." + exit 1 + fi + if (( CONFIG_SHARD_ID >= CONFIG_NUM_SHARDS )); then + echo "CONFIG_SHARD_ID=${CONFIG_SHARD_ID} must be less than CONFIG_NUM_SHARDS=${CONFIG_NUM_SHARDS}." + exit 1 + fi + + shard_configs=() + for idx in "${!configs[@]}"; do + if (( idx % CONFIG_NUM_SHARDS == CONFIG_SHARD_ID )); then + shard_configs+=("${configs[$idx]}") + fi + done + if (( ${#shard_configs[@]} == 0 )); then + echo "Shard ${CONFIG_SHARD_ID}/${CONFIG_NUM_SHARDS} selected no configs." + exit 1 + fi + configs=("${shard_configs[@]}") + echo "Shard ${CONFIG_SHARD_ID}/${CONFIG_NUM_SHARDS}: running ${#configs[@]} configs" +fi + run_tests() { local label=$1 local extra_args=$2