opengris-scaler 1.12.28__cp313-cp313-musllinux_1_2_x86_64.whl

scaler/scheduler/scheduler.py

@@ -0,0 +1,251 @@
+import asyncio
+import functools
+import logging
+
+import zmq.asyncio
+
+from scaler.config.defaults import CLEANUP_INTERVAL_SECONDS, STATUS_REPORT_INTERVAL_SECONDS
+from scaler.config.section.scheduler import SchedulerConfig
+from scaler.config.types.zmq import ZMQConfig, ZMQType
+from scaler.io.async_binder import ZMQAsyncBinder
+from scaler.io.async_connector import ZMQAsyncConnector
+from scaler.io.mixins import AsyncBinder, AsyncConnector, AsyncObjectStorageConnector
+from scaler.io.utility import create_async_object_storage_connector
+from scaler.protocol.python.common import ObjectStorageAddress
+from scaler.protocol.python.message import (
+    ClientDisconnect,
+    ClientHeartbeat,
+    DisconnectRequest,
+    GraphTask,
+    InformationRequest,
+    ObjectInstruction,
+    Task,
+    TaskCancel,
+    TaskCancelConfirm,
+    TaskLog,
+    TaskResult,
+    WorkerHeartbeat,
+)
+from scaler.protocol.python.mixins import Message
+from scaler.scheduler.controllers.balance_controller import VanillaBalanceController
+from scaler.scheduler.controllers.client_controller import VanillaClientController
+from scaler.scheduler.controllers.config_controller import VanillaConfigController
+from scaler.scheduler.controllers.graph_controller import VanillaGraphTaskController
+from scaler.scheduler.controllers.information_controller import VanillaInformationController
+from scaler.scheduler.controllers.object_controller import VanillaObjectController
+from scaler.scheduler.controllers.scaling_policies.utility import create_scaling_controller
+from scaler.scheduler.controllers.task_controller import VanillaTaskController
+from scaler.scheduler.controllers.worker_controller import VanillaWorkerController
+from scaler.utility.event_loop import create_async_loop_routine
+from scaler.utility.exceptions import ClientShutdownException
+from scaler.utility.identifiers import ClientID, WorkerID
+
+
+class Scheduler:
+    def __init__(self, config: SchedulerConfig):
+        self._config_controller = VanillaConfigController(config)
+
+        if config.scheduler_address.type != ZMQType.tcp:
+            raise TypeError(
+                f"{self.__class__.__name__}: scheduler address must be tcp type: \
+                    {config.scheduler_address.to_address()}"
+            )
+
+        if config.object_storage_address is None:
+            object_storage_address = ObjectStorageAddress.new_msg(
+                host=config.scheduler_address.host, port=config.scheduler_address.port + 1
+            )
+        else:
+            object_storage_address = ObjectStorageAddress.new_msg(
+                host=config.object_storage_address.host, port=config.object_storage_address.port
+            )
+        self._config_controller.update_config("object_storage_address", object_storage_address)
+
+        if config.monitor_address is None:
+            monitor_address = ZMQConfig(
+                type=ZMQType.tcp, host=config.scheduler_address.host, port=config.scheduler_address.port + 2
+            )
+        else:
+            monitor_address = config.monitor_address
+        self._config_controller.update_config("monitor_address", monitor_address)
+
+        self._context = zmq.asyncio.Context(io_threads=config.io_threads)
+
+        self._binder: AsyncBinder = ZMQAsyncBinder(
+            context=self._context, name="scheduler", address=config.scheduler_address
+        )
+        logging.info(f"{self.__class__.__name__}: listen to scheduler address {config.scheduler_address}")
+
+        self._connector_storage: AsyncObjectStorageConnector = create_async_object_storage_connector()
+        logging.info(f"{self.__class__.__name__}: connect to object storage server {object_storage_address!r}")
+
+        self._binder_monitor: AsyncConnector = ZMQAsyncConnector(
+            context=self._context,
+            name="scheduler_monitor",
+            socket_type=zmq.PUB,
+            address=monitor_address,
+            bind_or_connect="bind",
+            callback=None,
+            identity=None,
+        )
+        logging.info(f"{self.__class__.__name__}: listen to scheduler monitor address {monitor_address.to_address()}")
+
+        self._task_allocate_policy = config.allocate_policy.value()
+
+        self._client_manager = VanillaClientController(config_controller=self._config_controller)
+        self._object_controller = VanillaObjectController(config_controller=self._config_controller)
+        self._graph_controller = VanillaGraphTaskController(config_controller=self._config_controller)
+        self._task_controller = VanillaTaskController(config_controller=self._config_controller)
+        self._worker_controller = VanillaWorkerController(
+            config_controller=self._config_controller, task_allocate_policy=self._task_allocate_policy
+        )
+        self._balance_controller = VanillaBalanceController(
+            config_controller=self._config_controller, task_allocate_policy=self._task_allocate_policy
+        )
+        self._information_controller = VanillaInformationController(config_controller=self._config_controller)
+        self._scaling_controller = create_scaling_controller(
+            config.scaling_controller_strategy, config.adapter_webhook_urls
+        )
+
+        # register
+        self._binder.register(self.on_receive_message)
+        self._client_manager.register(
+            self._binder, self._binder_monitor, self._object_controller, self._task_controller, self._worker_controller
+        )
+        self._object_controller.register(
+            self._binder, self._binder_monitor, self._connector_storage, self._client_manager, self._worker_controller
+        )
+        self._graph_controller.register(
+            self._binder,
+            self._binder_monitor,
+            self._connector_storage,
+            self._client_manager,
+            self._task_controller,
+            self._object_controller,
+        )
+        self._task_controller.register(
+            self._binder,
+            self._binder_monitor,
+            self._client_manager,
+            self._object_controller,
+            self._worker_controller,
+            self._graph_controller,
+        )
+        self._worker_controller.register(self._binder, self._binder_monitor, self._task_controller)
+        self._balance_controller.register(self._binder, self._binder_monitor, self._task_controller)
+
+        self._information_controller.register_managers(
+            self._binder_monitor,
+            self._binder,
+            self._client_manager,
+            self._object_controller,
+            self._task_controller,
+            self._worker_controller,
+            self._scaling_controller,
+        )
+
+    async def connect_to_storage(self):
+        object_storage_address = self._config_controller.get_config("object_storage_address")
+        await self._connector_storage.connect(object_storage_address.host, object_storage_address.port)
+
+    async def on_receive_message(self, source: bytes, message: Message):
+        # =====================================================================================
+        # client manager
+        if isinstance(message, ClientHeartbeat):
+            await self._client_manager.on_heartbeat(ClientID(source), message)
+            return
+
+        # scheduler receives client shutdown request from upstream
+        if isinstance(message, ClientDisconnect):
+            await self._client_manager.on_client_disconnect(ClientID(source), message)
+            return
+
+        # =====================================================================================
+        # graph manager
+        if isinstance(message, GraphTask):
+            await self._graph_controller.on_graph_task(ClientID(source), message)
+            return
+
+        # =====================================================================================
+        # task manager
+        if isinstance(message, Task):
+            await self._task_controller.on_task_new(message)
+            return
+
+        if isinstance(message, TaskCancel):
+            if self._graph_controller.is_graph_subtask(message.task_id):
+                await self._graph_controller.on_graph_task_cancel(message)
+            else:
+                await self._task_controller.on_task_cancel(ClientID(source), message)
+            return
+
+        if isinstance(message, TaskCancelConfirm):
+            await self._task_controller.on_task_cancel_confirm(message)
+            return
+
+        if isinstance(message, TaskResult):
+            await self._task_controller.on_task_result(message)
+            return
+
+        if isinstance(message, TaskLog):
+            client = self._client_manager.get_client_id(message.task_id)
+            if client is not None:
+                await self._binder.send(client, message)
+            return
+
+        # =====================================================================================
+        # worker manager
+        if isinstance(message, WorkerHeartbeat):
+            await self._worker_controller.on_heartbeat(WorkerID(source), message)
+            return
+
+        # scheduler receives worker disconnect request from downstream
+        if isinstance(message, DisconnectRequest):
+            await self._worker_controller.on_disconnect(WorkerID(source), message)
+            return
+
+        # =====================================================================================
+        # object manager
+        if isinstance(message, ObjectInstruction):
+            await self._object_controller.on_object_instruction(source, message)
+            return
+
+        # =====================================================================================
+        # information manager
+        if isinstance(message, InformationRequest):
+            await self._information_controller.on_request(message)
+
+        logging.error(f"{self.__class__.__name__}: unknown message from {source=}: {message}")
+
+    async def get_loops(self):
+        await self.connect_to_storage()
+
+        loops = [
+            create_async_loop_routine(self._binder.routine, 0),
+            create_async_loop_routine(self._connector_storage.routine, 0),
+            create_async_loop_routine(self._graph_controller.routine, 0),
+            create_async_loop_routine(
+                self._balance_controller.routine, self._config_controller.get_config("load_balance_seconds")
+            ),
+            create_async_loop_routine(self._client_manager.routine, CLEANUP_INTERVAL_SECONDS),
+            create_async_loop_routine(self._object_controller.routine, CLEANUP_INTERVAL_SECONDS),
+            create_async_loop_routine(self._worker_controller.routine, CLEANUP_INTERVAL_SECONDS),
+            create_async_loop_routine(self._information_controller.routine, STATUS_REPORT_INTERVAL_SECONDS),
+        ]
+
+        try:
+            await asyncio.gather(*loops)
+        except asyncio.CancelledError:
+            pass
+        except ClientShutdownException as e:
+            logging.info(f"{self.__class__.__name__}: {e}")
+            pass
+
+        self._binder.destroy()
+        self._binder_monitor.destroy()
+
+
+@functools.wraps(Scheduler)
+async def scheduler_main(*args, **kwargs):
+    scheduler = Scheduler(*args, **kwargs)
+    await scheduler.get_loops()

scaler/scheduler/task/task_state_machine.py

@@ -0,0 +1,92 @@
+from typing import Dict, Optional
+
+from scaler.protocol.python.common import TaskState, TaskTransition
+
+
+class TaskStateMachine:
+    # see https://github.com/finos/opengris-scaler/issues/56
+    TRANSITION_MAP: Dict[TaskState, Dict[TaskTransition, TaskState]] = {
+        TaskState.Inactive: {
+            TaskTransition.HasCapacity: TaskState.Running,
+            TaskTransition.TaskCancel: TaskState.Canceled,
+        },
+        TaskState.Canceling: {
+            TaskTransition.TaskCancelConfirmCanceled: TaskState.Canceled,
+            TaskTransition.WorkerDisconnect: TaskState.Canceled,
+            TaskTransition.TaskCancelConfirmFailed: TaskState.Running,
+            TaskTransition.TaskCancelConfirmNotFound: TaskState.CanceledNotFound,
+        },
+        TaskState.Running: {
+            TaskTransition.TaskResultSuccess: TaskState.Success,
+            TaskTransition.TaskResultFailed: TaskState.Failed,
+            TaskTransition.TaskResultWorkerDied: TaskState.FailedWorkerDied,
+            TaskTransition.TaskCancel: TaskState.Canceling,
+            TaskTransition.BalanceTaskCancel: TaskState.BalanceCanceling,
+            TaskTransition.WorkerDisconnect: TaskState.WorkerDisconnecting,
+        },
+        TaskState.BalanceCanceling: {
+            TaskTransition.TaskResultSuccess: TaskState.Success,
+            TaskTransition.TaskResultFailed: TaskState.Failed,
+            TaskTransition.TaskResultWorkerDied: TaskState.FailedWorkerDied,
+            TaskTransition.TaskCancel: TaskState.Canceling,
+            TaskTransition.TaskCancelConfirmCanceled: TaskState.Inactive,
+            TaskTransition.TaskCancelConfirmFailed: TaskState.Running,
+            TaskTransition.WorkerDisconnect: TaskState.WorkerDisconnecting,
+        },
+        TaskState.WorkerDisconnecting: {
+            TaskTransition.SchedulerHasTask: TaskState.Inactive,
+            TaskTransition.SchedulerHasNoTask: TaskState.FailedWorkerDied,
+        },
+    }
+
+    def __init__(self, debug):
+        self._debug = debug
+        self._paths = list()
+
+        self._previous_state = None
+        self._state = TaskState.Inactive
+
+    def __repr__(self):
+        return f"TaskStateMachine(previous_state={self._previous_state}, state={self._state})"
+
+    def get_path(self):
+        return (
+            " ".join(f"[{state.name}] -{transition.name}->" for state, transition in self._paths)
+            + f" [{self._state.name}]"
+        )
+
+    def previous_state(self) -> Optional[TaskState]:
+        return self._previous_state
+
+    def current_state(self) -> TaskState:
+        return self._state
+
+    def is_running(self) -> bool:
+        return self._state == TaskState.Running
+
+    def is_canceling(self) -> bool:
+        return self._state == TaskState.Canceling
+
+    def is_finished(self) -> bool:
+        return self._state in {TaskState.Success, TaskState.Failed, TaskState.FailedWorkerDied}
+
+    def is_canceled(self) -> bool:
+        return self._state in {TaskState.Canceled, TaskState.CanceledNotFound}
+
+    def is_done(self) -> bool:
+        return self.is_finished() or self.is_canceled()
+
+    def on_transition(self, transition: TaskTransition) -> bool:
+        if self._state not in TaskStateMachine.TRANSITION_MAP:
+            return False
+
+        options = TaskStateMachine.TRANSITION_MAP[self._state]
+        if transition not in options:
+            return False
+
+        if self._debug:
+            self._paths.append((self._state, transition))
+
+        self._previous_state = self._state
+        self._state = options[transition]
+        return True

scaler/scheduler/task/task_state_manager.py

@@ -0,0 +1,61 @@
+import logging
+from typing import Dict, Optional
+
+from scaler.protocol.python.common import TaskState, TaskTransition
+from scaler.scheduler.task.task_state_machine import TaskStateMachine
+from scaler.utility.identifiers import TaskID
+
+
+class TaskStateManager:
+    def __init__(self, debug: bool):
+        self._debug = debug
+        self._task_id_to_state_machine: Dict[TaskID, TaskStateMachine] = dict()
+        self._statistics: Dict[TaskState, int] = {state: 0 for state in TaskState}
+
+    def add_state_machine(self, task_id: TaskID) -> TaskStateMachine:
+        """Create new task state machine, return True if success, False otherwise"""
+        assert task_id not in self._task_id_to_state_machine
+
+        state_machine = TaskStateMachine(self._debug)
+        self._task_id_to_state_machine[task_id] = state_machine
+        self._statistics[state_machine.current_state()] += 1
+        return state_machine
+
+    def remove_state_machine(self, task_id: TaskID):
+        self._task_id_to_state_machine.pop(task_id)
+
+    def get_state_machine(self, task_id: TaskID) -> Optional[TaskStateMachine]:
+        return self._task_id_to_state_machine.get(task_id, None)
+
+    def on_transition(self, task_id: TaskID, transition: TaskTransition) -> Optional[TaskStateMachine]:
+        """if adjust task state machine is successful, then return TaskStateFlags object associate with the task_id,
+        return None otherwise
+
+        This should be a central place to synchronize task state machine, if any unexpected event happened, it will not
+        return the TaskStateFlags
+        """
+
+        task_state_machine = self._task_id_to_state_machine.get(task_id, None)
+        if task_state_machine is None:
+            logging.error(f"{task_id!r}: unknown {transition=} for non-existed state machine")
+            return None
+
+        transit_success = task_state_machine.on_transition(transition)
+        if transit_success:
+            self._statistics[task_state_machine.previous_state()] -= 1
+            self._statistics[task_state_machine.current_state()] += 1
+        else:
+            logging.error(
+                f"{task_id!r}: cannot apply {transition} to current state" f" {task_state_machine.current_state()}"
+            )
+
+        return task_state_machine if transit_success else None
+
+    def get_statistics(self) -> Dict[TaskState, int]:
+        return self._statistics
+
+    def get_debug_paths(self):
+        return "\n".join(
+            f"{task_id!r}: {state_machine.get_path()}"
+            for task_id, state_machine in self._task_id_to_state_machine.items()
+        )

scaler/ui/constants.py

@@ -0,0 +1,9 @@
+# times are in seconds
+
+TASK_LOG_REFRESH_INTERVAL = 0.5
+
+WORKER_PROCESSORS_REFRESH_INTERVAL = 2
+
+TASK_STREAM_UPDATE_INTERVAL = 0.1
+
+MEMORY_USAGE_UPDATE_INTERVAL = 0.1

scaler/ui/live_display.py

@@ -0,0 +1,147 @@
+import dataclasses
+from collections import defaultdict
+from typing import Dict, List, Optional, Set
+
+from nicegui import ui
+from nicegui.element import Element
+
+from scaler.protocol.python.common import WorkerState
+from scaler.protocol.python.message import StateTask, StateWorker
+from scaler.protocol.python.status import WorkerStatus
+from scaler.ui.utility import display_capabilities, format_worker_name
+from scaler.utility.formatter import format_microseconds, format_seconds
+
+
+@dataclasses.dataclass
+class SchedulerSection:
+    cpu: str = dataclasses.field(default="")
+    rss: str = dataclasses.field(default="")
+    rss_free: str = dataclasses.field(default="")
+
+    handler: Optional[Element] = dataclasses.field(default=None)
+
+    def handle_task_state(self, _: StateTask):
+        return
+
+    def handle_worker_state(self, _: StateWorker):
+        return
+
+    def draw_section(self):
+        with ui.card().classes("w-full"), ui.row() as handler:
+            self.handler = handler
+            ui.label("Scheduler")
+            ui.label()
+            ui.label("CPU:")
+            ui.label().bind_text_from(self, "cpu")
+            ui.label()
+            ui.label("RSS:")
+            ui.label().bind_text_from(self, "rss")
+            ui.label()
+            ui.label("RSS Free:")
+            ui.label().bind_text_from(self, "rss_free")
+
+    def delete_section(self):
+        self.handler.clear()
+        self.handler.delete()
+
+
+@dataclasses.dataclass
+class WorkerRow:
+    worker: str = dataclasses.field(default="")
+    agt_cpu: float = dataclasses.field(default=0)
+    agt_rss: int = dataclasses.field(default=0)
+    cpu: float = dataclasses.field(default=0)
+    rss: int = dataclasses.field(default=0)
+    rss_free: int = dataclasses.field(default=0)
+    free: int = dataclasses.field(default=0)
+    sent: int = dataclasses.field(default=0)
+    queued: int = dataclasses.field(default=0)
+    suspended: int = dataclasses.field(default=0)
+    lag: str = dataclasses.field(default="")
+    itl: str = dataclasses.field(default="")
+    last_seen: str = dataclasses.field(default="")
+    capabilities: Set[str] = dataclasses.field(default_factory=set)
+    display_capabilities: str = dataclasses.field(default="")
+
+    handlers: List[Element] = dataclasses.field(default_factory=list)
+
+    def populate(self, data: WorkerStatus):
+        self.worker = data.worker_id.decode()
+        self.agt_cpu = data.agent.cpu / 10
+        self.agt_rss = int(data.agent.rss / 1e6)
+        self.cpu = sum(p.resource.cpu for p in data.processor_statuses) / 10
+        self.rss = int(sum(p.resource.rss for p in data.processor_statuses) / 1e6)
+        self.rss_free = int(data.rss_free / 1e6)
+        self.free = data.free
+        self.sent = data.sent
+        self.queued = data.queued
+        self.suspended = data.suspended
+        self.lag = format_microseconds(data.lag_us)
+        self.itl = data.itl
+        self.last_seen = format_seconds(data.last_s)
+
+    def set_capabilities(self, capabilities: Set[str]):
+        self.capabilities = capabilities
+        self.display_capabilities = display_capabilities(self.capabilities)
+
+    def draw_row(self):
+        total_rss = self.rss + self.rss_free
+
+        ui.label(format_worker_name(self.worker))
+        ui.knob(track_color="grey-2", show_value=True, min=0, max=100).bind_value_from(self, "agt_cpu")
+        ui.knob(track_color="grey-2", show_value=True, min=0, max=total_rss).bind_value_from(self, "agt_rss")
+        ui.knob(track_color="grey-2", show_value=True, min=0, max=100).bind_value_from(self, "cpu")
+        ui.knob(track_color="grey-2", show_value=True, min=0, max=total_rss).bind_value_from(self, "rss")
+        ui.label().bind_text_from(self, "free")
+        ui.label().bind_text_from(self, "sent")
+        ui.label().bind_text_from(self, "queued")
+        ui.label().bind_text_from(self, "suspended")
+        ui.label().bind_text_from(self, "lag")
+        ui.label().bind_text_from(self, "ITL")
+        ui.label().bind_text_from(self, "last_seen")
+        ui.label().bind_text_from(self, "display_capabilities")
+
+    def delete_row(self):
+        for element in self.handlers:
+            element.delete()
+
+
+@dataclasses.dataclass
+class WorkersSection:
+    workers: Dict[str, WorkerRow] = dataclasses.field(default_factory=lambda: defaultdict(WorkerRow))
+
+    @ui.refreshable
+    def draw_section(self):
+        with ui.row().classes("h-max"), ui.card().classes("w-full"), ui.grid(columns=13):
+            self.__draw_titles()
+            for worker_row in self.workers.values():
+                worker_row.draw_row()
+
+    def handle_task_state(self, _: StateTask):
+        return
+
+    def handle_worker_state(self, state_worker: StateWorker):
+        worker_id = state_worker.worker_id.decode()
+        state = state_worker.state
+
+        if state == WorkerState.Connected:
+            self.workers[worker_id].set_capabilities(set(state_worker.capabilities.keys()))
+        if state == WorkerState.Disconnected:
+            self.workers.pop(worker_id, None)
+            self.draw_section.refresh()
+
+    @staticmethod
+    def __draw_titles():
+        ui.label("Worker")
+        ui.label("Agt CPU %")
+        ui.label("Agt RSS (in MB)")
+        ui.label("Processors CPU %")
+        ui.label("Processors RSS (in MB)")
+        ui.label("Queue Capacity")
+        ui.label("Tasks Sent")
+        ui.label("Tasks Queued")
+        ui.label("Tasks Suspended")
+        ui.label("Lag")
+        ui.label("ITL")
+        ui.label("Last Seen")
+        ui.label("Capabilities")