diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 0c89677090..d22afc09f7 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -2714,7 +2714,7 @@ steps: dind: false agent_pool: mi300_2 num_gpus: 2 - working_dir: "/vllm-workspace/tests" + working_dir: "/vllm-workspace/" source_file_dependencies: - vllm/config/speculative.py - vllm/distributed/kv_transfer/kv_connector/ @@ -2738,7 +2738,7 @@ steps: - vllm/platforms/rocm.py commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s v1/kv_connector/extract_hidden_states_integration + - pytest -v -s tests/v1/kv_connector/extract_hidden_states_integration - label: ":amd: (MI300) V1 Attention Shard %N" # TBD timeout_in_minutes: 180 diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index b564896fbd..1d952b028c 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -136,6 +136,7 @@ steps: key: extract-hidden-states-integration timeout_in_minutes: 20 device: h200_18gb + working_dir: "/vllm-workspace" source_file_dependencies: - vllm/v1/spec_decode/extract_hidden_states.py - vllm/model_executor/models/extract_hidden_states.py @@ -144,12 +145,13 @@ steps: - tests/v1/kv_connector/extract_hidden_states_integration commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s v1/kv_connector/extract_hidden_states_integration + - pytest -v -s tests/v1/kv_connector/extract_hidden_states_integration - label: ":nvidia: (L4) Extract Hidden States Integration" key: extract-hidden-states-integration-2-gpus timeout_in_minutes: 20 num_devices: 2 + working_dir: "/vllm-workspace" source_file_dependencies: - vllm/v1/spec_decode/extract_hidden_states.py - vllm/model_executor/models/extract_hidden_states.py @@ -158,7 +160,7 @@ steps: - tests/v1/kv_connector/extract_hidden_states_integration commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -v -s -m 'distributed' v1/kv_connector/extract_hidden_states_integration + - pytest -v -s -m 'distributed' tests/v1/kv_connector/extract_hidden_states_integration - label: ":nvidia: (H200) Regression" key: regression diff --git a/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py b/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py index 82e0eb9d8f..9c43a36f92 100644 --- a/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py +++ b/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py @@ -8,7 +8,7 @@ import pytest import torch from tests.utils import create_new_process_for_each_test, multi_gpu_test -from vllm import LLM, ModelRegistry, SamplingParams +from vllm import LLM, SamplingParams from vllm.distributed.kv_transfer.kv_connector.v1 import ( example_hidden_states_connector, ) @@ -67,16 +67,12 @@ def predictable_llama_config_path(tmp_path_factory): return str(config_dir) -@pytest.fixture(scope="module", autouse=True) -def register_predictable_model(): - """Register the PredictableLlamaForCausalLM model.""" - from .predictable_llama import PredictableLlamaForCausalLM - - if "PredictableLlamaForCausalLM" not in ModelRegistry.get_supported_archs(): - ModelRegistry.register_model( - "PredictableLlamaForCausalLM", PredictableLlamaForCausalLM - ) - yield +# Need for Python 3.14 compatibility, vLLM with force spawn due to CUDA +# reinitialization, and register will not persist. +PREDICTABLE_LLAMA_MODEL_CLS = ( + "tests.v1.kv_connector.extract_hidden_states_integration" + ".predictable_llama:PredictableLlamaForCausalLM" +) def test_extract_hidden_states_with_predictable_dummy_model( @@ -104,6 +100,9 @@ def test_extract_hidden_states_with_predictable_dummy_model( llm = LLM( model=predictable_llama_config_path, + model_class_overrides={ + "PredictableLlamaForCausalLM": PREDICTABLE_LLAMA_MODEL_CLS, + }, speculative_config={ "method": "extract_hidden_states", "num_speculative_tokens": 1,