fix how we find num replicas for endpoints and collectors (#340)

launchflow · Dec 22, 2023 · a899c7b · a899c7b
1 parent 5e959fe
commit a899c7b
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/buildflow/core/app/runtime/actors/collector_pattern/receive_process_push_ack.py b/buildflow/core/app/runtime/actors/collector_pattern/receive_process_push_ack.py
@@ -158,7 +158,12 @@ async def snapshot(self) -> Snapshot:
                 avg_process_time_millis=0,
             )
         if self.collector_deployment is not None:
-            num_replicas = self.collector_deployment.num_replicas
+            num_replicas = (
+                serve.status()
+                .applications.get(self.processor_group.group_id, {})
+                .deployments.get(self.endpoint_deployment.name, {})
+                .replica_states.get("RUNNING", 0)
+            )
         else:
             num_replicas = 0
 

diff --git a/buildflow/core/app/runtime/actors/endpoint_pattern/receive_process_respond.py b/buildflow/core/app/runtime/actors/endpoint_pattern/receive_process_respond.py
@@ -146,7 +146,12 @@ async def snapshot(self) -> Snapshot:
                 avg_process_time_millis=0,
             )
         if self.endpoint_deployment is not None:
-            num_replicas = self.endpoint_deployment.num_replicas
+            num_replicas = (
+                serve.status()
+                .applications.get(self.processor_group.group_id, {})
+                .deployments.get(self.endpoint_deployment.name, {})
+                .replica_states.get("RUNNING", 0)
+            )
         else:
             num_replicas = 0