Skip to content

Commit b2996e3

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

File tree

1 file changed

+6
-0
lines changed

1 file changed

+6
-0
lines changed

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)