add tests defines

Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>

Author: reasonsolo

tensorrt_llm/llmapi/disagg_utils.py

@@ -43,6 +43,21 @@ class ConditionalDisaggConfig():
     max_local_prefill_length: int = 0
 
 
+@dataclass
+class MinimalInstances:
+    context_servers: int = 1
+    generation_servers: int = 1
+
+
+@dataclass
+class DisaggClusterConfig:
+    cluster_uri: str
+    cluster_name: str = ""
+    minimal_instances: Optional[MinimalInstances] = None
+    heartbeat_interval: int = 5
+    inactive_timeout: int = 10
+
+
 @dataclass
 class DisaggServerConfig():
     server_configs: List[CtxGenServerConfig]

tensorrt_llm/serve/auto_scaling.py

@@ -105,9 +105,6 @@ async def get_worker_events(
         worker_events = []
         for event in events:
             try:
-                print(
-                    f"Processing event: {event.event_type} for key: {event.storage_item.key} value {event.storage_item.value}"
-                )
                 worker_info = self._parse_worker_info(event)
                 worker_events.append((worker_info, event.event_type))
             except Exception as e:
@@ -192,6 +189,11 @@ def __init__(self, role: ServerRole, host: str, port: int,
         self._last_heartbeat = 0
         self._worker_id = f"{role.name}-{host}:{port}-{int(time.time()*1000)}-{os.getpid()}-{random.randint(0, 1000):03}"
 
+    def __del__(self):
+        if asyncio.get_event_loop():
+            asyncio.run_coroutine_threadsafe(self.deregister_worker(),
+                                             asyncio.get_event_loop())
+
     @property
     def worker_id(self) -> str:
         return self._worker_id
@@ -230,8 +232,8 @@ async def register_worker(self, validator=None, retry_interval=5):
                 logger.warning(
                     f"Worker {self.worker_info.worker_id} registration failed, retry in {retry_interval} seconds"
                 )
-                await asyncio.sleep(max(10, retry_interval))
-                return await self.register_worker(validator, retry_interval + 1)
+                await asyncio.sleep(retry_interval)
+                return await self.register_worker(validator, retry_interval)
         else:
             logger.info(
                 f"Worker {self.worker_info.worker_id} registration successful")
@@ -248,8 +250,9 @@ async def register_worker(self, validator=None, retry_interval=5):
 
     async def deregister_worker(self):
         self._stop = True
-        self._heartbeat_task.cancel()
-        self._heartbeat_task = None
+        if self._heartbeat_task:
+            self._heartbeat_task.cancel()
+            self._heartbeat_task = None
         await self._cluster_storage.stop()
         success = await self._cluster_storage.delete(self.worker_key)
         if not success:

tensorrt_llm/serve/cluster_storage.py

@@ -99,18 +99,12 @@ async def get_prefix(self,
 def create_cluster_storage(cluster_uri, cluster_name, **kwargs):
     if cluster_uri.startswith("http"):
         return HttpClusterStorageServer(cluster_uri, cluster_name, **kwargs)
-    elif cluster_uri.startswith("etcd"):
-        from tensorrt_llm.serve.cluster_storage_etcd import Etcd3ClusterStorage
-        return Etcd3ClusterStorage(cluster_uri, cluster_name, **kwargs)
     raise ValueError(f"Invalid cluster storage URI: {cluster_uri}")
 
 
 def create_cluster_storage_client(cluster_uri, cluster_name):
     if cluster_uri.startswith("http"):
         return HttpClusterStorageClient(cluster_uri, cluster_name)
-    elif cluster_uri.startswith("etcd"):
-        from tensorrt_llm.serve.cluster_storage_etcd import Etcd3ClusterStorage
-        return Etcd3ClusterStorage(cluster_uri, cluster_name)
     raise ValueError(f"Invalid cluster storage URI: {cluster_uri}")
 
 
@@ -356,7 +350,7 @@ async def get_prefix(self,
                          keys_only: bool = False) -> Dict[str, str]:
         return await self._get("get_prefix",
                                key_prefix=key_prefix,
-                               keys_only=keys_only)
+                               keys_only=int(keys_only))
 
     async def delete(self, key: str) -> bool:
         try:

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

@@ -35,6 +35,8 @@ l0_h100:
   - unittest/disaggregated/test_disagg_utils.py
   - unittest/disaggregated/test_router.py
   - unittest/disaggregated/test_remoteDictionary.py
+  - unittest/disaggregated/test_cluster_manager_worker.py
+  - unittest/disaggregated/test_cluster_storage.py
   - accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::test_auto_dtype
   - accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::test_auto_dtype_vswa
   - accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::test_auto_dtype_chunked_prefill

tests/unittest/disaggregated/test_cluster_manager_worker.py

@@ -0,0 +1,227 @@
+import asyncio
+import subprocess
+import tempfile
+import time
+
+import pytest
+
+from tensorrt_llm.llmapi.disagg_utils import (DisaggClusterConfig,
+                                              MinimalInstances, ServerRole)
+from tensorrt_llm.serve.auto_scaling import ClusterManager, ClusterWorker
+from tensorrt_llm.serve.cluster_storage import (WatchEventType,
+                                                create_cluster_storage,
+                                                create_cluster_storage_client)
+
+from .test_cluster_storage import http_server_storage, pytest_async_fixture
+
+INACTIVE_TIMEOUT = 4
+HEARTBEAT_INTERVAL = 2
+
+storage_types = ["http"]
+
+
+def get_uri(storage_type):
+    if storage_type == "http":
+        return f"http://localhost:18000"
+    elif storage_type == "etcd":
+        return f"etcd://localhost:2379"
+    else:
+        raise ValueError(f"Invalid storage type: {storage_type}")
+
+
+@pytest.fixture(scope="module")
+def config(request):
+    cluster_uri = get_uri(request.param)
+    return DisaggClusterConfig(cluster_uri=cluster_uri,
+                               cluster_name="test",
+                               minimal_instances=MinimalInstances(
+                                   context_servers=1, generation_servers=1),
+                               inactive_timeout=INACTIVE_TIMEOUT,
+                               heartbeat_interval=HEARTBEAT_INTERVAL)
+
+
+@pytest.fixture(scope="module")
+def storage_server(config):
+    if config.cluster_uri.startswith("http"):
+        port = 18000
+        server, cluster_storage = http_server_storage(port)
+        with server.run_in_thread():
+            yield cluster_storage, config.cluster_uri
+    elif config.cluster_uri.startswith("etcd"):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            etcd = subprocess.Popen(
+                ["etcd", "--data-dir", temp_dir, "--log-level", "debug"])
+            time.sleep(2)  # wait for etcd to start
+            yield create_cluster_storage(
+                config.cluster_uri, config.cluster_name), config.cluster_uri
+        etcd.kill()
+        etcd.wait()
+    else:
+        raise ValueError(f"Invalid cluster storage URI: {config.cluster_uri}")
+
+
+@pytest_async_fixture(scope="module")
+async def storage_client(storage_server):
+    _, cluster_uri = storage_server
+    return create_cluster_storage_client(cluster_uri, "test")
+
+
+@pytest_async_fixture(scope="module")
+async def cluster_manager(config, storage_server):
+    storage, cluster_uri = storage_server
+    manager = ClusterManager(config, storage)
+    await manager.start()
+    yield manager
+    await manager.stop()
+
+
+@pytest.mark.parametrize("config", storage_types, indirect=True)
+@pytest.mark.threadleak(enabled=False)
+@pytest.mark.asyncio(scope="module")
+async def test_init_workers_first(config, storage_server):
+    try:
+        # init workers before initializing the manager, so the manager should be able to
+        # get the pre-registered workers
+        server, storage_uri = storage_server
+        storage_client = create_cluster_storage_client(storage_uri, "test")
+        ctx_worker = ClusterWorker(ServerRole.CONTEXT, "127.0.0.1", 8001,
+                                   config, storage_client)
+        gen_worker = ClusterWorker(ServerRole.GENERATION, "127.0.0.1", 8002,
+                                   config, storage_client)
+        await ctx_worker.register_worker()
+        await gen_worker.register_worker()
+
+        cluster_manager = ClusterManager(config, server)
+        await cluster_manager.start()
+        existing_workers = await cluster_manager.watch_workers(
+            get_existing_first=True)
+        assert set([worker.worker_id for worker in existing_workers]) == {
+            ctx_worker.worker_id,
+            gen_worker.worker_id,
+        }
+
+        assert await cluster_manager.is_ready() == True
+    finally:
+        await ctx_worker.deregister_worker()
+        await gen_worker.deregister_worker()
+
+
+@pytest.mark.parametrize("config", storage_types, indirect=True)
+@pytest.mark.threadleak(enabled=False)
+@pytest.mark.timeout(20)
+@pytest.mark.asyncio(scope="module")
+async def test_cluster_manager(cluster_manager, storage_client, config):
+    try:
+        cluster_manager.current_ctx_worker_num == 0
+        cluster_manager.current_gen_worker_num == 0
+        await cluster_manager.watch_workers()
+        try:
+            await asyncio.wait_for(cluster_manager.get_worker_events(),
+                                   timeout=1)
+        except asyncio.TimeoutError:
+            pass
+        assert await cluster_manager.is_ready() == False
+
+        ctx_worker = ClusterWorker(ServerRole.CONTEXT, "127.0.0.1", 8001,
+                                   config, storage_client)
+        await cluster_manager.watch_workers()
+        await ctx_worker.register_worker()
+        worker_events = await cluster_manager.get_worker_events()
+        assert worker_events == [(ctx_worker.worker_info, WatchEventType.SET)]
+        assert cluster_manager.current_ctx_worker_num == 1
+        assert cluster_manager.current_gen_worker_num == 0
+        assert await cluster_manager.is_ready() == False
+
+        gen_worker = ClusterWorker(ServerRole.GENERATION, "127.0.0.1", 8002,
+                                   config, storage_client)
+        await gen_worker.register_worker()
+        worker_events = await cluster_manager.get_worker_events()
+        assert worker_events == [(gen_worker.worker_info, WatchEventType.SET)]
+        assert cluster_manager.current_ctx_worker_num == 1
+        assert cluster_manager.current_gen_worker_num == 1
+        assert await cluster_manager.is_ready() == True
+
+        await ctx_worker.deregister_worker()
+        worker_events = await cluster_manager.get_worker_events()
+        assert worker_events == [(ctx_worker.worker_info, WatchEventType.DELETE)
+                                 ]
+        assert cluster_manager.current_ctx_worker_num == 0
+        assert cluster_manager.current_gen_worker_num == 1
+        assert await cluster_manager.is_ready() == False
+
+        await gen_worker.deregister_worker()
+        worker_events = await cluster_manager.get_worker_events()
+        assert worker_events == [(gen_worker.worker_info, WatchEventType.DELETE)
+                                 ]
+        assert cluster_manager.current_ctx_worker_num == 0
+        assert cluster_manager.current_gen_worker_num == 0
+        assert await cluster_manager.is_ready() == False
+    finally:
+        await ctx_worker.deregister_worker()
+        await gen_worker.deregister_worker()
+
+
+@pytest.mark.timeout(20)
+@pytest.mark.parametrize("config", storage_types, indirect=True)
+@pytest.mark.threadleak(enabled=False)
+@pytest.mark.asyncio(scope="module")
+async def test_cluster_worker(cluster_manager, storage_client, config):
+
+    async def wait_for_worker_events(expected_new_event_num,
+                                     expected_dead_event_num):
+        new_worker_ids = []
+        dead_workers_ids = []
+        while len(new_worker_ids) < expected_new_event_num or len(
+                dead_workers_ids) < expected_dead_event_num:
+            try:
+                worker_events = await asyncio.wait_for(
+                    cluster_manager.get_worker_events(), timeout=2)
+                new_workers = [
+                    worker_info.worker_id
+                    for worker_info, event_type in worker_events
+                    if event_type == WatchEventType.SET
+                ]
+                dead_workers = [
+                    worker_info.worker_id
+                    for worker_info, event_type in worker_events
+                    if event_type == WatchEventType.DELETE
+                ]
+                print(f"Worker events: {worker_events} {time.time()}")
+                new_worker_ids += new_workers
+                dead_workers_ids += dead_workers
+            except asyncio.TimeoutError:
+                pass
+        return new_worker_ids, dead_workers_ids
+
+    try:
+        await cluster_manager.start()
+        await cluster_manager.watch_workers()
+        ctx_worker = ClusterWorker(ServerRole.CONTEXT, "127.0.0.1", 8001,
+                                   config, storage_client)
+        gen_worker = ClusterWorker(ServerRole.GENERATION, "127.0.0.1", 8002,
+                                   config, storage_client)
+
+        keep_heartbeat = True
+        assert await ctx_worker.register_worker(validator=lambda: keep_heartbeat
+                                                )
+        assert await gen_worker.register_worker(validator=lambda: keep_heartbeat
+                                                )
+        worker_ids = set([ctx_worker.worker_id, gen_worker.worker_id])
+        new_worker_ids, dead_workers_ids = await wait_for_worker_events(2, 0)
+        assert set(new_worker_ids) == worker_ids
+        assert len(dead_workers_ids) == 0
+        assert await cluster_manager.is_ready() == True
+
+        await asyncio.sleep(config.inactive_timeout + 1)
+        assert await cluster_manager.is_ready() == True
+
+        # stop heartbeat, then we should see two workers deleted
+        keep_heartbeat = False
+        new_worker_ids, dead_workers_ids = await wait_for_worker_events(0, 2)
+        assert len(new_worker_ids) == 0
+        assert len(dead_workers_ids) == 2
+        assert set(dead_workers_ids) == worker_ids
+        assert await cluster_manager.is_ready() == False
+    finally:
+        await ctx_worker.deregister_worker()
+        await gen_worker.deregister_worker()

tests/unittest/serve/test_cluster_storage.py renamed to tests/unittest/disaggregated/test_cluster_storage.py

@@ -96,19 +96,22 @@ async def test_expire(self, storage_server, storage_client):
 
     @timeout(5)
     @pytest_async_module
-    async def test_get_prefix(self, storage_server, storage_client):
+    async def test_get_keys(self, storage_server, storage_client):
         keys = [gen_key("test_key_unique") for _ in range(3)]
-        for key in keys:
+        values = [f"test_value{i}" for i in range(3)]
+        for key, value in zip(keys, values):
             assert await storage_client.set(key,
-                                            "test_value1",
+                                            value,
                                             overwrite_if_exists=True)
 
-        answer_keys = await storage_client.get_prefix("test_key_unique")
-        assert set(keys) == set(answer_keys)
-        answer_keys = await storage_client.get_prefix(keys[0])
-        assert answer_keys == [keys[0]]
-        answer_keys = await storage_client.get_prefix(keys[1])
-        assert answer_keys == [keys[1]]
+        answer_keys = await storage_client.get_prefix("test_key_unique",
+                                                      keys_only=False)
+        assert set(keys) == set(answer_keys.keys())
+        assert set(values) == set(answer_keys.values())
+        answer_keys = await storage_client.get_prefix(keys[0], keys_only=True)
+        assert answer_keys == {keys[0]: ""}
+        answer_keys = await storage_client.get_prefix(keys[1], keys_only=True)
+        assert answer_keys == {keys[1]: ""}
 
     @pytest_ignore_tleak
     @pytest_async_module
@@ -199,7 +202,8 @@ def storage_server(self):
 
 
 class TestEtcdClusterStorage(TestClusterStorage):
-    __test__ = True
+    # Disable this test until Etcd functionality is ready.
+    __test__ = False
 
     @pytest.fixture(scope="class")
     def storage_server(self):