Skip to content

Commit 3946c3d

Browse files
authored
Merge branch 'main' into rupei/hashes-to-texts
2 parents de04c4f + 3c85cfd commit 3946c3d

File tree

3 files changed

+9
-1
lines changed

3 files changed

+9
-1
lines changed

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ async def add_component(self, component_name: str, blocking: bool = True):
4343
self._get_graph_deployment_name(deployment)
4444
)
4545

46-
async def remove_component(self, component_name: str):
46+
async def remove_component(self, component_name: str, blocking: bool = True):
4747
"""Remove a component by decreasing its replica count by 1"""
4848
deployment = await self.kube_api.get_graph_deployment(
4949
component_name, self.namespace
@@ -60,6 +60,10 @@ async def remove_component(self, component_name: str):
6060
component_name,
6161
current_replicas - 1,
6262
)
63+
if blocking:
64+
await self.kube_api.wait_for_graph_deployment_ready(
65+
self._get_graph_deployment_name(deployment)
66+
)
6367

6468
def _get_current_replicas(self, deployment: dict, component_name: str) -> int:
6569
"""Get the current replicas for a component in a graph deployment"""

components/planner/test/kubernetes_connector.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,7 @@ async def test_remove_component_decreases_replicas(kubernetes_connector, mock_ku
123123
mock_kube_api.update_graph_replicas.assert_called_once_with(
124124
"test-graph", component_name, 1
125125
)
126+
mock_kube_api.wait_for_graph_deployment_ready.assert_called_once_with("test-graph")
126127

127128

128129
@pytest.mark.asyncio
@@ -140,3 +141,4 @@ async def test_remove_component_with_zero_replicas(kubernetes_connector, mock_ku
140141

141142
# Assert
142143
mock_kube_api.update_graph_replicas.assert_not_called()
144+
mock_kube_api.wait_for_graph_deployment_ready.assert_not_called()

pyproject.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@ dependencies = [
3131
"ai-dynamo-runtime==0.3.0",
3232
"fastapi==0.115.6",
3333
"distro",
34+
# filelock: required by planner
35+
"filelock",
3436
"typer",
3537
"circus==0.19.0",
3638
"click<8.2.0",

0 commit comments

Comments
 (0)