fix: add blocking mode for k8s connector in planner (#1446)

components/planner/src/dynamo/planner/kubernetes_connector.py

@@ -43,7 +43,7 @@ class KubernetesConnector(PlannerConnector):
                 self._get_graph_deployment_name(deployment)
             )
 
-    async def remove_component(self, component_name: str):
+    async def remove_component(self, component_name: str, blocking: bool = True):
         """Remove a component by decreasing its replica count by 1"""
         deployment = await self.kube_api.get_graph_deployment(
             component_name, self.namespace
@@ -60,6 +60,10 @@ class KubernetesConnector(PlannerConnector):
                 component_name,
                 current_replicas - 1,
             )
+            if blocking:
+                await self.kube_api.wait_for_graph_deployment_ready(
+                    self._get_graph_deployment_name(deployment)
+                )
 
     def _get_current_replicas(self, deployment: dict, component_name: str) -> int:
         """Get the current replicas for a component in a graph deployment"""

components/planner/test/kubernetes_connector.py

@@ -123,6 +123,7 @@ async def test_remove_component_decreases_replicas(kubernetes_connector, mock_ku
     mock_kube_api.update_graph_replicas.assert_called_once_with(
         "test-graph", component_name, 1
     )
+    mock_kube_api.wait_for_graph_deployment_ready.assert_called_once_with("test-graph")
 
 
 @pytest.mark.asyncio
@@ -140,3 +141,4 @@ async def test_remove_component_with_zero_replicas(kubernetes_connector, mock_ku
 
     # Assert
     mock_kube_api.update_graph_replicas.assert_not_called()
+    mock_kube_api.wait_for_graph_deployment_ready.assert_not_called()