Skip to content

Commit 33eda74

Browse files
authored
fix: cleanup after scale_in (#302)
When we do scale in, we remove some workers. We need to do cleanup for these workers in the worker status data structure for job state to correctly evaluate the status based on worker statuses.
1 parent d7f9bfb commit 33eda74

File tree

1 file changed

+8
-8
lines changed

1 file changed

+8
-8
lines changed

infscale/controller/job_context.py

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -474,11 +474,7 @@ async def _remove_pipeline_n_update(
474474
"""Remove pipeline and update job."""
475475
updated_cfg = JobConfig.remove_pipeline(cfg, failed_wrk_ids)
476476

477-
worker_diff = JobConfig.get_workers_diff(cfg, updated_cfg)
478-
479-
# since we remove a pipeline,
480-
# we need to cleanup worker status data structure
481-
self.context.remove_wrk_status(worker_diff)
477+
self.context._reconcile_wrk_status(cfg, updated_cfg)
482478

483479
# checker setup with updated config
484480
self.context.job_checker.setup(updated_cfg)
@@ -808,8 +804,12 @@ def set_wrkr_metrics(self, wrkr_id: str, metrics: PerfMetrics) -> None:
808804
"""Set worker's performance metrics."""
809805
self.wrkr_metrics[wrkr_id] = metrics
810806

811-
def _init_wrk_status(self) -> None:
812-
"""Create worker status dict with default status."""
807+
def _reconcile_wrk_status(self, cur_cfg: JobConfig, new_cfg: JobConfig) -> None:
808+
"""Reconcile worker status dict by adding or removing entries."""
809+
if cur_cfg:
810+
worker_diff = JobConfig.get_workers_diff(cur_cfg, new_cfg)
811+
self.remove_wrk_status(worker_diff)
812+
813813
for w in self._new_cfg.workers:
814814
if w.id not in self.wrk_status:
815815
self.wrk_status[w.id] = WorkerStatus.READY
@@ -829,7 +829,7 @@ def process_cfg(self) -> None:
829829
if JobConfig.is_identical(self._cur_cfg, self._new_cfg):
830830
raise InvalidConfig("current and new configs are identical")
831831

832-
self._init_wrk_status()
832+
self._reconcile_wrk_status(self._cur_cfg, self._new_cfg)
833833

834834
self._new_cfg.reqgen_config = self.ctrl.reqgen_config
835835

0 commit comments

Comments
 (0)