Add disaggregated unittest (NVIDIA#4899)

Signed-off-by: Shunkang <[email protected]>

Author: Shunkangz

tensorrt_llm/commands/serve.py

@@ -292,6 +292,7 @@ def serve(model: str, tokenizer: Optional[str], host: str, port: int,
         metadata_server_config_file)
 
     if metadata_server_cfg is not None:
+        assert server_role is not None, "server_role is required when metadata_server_cfg is provided"
         try:
             server_role = ServerRole[server_role.upper()]
         except ValueError:

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

@@ -19,6 +19,9 @@ l0_h100:
   - unittest/_torch -k "modeling_llama"
   - unittest/_torch/modeling -k "modeling_mixtral"
   - unittest/_torch/modeling -k "modeling_nemotron"
+  - unittest/disaggregated/test_disagg_utils.py
+  - unittest/disaggregated/test_router.py
+  - unittest/disaggregated/test_remoteDictionary.py
   - accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::test_auto_dtype
   - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16[attn_backend=TRTLLM-torch_compile=False]
   - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_bfloat16[attn_backend=TRTLLM-torch_compile=True]

tests/unittest/disaggregated/test_disagg_utils.py

@@ -57,7 +57,6 @@ def verify_disagg_config(config: DisaggServerConfig):
     assert config.ctx_router_config.type == "round_robin"
     assert config.gen_router_config.type == "load_balancing"
     assert len(config.server_configs) == 3
-    assert config.condition is None
 
 
 def test_parse_disagg_config_file(sample_yaml_file):

tests/unittest/disaggregated/test_remoteDictionary.py

@@ -1,13 +1,46 @@
+import os
+import signal
+import subprocess
+import time
 import unittest
 
 import etcd3
 
 from tensorrt_llm.serve.metadata_server import EtcdDictionary
 
 
+def start_etcd_server():
+    # Command to start etcd
+    etcd_cmd = ["etcd"]
+
+    # Start etcd in background
+    process = subprocess.Popen(
+        etcd_cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        preexec_fn=os.setsid)  # This makes it run in a new process group
+
+    # Wait a bit for etcd to start
+    time.sleep(5)
+
+    return process
+
+
+def stop_etcd_server(process):
+    # Kill the process group
+    os.killpg(os.getpgid(process.pid), signal.SIGTERM)
+    process.wait()
+
+
 class TestEtcdDictionary(unittest.TestCase):
 
     def setUp(self):
+        # Set the protocol buffers implementation to python
+        os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
+
+        # Start etcd server
+        self.etcd_process = start_etcd_server()
+
         # Setup etcd connection parameters
         self.host = "localhost"
         self.port = 2379
@@ -25,6 +58,12 @@ def tearDown(self):
         # Clean up test keys after each test
         self._cleanup_test_keys()
 
+        # Stop etcd server
+        stop_etcd_server(self.etcd_process)
+
+        # Unset the protocol buffers implementation
+        del os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"]
+
     def _cleanup_test_keys(self):
         # Helper method to remove test keys
         test_keys = [
@@ -47,6 +86,8 @@ def test_put_and_get(self):
         # Assert
         self.assertEqual(value.decode('utf-8'), test_value)
 
+        self._cleanup_test_keys()
+
     def test_remove(self):
         # Test removing a key
         test_key = "trtllm/1/test_key2"
@@ -63,6 +104,8 @@ def test_remove(self):
         self.assertIsNone(
             result[0])  # etcd3 returns (None, None) when key doesn't exist
 
+        self._cleanup_test_keys()
+
     def test_keys(self):
         # Test listing all keys
         test_data = {
@@ -85,6 +128,8 @@ def test_keys(self):
         extract_keys = set(keys)
         self.assertEqual(prefix_keys, extract_keys)
 
+        self._cleanup_test_keys()
+
     def test_get_nonexistent_key(self):
         # Test getting a key that doesn't exist
         result, _ = self.etcd_dict.get("nonexistent_key")
@@ -108,6 +153,8 @@ def test_put_update_existing(self):
         # Assert
         self.assertEqual(value.decode('utf-8'), updated_value)
 
+        self._cleanup_test_keys()
+
 
 if __name__ == '__main__':
     unittest.main()