avoid node migration overlaps for a pool

Author: piax93

clusterman/batch/node_migration.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import time
+from collections import defaultdict
+from multiprocessing import Lock
 from typing import Callable
 from typing import Collection
 from typing import Dict
@@ -45,9 +47,6 @@
 from clusterman.util import setup_logging
 
 
-WorkerProcessLabel = Union[str, MigrationEvent]
-
-
 class NodeMigration(BatchDaemon, BatchLoggingMixin, BatchRunningSentinelMixin):
     notify_emails = ["compute-infra@yelp.com"]
 
@@ -56,6 +55,10 @@ class NodeMigration(BatchDaemon, BatchLoggingMixin, BatchRunningSentinelMixin):
     DEFAULT_MAX_WORKER_PROCESSES = 6
     DEFAULT_RUN_INTERVAL_SECONDS = 60
 
+    WORKER_LABEL_SEPARATOR = ":"
+    EVENT_WORKER_LABEL_PREFIX = "event"
+    UPTIME_WORKER_LABEL_PREFIX = "uptime"
+
     @batch_command_line_arguments
     def parse_args(self, parser):
         arg_group = parser.add_argument_group("NodeMigration batch options")
@@ -66,10 +69,11 @@ def parse_args(self, parser):
     def configure_initial(self):
         setup_config(self.options)
         self.logger = colorlog.getLogger(__name__)
-        self.migration_workers: Dict[WorkerProcessLabel, RestartableDaemonProcess] = {}
+        self.migration_workers: Dict[str, RestartableDaemonProcess] = {}
         self.migration_configs = {}
-        self.events_in_progress = set()
+        self.events_in_progress = {}
         self.pools_accepting_events = set()
+        self.worker_locks = defaultdict(Lock)
         self.cluster_connector = KubernetesClusterConnector(self.options.cluster, None, init_crd=True)
         self.run_interval = staticconf.read_int(
             "batches.node_migration.run_interval_seconds", self.DEFAULT_RUN_INTERVAL_SECONDS
@@ -101,7 +105,26 @@ def _get_worker_setup(self, pool: str) -> Optional[WorkerSetup]:
             self.logger.exception(f"Bad migration configuration for pool {pool}: {e}")
         return None
 
-    def _spawn_worker(self, label: WorkerProcessLabel, routine: Callable, *args, **kwargs) -> bool:
+    def _build_worker_label(self, pool: Optional[str] = None, event: Optional[MigrationEvent] = None) -> str:
+        """Composes label for worker process
+
+        :param Optional[str] pool: pool name in case of uptime worker label
+        :param Optional[MigrationEvent] event: event instance in case of event worker label
+        :return: label string
+        """
+        if event:
+            prefix, cluster, pool = self.EVENT_WORKER_LABEL_PREFIX, event.cluster, event.pool
+        elif pool:
+            prefix, cluster = self.UPTIME_WORKER_LABEL_PREFIX, self.options.cluster
+        else:
+            raise ValueError("Either 'pool' or 'event' must be provided as parameter")
+        return self.WORKER_LABEL_SEPARATOR.join((prefix, cluster, pool))
+
+    def _is_event_worker_label(self, label: str) -> bool:
+        """Check if process label is for an event worker"""
+        return label.startswith(self.EVENT_WORKER_LABEL_PREFIX)
+
+    def _spawn_worker(self, label: str, routine: Callable, *args, **kwargs) -> bool:
         """Start worker process
 
         :param Callable routine: worker method
@@ -113,10 +136,12 @@ def _spawn_worker(self, label: WorkerProcessLabel, routine: Callable, *args, **k
             self.logger.warning(f"Worker labelled {label} already running, skipping")
             return False
         running_workers = sum(proc.is_alive() for proc in self.migration_workers.values())
-        if isinstance(label, MigrationEvent) and running_workers >= self.available_worker_slots:
+        if self._is_event_worker_label(label) and running_workers >= self.available_worker_slots:
             # uptime workers are prioritized skipping this check
             self.logger.warning(f"Too many worker processes running already ({running_workers}), skipping")
             return False
+        lock_key = label.split(self.WORKER_LABEL_SEPARATOR, 1)[1]
+        kwargs["pool_lock"] = self.worker_locks[lock_key]
         proc = RestartableDaemonProcess(target=routine, args=args, kwargs=kwargs)
         self.migration_workers[label] = proc
         proc.start()
@@ -128,7 +153,7 @@ def fetch_event_crd(self) -> Collection[MigrationEvent]:
         events = self.cluster_connector.list_node_migration_resources(
             [MigrationStatus.PENDING, MigrationStatus.INPROGRESS]
         )
-        return set(events) - self.events_in_progress
+        return set(events) - set(self.events_in_progress.values())
 
     def mark_event(self, event: MigrationEvent, status: MigrationStatus = MigrationStatus.COMPLETED) -> None:
         """Set status for CRD event resource
@@ -153,14 +178,15 @@ def spawn_event_worker(self, event: MigrationEvent):
             self.mark_event(event, MigrationStatus.SKIPPED)
             return
         self.logger.info(f"Spawning migration worker for event: {event}")
+        worker_label = self._build_worker_label(event=event)
         if self._spawn_worker(
-            label=event,
+            label=worker_label,
             routine=event_migration_worker,
             migration_event=event,
             worker_setup=worker_setup,
         ):
             self.mark_event(event, MigrationStatus.INPROGRESS)
-            self.events_in_progress.add(event)
+            self.events_in_progress[worker_label] = event
 
     def spawn_uptime_worker(self, pool: str, uptime: Union[int, str]):
         """Start process monitoring pool node uptime, and recycling nodes accordingly
@@ -180,7 +206,7 @@ def spawn_uptime_worker(self, pool: str, uptime: Union[int, str]):
             return
         self.logger.info(f"Spawning uptime migration worker for {pool} pool")
         self._spawn_worker(
-            label=f"uptime-{self.options.cluster}-{pool}",
+            label=self._build_worker_label(pool=pool),
             routine=uptime_migration_worker,
             cluster=self.options.cluster,
             pool=pool,
@@ -198,9 +224,9 @@ def monitor_workers(self):
                 else:
                     torestart.append(label)
         for label in completed:
-            if isinstance(label, MigrationEvent):
-                self.mark_event(label, MigrationStatus.COMPLETED)
-                self.events_in_progress.discard(label)
+            if self._is_event_worker_label(label):
+                event = self.events_in_progress.pop(label)
+                self.mark_event(event, MigrationStatus.COMPLETED)
             del self.migration_workers[label]
         for label in torestart:
             self.migration_workers[label].restart()

clusterman/migration/worker.py

@@ -14,6 +14,7 @@
 import time
 from functools import partial
 from multiprocessing import Process
+from multiprocessing.synchronize import Lock as LockBase
 from statistics import mean
 from typing import Callable
 from typing import cast
@@ -122,7 +123,9 @@ def _drain_node_selection(
     return True
 
 
-def uptime_migration_worker(cluster: str, pool: str, uptime_seconds: int, worker_setup: WorkerSetup) -> None:
+def uptime_migration_worker(
+    cluster: str, pool: str, uptime_seconds: int, worker_setup: WorkerSetup, pool_lock: LockBase
+) -> None:
     """Worker monitoring and migrating nodes according to uptime
 
     :parma str cluster: cluster name
@@ -137,24 +140,28 @@ def uptime_migration_worker(cluster: str, pool: str, uptime_seconds: int, worker
         return
     while True:
         if manager.is_capacity_satisfied():
-            _drain_node_selection(manager, node_selector, worker_setup)
+            with pool_lock:
+                _drain_node_selection(manager, node_selector, worker_setup)
         else:
             logger.warning(f"Pool {cluster}:{pool} is currently underprovisioned, skipping uptime migration iteration")
         time.sleep(UPTIME_CHECK_INTERVAL_SECONDS)
         manager.reload_state()
 
 
-def event_migration_worker(migration_event: MigrationEvent, worker_setup: WorkerSetup) -> None:
+def event_migration_worker(migration_event: MigrationEvent, worker_setup: WorkerSetup, pool_lock: LockBase) -> None:
     """Worker migrating nodes according to event configuration
 
     :param MigrationEvent migration_event: event instance
     :param WorkerSetup worker_setup: migration setup
     """
+    pool_lock_acquired = False
     manager = PoolManager(migration_event.cluster, migration_event.pool, SUPPORTED_POOL_SCHEDULER, fetch_state=False)
     connector = cast(KubernetesClusterConnector, manager.cluster_connector)
     connector.set_label_selectors(migration_event.label_selectors, add_to_existing=True)
     manager.reload_state()
     try:
+        pool_lock.acquire(timeout=worker_setup.expected_duration)
+        pool_lock_acquired = True
         if worker_setup.disable_autoscaling:
             logger.info(f"Disabling autoscaling for {migration_event.cluster}:{migration_event.pool}")
             disable_autoscaling(
@@ -181,6 +188,8 @@ def event_migration_worker(migration_event: MigrationEvent, worker_setup: Worker
         logger.error(f"Issue while processing migration event {migration_event}: {e}")
         raise
     finally:
+        if pool_lock_acquired:
+            pool_lock.release()
         # we do not reset the pool target capacity in case of pre-scaling as we
         # trust the autoscaler to readjust that in a short time eventually
         if worker_setup.disable_autoscaling:

tests/batch/node_migration_test.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from argparse import Namespace
+from collections import defaultdict
 from unittest.mock import MagicMock
 from unittest.mock import patch
 
@@ -45,7 +46,7 @@ def migration_batch():
 
 def test_fetch_event_crd(migration_batch: NodeMigration):
     migration_batch.events_in_progress = {
-        MigrationEvent(
+        "event:mesos-test:bar": MigrationEvent(
             resource_name="mesos-test-bar-220912-1",
             cluster="mesos-test",
             pool="bar",
@@ -113,26 +114,16 @@ def test_get_worker_setup(migration_batch):
     )
 
 
-@pytest.mark.parametrize(
-    "worker_label",
-    (
-        (
-            MigrationEvent(
-                resource_name="mesos-test-bar-220912-1",
-                cluster="mesos-test",
-                pool="bar",
-                label_selectors=[],
-                condition=MigrationCondition(ConditionTrait.KERNEL, ConditionOperator.GE, "3.2.1"),
-            ),
-        ),
-        ("foobar",),
-    ),
-)
 @patch("clusterman.batch.node_migration.RestartableDaemonProcess")
-def test_spawn_worker(mock_process, migration_batch, worker_label):
+def test_spawn_worker(mock_process, migration_batch):
+    mock_lock = MagicMock()
     mock_routine = MagicMock()
+    worker_label = "foobar:123:456"
+    migration_batch.worker_locks = defaultdict(mock_lock)
     assert migration_batch._spawn_worker(worker_label, mock_routine, 1, x=2) is True
-    mock_process.assert_called_once_with(target=mock_routine, args=(1,), kwargs={"x": 2})
+    mock_process.assert_called_once_with(
+        target=mock_routine, args=(1,), kwargs={"x": 2, "pool_lock": mock_lock.return_value}
+    )
     assert migration_batch.migration_workers == {worker_label: mock_process.return_value}
 
 
@@ -146,7 +137,7 @@ def test_spawn_worker_existing(mock_process, migration_batch):
 @patch("clusterman.batch.node_migration.RestartableDaemonProcess")
 def test_spawn_worker_over_capacity(mock_process, migration_batch):
     migration_batch.migration_workers = {f"foobar{i}": MagicMock(is_alive=lambda: True) for i in range(6)}
-    assert migration_batch._spawn_worker(MigrationEvent(None, None, None, None, None), MagicMock(), 1, x=2) is False
+    assert migration_batch._spawn_worker("event:foo:bar", MagicMock(), 1, x=2) is False
     mock_process.assert_not_called()
 
 
@@ -197,7 +188,7 @@ def test_spawn_event_worker(mock_worker_routine, migration_batch, event, worker_
         migration_batch.spawn_event_worker(event)
         if is_spawned:
             mock_spawn.assert_called_once_with(
-                label=event,
+                label=f"event:{event.cluster}:{event.pool}",
                 routine=mock_worker_routine,
                 migration_event=event,
                 worker_setup=worker_setup,
@@ -225,7 +216,7 @@ def test_spawn_uptime_worker(mock_worker_routine, migration_batch, uptime, worke
         migration_batch.spawn_uptime_worker("bar", uptime)
         if expected_uptime_spawn:
             mock_spawn.assert_called_once_with(
-                label="uptime-mesos-test-bar",
+                label="uptime:mesos-test:bar",
                 routine=mock_worker_routine,
                 cluster="mesos-test",
                 pool="bar",
@@ -244,10 +235,12 @@ def test_monitor_workers(migration_batch):
         "foobar": mock_ok_worker,
         "buzz": mock_to_restart,
         "some": MagicMock(is_alive=lambda: False, exitcode=0),
-        mock_event: MagicMock(is_alive=lambda: False, exitcode=0),
+        "event:123:456": MagicMock(is_alive=lambda: False, exitcode=0),
     }
+    migration_batch.events_in_progress = {"event:123:456": mock_event}
     with patch.object(migration_batch, "mark_event") as mock_mark:
         migration_batch.monitor_workers()
         mock_mark.assert_called_once_with(mock_event, MigrationStatus.COMPLETED)
     mock_to_restart.restart.assert_called_once_with()
     assert migration_batch.migration_workers == {"foobar": mock_ok_worker, "buzz": mock_to_restart}
+    assert not migration_batch.events_in_progress

tests/migration/migration_worker_test.py

@@ -128,7 +128,7 @@ def test_uptime_migration_worker(mock_drain_selection, mock_manager_class, mock_
     mock_manager = mock_manager_class.return_value
     mock_manager.is_capacity_satisfied.side_effect = [True, False, True]
     with pytest.raises(StopIteration):  # using end of mock side-effect to get out of forever looop
-        uptime_migration_worker("mesos-test", "bar", 10000, mock_setup)
+        uptime_migration_worker("mesos-test", "bar", 10000, mock_setup, pool_lock=MagicMock())
     assert mock_drain_selection.call_count == 2
     selector = mock_drain_selection.call_args_list[0][0][1]
     assert selector(ClusterNodeMetadata(None, InstanceMetadata(None, None, uptime=timedelta(seconds=10001)))) is True
@@ -160,7 +160,7 @@ def test_event_migration_worker(
         for i in range(1, 6)
     ]
     mock_manager.target_capacity = 19
-    event_migration_worker(mock_migration_event, event_worker_setup)
+    event_migration_worker(mock_migration_event, event_worker_setup, pool_lock=MagicMock())
     mock_manager.modify_target_capacity.assert_called_once_with(23)
     mock_disable_scaling.assert_called_once_with("mesos-test", "bar", "kubernetes", 3)
     mock_enable_scaling.assert_called_once_with("mesos-test", "bar", "kubernetes")
@@ -192,7 +192,7 @@ def test_event_migration_worker_error(
     mock_manager = mock_manager_class.return_value
     mock_manager.get_node_metadatas.side_effect = Exception(123)
     with pytest.raises(Exception):
-        event_migration_worker(mock_migration_event, event_worker_setup)
+        event_migration_worker(mock_migration_event, event_worker_setup, pool_lock=MagicMock())
     mock_disable_scaling.assert_called_once_with("mesos-test", "bar", "kubernetes", 3)
     mock_enable_scaling.assert_called_once_with("mesos-test", "bar", "kubernetes")
     mock_manager.modify_target_capacity.assert_not_called()