Fix test

Signed-off-by: Jee Jee Li <[email protected]>

Author: jeejeelee

tests/lora/test_llama_tp.py

@@ -11,19 +11,27 @@
 from vllm.lora.request import LoRARequest
 from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
 
-from ..utils import VLLM_PATH, create_new_process_for_each_test, multi_gpu_test
-
-MODEL_PATH = "meta-llama/Llama-2-7b-hf"
+PROMPT_TEMPLATE = """<|eot_id|><|start_header_id|>user<|end_header_id|>
+I want you to act as a SQL terminal in front of an example database, you need only to return the sql command to me.Below is an instruction that describes a task, Write a response that appropriately completes the request.
+"
+##Instruction:
+candidate_poll contains tables such as candidate, people. Table candidate has columns such as Candidate_ID, People_ID, Poll_Source, Date, Support_rate, Consider_rate, Oppose_rate, Unsure_rate. Candidate_ID is the primary key.
+Table people has columns such as People_ID, Sex, Name, Date_of_Birth, Height, Weight. People_ID is the primary key.
+The People_ID of candidate is the foreign key of People_ID of people.
+###Input:
+{context}
+###Response:<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+"""  # noqa: E501
 
 EXPECTED_LORA_OUTPUT = [
-    "  SELECT icao FROM table_name_74 WHERE airport = 'lilongwe international airport' ",  # noqa: E501
-    "  SELECT nationality FROM table_name_11 WHERE elector = 'anchero pantaleone' ",
-    "  SELECT one_mora FROM table_name_95 WHERE gloss = 'low tone mora with a gloss of /˩okiru/' [òkìɽɯ́] AND accented_mora = 'low tone mora with a gloss of /˩okiru/' [òkìɽɯ́] ",  # noqa: E501
-    "  SELECT sex FROM people WHERE people_id IN (SELECT people_id FROM candidate GROUP BY sex ORDER BY COUNT(people_id) DESC LIMIT 1) ",  # noqa: E501
-    "  SELECT pick FROM table_name_60 WHERE former_wnba_team = 'Minnesota Lynx' ",
-    "  SELECT womens_doubles FROM table_28138035_4 WHERE mens_singles = 'Werner Schlager' ",  # noqa: E501
+    "SELECT count(*) FROM candidate",
+    "SELECT count(*) FROM candidate",
+    "SELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1",  # noqa: E501
+    "SELECT poll_source FROM candidate GROUP BY poll_source ORDER BY count(*) DESC LIMIT 1",  # noqa: E501
 ]
 
+MODEL_PATH = "meta-llama/Llama-3.2-3B-Instruct"
+
 
 def do_sample(
     llm: vllm.LLM,
@@ -32,18 +40,19 @@ def do_sample(
     tensorizer_config_dict: dict | None = None,
 ) -> list[str]:
     prompts = [
-        "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_74 (icao VARCHAR, airport VARCHAR)\n\n question: Name the ICAO for lilongwe international airport [/user] [assistant]",  # noqa: E501
-        "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_11 (nationality VARCHAR, elector VARCHAR)\n\n question: When Anchero Pantaleone was the elector what is under nationality? [/user] [assistant]",  # noqa: E501
-        "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_95 (one_mora VARCHAR, gloss VARCHAR, accented_mora VARCHAR)\n\n question: What is the one mora for a low tone mora with a gloss of /˩okiru/ [òkìɽɯ́]? [/user] [assistant]",  # noqa: E501
-        "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE candidate (people_id VARCHAR, unsure_rate INTEGER); CREATE TABLE people (sex VARCHAR, people_id VARCHAR)\n\n question: which gender got the highest average uncertain ratio. [/user] [assistant]",  # noqa: E501
-        "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_60 (pick INTEGER, former_wnba_team VARCHAR)\n\n question: What pick was a player that previously played for the Minnesota Lynx? [/user] [assistant]",  # noqa: E501
-        "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_28138035_4 (womens_doubles VARCHAR, mens_singles VARCHAR)\n\n question: Name the women's doubles for werner schlager [/user] [assistant]",  # noqa: E501
+        PROMPT_TEMPLATE.format(context="How many candidates are there?"),
+        PROMPT_TEMPLATE.format(context="Count the number of candidates."),
+        PROMPT_TEMPLATE.format(
+            context="Which poll resource provided the most number of candidate information?"  # noqa: E501
+        ),
+        PROMPT_TEMPLATE.format(
+            context="Return the poll resource associated with the most candidates."
+        ),
     ]
 
     sampling_params = vllm.SamplingParams(
-        temperature=0, max_tokens=256, skip_special_tokens=False, stop=["[/assistant]"]
+        temperature=0, max_tokens=64, stop=["<|im_end|>"]
     )
-
     if tensorizer_config_dict is not None:
         outputs = llm.generate(
             prompts,
@@ -75,13 +84,15 @@ def do_sample(
     return generated_texts
 
 
-def generate_and_test(llm, sql_lora_files, tensorizer_config_dict: dict | None = None):
+def generate_and_test(
+    llm, llama32_lora_files, tensorizer_config_dict: dict | None = None
+):
     print("lora adapter created")
     print("lora 1")
     assert (
         do_sample(
             llm,
-            sql_lora_files,
+            llama32_lora_files,
             tensorizer_config_dict=tensorizer_config_dict,
             lora_id=1,
         )
@@ -92,7 +103,7 @@ def generate_and_test(llm, sql_lora_files, tensorizer_config_dict: dict | None =
     assert (
         do_sample(
             llm,
-            sql_lora_files,
+            llama32_lora_files,
             tensorizer_config_dict=tensorizer_config_dict,
             lora_id=2,
         )
@@ -102,12 +113,10 @@ def generate_and_test(llm, sql_lora_files, tensorizer_config_dict: dict | None =
     print("removing lora")
 
 
-@create_new_process_for_each_test()
 @pytest.mark.parametrize("cudagraph_specialize_lora", [True, False])
-def test_llama_lora(sql_lora_files, cudagraph_specialize_lora: bool):
+def test_llama_lora(llama32_lora_files, cudagraph_specialize_lora: bool):
     llm = vllm.LLM(
         MODEL_PATH,
-        tokenizer=sql_lora_files,
         enable_lora=True,
         # also test odd max_num_seqs
         max_num_seqs=13,
@@ -116,39 +125,35 @@ def test_llama_lora(sql_lora_files, cudagraph_specialize_lora: bool):
             cudagraph_specialize_lora=cudagraph_specialize_lora,
         ),
     )
-    generate_and_test(llm, sql_lora_files)
+    generate_and_test(llm, llama32_lora_files)
 
 
-@multi_gpu_test(num_gpus=4)
-def test_llama_lora_tp4(sql_lora_files):
+def test_llama_lora_tp4(llama32_lora_files):
     llm = vllm.LLM(
         MODEL_PATH,
-        tokenizer=sql_lora_files,
         enable_lora=True,
         max_num_seqs=16,
         max_loras=4,
         tensor_parallel_size=4,
     )
-    generate_and_test(llm, sql_lora_files)
+    generate_and_test(llm, llama32_lora_files)
 
 
-@multi_gpu_test(num_gpus=4)
-def test_llama_lora_tp4_fully_sharded_loras(sql_lora_files):
+def test_llama_lora_tp4_fully_sharded_loras(llama32_lora_files):
     llm = vllm.LLM(
         MODEL_PATH,
-        tokenizer=sql_lora_files,
         enable_lora=True,
         max_num_seqs=16,
         max_loras=4,
         tensor_parallel_size=4,
         fully_sharded_loras=True,
     )
-    generate_and_test(llm, sql_lora_files)
+    generate_and_test(llm, llama32_lora_files)
 
 
-@multi_gpu_test(num_gpus=2)
 def test_tp2_serialize_and_deserialize_lora(
-    tmp_path, sql_lora_files, sql_lora_huggingface_id
+    tmp_path,
+    llama32_lora_files,
 ):
     # Run the tensorizing of the LoRA adapter and the model in a subprocess
     # to guarantee cleanup
@@ -157,7 +162,7 @@ def test_tp2_serialize_and_deserialize_lora(
     model_name = "model-rank-%03d.tensors"
 
     model_ref = MODEL_PATH
-    lora_path = sql_lora_huggingface_id
+    lora_path = llama32_lora_files
     suffix = "test"
     try:
         result = subprocess.run(
@@ -195,7 +200,6 @@ def test_tp2_serialize_and_deserialize_lora(
 
     loaded_llm = LLM(
         model=model_ref,
-        tokenizer=sql_lora_files,
         load_format="tensorizer",
         enable_lora=True,
         enforce_eager=True,
@@ -211,7 +215,7 @@ def test_tp2_serialize_and_deserialize_lora(
     print("lora 1")
     assert (
         do_sample(
-            loaded_llm, sql_lora_files, tensorizer_config_dict=tc_as_dict, lora_id=1
+            loaded_llm, llama32_lora_files, tensorizer_config_dict=tc_as_dict, lora_id=1
         )
         == EXPECTED_LORA_OUTPUT
     )