@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 156 lines 5.72 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Dynamic registration and liveness for the EPD proxy's instance registry."""
4
5 from unittest.mock import patch
6
7 import pytest
8
9 from examples.disaggregated.disaggregated_encoder.disagg_epd_proxy import (
10 InstanceRecord,
11 InstanceRegistry,
12 InstanceRole,
13 )
14
15 ENCODE = InstanceRole.ENCODE
16 DECODE = InstanceRole.DECODE
17
18
19 @pytest.fixture
20 def registry():
21 return InstanceRegistry(probe_interval=0, fail_threshold=3, evicted_ttl=60.0)
22
23
24 def _fill(registry, role, count, prefix="http://h"):
25 for index in range(count):
26 registry.register(InstanceRecord(role, f"{prefix}{index}:8000"))
27
28
29 async def _probe_round(registry, healthy: set[str], now: float = 0.0):
30 """Run one probe pass where only `healthy` URLs answer."""
31
32 async def fake_probe(self, session, url):
33 return url in healthy
34
35 with (
36 patch.object(InstanceRegistry, "_probe", fake_probe),
37 patch(
38 "examples.disaggregated.disaggregated_encoder.disagg_epd_proxy.time.monotonic",
39 return_value=now,
40 ),
41 ):
42 await registry._probe_once(session=None)
43
44
45 class TestRegistration:
46 def test_registering_is_idempotent_and_reports_novelty(self, registry):
47 record = InstanceRecord(ENCODE, "http://e0:8000")
48 assert registry.register(record) is True
49 assert registry.register(record) is False
50 assert registry.urls(ENCODE) == ["http://e0:8000"]
51
52 def test_transfer_addresses_stay_with_the_registered_consumer(self, registry):
53 """Roster updates must not shift a consumer's transfer addresses."""
54 registry.register(InstanceRecord(DECODE, "http://d0:8000"))
55 registry.register(
56 InstanceRecord(
57 DECODE,
58 "http://d0:8000",
59 ec_zmq_addrs=["tcp://d0:20001", "tcp://d0:20002"],
60 dp_size=2,
61 )
62 )
63 picked = registry.pick(DECODE)
64 assert picked.ec_zmq_addrs == ["tcp://d0:20001", "tcp://d0:20002"]
65 assert picked.dp_size == 2
66
67 def test_unregister_is_terminal(self, registry):
68 registry.register(InstanceRecord(ENCODE, "http://e0:8000"))
69 assert registry.unregister("http://e0:8000") is True
70 assert registry.unregister("http://e0:8000") is False
71 assert registry.pick(ENCODE) is None
72
73 def test_round_robin_survives_registration(self, registry):
74 _fill(registry, ENCODE, 2)
75 assert [record.url for record in registry.pick_many(ENCODE, 2)] == [
76 "http://h0:8000",
77 "http://h1:8000",
78 ]
79 registry.register(InstanceRecord(ENCODE, "http://h2:8000"))
80 assert registry.pick(ENCODE).url == "http://h2:8000"
81
82
83 class TestLiveness:
84 @pytest.mark.asyncio
85 async def test_identical_reregistration_preserves_inflight_probe_results(
86 self, registry
87 ):
88 """Repeated registrations must not suppress eviction or recovery."""
89 url = "http://e0:8000"
90 registry.register(InstanceRecord(ENCODE, url))
91 healthy = False
92
93 async def probe(self, session, url):
94 registry.register(InstanceRecord(ENCODE, url))
95 return healthy
96
97 with patch.object(InstanceRegistry, "_probe", probe):
98 for _ in range(registry._fail_threshold):
99 await registry._probe_once(None)
100 assert registry.status()["encode"] == {"live": [], "evicted": [url]}
101
102 healthy = True
103 await registry._probe_once(None)
104 assert registry.status()["encode"] == {"live": [url], "evicted": []}
105
106 @pytest.mark.asyncio
107 async def test_consecutive_failures_stop_routing(self, registry):
108 registry.register(InstanceRecord(ENCODE, "http://e0:8000"))
109 registry.register(InstanceRecord(ENCODE, "http://e1:8000"))
110 for _ in range(registry._fail_threshold):
111 await _probe_round(registry, healthy={"http://e1:8000"})
112 assert registry.urls(ENCODE) == ["http://e1:8000"]
113 assert registry.status()["encode"]["evicted"] == ["http://e0:8000"]
114
115 @pytest.mark.asyncio
116 @pytest.mark.parametrize("role", [ENCODE, DECODE])
117 async def test_a_recovered_instance_rejoins_without_re_registering(
118 self, registry, role
119 ):
120 """Restarting every encoder to recover from a blip is not acceptable."""
121 registry.register(InstanceRecord(role, "http://e0:8000"))
122 for _ in range(registry._fail_threshold):
123 await _probe_round(registry, healthy=set())
124 assert registry.urls(role) == []
125
126 await _probe_round(registry, healthy={"http://e0:8000"}, now=1)
127 assert registry.urls(role) == ["http://e0:8000"]
128
129 @pytest.mark.asyncio
130 async def test_re_registration_does_not_override_failed_health_checks(
131 self, registry
132 ):
133 record = InstanceRecord(ENCODE, "http://e0:8000")
134 registry.register(record)
135 for _ in range(registry._fail_threshold):
136 await _probe_round(registry, healthy=set())
137 registry.register(record)
138 assert registry.urls(ENCODE) == []
139
140 @pytest.mark.asyncio
141 @pytest.mark.parametrize("healthy", [False, True])
142 async def test_inflight_probe_cannot_restore_unregistered_instance(
143 self, registry, healthy
144 ):
145 record = InstanceRecord(ENCODE, "http://e0:8000")
146 registry.register(record)
147 for _ in range(registry._fail_threshold):
148 await _probe_round(registry, healthy=set())
149
150 async def probe(self, session, url):
151 registry.unregister(url)
152 return healthy
153
154 with patch.object(InstanceRegistry, "_probe", probe):
155 await registry._probe_once(None)
156 assert registry.status()["encode"] == {"live": [], "evicted": []}