[CI/Build][The Rock] Use model_class_overrides so spawned worker can use test PredictableLlamaForCausalLM class when worker spawned using Python 3.14 (#49218)
Signed-off-by: Randall Smith <Randall.Smith@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
rasmith committed
Aug 26, 2026 at 02:55 UTC
cde7ba92da0e1e94bad2a410f9ec2e68bda1c7f7
3 files changed
+16
-15
.buildkite/test-amd.yaml
+2
-2
@@ -2714,7 +2714,7 @@ steps:
2714
dind: false
2715
agent_pool: mi300_2
2716
num_gpus: 2
2717
- working_dir: "/vllm-workspace/tests"
2717
+ working_dir: "/vllm-workspace/"
2718
source_file_dependencies:
2719
- vllm/config/speculative.py
2720
- vllm/distributed/kv_transfer/kv_connector/
@@ -2738,7 +2738,7 @@ steps:
2738
- vllm/platforms/rocm.py
2739
commands:
2740
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
2741
- - pytest -v -s v1/kv_connector/extract_hidden_states_integration
2741
+ - pytest -v -s tests/v1/kv_connector/extract_hidden_states_integration
2742
2743
- label: ":amd: (MI300) V1 Attention Shard %N" # TBD
2744
timeout_in_minutes: 180
.buildkite/test_areas/misc.yaml
+4
-2
@@ -136,6 +136,7 @@ steps:
136
key: extract-hidden-states-integration
137
timeout_in_minutes: 20
138
device: h200_18gb
139
+ working_dir: "/vllm-workspace"
140
source_file_dependencies:
141
- vllm/v1/spec_decode/extract_hidden_states.py
142
- vllm/model_executor/models/extract_hidden_states.py
@@ -144,12 +145,13 @@ steps:
145
- tests/v1/kv_connector/extract_hidden_states_integration
146
commands:
147
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
147
- - pytest -v -s v1/kv_connector/extract_hidden_states_integration
148
+ - pytest -v -s tests/v1/kv_connector/extract_hidden_states_integration
149
150
- label: ":nvidia: (L4) Extract Hidden States Integration"
151
key: extract-hidden-states-integration-2-gpus
152
timeout_in_minutes: 20
153
num_devices: 2
154
+ working_dir: "/vllm-workspace"
155
source_file_dependencies:
156
- vllm/v1/spec_decode/extract_hidden_states.py
157
- vllm/model_executor/models/extract_hidden_states.py
@@ -158,7 +160,7 @@ steps:
160
- tests/v1/kv_connector/extract_hidden_states_integration
161
commands:
162
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
161
- - pytest -v -s -m 'distributed' v1/kv_connector/extract_hidden_states_integration
163
+ - pytest -v -s -m 'distributed' tests/v1/kv_connector/extract_hidden_states_integration
164
165
- label: ":nvidia: (H200) Regression"
166
key: regression
tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py
+10
-11
@@ -8,7 +8,7 @@ import pytest
8
import torch
9
10
from tests.utils import create_new_process_for_each_test, multi_gpu_test
11
-from vllm import LLM, ModelRegistry, SamplingParams
11
+from vllm import LLM, SamplingParams
12
from vllm.distributed.kv_transfer.kv_connector.v1 import (
13
example_hidden_states_connector,
14
)
@@ -67,16 +67,12 @@ def predictable_llama_config_path(tmp_path_factory):
67
return str(config_dir)
68
69
70
-@pytest.fixture(scope="module", autouse=True)
71
-def register_predictable_model():
72
- """Register the PredictableLlamaForCausalLM model."""
73
- from .predictable_llama import PredictableLlamaForCausalLM
74
-
75
- if "PredictableLlamaForCausalLM" not in ModelRegistry.get_supported_archs():
76
- ModelRegistry.register_model(
77
- "PredictableLlamaForCausalLM", PredictableLlamaForCausalLM
78
- )
79
- yield
70
+# Need for Python 3.14 compatibility, vLLM with force spawn due to CUDA
71
+# reinitialization, and register will not persist.
72
+PREDICTABLE_LLAMA_MODEL_CLS = (
73
+ "tests.v1.kv_connector.extract_hidden_states_integration"
74
+ ".predictable_llama:PredictableLlamaForCausalLM"
75
+)
76
77
78
def test_extract_hidden_states_with_predictable_dummy_model(
@@ -104,6 +100,9 @@ def test_extract_hidden_states_with_predictable_dummy_model(
100
101
llm = LLM(
102
model=predictable_llama_config_path,
103
+ model_class_overrides={
104
+ "PredictableLlamaForCausalLM": PREDICTABLE_LLAMA_MODEL_CLS,
105
+ },
106
speculative_config={
107
"method": "extract_hidden_states",
108
"num_speculative_tokens": 1,