opengris-scaler 1.12.37__cp38-cp38-musllinux_1_2_x86_64.whl

scaler/scheduler/allocate_policy/capability_allocate_policy.py

@@ -0,0 +1,280 @@
+import dataclasses
+import logging
+import typing
+from collections import OrderedDict, defaultdict
+from itertools import takewhile
+from typing import Dict, Iterable, List, Optional, Set
+
+from sortedcontainers import SortedList
+
+from scaler.protocol.python.message import Task
+from scaler.scheduler.allocate_policy.mixins import TaskAllocatePolicy
+from scaler.utility.identifiers import TaskID, WorkerID
+
+
+@dataclasses.dataclass(frozen=True)
+class _TaskHolder:
+    task_id: TaskID = dataclasses.field()
+    capabilities: Set[str] = dataclasses.field()
+
+
+@dataclasses.dataclass(frozen=True)
+class _WorkerHolder:
+    worker_id: WorkerID = dataclasses.field()
+
+    capabilities: Set[str] = dataclasses.field()
+    queue_size: int = dataclasses.field()
+
+    # Queued tasks, ordered from oldest to youngest tasks.
+    task_id_to_task: typing.OrderedDict[TaskID, _TaskHolder] = dataclasses.field(default_factory=OrderedDict)
+
+    def n_tasks(self) -> int:
+        return len(self.task_id_to_task)
+
+    def n_free(self) -> int:
+        return self.queue_size - self.n_tasks()
+
+    def copy(self) -> "_WorkerHolder":
+        return _WorkerHolder(self.worker_id, self.capabilities, self.queue_size, self.task_id_to_task.copy())
+
+
+class CapabilityAllocatePolicy(TaskAllocatePolicy):
+    """
+    This allocator policy assigns the tasks to workers supporting the requested task capabilities, trying to make all
+    workers load as equal as possible.
+    """
+
+    def __init__(self):
+        self._worker_id_to_worker: Dict[WorkerID, _WorkerHolder] = {}
+
+        self._task_id_to_worker_id: Dict[TaskID, WorkerID] = {}
+        self._capability_to_worker_ids: Dict[str, Set[WorkerID]] = {}
+
+    def add_worker(self, worker: WorkerID, capabilities: Dict[str, int], queue_size: int) -> bool:
+        if any(capability_value != -1 for capability_value in capabilities.values()):
+            logging.warning(f"allocate policy ignores non-infinite worker capabilities: {capabilities!r}.")
+
+        if worker in self._worker_id_to_worker:
+            return False
+
+        worker_holder = _WorkerHolder(worker_id=worker, capabilities=set(capabilities.keys()), queue_size=queue_size)
+        self._worker_id_to_worker[worker] = worker_holder
+
+        for capability in worker_holder.capabilities:
+            if capability not in self._capability_to_worker_ids:
+                self._capability_to_worker_ids[capability] = set()
+
+            self._capability_to_worker_ids[capability].add(worker)
+
+        return True
+
+    def remove_worker(self, worker: WorkerID) -> List[TaskID]:
+        worker_holder = self._worker_id_to_worker.pop(worker, None)
+
+        if worker_holder is None:
+            return []
+
+        for capability in worker_holder.capabilities:
+            self._capability_to_worker_ids[capability].discard(worker)
+            if len(self._capability_to_worker_ids[capability]) == 0:
+                self._capability_to_worker_ids.pop(capability)
+
+        task_ids = list(worker_holder.task_id_to_task.keys())
+        for task_id in task_ids:
+            self._task_id_to_worker_id.pop(task_id)
+
+        return task_ids
+
+    def get_worker_ids(self) -> Set[WorkerID]:
+        return set(self._worker_id_to_worker.keys())
+
+    def get_worker_by_task_id(self, task_id: TaskID) -> WorkerID:
+        return self._task_id_to_worker_id.get(task_id, WorkerID.invalid_worker_id())
+
+    def balance(self) -> Dict[WorkerID, List[TaskID]]:
+        """Returns, for every worker id, the list of task ids to balance out."""
+
+        has_idle_workers = any(worker.n_tasks() == 0 for worker in self._worker_id_to_worker.values())
+
+        if not has_idle_workers:
+            return {}
+
+        # The balancing algorithm works by trying to move tasks from workers that have more queued tasks than the
+        # average (high-load workers) to workers that have less tasks than the average (low-load workers).
+        #
+        # The overall worst-case time complexity of the balancing algorithm is:
+        #
+        #     O(n_workers * log(n_workers) + n_tasks * n_workers * n_capabilities)
+        #
+        # However, if the cluster does not use any capability, time complexity is always:
+        #
+        #     O(n_workers * log(n_workers) + n_tasks * log(n_workers))
+        #
+        # If capability constraints are used, this might result in less than optimal balancing. That's because, in some
+        # cases, the optimal balancing might require to move tasks between more than two workers. Consider this
+        # cluster's state:
+        #
+        #   Worker 1
+        #       Supported capabilities: {Linux, GPU}
+        #       Tasks:
+        #           Task 1: {Linux}
+        #
+        #   Worker 2
+        #       Supported capabilities: {Linux}
+        #       Tasks: None
+        #
+        #   Worker 3:
+        #       Supported capabilities: {GPU}
+        #       Tasks:
+        #           Task 2: {GPU}
+        #           Task 3: {GPU}
+        #
+        # Here, the algorithm will not be able to rebalance the cluster, while ideally we could move Task 1 to Worker 2
+        # and then Task 3 to Worker 1.
+        #
+        # Balancing algorithms that can find this optimal balancing exist (assignment problem), but these are complex
+        # and slow. These might also cause a lot of messages to be propagated through the cluster.
+        #
+        # See <https://github.com/finos/opengris-scaler/issues/32#issuecomment-2541897645> for more details.
+
+        n_tasks = sum(worker.n_tasks() for worker in self._worker_id_to_worker.values())
+        avg_tasks_per_worker = n_tasks / len(self._worker_id_to_worker)
+
+        def is_balanced(worker: _WorkerHolder) -> bool:
+            return abs(worker.n_tasks() - avg_tasks_per_worker) < 1
+
+        # First, we create a copy of the current workers objects so that we can modify their respective task queues.
+        # We also filter out workers that are already balanced as we will not touch these.
+        #
+        # Time complexity is O(n_workers + n_tasks)
+
+        workers = [worker.copy() for worker in self._worker_id_to_worker.values() if not is_balanced(worker)]
+
+        # Then, we sort the remaining workers by the number of queued tasks.
+        #
+        # Time complexity is O(n_workers * log(n_workers))
+
+        sorted_workers: SortedList[_WorkerHolder] = SortedList(workers, key=lambda worker: worker.n_tasks())
+
+        # Finally, we repeatedly remove one task from the most loaded worker until either:
+        #
+        # - all workers are balanced;
+        # - we cannot find a low-load worker than can accept tasks from a high-load worker.
+        #
+        # Worst-case time complexity is O(n_tasks * n_workers * n_capabilities).
+        # If no tag is used in the cluster, complexity is always O(n_tasks * log(n_workers))
+
+        balancing_advice: Dict[WorkerID, List[TaskID]] = defaultdict(list)
+        unbalanceable_tasks: Set[bytes] = set()
+
+        while len(sorted_workers) >= 2:
+            most_loaded_worker: _WorkerHolder = sorted_workers.pop(-1)
+
+            if is_balanced(most_loaded_worker):
+                # Most loaded worker is not high-load, stop
+                break
+
+            # Go through all of the most loaded worker's tasks, trying to find a low-load worker that can accept it.
+
+            receiving_worker: Optional[_WorkerHolder] = None
+            moved_task: Optional[_TaskHolder] = None
+
+            for task in reversed(most_loaded_worker.task_id_to_task.values()):  # Try to balance youngest tasks first.
+                if task.task_id in unbalanceable_tasks:
+                    continue
+
+                worker_candidates = takewhile(lambda worker: worker.n_tasks() < avg_tasks_per_worker, sorted_workers)
+
+                receiving_worker_index = self.__balance_try_reassign_task(task, worker_candidates)
+
+                if receiving_worker_index is not None:
+                    receiving_worker = sorted_workers.pop(receiving_worker_index)
+                    moved_task = task
+                    break
+                else:
+                    # We could not find a receiving worker for this task, remember the task as unbalanceable in case the
+                    # worker pops-up again. This greatly reduces the worst-case big-O complexity of the algorithm.
+                    unbalanceable_tasks.add(task.task_id)
+
+            # Re-inserts the workers in the sorted list if these can be balanced more.
+
+            if moved_task is not None:
+                assert receiving_worker is not None
+
+                balancing_advice[most_loaded_worker.worker_id].append(moved_task.task_id)
+
+                most_loaded_worker.task_id_to_task.pop(moved_task.task_id)
+                receiving_worker.task_id_to_task[moved_task.task_id] = moved_task
+
+                if not is_balanced(most_loaded_worker):
+                    sorted_workers.add(most_loaded_worker)
+
+                if not is_balanced(receiving_worker):
+                    sorted_workers.add(receiving_worker)
+
+        return balancing_advice
+
+    @staticmethod
+    def __balance_try_reassign_task(task: _TaskHolder, worker_candidates: Iterable[_WorkerHolder]) -> Optional[int]:
+        """Returns the index of the first worker that can accept the task."""
+
+        # Time complexity is O(n_workers * len(task.capabilities))
+
+        for worker_index, worker in enumerate(worker_candidates):
+            if task.capabilities.issubset(worker.capabilities):
+                return worker_index
+
+        return None
+
+    def assign_task(self, task: Task) -> WorkerID:
+        # Worst-case time complexity is O(n_workers * len(task.capabilities))
+
+        available_workers = self.__get_available_workers_for_capabilities(task.capabilities)
+
+        if len(available_workers) == 0:
+            return WorkerID.invalid_worker_id()
+
+        # Selects the worker that has the least amount of queued tasks. We could select the worker that has the most
+        # free queue task slots, but that might needlessly idle workers that have a smaller queue.
+
+        min_loaded_worker = min(available_workers, key=lambda worker: worker.n_tasks())
+        min_loaded_worker.task_id_to_task[task.task_id] = _TaskHolder(task.task_id, set(task.capabilities.keys()))
+
+        self._task_id_to_worker_id[task.task_id] = min_loaded_worker.worker_id
+
+        return min_loaded_worker.worker_id
+
+    def remove_task(self, task_id: TaskID) -> WorkerID:
+        worker_id = self._task_id_to_worker_id.pop(task_id, None)
+
+        if worker_id is None:
+            return WorkerID.invalid_worker_id()
+
+        worker = self._worker_id_to_worker[worker_id]
+        worker.task_id_to_task.pop(task_id)
+
+        return worker_id
+
+    def has_available_worker(self, capabilities: Optional[Dict[str, int]] = None) -> bool:
+        return len(self.__get_available_workers_for_capabilities(capabilities or {})) > 0
+
+    def statistics(self) -> Dict:
+        return {
+            worker.worker_id: {"free": worker.n_free(), "sent": worker.n_tasks(), "capabilities": worker.capabilities}
+            for worker in self._worker_id_to_worker.values()
+        }
+
+    def __get_available_workers_for_capabilities(self, capabilities: Dict[str, int]) -> List[_WorkerHolder]:
+        # Worst-case time complexity is O(n_workers * len(capabilities))
+
+        if any(capability not in self._capability_to_worker_ids for capability in capabilities.keys()):
+            return []
+
+        matching_worker_ids = set(self._worker_id_to_worker.keys())
+
+        for capability in capabilities.keys():
+            matching_worker_ids.intersection_update(self._capability_to_worker_ids[capability])
+
+        matching_workers = [self._worker_id_to_worker[worker_id] for worker_id in matching_worker_ids]
+
+        return [worker for worker in matching_workers if worker.n_free() > 0]

scaler/scheduler/allocate_policy/even_load_allocate_policy.py

@@ -0,0 +1,159 @@
+import logging
+import math
+from typing import Dict, List, Optional, Set
+
+from scaler.protocol.python.message import Task
+from scaler.scheduler.allocate_policy.mixins import TaskAllocatePolicy
+from scaler.utility.identifiers import TaskID, WorkerID
+from scaler.utility.queues.async_priority_queue import AsyncPriorityQueue
+from scaler.utility.queues.indexed_queue import IndexedQueue
+
+
+class EvenLoadAllocatePolicy(TaskAllocatePolicy):
+    """This Allocator policy is trying to make all workers load as equal as possible"""
+
+    def __init__(self):
+        self._workers_to_queue_size: Dict[bytes, int] = dict()
+        self._workers_to_task_ids: Dict[WorkerID, IndexedQueue] = dict()
+        self._task_id_to_worker: Dict[TaskID, WorkerID] = {}
+
+        self._worker_queue: AsyncPriorityQueue = AsyncPriorityQueue()
+
+    def add_worker(self, worker: WorkerID, capabilities: Dict[str, int], queue_size: int) -> bool:
+        if len(capabilities) > 0:
+            logging.warning(f"allocate policy ignores worker capabilities: {capabilities!r}.")
+
+        # TODO: handle uneven queue size for each worker
+        if worker in self._workers_to_task_ids:
+            return False
+
+        self._workers_to_task_ids[worker] = IndexedQueue()
+        self._workers_to_queue_size[worker] = queue_size
+
+        self._worker_queue.put_nowait([0, worker])
+        return True
+
+    def remove_worker(self, worker: WorkerID) -> List[TaskID]:
+        if worker not in self._workers_to_task_ids:
+            return []
+
+        self._worker_queue.remove(worker)
+
+        task_ids = list(self._workers_to_task_ids.pop(worker))
+        for task_id in task_ids:
+            self._task_id_to_worker.pop(task_id)
+        return task_ids
+
+    def get_worker_ids(self) -> Set[WorkerID]:
+        return set(self._workers_to_task_ids.keys())
+
+    def get_worker_by_task_id(self, task_id: TaskID) -> WorkerID:
+        return self._task_id_to_worker.get(task_id, WorkerID.invalid_worker_id())
+
+    def balance(self) -> Dict[WorkerID, List[TaskID]]:
+        """Returns, for every worker, the list of tasks to balance out."""
+
+        # TODO: handle uneven queue size for each worker
+        balance_count = self.__get_balance_count_by_worker()
+
+        balance_result = {}
+
+        for worker, count in balance_count.items():
+            if count == 0:
+                continue
+
+            tasks = list(self._workers_to_task_ids[worker])
+            balance_result[worker] = tasks[-count:]  # balance out the most recently queued tasks
+
+        return balance_result
+
+    def __get_balance_count_by_worker(self) -> Dict[WorkerID, int]:
+        """Returns, for every worker, the number of tasks to balance out."""
+
+        queued_tasks_per_worker = {
+            worker: max(0, len(tasks) - 1) for worker, tasks in self._workers_to_task_ids.items()
+        }
+
+        any_worker_has_queued_task = any(queued_tasks_per_worker.values())
+
+        if not any_worker_has_queued_task:
+            return {}
+
+        number_of_idle_workers = sum(1 for tasks in self._workers_to_task_ids.values() if len(tasks) == 0)
+
+        if number_of_idle_workers == 0:
+            return {}
+
+        mean_queued = math.ceil(sum(queued_tasks_per_worker.values()) / len(queued_tasks_per_worker))
+
+        balance_count = {worker: max(0, count - mean_queued) for worker, count in queued_tasks_per_worker.items()}
+
+        over_mean_advice_total = sum(balance_count.values())
+        minimal_allocate = min(number_of_idle_workers, sum(queued_tasks_per_worker.values()))
+
+        if over_mean_advice_total >= minimal_allocate:
+            return balance_count
+
+        total_to_be_balance = minimal_allocate - over_mean_advice_total
+        for worker, count in queued_tasks_per_worker.items():
+            assert total_to_be_balance >= 0, "total_to_be_balance must be positive"
+            if total_to_be_balance == 0:
+                break
+
+            leftover = count - balance_count[worker]
+            if leftover < 1:
+                continue
+
+            to_to_balance = min(leftover, total_to_be_balance)
+            balance_count[worker] += to_to_balance
+            total_to_be_balance -= to_to_balance
+
+        return balance_count
+
+    def assign_task(self, task: Task) -> WorkerID:
+        if len(task.capabilities) > 0:
+            logging.warning(f"allocate policy ignores task capabilities: {task.capabilities!r}.")
+
+        task_id = task.task_id
+
+        if task_id in self._task_id_to_worker:
+            return self._task_id_to_worker[task_id]
+
+        if self._worker_queue.empty():
+            return WorkerID.invalid_worker_id()
+
+        count, worker = self._worker_queue.get_nowait()
+        if count == self._workers_to_queue_size[worker]:
+            self._worker_queue.put_nowait([count, worker])
+            return WorkerID.invalid_worker_id()
+
+        self._workers_to_task_ids[worker].put(task_id)
+        self._task_id_to_worker[task_id] = worker
+        self._worker_queue.put_nowait([count + 1, worker])
+        return worker
+
+    def remove_task(self, task_id: TaskID) -> WorkerID:
+        if task_id not in self._task_id_to_worker:
+            return WorkerID.invalid_worker_id()
+
+        worker = self._task_id_to_worker.pop(task_id)
+        self._workers_to_task_ids[worker].remove(task_id)
+
+        self._worker_queue.decrease_priority(worker)
+        return worker
+
+    def has_available_worker(self, capabilities: Optional[Dict[str, int]] = None) -> bool:
+        if not len(self._worker_queue):
+            return False
+
+        count, worker = self._worker_queue.max_priority_item()
+        if count == self._workers_to_queue_size[worker]:
+            return False
+
+        return True
+
+    def statistics(self) -> Dict:
+        return {
+            worker: {"free": self._workers_to_queue_size[worker] - len(tasks), "sent": len(tasks)}
+            for worker, tasks in self._workers_to_task_ids.items()
+        }

scaler/scheduler/allocate_policy/mixins.py

@@ -0,0 +1,55 @@
+import abc
+from typing import Dict, List, Optional, Set
+
+from scaler.protocol.python.message import Task
+from scaler.utility.identifiers import TaskID, WorkerID
+
+
+class TaskAllocatePolicy(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def add_worker(self, worker: WorkerID, capabilities: Dict[str, int], queue_size: int) -> bool:
+        """add worker to worker collection"""
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def remove_worker(self, worker: WorkerID) -> List[TaskID]:
+        """remove worker to worker collection, and return list of task_ids of removed worker"""
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_worker_ids(self) -> Set[WorkerID]:
+        """get all worker ids as list"""
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_worker_by_task_id(self, task_id: TaskID) -> WorkerID:
+        """get worker that been assigned to this task_id, return an invalid worker ID if it cannot find the worker
+        assigned to this task id"""
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def balance(self) -> Dict[WorkerID, List[TaskID]]:
+        """balance worker, it should return list of task ids for over burdened worker, represented as worker
+        identity to list of task ids dictionary"""
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def assign_task(self, task: Task) -> WorkerID:
+        """assign task in allocator, return an invalid worker ID if available worker, otherwise will return worker been
+        assigned to"""
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def remove_task(self, task_id: TaskID) -> WorkerID:
+        """remove task in allocator, return an invalid worker ID if it did not find any worker, otherwise will return
+        worker associate with the removed task_id"""
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def has_available_worker(self, capabilities: Optional[Dict[str, int]] = None) -> bool:
+        """has available worker or not, possibly constrained to the requested task capabilities"""
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def statistics(self) -> Dict:
+        raise NotImplementedError()

scaler/scheduler/controllers/balance_controller.py

@@ -0,0 +1,65 @@
+import logging
+from typing import Dict, List, Optional
+
+from scaler.io.mixins import AsyncBinder, AsyncConnector
+from scaler.protocol.python.message import StateBalanceAdvice
+from scaler.scheduler.allocate_policy.mixins import TaskAllocatePolicy
+from scaler.scheduler.controllers.config_controller import VanillaConfigController
+from scaler.scheduler.controllers.mixins import TaskController
+from scaler.utility.identifiers import TaskID, WorkerID
+from scaler.utility.mixins import Looper
+
+
+class VanillaBalanceController(Looper):
+    def __init__(self, config_controller: VanillaConfigController, task_allocate_policy: TaskAllocatePolicy):
+        self._config_controller = config_controller
+
+        self._task_allocate_policy = task_allocate_policy
+
+        self._last_balance_advice: Dict[WorkerID, List[TaskID]] = dict()
+        self._same_load_balance_advice_count = 0
+
+        self._binder: Optional[AsyncBinder] = None
+        self._binder_monitor: Optional[AsyncConnector] = None
+
+        self._task_controller: Optional[TaskController] = None
+
+    def register(self, binder: AsyncBinder, binder_monitor: AsyncConnector, task_controller: TaskController):
+        self._binder = binder
+        self._binder_monitor = binder_monitor
+
+        self._task_controller = task_controller
+
+    async def routine(self):
+        current_advice = self._task_allocate_policy.balance()
+        if not self.__should_balance(current_advice):
+            return
+
+        worker_to_num_tasks = {worker: len(task_ids) for worker, task_ids in current_advice.items()}
+        logging.info(f"balancing task: {worker_to_num_tasks}")
+        for worker, task_ids in current_advice.items():
+            await self._binder_monitor.send(StateBalanceAdvice.new_msg(worker, task_ids))
+
+        self._last_balance_advice = current_advice
+        for worker, task_ids in current_advice.items():
+            for task_id in task_ids:
+                await self._task_controller.on_task_balance_cancel(task_id)
+
+    def __should_balance(self, current_advice: Dict[WorkerID, List[TaskID]]) -> bool:
+        # 1. if this is the same advise as last time, then we +1 on same advice count
+        # 2. if there is another different advice come in, then we reset same advice count to 0
+        if self._last_balance_advice == current_advice:
+            self._same_load_balance_advice_count += 1
+        else:
+            self._last_balance_advice = current_advice
+            self._same_load_balance_advice_count = 0
+
+        # if we have same advice for more than trigger times, then we start doing the balancing
+        if 0 < self._same_load_balance_advice_count < self._config_controller.get_config("load_balance_trigger_times"):
+            return False
+
+        # if current advice is empty, then we skip
+        if not current_advice:
+            return False
+
+        return True

scaler/scheduler/controllers/client_controller.py

@@ -0,0 +1,131 @@
+import logging
+import time
+from typing import Dict, Optional, Set, Tuple
+
+from scaler.io.mixins import AsyncBinder, AsyncConnector
+from scaler.protocol.python.message import (
+    ClientDisconnect,
+    ClientHeartbeat,
+    ClientHeartbeatEcho,
+    ClientShutdownResponse,
+    TaskCancel,
+)
+from scaler.protocol.python.status import ClientManagerStatus
+from scaler.scheduler.controllers.config_controller import VanillaConfigController
+from scaler.scheduler.controllers.mixins import ClientController, ObjectController, TaskController, WorkerController
+from scaler.utility.exceptions import ClientShutdownException
+from scaler.utility.identifiers import ClientID, TaskID
+from scaler.utility.mixins import Looper, Reporter
+from scaler.utility.one_to_many_dict import OneToManyDict
+
+
+class VanillaClientController(ClientController, Looper, Reporter):
+    def __init__(self, config_controller: VanillaConfigController):
+        self._config_controller = config_controller
+
+        self._client_to_task_ids: OneToManyDict[ClientID, TaskID] = OneToManyDict()
+
+        self._binder: Optional[AsyncBinder] = None
+        self._binder_monitor: Optional[AsyncConnector] = None
+        self._object_controller: Optional[ObjectController] = None
+        self._task_controller: Optional[TaskController] = None
+        self._worker_controller: Optional[WorkerController] = None
+
+        self._client_last_seen: Dict[ClientID, Tuple[float, ClientHeartbeat]] = dict()
+
+    def register(
+        self,
+        binder: AsyncBinder,
+        binder_monitor: AsyncConnector,
+        object_controller: ObjectController,
+        task_controller: TaskController,
+        worker_controller: WorkerController,
+    ):
+        self._binder = binder
+        self._binder_monitor = binder_monitor
+        self._object_controller = object_controller
+        self._task_controller = task_controller
+        self._worker_controller = worker_controller
+
+    def get_client_task_ids(self, client_id: ClientID) -> Set[TaskID]:
+        return self._client_to_task_ids.get_values(client_id)
+
+    def has_client_id(self, client_id: ClientID) -> bool:
+        return client_id in self._client_last_seen
+
+    def get_client_id(self, task_id: TaskID) -> Optional[ClientID]:
+        return self._client_to_task_ids.get_key(task_id)
+
+    def on_task_begin(self, client_id: ClientID, task_id: TaskID):
+        self._client_to_task_ids.add(client_id, task_id)
+
+    def on_task_finish(self, task_id: TaskID) -> ClientID:
+        return self._client_to_task_ids.remove_value(task_id)
+
+    async def on_heartbeat(self, client_id: ClientID, info: ClientHeartbeat):
+        await self._binder.send(
+            client_id,
+            ClientHeartbeatEcho.new_msg(
+                object_storage_address=self._config_controller.get_config("object_storage_address")
+            ),
+        )
+        if client_id not in self._client_last_seen:
+            logging.info(f"{client_id!r} connected")
+
+        self._client_last_seen[client_id] = (time.time(), info)
+
+    async def on_client_disconnect(self, client_id: ClientID, request: ClientDisconnect):
+        if request.disconnect_type == ClientDisconnect.DisconnectType.Disconnect:
+            await self.__on_client_disconnect(client_id)
+            return
+
+        if self._config_controller.get_config("protected"):
+            logging.warning("cannot shutdown clusters as scheduler is running in protected mode")
+            accepted = False
+        else:
+            logging.info(f"shutdown scheduler and all clusters as received signal from {client_id!r}")
+            accepted = True
+
+        await self._binder.send(client_id, ClientShutdownResponse.new_msg(accepted=accepted))
+
+        if self._config_controller.get_config("protected"):
+            return
+
+        await self._worker_controller.on_client_shutdown(client_id)
+
+        raise ClientShutdownException(f"received client shutdown from {client_id!r}, quitting")
+
+    async def routine(self):
+        await self.__routine_cleanup_clients()
+
+    def get_status(self) -> ClientManagerStatus:
+        return ClientManagerStatus.new_msg(
+            {client: len(task_ids) for client, task_ids in self._client_to_task_ids.items()}
+        )
+
+    async def __routine_cleanup_clients(self):
+        now = time.time()
+        dead_clients = {
+            client
+            for client, (last_seen, info) in self._client_last_seen.items()
+            if now - last_seen > self._config_controller.get_config("client_timeout_seconds")
+        }
+
+        for client in dead_clients:
+            await self.__on_client_disconnect(client)
+
+    async def __on_client_disconnect(self, client_id: ClientID):
+        logging.info(f"{client_id!r} disconnected")
+        if client_id in self._client_last_seen:
+            self._client_last_seen.pop(client_id)
+
+        await self.__cancel_client_all_tasks(client_id)
+        self._object_controller.clean_client(client_id)
+
+    async def __cancel_client_all_tasks(self, client_id: ClientID):
+        if client_id not in self._client_to_task_ids.keys():
+            return
+
+        tasks = self._client_to_task_ids.get_values(client_id).copy()
+        for task in tasks:
+            await self._task_controller.on_task_cancel(client_id, TaskCancel.new_msg(task))