Skip to content

Commit 9f5e8a1

Browse files
committed
add runai model streamer e2e test for GCS
1 parent 2d977a7 commit 9f5e8a1

File tree

2 files changed

+18
-0
lines changed

2 files changed

+18
-0
lines changed

.buildkite/test-pipeline.yaml

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -541,6 +541,18 @@ steps:
541541
commands:
542542
- pytest -v -s kernels/mamba
543543

544+
- label: RunAI Model Streamer Test # 77s
545+
# Cross continent streaming could take much take much longer.
546+
timeout_in_minutes: 10
547+
torch_nightly: true
548+
mirror_hardwares: [amdexperimental]
549+
source_file_dependencies:
550+
- vllm/engine
551+
- vllm/model_executor/model_loader
552+
- tests/model_executor/runai_model_streamer
553+
commands:
554+
- pytest -v -s model_executor/model_loader/runai_model_streamer
555+
544556
- label: Model Executor Test # 23min
545557
timeout_in_minutes: 35
546558
mirror_hardwares: [amdexperimental]

tests/model_executor/model_loader/runai_model_streamer/test_runai_model_streamer_loader.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
load_format = "runai_streamer"
99
test_model = "openai-community/gpt2"
10+
test_gcs_model = "gs://vertex-model-garden-public-us/codegemma/codegemma-2b/"
1011

1112
prompts = [
1213
"Hello, my name is",
@@ -32,3 +33,8 @@ def test_runai_model_loader_download_files(vllm_runner):
3233
with vllm_runner(test_model, load_format=load_format) as llm:
3334
deserialized_outputs = llm.generate(prompts, sampling_params)
3435
assert deserialized_outputs
36+
37+
def test_runai_model_loader_download_files_gcs(vllm_runner):
38+
with vllm_runner(test_gcs_model, load_format=load_format) as llm:
39+
deserialized_outputs = llm.generate(prompts, sampling_params)
40+
assert deserialized_outputs

0 commit comments

Comments
 (0)