initial lb changes

microsoft · Jun 29, 2023 · 4d4e0d8 · 4d4e0d8
1 parent 38e270e
commit 4d4e0d8
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/mii/grpc_related/modelresponse_server.py b/mii/grpc_related/modelresponse_server.py
@@ -168,11 +168,12 @@ def __init__(self, task_name, replica_configs):
         super().__init__()
         self.asyncio_loop = asyncio.get_event_loop()
 
-        self.stubs = [
-            ParallelStubInvoker(replica.hostname,
-                                replica.tensor_parallel_ports)
-            for replica in replica_configs
-        ]
+        self.stubs = {}
+        for repl in replica_configs:
+            stubs[repl.deployment_name] = [ParallelStubInvoker(replica.hostname,
+                                                               replica.tensor_parallel_ports)
+                                                    for replica in replica_configs
+                                          ]
         print(self.stubs)
         self.counter = AtomicCounter()
         self.task = get_task(task_name)