Add disagg test

Signed-off-by: Pengyun Lin <[email protected]>

Author: LinPoly

examples/disaggregated/clients/long_prompts.json

@@ -0,0 +1,3 @@
+[
+    "Place of birth\nThe place of birth (POB) or birthplace is the place where a person was born. This place is often used in legal documents, together with name and date of birth, to uniquely identify a person. Practice regarding whether this place should be a country, a territory or a city/town/locality differs in different countries, but often city or territory is used for native-born citizen passports and countries for foreign-born ones.\nAs a general rule with respect to passports, if the place of birth is to be a country, it's determined to be the country that currently has sovereignty over the actual place of birth, regardless of when the birth actually occurred. The place of birth is not necessarily the place where the parents of the new baby live. If the baby is born in a hospital in another place, that place is the place of birth. In many countries, this also means that the government requires that the birth of the new baby is registered in the place of birth.\nSome countries place less or no importance on the place of birth, instead using alternative geographical characteristics for the purpose of identity documents. For example, Sweden has used the concept of födelsehemort (\"domicile of birth\") since 1947. This means that the domicile of the baby's mother is the registered place of birth.\nSimilarly, Switzerland uses the concept of place of origin. A child born to Swiss parents is automatically assigned the place of origin of the parent with the same last name, so the child either gets their mother's or father's place of origin. A child born to one Swiss parent and one foreign parent acquires the place of origin of their Swiss parent. In a Swiss passport and identity card, the holder's place of origin is stated, not their place of birth. In Japan, the registered domicile is a similar concept.\nIn some countries (primarily in the Americas), the place of birth automatically determines the nationality of the baby, a practice often referred to by the Latin phrase jus soli."
+]

tests/integration/defs/disaggregated/test_configs/disagg_config_diff_max_tokens.yaml

@@ -0,0 +1,23 @@
+hostname: localhost
+port: 8000
+model: TinyLlama/TinyLlama-1.1B-Chat-v1.0
+free_gpu_memory_fraction: 0.5
+backend: "pytorch"
+cuda_graph_config: null
+disable_overlap_scheduler: True
+context_servers:
+  num_instances: 1
+  max_num_tokens: 512
+  max_batch_size: 256
+  cache_transceiver_config:
+    backend: default
+  urls:
+      - "localhost:8001"
+generation_servers:
+  num_instances: 1
+  max_num_tokens: 256
+  max_batch_size: 128
+  cache_transceiver_config:
+    backend: default
+  urls:
+      - "localhost:8002"

tests/integration/defs/disaggregated/test_disaggregated.py

@@ -36,6 +36,8 @@ def get_test_config(test_desc, example_dir, test_root):
     """Get test configuration based on test description."""
     test_configs_root = f"{test_root}/test_configs"
     config_map = {
+        "2_ranks_diff_max_tokens":
+        (2, f"{test_configs_root}/disagg_config_diff_max_tokens.yaml"),
         "2_ranks": (2, f"{example_dir}/disagg_config.yaml"),
         "2_ranks_trt_backend":
         (2, f"{test_configs_root}/disagg_config_trt_backend.yaml"),
@@ -134,7 +136,8 @@ def run_disaggregated_test(example_dir,
                            test_desc,
                            num_iters=5,
                            env=None,
-                           cwd=None):
+                           cwd=None,
+                           prompt_file="prompts.json"):
     """Run disaggregated test with given configuration."""
     cleanup_output_files()
     run_env = env.copy()
@@ -175,10 +178,13 @@ def run_disaggregated_test(example_dir,
                 client_cmd = [
                     'python3', f'{client_dir}/disagg_client.py', '-c',
                     f'{example_dir}/disagg_config.yaml', '-p',
-                    f'{client_dir}/prompts.json', '--ignore-eos',
+                    f'{client_dir}/{prompt_file}', '--ignore-eos',
                     '--server-start-timeout',
                     str(server_start_timeout)
                 ]
+                if prompt_file == "long_prompts.json":
+                    # Use max_tokens 4 for long prompts to reduce test time
+                    client_cmd.extend(['--max-tokens', '4'])
                 check_call(client_cmd,
                            env=env,
                            poll_procs=[workers_proc, server_proc])
@@ -207,6 +213,10 @@ def run_disaggregated_test(example_dir,
                                env=env,
                                poll_procs=[workers_proc, server_proc])
 
+                # Skip output verification for long prompts test
+                if prompt_file == "long_prompts.json":
+                    continue
+
                 # Verify outputs
                 not_expected_strings = ["Berlin Berlin"]
 
@@ -250,6 +260,27 @@ def run_disaggregated_test(example_dir,
         workers_proc.wait()
 
 
+@pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'],
+                         indirect=True)
+def test_disaggregated_diff_max_tokens(disaggregated_test_root,
+                                       disaggregated_example_root, llm_venv,
+                                       llama_model_root):
+    src_dst_dict = {
+        llama_model_root:
+        f"{llm_venv.get_working_directory()}/TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+    }
+    for src, dst in src_dst_dict.items():
+        if not os.path.islink(dst):
+            os.makedirs(os.path.dirname(dst), exist_ok=True)
+            os.symlink(src, dst, target_is_directory=True)
+
+    run_disaggregated_test(disaggregated_example_root,
+                           "2_ranks_diff_max_tokens",
+                           env=llm_venv._new_env,
+                           cwd=llm_venv.get_working_directory(),
+                           prompt_file="long_prompts.json")
+
+
 @pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'],
                          indirect=True)
 def test_disaggregated_single_gpu_with_mpirun(disaggregated_test_root,

tests/integration/test_lists/test-db/l0_a10.yml

@@ -21,6 +21,7 @@ l0_a10:
   - disaggregated/test_disaggregated.py::test_disaggregated_cuda_graph[TinyLlama-1.1B-Chat-v1.0]
   - disaggregated/test_disaggregated.py::test_disaggregated_mixed[TinyLlama-1.1B-Chat-v1.0]
   - disaggregated/test_disaggregated.py::test_disaggregated_overlap[TinyLlama-1.1B-Chat-v1.0]
+  - disaggregated/test_disaggregated.py::test_disaggregated_diff_max_tokens[TinyLlama-1.1B-Chat-v1.0]
   - test_e2e.py::test_openai_chat_structural_tag_example
   - test_e2e.py::test_openai_chat_multimodal_example
   - test_e2e.py::test_openai_lora

tests/unittest/llmapi/test_llm_pytorch.py

@@ -6,13 +6,13 @@
 
 # isort: off
 from .lora_test_utils import check_llama_7b_multi_unique_lora_adapters_from_request, create_mock_nemo_lora_checkpoint
-from .test_llm import (get_model_path, global_kvcache_config, llama_model_path,
+from .test_llm import (_test_llm_capture_request_error, get_model_path,
+                       global_kvcache_config, llama_model_path,
                        llm_get_stats_async_test_harness,
                        llm_get_stats_test_harness, prompts,
                        run_llm_abort_request,
                        run_llm_with_postprocess_parallel_and_result_handler,
-                       tinyllama_logits_processor_test_harness,
-                       _test_llm_capture_request_error)
+                       tinyllama_logits_processor_test_harness)
 from utils.util import (EnvVarsContextManager, force_ampere,
                         run_function_in_sub_process, similar,
                         skip_gpu_memory_less_than_40gb,