opengris-scaler 1.12.37__cp38-cp38-musllinux_1_2_x86_64.whl

scaler/scheduler/controllers/task_controller.py

@@ -0,0 +1,376 @@
+import asyncio
+import logging
+from collections import deque
+from typing import Any, Awaitable, Callable, Deque, Dict, List, Optional, Tuple
+
+from scaler.io.mixins import AsyncBinder, AsyncConnector
+from scaler.protocol.python.common import TaskCancelConfirmType, TaskResultType, TaskState, TaskTransition
+from scaler.protocol.python.message import StateTask, Task, TaskCancel, TaskCancelConfirm, TaskResult
+from scaler.protocol.python.status import TaskManagerStatus
+from scaler.scheduler.controllers.config_controller import VanillaConfigController
+from scaler.scheduler.controllers.mixins import (
+    ClientController,
+    GraphTaskController,
+    ObjectController,
+    TaskController,
+    WorkerController,
+)
+from scaler.scheduler.task.task_state_machine import TaskStateMachine
+from scaler.scheduler.task.task_state_manager import TaskStateManager
+from scaler.utility.identifiers import ClientID, TaskID, WorkerID
+from scaler.utility.mixins import Looper, Reporter
+
+
+class VanillaTaskController(TaskController, Looper, Reporter):
+    def __init__(self, config_controller: VanillaConfigController):
+        self._config_controller = config_controller
+        self._binder: Optional[AsyncBinder] = None
+        self._binder_monitor: Optional[AsyncConnector] = None
+
+        self._client_controller: Optional[ClientController] = None
+        self._object_controller: Optional[ObjectController] = None
+        self._worker_controller: Optional[WorkerController] = None
+
+        self._graph_controller: Optional[GraphTaskController] = None
+
+        self._task_id_to_task: Dict[TaskID, Task] = dict()
+        self._task_state_manager: TaskStateManager = TaskStateManager(debug=True)
+
+        self._unassigned: Deque[TaskID] = deque()  # type: ignore[misc]
+
+        self._state_functions: Dict[TaskState, Callable[[*Tuple[Any, ...]], Awaitable[None]]] = {
+            TaskState.Inactive: self.__state_inactive,  # type: ignore[dict-item]
+            TaskState.Running: self.__state_running,  # type: ignore[dict-item]
+            TaskState.Canceling: self.__state_canceling,  # type: ignore[dict-item]
+            TaskState.BalanceCanceling: self.__state_balance_canceling,  # type: ignore[dict-item]
+            TaskState.WorkerDisconnecting: self.__state_worker_disconnecting,  # type: ignore[dict-item]
+            TaskState.Canceled: self.__state_canceled,  # type: ignore[dict-item]
+            TaskState.CanceledNotFound: self.__state_canceled_not_found,  # type: ignore[dict-item]
+            TaskState.Success: self.__state_success,  # type: ignore[dict-item]
+            TaskState.Failed: self.__state_failed,  # type: ignore[dict-item]
+            TaskState.FailedWorkerDied: self.__state_failed_worker_died,  # type: ignore[dict-item]
+        }
+        self._task_result_transition_map = {
+            TaskResultType.Success: TaskTransition.TaskResultSuccess,
+            TaskResultType.Failed: TaskTransition.TaskResultFailed,
+            TaskResultType.FailedWorkerDied: TaskTransition.TaskResultWorkerDied,
+        }
+        self._task_cancel_confirm_transition_map = {
+            TaskCancelConfirmType.Canceled: TaskTransition.TaskCancelConfirmCanceled,
+            TaskCancelConfirmType.CancelNotFound: TaskTransition.TaskCancelConfirmNotFound,
+            TaskCancelConfirmType.CancelFailed: TaskTransition.TaskCancelConfirmFailed,
+        }
+
+    def register(
+        self,
+        binder: AsyncBinder,
+        binder_monitor: AsyncConnector,
+        client_controller: ClientController,
+        object_controller: ObjectController,
+        worker_controller: WorkerController,
+        graph_controller: GraphTaskController,
+    ):
+        self._binder = binder
+        self._binder_monitor = binder_monitor
+
+        self._client_controller = client_controller
+        self._object_controller = object_controller
+        self._worker_controller = worker_controller
+        self._graph_controller = graph_controller
+
+    async def routine(self):
+        # TODO: we don't need loop task anymore, but I will leave this routine API here in case we need in the future
+        pass
+
+    async def on_task_new(self, task: Task):
+        if self._task_state_manager.get_state_machine(task.task_id) is not None:
+            logging.error(
+                f"{task.task_id!r}: state machine already exists: "
+                f"{self._task_state_manager.get_state_machine(task.task_id)}"
+            )
+            return
+
+        state_machine = self._task_state_manager.add_state_machine(task.task_id)
+        await self.__state_inactive(task_id=task.task_id, state_machine=state_machine, task=task)
+
+    async def on_task_cancel(self, client_id: ClientID, task_cancel: TaskCancel):
+        state_machine = self._task_state_manager.get_state_machine(task_cancel.task_id)
+        if state_machine is None:
+            logging.error(f"{task_cancel.task_id!r}: task not exists while received TaskCancel, send TaskCancelConfirm")
+
+            task_cancel_confirm = TaskCancelConfirm.new_msg(task_cancel.task_id, TaskCancelConfirmType.CancelNotFound)
+
+            if self._graph_controller.is_graph_subtask(task_cancel.task_id):
+                await self._graph_controller.on_graph_sub_task_cancel_confirm(task_cancel_confirm)
+
+            await self._binder.send(client_id, task_cancel_confirm)
+            return
+
+        if state_machine.current_state() == TaskState.Inactive:
+            await self.__routing(
+                task_cancel.task_id,
+                TaskTransition.TaskCancel,
+                task_cancel_confirm=TaskCancelConfirm.new_msg(
+                    task_id=task_cancel.task_id, cancel_confirm_type=TaskCancelConfirmType.Canceled
+                ),
+            )
+            return
+
+        await self.__routing(task_cancel.task_id, TaskTransition.TaskCancel, client=client_id, task_cancel=task_cancel)
+
+    async def on_task_balance_cancel(self, task_id: TaskID):
+        await self.__routing(task_id, TaskTransition.BalanceTaskCancel)
+
+    async def on_task_cancel_confirm(self, task_cancel_confirm: TaskCancelConfirm):
+        transition = self._task_cancel_confirm_transition_map.get(task_cancel_confirm.cancel_confirm_type, None)
+        if transition is None:
+            raise ValueError(f"unknown TaskCancelConfirmType: {task_cancel_confirm.cancel_confirm_type}")
+
+        state_machine = self._task_state_manager.get_state_machine(task_cancel_confirm.task_id)
+        if state_machine is None:
+            logging.error(
+                f"{task_cancel_confirm.task_id!r}: task not exists while received TaskCancelTaskCancelConfirm, ignore"
+            )
+            return
+
+        current_state = state_machine.current_state()
+        if (
+            current_state == TaskState.BalanceCanceling
+            and task_cancel_confirm.cancel_confirm_type == TaskCancelConfirmType.Canceled
+        ):
+            # if balance cancel success
+            task = self._task_id_to_task[task_cancel_confirm.task_id]
+            await self.__routing(task_cancel_confirm.task_id, transition, task=task)
+            return
+
+        if (
+            current_state == TaskState.BalanceCanceling
+            and task_cancel_confirm.cancel_confirm_type == TaskCancelConfirmType.CancelFailed
+        ):
+            # balance cancel failed
+            worker_id = self._worker_controller.get_worker_by_task_id(task_cancel_confirm.task_id)
+            await self.__routing(task_cancel_confirm.task_id, transition, worker_id=worker_id)
+            return
+
+        await self.__routing(task_cancel_confirm.task_id, transition, task_cancel_confirm=task_cancel_confirm)
+
+    async def on_task_result(self, task_result: TaskResult):
+        transition = self._task_result_transition_map.get(task_result.result_type, None)
+        if transition is None:
+            raise ValueError(f"unknown TaskResultType: {task_result.result_type}")
+
+        await self.__routing(task_result.task_id, transition, task_result=task_result)
+
+    async def on_worker_connect(self, worker_id: WorkerID):
+        await self.__retry_unassignable()
+
+    async def on_worker_disconnect(self, task_id: TaskID, worker_id: WorkerID):
+        await self.__routing(task_id, TaskTransition.WorkerDisconnect, worker_id=worker_id)
+
+    def get_status(self) -> TaskManagerStatus:
+        return TaskManagerStatus.new_msg(state_to_count=self._task_state_manager.get_statistics())
+
+    async def __state_inactive(self, task_id: TaskID, state_machine: TaskStateMachine, task: Task):
+        assert task_id == task.task_id
+        assert state_machine.current_state() == TaskState.Inactive
+
+        self._client_controller.on_task_begin(task.source, task.task_id)
+        self._task_id_to_task[task.task_id] = task
+
+        worker_id = self._worker_controller.acquire_worker(self._task_id_to_task[task_id])
+        if not worker_id.is_valid():
+            # put task on hold until there is worker is added or task is finished/canceled (means have capacity)
+            self._unassigned.append(task_id)
+            return
+
+        await self.__routing(task_id, TaskTransition.HasCapacity, worker_id=worker_id)
+        await self.__send_monitor(task.task_id, self._object_controller.get_object_name(task.func_object_id))
+
+    async def __state_running(self, task_id: TaskID, state_machine: TaskStateMachine, worker_id: WorkerID):
+        if state_machine.previous_state() in {TaskState.Canceling, TaskState.BalanceCanceling}:
+            # if cancel failed (task is ongoing), we should wait here for the result
+            return
+
+        assert state_machine.current_state() == TaskState.Running
+
+        task = self._task_id_to_task[task_id]
+        await self._binder.send(worker_id, task)
+        await self.__send_monitor(task_id, self._object_controller.get_object_name(task.func_object_id))
+
+    async def __state_canceling(
+        self, task_id: TaskID, state_machine: TaskStateMachine, client: ClientID, task_cancel: TaskCancel
+    ):
+        assert task_id == task_cancel.task_id
+        assert state_machine.current_state() == TaskState.Canceling
+
+        if state_machine.previous_state() == TaskState.BalanceCanceling:
+            # we don't need to send another TaskCancel as it's already sent in previous state
+            return
+
+        # in case if task trying to cancel doesn't have task in scheduler, so we know which client we can send
+        # confirm to
+        self._client_controller.on_task_begin(client, task_id)
+
+        if task_id not in self._unassigned:
+            await self.__send_task_cancel_to_worker(task_cancel)
+            return
+
+        # task is either in unassigned or unassignable
+        await self.__routing(
+            task_id,
+            TaskTransition.TaskCancelConfirmCanceled,
+            task_cancel_confirm=TaskCancelConfirm.new_msg(task_id, TaskCancelConfirmType.Canceled),
+        )
+
+        if task_id in self._unassigned:
+            self._unassigned.remove(task_id)
+
+    async def __state_balance_canceling(self, task_id: TaskID, state_machine: TaskStateMachine):
+        assert state_machine.current_state() == TaskState.BalanceCanceling
+        await self.__send_task_cancel_to_worker(
+            TaskCancel.new_msg(task_id=task_id, flags=TaskCancel.TaskCancelFlags(force=False))
+        )
+
+    async def __state_worker_disconnecting(self, task_id: TaskID, state_machine: TaskStateMachine, worker_id: WorkerID):
+        assert isinstance(worker_id, WorkerID)
+        assert state_machine.current_state() == TaskState.WorkerDisconnecting
+
+        # this is where we decide to reroute or just send fail
+        task = self._task_id_to_task.get(task_id)
+        if task is None:
+            await self.__routing(
+                task_id,
+                TaskTransition.SchedulerHasNoTask,
+                task_result=TaskResult.new_msg(task_id, TaskResultType.FailedWorkerDied),
+            )
+        else:
+            await self.__routing(task_id, TaskTransition.SchedulerHasTask, task=task)
+
+    async def __state_canceled(
+        self, task_id: TaskID, state_machine: TaskStateMachine, task_cancel_confirm: TaskCancelConfirm
+    ):
+        assert task_id == task_cancel_confirm.task_id
+        assert task_cancel_confirm.cancel_confirm_type == TaskCancelConfirmType.Canceled
+        assert state_machine.current_state() == TaskState.Canceled
+
+        if task_cancel_confirm.task_id in self._unassigned:
+            # if task is not assigned to any worker, we don't need to deal with worker manager
+            self._unassigned.remove(task_cancel_confirm.task_id)
+        else:
+            await self._worker_controller.on_task_done(task_cancel_confirm.task_id)
+
+        await self.__send_task_cancel_confirm_to_client(task_cancel_confirm)
+
+    async def __state_canceled_not_found(
+        self, task_id: TaskID, state_machine: TaskStateMachine, task_cancel_confirm: TaskCancelConfirm
+    ):
+        assert task_id == task_cancel_confirm.task_id
+        assert task_cancel_confirm.cancel_confirm_type == TaskCancelConfirmType.CancelNotFound
+        assert state_machine.current_state() == TaskState.CanceledNotFound
+
+        await self.__send_task_cancel_confirm_to_client(task_cancel_confirm)
+
+    async def __state_success(self, task_id: TaskID, state_machine: TaskStateMachine, task_result: TaskResult):
+        assert task_id == task_result.task_id
+        assert state_machine.current_state() == TaskState.Success
+        await self.__send_task_result_to_client(task_result)
+
+    async def __state_failed(self, task_id: TaskID, state_machine: TaskStateMachine, task_result: TaskResult):
+        assert task_id == task_result.task_id
+        assert state_machine.current_state() == TaskState.Failed
+        await self.__send_task_result_to_client(task_result)
+
+    async def __state_failed_worker_died(
+        self, task_id: TaskID, state_machine: TaskStateMachine, task_result: TaskResult
+    ):
+        assert task_id == task_result.task_id
+        assert state_machine.current_state() == TaskState.FailedWorkerDied
+        await self.__send_task_result_to_client(task_result)
+
+    async def __send_task_cancel_to_worker(self, task_cancel: TaskCancel):
+        worker = await self._worker_controller.on_task_cancel(task_cancel)
+        if not worker:
+            logging.error(f"{task_cancel.task_id!r}: cannot find task in worker to cancel")
+            await self.__routing(
+                task_cancel.task_id,
+                TaskTransition.TaskCancelConfirmNotFound,
+                task_cancel_confirm=TaskCancelConfirm.new_msg(
+                    task_cancel.task_id, TaskCancelConfirmType.CancelNotFound
+                ),
+            )
+            return
+
+        await self._binder.send(worker, TaskCancel.new_msg(task_cancel.task_id))
+        await self.__send_monitor(task_cancel.task_id, b"")
+
+    async def __send_task_result_to_client(self, task_result: TaskResult):
+        await self._worker_controller.on_task_done(task_result.task_id)
+        client = self._client_controller.on_task_finish(task_result.task_id)
+        await self._binder.send(client, task_result)
+        await self.__send_monitor(task_result.task_id, b"", task_result.metadata)
+
+        self._task_state_manager.remove_state_machine(task_result.task_id)
+        self._task_id_to_task.pop(task_result.task_id)
+
+        if self._graph_controller.is_graph_subtask(task_result.task_id):
+            await self._graph_controller.on_graph_sub_task_result(task_result)
+
+        await self.__retry_unassignable()
+
+    async def __send_task_cancel_confirm_to_client(self, task_cancel_confirm: TaskCancelConfirm):
+        client = self._client_controller.on_task_finish(task_cancel_confirm.task_id)
+        await self._binder.send(client, task_cancel_confirm)
+        await self.__send_monitor(task_cancel_confirm.task_id, b"")
+        self._task_state_manager.remove_state_machine(task_cancel_confirm.task_id)
+        self._task_id_to_task.pop(task_cancel_confirm.task_id)
+
+        if self._graph_controller.is_graph_subtask(task_cancel_confirm.task_id):
+            await self._graph_controller.on_graph_sub_task_cancel_confirm(task_cancel_confirm)
+
+        await self.__retry_unassignable()
+
+    async def __send_monitor(self, task_id: TaskID, function_name: bytes, metadata: bytes = b""):
+        worker = self._worker_controller.get_worker_by_task_id(task_id)
+        task_state = self._task_state_manager.get_state_machine(task_id).current_state()
+        capabilities = self._task_id_to_task[task_id].capabilities if task_id in self._task_id_to_task else {}
+        await self._binder_monitor.send(
+            StateTask.new_msg(task_id, function_name, task_state, worker, capabilities, metadata)
+        )
+
+    async def __routing(self, task_id: TaskID, transition: TaskTransition, **kwargs):
+        state_machine = self._task_state_manager.on_transition(task_id, transition)
+        if state_machine is None:
+            logging.info(f"{task_id!r}: unknown transition: {transition}")
+            return
+
+        try:
+            await self._state_functions[state_machine.current_state()](task_id, state_machine, **kwargs)  # noqa
+        except Exception as e:
+            logging.exception(
+                f"{task_id!r}: exception happened, transition: {transition} path: {state_machine.get_path()}"
+            )
+            raise e
+
+    async def __retry_unassignable(self):
+        futures = [
+            self.__routing(task_id, TaskTransition.HasCapacity, worker_id=worker_id)
+            for task_id, worker_id in self.__acquire_workers()
+        ]
+
+        await asyncio.gather(*futures)
+
+    def __acquire_workers(self) -> List[Tuple[TaskID, WorkerID]]:
+        """please note this function has to be atomic, means no async decorated in order to make unassigned queue to be
+        synced, also this function should return as list not generator because of atomic
+        """
+
+        ready_to_assign = list()
+        while len(self._unassigned) > 0:
+            worker_id = self._worker_controller.acquire_worker(self._task_id_to_task[self._unassigned[0]])
+            if not worker_id.is_valid():
+                break
+
+            task_id = self._unassigned.popleft()
+            ready_to_assign.append((task_id, worker_id))
+
+        return ready_to_assign

scaler/scheduler/controllers/worker_controller.py

@@ -0,0 +1,169 @@
+import logging
+import time
+from typing import Dict, Optional, Set, Tuple
+
+from scaler.io.mixins import AsyncBinder, AsyncConnector
+from scaler.protocol.python.common import WorkerState
+from scaler.protocol.python.message import (
+    ClientDisconnect,
+    DisconnectRequest,
+    DisconnectResponse,
+    StateWorker,
+    Task,
+    TaskCancel,
+    WorkerHeartbeat,
+    WorkerHeartbeatEcho,
+)
+from scaler.protocol.python.status import ProcessorStatus, Resource, WorkerManagerStatus, WorkerStatus
+from scaler.scheduler.allocate_policy.mixins import TaskAllocatePolicy
+from scaler.scheduler.controllers.config_controller import VanillaConfigController
+from scaler.scheduler.controllers.mixins import TaskController, WorkerController
+from scaler.utility.identifiers import ClientID, TaskID, WorkerID
+from scaler.utility.mixins import Looper, Reporter
+
+UINT8_MAX = 2**8 - 1
+
+
+class VanillaWorkerController(WorkerController, Looper, Reporter):
+    def __init__(self, config_controller: VanillaConfigController, task_allocate_policy: TaskAllocatePolicy):
+        self._config_controller = config_controller
+
+        self._binder: Optional[AsyncBinder] = None
+        self._binder_monitor: Optional[AsyncConnector] = None
+        self._task_controller: Optional[TaskController] = None
+
+        self._worker_alive_since: Dict[WorkerID, Tuple[float, WorkerHeartbeat]] = dict()
+        self._allocator_policy = task_allocate_policy
+
+    def register(self, binder: AsyncBinder, binder_monitor: AsyncConnector, task_controller: TaskController):
+        self._binder = binder
+        self._binder_monitor = binder_monitor
+        self._task_controller = task_controller
+
+    def acquire_worker(self, task: Task) -> WorkerID:
+        return self._allocator_policy.assign_task(task)
+
+    async def on_task_cancel(self, task_cancel: TaskCancel):
+        worker = self._allocator_policy.remove_task(task_cancel.task_id)
+        if not worker.is_valid():
+            logging.error(f"cannot find task_id={task_cancel.task_id.hex()} in task workers")
+            return
+
+        await self._binder.send(worker, task_cancel)
+
+    async def on_task_done(self, task_id: TaskID) -> WorkerID:
+        worker = self._allocator_policy.remove_task(task_id)
+        if not worker.is_valid():
+            logging.error(f"Cannot find task in worker queue: task_id={task_id.hex()}")
+
+        return worker
+
+    async def on_heartbeat(self, worker_id: WorkerID, info: WorkerHeartbeat):
+        if self._allocator_policy.add_worker(worker_id, info.capabilities, info.queue_size):
+            logging.info(f"worker {worker_id!r} connected")
+            await self._binder_monitor.send(StateWorker.new_msg(worker_id, WorkerState.Connected, info.capabilities))
+            await self._task_controller.on_worker_connect(worker_id)
+
+        self._worker_alive_since[worker_id] = (time.time(), info)
+        await self._binder.send(
+            worker_id,
+            WorkerHeartbeatEcho.new_msg(
+                object_storage_address=self._config_controller.get_config("object_storage_address")
+            ),
+        )
+
+    async def on_client_shutdown(self, client_id: ClientID):
+        for worker in self._allocator_policy.get_worker_ids():
+            await self.__shutdown_worker(worker)
+
+    async def on_disconnect(self, worker_id: WorkerID, request: DisconnectRequest):
+        await self.__disconnect_worker(request.worker)
+        await self._binder.send(worker_id, DisconnectResponse.new_msg(request.worker))
+
+    async def routine(self):
+        await self.__clean_workers()
+
+    def get_status(self) -> WorkerManagerStatus:
+        worker_to_task_numbers = self._allocator_policy.statistics()
+        return WorkerManagerStatus.new_msg(
+            [
+                self.__worker_status_from_heartbeat(worker, worker_to_task_numbers[worker], last, info)
+                for worker, (last, info) in self._worker_alive_since.items()
+            ]
+        )
+
+    @staticmethod
+    def __worker_status_from_heartbeat(
+        worker_id: WorkerID, worker_task_numbers: Dict, last: float, info: WorkerHeartbeat
+    ) -> WorkerStatus:
+        current_processor = next((p for p in info.processors if not p.suspended), None)
+        suspended = min(len([p for p in info.processors if p.suspended]), UINT8_MAX)
+        last_s = min(int(time.time() - last), UINT8_MAX)
+
+        if current_processor:
+            debug_info = f"{int(current_processor.initialized)}{int(current_processor.has_task)}{int(info.task_lock)}"
+        else:
+            debug_info = f"00{int(info.task_lock)}"
+
+        return WorkerStatus.new_msg(
+            worker_id=worker_id,
+            agent=info.agent,
+            rss_free=info.rss_free,
+            free=worker_task_numbers["free"],
+            sent=worker_task_numbers["sent"],
+            queued=info.queued_tasks,
+            suspended=suspended,
+            lag_us=info.latency_us,
+            last_s=last_s,
+            itl=debug_info,
+            processor_statuses=[
+                ProcessorStatus.new_msg(
+                    pid=p.pid,
+                    initialized=p.initialized,
+                    has_task=p.has_task,
+                    suspended=p.suspended,
+                    resource=Resource.new_msg(p.resource.cpu, p.resource.rss),
+                )
+                for p in info.processors
+            ],
+        )
+
+    def has_available_worker(self) -> bool:
+        return self._allocator_policy.has_available_worker()
+
+    def get_worker_by_task_id(self, task_id: TaskID) -> WorkerID:
+        return self._allocator_policy.get_worker_by_task_id(task_id)
+
+    def get_worker_ids(self) -> Set[WorkerID]:
+        return self._allocator_policy.get_worker_ids()
+
+    async def __clean_workers(self):
+        now = time.time()
+        dead_workers = [
+            dead_worker
+            for dead_worker, (alive_since, info) in self._worker_alive_since.items()
+            if now - alive_since > self._config_controller.get_config("worker_timeout_seconds")
+        ]
+        for dead_worker in dead_workers:
+            await self.__disconnect_worker(dead_worker)
+
+    async def __disconnect_worker(self, worker_id: WorkerID):
+        """return True if disconnect worker success"""
+        if worker_id not in self._worker_alive_since:
+            return
+
+        logging.info(f"{worker_id!r} disconnected")
+        await self._binder_monitor.send(StateWorker.new_msg(worker_id, WorkerState.Disconnected, {}))
+        self._worker_alive_since.pop(worker_id)
+
+        task_ids = self._allocator_policy.remove_worker(worker_id)
+        if not task_ids:
+            return
+
+        logging.info(f"{len(task_ids)} task(s) failed due to worker {worker_id!r} disconnected")
+        for task_id in task_ids:
+            await self._task_controller.on_worker_disconnect(task_id, worker_id)
+
+    async def __shutdown_worker(self, worker_id: WorkerID):
+        await self._binder.send(worker_id, ClientDisconnect.new_msg(ClientDisconnect.DisconnectType.Shutdown))
+        await self.__disconnect_worker(worker_id)

scaler/scheduler/object_usage/object_tracker.py

@@ -0,0 +1,131 @@
+import abc
+from typing import Callable, Dict, Generator, Generic, Optional, Set, Tuple, TypeVar
+
+from scaler.utility.many_to_many_dict import ManyToManyDict
+
+BlockType = TypeVar("BlockType")
+ObjectKeyType = TypeVar("ObjectKeyType")
+
+
+class ObjectUsage(Generic[ObjectKeyType], metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def get_object_key(self) -> ObjectKeyType:
+        raise NotImplementedError()
+
+
+ObjectType = TypeVar("ObjectType", bound=ObjectUsage)
+
+
+class ObjectTracker(Generic[BlockType, ObjectKeyType, ObjectType]):
+    def __init__(self, prefix: str, callback: Callable[[ObjectType], None]):
+        self._prefix = prefix
+        self._callback = callback
+
+        self._current_blocks: Set[BlockType] = set()
+        self._object_key_to_block: ManyToManyDict[ObjectKeyType, BlockType] = ManyToManyDict()
+        self._object_key_to_object: Dict[ObjectKeyType, ObjectType] = dict()
+
+    def object_count(self):
+        return len(self._object_key_to_object)
+
+    def items(self):
+        return self._object_key_to_object.items()
+
+    def get_all_object_keys(self) -> Set[ObjectKeyType]:
+        return set(self._object_key_to_object.keys())
+
+    def has_object(self, key: ObjectKeyType) -> bool:
+        return key in self._object_key_to_object
+
+    def get_object(self, key: ObjectKeyType) -> ObjectType:
+        return self._object_key_to_object[key]
+
+    def add_object(self, obj: ObjectType):
+        self._object_key_to_object[obj.get_object_key()] = obj
+
+    def get_object_block_pairs(self, blocks: Set[BlockType]) -> Generator[Tuple[ObjectKeyType, BlockType], None, None]:
+        for block in blocks:
+            if not self._object_key_to_block.has_right_key(block):
+                continue
+
+            for object_key in self._object_key_to_block.get_left_items(block):
+                yield object_key, block
+
+    def add_blocks_for_one_object(self, object_key: ObjectKeyType, blocks: Set[BlockType]):
+        if object_key not in self._object_key_to_object:
+            raise KeyError(f"cannot find key={object_key} in ObjectTracker")
+
+        for block in blocks:
+            self._object_key_to_block.add(object_key, block)
+
+        self._current_blocks.update(blocks)
+
+    def remove_blocks_for_one_object(self, object_key: ObjectKeyType, blocks: Set[BlockType]):
+        ready_objects = []
+        for block in blocks:
+            obj = self.__remove_block_for_object(object_key, block)
+            if obj is None:
+                continue
+
+            ready_objects.append(obj)
+
+        for obj in ready_objects:
+            self._callback(obj)
+
+    def add_one_block_for_objects(self, object_keys: Set[ObjectKeyType], block: BlockType):
+        for object_key in object_keys:
+            if object_key not in self._object_key_to_object:
+                raise KeyError(f"cannot find key={object_key} in ObjectTracker")
+
+            self._object_key_to_block.add(object_key, block)
+
+        self._current_blocks.add(block)
+
+    def remove_one_block_for_objects(self, object_keys: Set[ObjectKeyType], block: BlockType):
+        ready_objects = []
+        for object_key in object_keys:
+            obj = self.__remove_block_for_object(object_key, block)
+            if obj is None:
+                continue
+
+            ready_objects.append(obj)
+
+        for obj in ready_objects:
+            self._callback(obj)
+
+    def remove_blocks(self, blocks: Set[BlockType]):
+        ready_objects = []
+        for block in blocks:
+            if not self._object_key_to_block.has_right_key(block):
+                continue
+
+            object_keys = self._object_key_to_block.get_left_items(block).copy()
+            for object_key in object_keys:
+                obj = self.__remove_block_for_object(object_key, block)
+                if obj is None:
+                    continue
+
+                ready_objects.append(obj)
+
+        for obj in ready_objects:
+            self._callback(obj)
+
+    def __remove_block_for_object(self, object_key: ObjectKeyType, block: BlockType) -> Optional[ObjectType]:
+        if block not in self._current_blocks:
+            return None
+
+        if object_key not in self._object_key_to_object:
+            return None
+
+        if not self._object_key_to_block.has_key_pair(object_key, block):
+            return None
+
+        self._object_key_to_block.remove(object_key, block)
+
+        if not self._object_key_to_block.has_right_key(block):
+            self._current_blocks.remove(block)
+
+        if self._object_key_to_block.has_left_key(object_key):
+            return None
+
+        return self._object_key_to_object.pop(object_key)