opengris-scaler 1.12.7__cp312-cp312-musllinux_1_2_x86_64.whl

scaler/scheduler/controllers/mixins.py

@@ -0,0 +1,201 @@
+import abc
+from typing import Any, Optional, Set
+
+from scaler.protocol.python.common import ObjectMetadata
+from scaler.protocol.python.message import (
+    ClientDisconnect,
+    ClientHeartbeat,
+    DisconnectRequest,
+    GraphTask,
+    InformationRequest,
+    InformationSnapshot,
+    ObjectInstruction,
+    Task,
+    TaskCancel,
+    TaskCancelConfirm,
+    TaskResult,
+    WorkerHeartbeat,
+)
+from scaler.utility.identifiers import ClientID, ObjectID, TaskID, WorkerID
+from scaler.utility.mixins import Reporter
+
+
+class ConfigController(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def get_config(self, path: str) -> Any:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def update_config(self, path: str, value: Any):
+        raise NotImplementedError()
+
+
+class ObjectController(Reporter):
+    @abc.abstractmethod
+    async def on_object_instruction(self, source: bytes, request: ObjectInstruction):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def on_add_object(
+        self,
+        client_id: ClientID,
+        object_id: ObjectID,
+        object_type: ObjectMetadata.ObjectContentType,
+        object_name: bytes,
+    ):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def on_del_objects(self, client_id: ClientID, object_ids: Set[ObjectID]):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def clean_client(self, client_id: ClientID):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def has_object(self, object_id: ObjectID) -> bool:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_object_name(self, object_id: ObjectID) -> bytes:
+        raise NotImplementedError()
+
+
+class ClientController(Reporter):
+    @abc.abstractmethod
+    def get_client_task_ids(self, client_id: ClientID) -> Set[TaskID]:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def has_client_id(self, client_id: ClientID) -> bool:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_client_id(self, task_id: TaskID) -> Optional[ClientID]:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def on_task_begin(self, client_id: ClientID, task_id: TaskID):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def on_task_finish(self, task_id: TaskID) -> bytes:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_heartbeat(self, client_id: ClientID, info: ClientHeartbeat):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_client_disconnect(self, client_id: ClientID, request: ClientDisconnect):
+        raise NotImplementedError()
+
+
+class GraphTaskController(Reporter):
+    @abc.abstractmethod
+    async def on_graph_task(self, client_id: ClientID, graph_task: GraphTask):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_graph_task_cancel(self, graph_task_cancel: TaskCancel):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_graph_sub_task_cancel_confirm(self, task_cancel_confirm: TaskCancelConfirm):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_graph_sub_task_result(self, result: TaskResult) -> bool:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def is_graph_subtask(self, task_id: TaskID) -> bool:
+        raise NotImplementedError()
+
+
+class TaskController(Reporter):
+    @abc.abstractmethod
+    async def on_task_new(self, task: Task):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_task_cancel(self, client_id: ClientID, task_cancel: TaskCancel):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_task_balance_cancel(self, task_id: TaskID):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_task_cancel_confirm(self, task_cancel_confirm: TaskCancelConfirm):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_task_result(self, result: TaskResult):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_worker_connect(self, worker_id: WorkerID):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_worker_disconnect(self, task_id: TaskID, worker_id: WorkerID):
+        raise NotImplementedError()
+
+
+class WorkerController(Reporter):
+    @abc.abstractmethod
+    def acquire_worker(self, task: Task) -> Optional[WorkerID]:
+        """this acquires worker should be atomic, means it cannot be async decorated, otherwise it will create gap that
+        get worker but task is not send to worker, and cannot find task in the worker state"""
+
+        # TODO: this function should return things that expose 3 kinds of information:
+        # TODO: 1. worker id as bytes if have capacity and able to assign to worker id
+        # TODO: 2. capacity is full, and unable to add new task
+        # TODO: 3. capacity is not full, but all the workers are busy right now, so tasks will be queued
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_task_cancel(self, task_cancel: TaskCancel) -> bytes:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_task_done(self, task_id: TaskID):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_heartbeat(self, worker_id: WorkerID, info: WorkerHeartbeat):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_client_shutdown(self, client_id: ClientID):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_disconnect(self, worker_id: WorkerID, request: DisconnectRequest):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def has_available_worker(self) -> bool:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_worker_by_task_id(self, task_id: TaskID) -> WorkerID:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_worker_ids(self) -> Set[WorkerID]:
+        raise NotImplementedError()
+
+
+class ScalingController(Reporter):
+    @abc.abstractmethod
+    async def on_snapshot(self, snapshot: InformationSnapshot):
+        raise NotImplementedError()
+
+
+class InformationController(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    async def on_request(self, request: InformationRequest):
+        raise NotImplementedError()

scaler/scheduler/controllers/object_controller.py

@@ -0,0 +1,147 @@
+import dataclasses
+import logging
+from asyncio import Queue
+from typing import Optional, Set
+
+from scaler.io.mixins import AsyncBinder, AsyncConnector, AsyncObjectStorageConnector
+from scaler.protocol.python.common import ObjectMetadata
+from scaler.protocol.python.message import ObjectInstruction
+from scaler.protocol.python.status import ObjectManagerStatus
+from scaler.scheduler.controllers.config_controller import VanillaConfigController
+from scaler.scheduler.controllers.mixins import ClientController, ObjectController, WorkerController
+from scaler.scheduler.object_usage.object_tracker import ObjectTracker, ObjectUsage
+from scaler.utility.identifiers import ClientID, ObjectID
+from scaler.utility.mixins import Looper, Reporter
+
+
+@dataclasses.dataclass
+class _ObjectCreation(ObjectUsage):
+    object_id: ObjectID
+    object_creator: ClientID
+    object_type: ObjectMetadata.ObjectContentType
+    object_name: bytes
+
+    def get_object_key(self) -> ObjectID:
+        return self.object_id
+
+
+class VanillaObjectController(ObjectController, Looper, Reporter):
+    def __init__(self, config_controller: VanillaConfigController):
+        self._config_controller = config_controller
+
+        self._object_tracker: ObjectTracker[ClientID, ObjectID, _ObjectCreation] = ObjectTracker(
+            "object_usage", self.__finished_object_storage
+        )
+
+        self._queue_deleted_object_ids: Queue[ObjectID] = Queue()
+
+        self._binder: Optional[AsyncBinder] = None
+        self._binder_monitor: Optional[AsyncConnector] = None
+        self._connector_storage: Optional[AsyncObjectStorageConnector] = None
+
+        self._client_manager: Optional[ClientController] = None
+        self._worker_manager: Optional[WorkerController] = None
+
+    def register(
+        self,
+        binder: AsyncBinder,
+        binder_monitor: AsyncConnector,
+        connector_storage: AsyncObjectStorageConnector,
+        client_manager: ClientController,
+        worker_manager: WorkerController,
+    ):
+        self._binder = binder
+        self._binder_monitor = binder_monitor
+        self._connector_storage = connector_storage
+        self._client_manager = client_manager
+        self._worker_manager = worker_manager
+
+    async def on_object_instruction(self, source: bytes, instruction: ObjectInstruction):
+        if instruction.instruction_type == ObjectInstruction.ObjectInstructionType.Create:
+            self.__on_object_create(source, instruction)
+            return
+
+        if instruction.instruction_type == ObjectInstruction.ObjectInstructionType.Delete:
+            self.on_del_objects(instruction.object_user, set(instruction.object_metadata.object_ids))
+            return
+
+        logging.error(f"received unknown object instruction_type={instruction.instruction_type} from {source=}")
+
+    def on_add_object(
+        self,
+        client_id: ClientID,
+        object_id: ObjectID,
+        object_type: ObjectMetadata.ObjectContentType,
+        object_name: bytes,
+    ):
+        creation = _ObjectCreation(object_id, client_id, object_type, object_name)
+        logging.debug(
+            f"add object cache "
+            f"object_name={creation.object_name!r}, "
+            f"object_type={creation.object_type}, "
+            f"object_id={creation.object_id!r}"
+        )
+
+        self._object_tracker.add_object(creation)
+        self._object_tracker.add_blocks_for_one_object(creation.get_object_key(), {creation.object_creator})
+
+    def on_del_objects(self, client_id: ClientID, object_ids: Set[ObjectID]):
+        for object_id in object_ids:
+            self._object_tracker.remove_one_block_for_objects({object_id}, client_id)
+
+    def clean_client(self, client_id: ClientID):
+        self._object_tracker.remove_blocks({client_id})
+
+    async def routine(self):
+        await self.__routine_send_objects_deletions()
+
+    def has_object(self, object_id: ObjectID) -> bool:
+        return self._object_tracker.has_object(object_id)
+
+    def get_object_name(self, object_id: ObjectID) -> bytes:
+        if not self.has_object(object_id):
+            return b"<Unknown>"
+
+        return self._object_tracker.get_object(object_id).object_name
+
+    def get_status(self) -> ObjectManagerStatus:
+        return ObjectManagerStatus.new_msg(self._object_tracker.object_count())
+
+    async def __routine_send_objects_deletions(self):
+        deleted_object_ids = [await self._queue_deleted_object_ids.get()]
+        self._queue_deleted_object_ids.task_done()
+
+        while not self._queue_deleted_object_ids.empty():
+            deleted_object_ids.append(self._queue_deleted_object_ids.get_nowait())
+            self._queue_deleted_object_ids.task_done()
+
+        for worker in self._worker_manager.get_worker_ids():
+            await self._binder.send(
+                worker,
+                ObjectInstruction.new_msg(
+                    ObjectInstruction.ObjectInstructionType.Delete,
+                    # TODO: ideally object_user should be set to the owning client ID, but then we cannot batch these
+                    # Delete instructions.
+                    None,
+                    ObjectMetadata.new_msg(tuple(deleted_object_ids)),
+                ),
+            )
+
+        for object_id in deleted_object_ids:
+            await self._connector_storage.delete_object(object_id)
+
+    def __on_object_create(self, source: bytes, instruction: ObjectInstruction):
+        if not self._client_manager.has_client_id(instruction.object_user):
+            logging.error(f"received object creation from {source!r} for unknown client {instruction.object_user!r}")
+            return
+
+        for object_id, object_type, object_name in zip(
+            instruction.object_metadata.object_ids,
+            instruction.object_metadata.object_types,
+            instruction.object_metadata.object_names,
+        ):
+            self.on_add_object(instruction.object_user, object_id, object_type, object_name)
+
+    def __finished_object_storage(self, creation: _ObjectCreation):
+        logging.debug(f"del object cache object_name={creation.object_name!r}, object_id={creation.object_id!r}")
+        self._queue_deleted_object_ids.put_nowait(creation.object_id)

scaler/scheduler/controllers/scaling_controller.py

@@ -0,0 +1,86 @@
+import logging
+from typing import Dict, List
+
+import aiohttp
+from aiohttp import web
+
+from scaler.protocol.python.message import InformationSnapshot
+from scaler.protocol.python.status import ScalingManagerStatus
+from scaler.scheduler.controllers.mixins import ScalingController
+from scaler.utility.identifiers import WorkerID
+
+WorkerGroupID = bytes
+
+
+class NullScalingController(ScalingController):
+    def get_status(self):
+        return ScalingManagerStatus.new_msg(worker_groups={})
+
+    async def on_snapshot(self, information_snapshot: InformationSnapshot):
+        pass
+
+
+class VanillaScalingController(ScalingController):
+    def __init__(self, adapter_webhook_url: str, lower_task_ratio: float = 1, upper_task_ratio: float = 10):
+        self._adapter_webhook_url = adapter_webhook_url
+        self._lower_task_ratio = lower_task_ratio
+        self._upper_task_ratio = upper_task_ratio
+        assert upper_task_ratio >= lower_task_ratio
+
+        self._worker_groups: Dict[WorkerGroupID, List[WorkerID]] = {}
+
+    def get_status(self):
+        return ScalingManagerStatus.new_msg(worker_groups=self._worker_groups)
+
+    async def on_snapshot(self, information_snapshot: InformationSnapshot):
+        if not information_snapshot.workers:
+            if information_snapshot.tasks:
+                await self._start_worker_group()
+            return
+
+        task_ratio = len(information_snapshot.tasks) / len(information_snapshot.workers)
+        if task_ratio > self._upper_task_ratio:
+            await self._start_worker_group()
+        elif task_ratio < self._lower_task_ratio:
+            worker_group_task_counts = {
+                worker_group_id: sum(information_snapshot.workers[worker_id].queued_tasks for worker_id in worker_ids)
+                for worker_group_id, worker_ids in self._worker_groups.items()
+            }
+            worker_group_id = min(worker_group_task_counts, key=worker_group_task_counts.get)
+            await self._shutdown_worker_group(worker_group_id)
+
+    async def _start_worker_group(self):
+        response, status = await self._make_request({"action": "start_worker_group"})
+        if status == web.HTTPTooManyRequests.status_code:
+            logging.warning("Capacity exceeded, cannot start new worker group.")
+            return
+        if status == web.HTTPInternalServerError.status_code:
+            logging.error(f"Failed to start worker group: {response.get('error', 'Unknown error')}")
+            return
+
+        worker_group_id = response["worker_group_id"].encode()
+        self._worker_groups[worker_group_id] = [WorkerID(worker_id.encode()) for worker_id in response["worker_ids"]]
+        logging.info(f"Started worker group: {worker_group_id.decode()}")
+
+    async def _shutdown_worker_group(self, worker_group_id: WorkerGroupID):
+        if worker_group_id not in self._worker_groups:
+            logging.error(f"Worker group with ID {worker_group_id.decode()} does not exist.")
+            return
+
+        response, status = await self._make_request(
+            {"action": "shutdown_worker_group", "worker_group_id": worker_group_id.decode()}
+        )
+        if status == web.HTTPNotFound.status_code:
+            logging.error(f"Worker group with ID {worker_group_id.decode()} not found in adapter.")
+            return
+        if status == web.HTTPInternalServerError.status_code:
+            logging.error(f"Failed to shutdown worker group: {response.get('error', 'Unknown error')}")
+            return
+
+        self._worker_groups.pop(worker_group_id)
+        logging.info(f"Shutdown worker group: {worker_group_id.decode()}")
+
+    async def _make_request(self, payload):
+        async with aiohttp.ClientSession() as session:
+            async with session.post(self._adapter_webhook_url, json=payload) as response:
+                return await response.json(), response.status