fix test

richardhuo-nv · richardhuo-nv · commit 4da12cc4a50a · 2025-03-06T11:04:01.000-08:00
diff --git a/src/triton_cli/parser.py b/src/triton_cli/parser.py
@@ -68,6 +68,7 @@
     "opt125m": "hf:facebook/opt-125m",
     "mistral-7b": "hf:mistralai/Mistral-7B-v0.1",
     "falcon-7b": "hf:tiiuae/falcon-7b",
+    "tinyllama-1.1b-chat-v1.0": "hf:TinyLlama/TinyLlama-1.1B-Chat-v1.0",
 }
 
 
diff --git a/tests/test_e2e.py b/tests/test_e2e.py
@@ -149,6 +149,10 @@ def test_vllm_openai_e2e(self, vllm_openai_server):
         os.environ.get("IMAGE_KIND") != "TRTLLM",
         reason="Only run for TRT-LLM image with LLM-API",
     )
+    @pytest.mark.skipif(
+        os.environ.get("TRTLLM_MODEL") == "gpt2",
+        reason="LLM API doesn't support gpt2's model architecture yet",
+    )
     @pytest.mark.parametrize(
         "protocol",
         [

Original file line number	Diff line number	Diff line change
`@@ -68,6 +68,7 @@`
`68`	`68`	`"opt125m": "hf:facebook/opt-125m",`
`69`	`69`	`"mistral-7b": "hf:mistralai/Mistral-7B-v0.1",`
`70`	`70`	`"falcon-7b": "hf:tiiuae/falcon-7b",`
	`71`	`+ "tinyllama-1.1b-chat-v1.0": "hf:TinyLlama/TinyLlama-1.1B-Chat-v1.0",`
`71`	`72`	`}`
`72`	`73`
`73`	`74`
Original file line number	Diff line number	Diff line change
`@@ -149,6 +149,10 @@ def test_vllm_openai_e2e(self, vllm_openai_server):`
`149`	`149`	`os.environ.get("IMAGE_KIND") != "TRTLLM",`
`150`	`150`	`reason="Only run for TRT-LLM image with LLM-API",`
`151`	`151`	`)`
	`152`	`+ @pytest.mark.skipif(`
	`153`	`+ os.environ.get("TRTLLM_MODEL") == "gpt2",`
	`154`	`+ reason="LLM API doesn't support gpt2's model architecture yet",`
	`155`	`+ )`
`152`	`156`	`@pytest.mark.parametrize(`
`153`	`157`	`"protocol",`
`154`	`158`	`[`