[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>

.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

.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

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,