opengris-scaler 1.12.37__cp38-cp38-musllinux_1_2_x86_64.whl

scaler/worker_adapter/symphony/worker.py

@@ -0,0 +1,204 @@
+import asyncio
+import logging
+import multiprocessing
+import signal
+from collections import deque
+from typing import Dict, Optional
+
+import zmq
+
+from scaler.config.types.object_storage_server import ObjectStorageAddressConfig
+from scaler.config.types.zmq import ZMQConfig
+from scaler.io.async_connector import ZMQAsyncConnector
+from scaler.io.mixins import AsyncConnector, AsyncObjectStorageConnector
+from scaler.io.utility import create_async_object_storage_connector
+from scaler.protocol.python.message import (
+    ClientDisconnect,
+    DisconnectRequest,
+    ObjectInstruction,
+    Task,
+    TaskCancel,
+    WorkerHeartbeatEcho,
+)
+from scaler.protocol.python.mixins import Message
+from scaler.utility.event_loop import create_async_loop_routine, register_event_loop
+from scaler.utility.exceptions import ClientShutdownException
+from scaler.utility.identifiers import WorkerID
+from scaler.utility.logging.utility import setup_logger
+from scaler.worker.agent.timeout_manager import VanillaTimeoutManager
+from scaler.worker_adapter.symphony.heartbeat_manager import SymphonyHeartbeatManager
+from scaler.worker_adapter.symphony.task_manager import SymphonyTaskManager
+
+
+class SymphonyWorker(multiprocessing.get_context("spawn").Process):  # type: ignore
+    """
+    SymphonyWorker is an implementation of a worker that can handle multiple tasks concurrently.
+    Most of the task execution logic is handled by SymphonyTaskManager.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        address: ZMQConfig,
+        object_storage_address: Optional[ObjectStorageAddressConfig],
+        service_name: str,
+        capabilities: Dict[str, int],
+        base_concurrency: int,
+        heartbeat_interval_seconds: int,
+        death_timeout_seconds: int,
+        task_queue_size: int,
+        io_threads: int,
+        event_loop: str,
+    ):
+        multiprocessing.Process.__init__(self, name="Agent")
+
+        self._event_loop = event_loop
+        self._name = name
+        self._address = address
+        self._object_storage_address = object_storage_address
+        self._capabilities = capabilities
+        self._io_threads = io_threads
+
+        self._ident = WorkerID.generate_worker_id(name)  # _identity is internal to multiprocessing.Process
+
+        self._service_name = service_name
+        self._base_concurrency = base_concurrency
+
+        self._heartbeat_interval_seconds = heartbeat_interval_seconds
+        self._death_timeout_seconds = death_timeout_seconds
+        self._task_queue_size = task_queue_size
+
+        self._context: Optional[zmq.asyncio.Context] = None
+        self._connector_external: Optional[AsyncConnector] = None
+        self._connector_storage: Optional[AsyncObjectStorageConnector] = None
+        self._task_manager: Optional[SymphonyTaskManager] = None
+        self._heartbeat_manager: Optional[SymphonyHeartbeatManager] = None
+
+        """
+        Sometimes the first message received is not a heartbeat echo, so we need to backoff processing other tasks
+        until we receive the first heartbeat echo.
+        """
+        self._heartbeat_received: bool = False
+        self._backoff_message_queue: deque = deque()
+
+    @property
+    def identity(self) -> WorkerID:
+        return self._ident
+
+    def run(self) -> None:
+        self.__initialize()
+        self.__run_forever()
+
+    def __initialize(self):
+        setup_logger()
+        register_event_loop(self._event_loop)
+
+        self._context = zmq.asyncio.Context()
+        self._connector_external = ZMQAsyncConnector(
+            context=self._context,
+            name=self.name,
+            socket_type=zmq.DEALER,
+            address=self._address,
+            bind_or_connect="connect",
+            callback=self.__on_receive_external,
+            identity=self._ident,
+        )
+
+        self._connector_storage = create_async_object_storage_connector()
+
+        self._heartbeat_manager = SymphonyHeartbeatManager(
+            object_storage_address=self._object_storage_address,
+            capabilities=self._capabilities,
+            task_queue_size=self._task_queue_size,
+        )
+        self._task_manager = SymphonyTaskManager(
+            base_concurrency=self._base_concurrency, service_name=self._service_name
+        )
+        self._timeout_manager = VanillaTimeoutManager(death_timeout_seconds=self._death_timeout_seconds)
+
+        # register
+        self._heartbeat_manager.register(
+            connector_external=self._connector_external,
+            connector_storage=self._connector_storage,
+            worker_task_manager=self._task_manager,
+            timeout_manager=self._timeout_manager,
+        )
+        self._task_manager.register(
+            connector_external=self._connector_external,
+            connector_storage=self._connector_storage,
+            heartbeat_manager=self._heartbeat_manager,
+        )
+
+        self._loop = asyncio.get_event_loop()
+        self.__register_signal()
+        self._task = self._loop.create_task(self.__get_loops())
+
+    async def __on_receive_external(self, message: Message):
+        if not self._heartbeat_received and not isinstance(message, WorkerHeartbeatEcho):
+            self._backoff_message_queue.append(message)
+            return
+
+        if isinstance(message, WorkerHeartbeatEcho):
+            await self._heartbeat_manager.on_heartbeat_echo(message)
+            self._heartbeat_received = True
+
+            while self._backoff_message_queue:
+                backoff_message = self._backoff_message_queue.popleft()
+                await self.__on_receive_external(backoff_message)
+
+            return
+
+        if isinstance(message, Task):
+            await self._task_manager.on_task_new(message)
+            return
+
+        if isinstance(message, TaskCancel):
+            await self._task_manager.on_cancel_task(message)
+            return
+
+        if isinstance(message, ObjectInstruction):
+            await self._task_manager.on_object_instruction(message)
+            return
+
+        if isinstance(message, ClientDisconnect):
+            if message.disconnect_type == ClientDisconnect.DisconnectType.Shutdown:
+                raise ClientShutdownException("received client shutdown, quitting")
+            logging.error(f"Worker received invalid ClientDisconnect type, ignoring {message=}")
+            return
+
+        raise TypeError(f"Unknown {message=}")
+
+    async def __get_loops(self):
+        if self._object_storage_address is not None:
+            # With a manually set storage address, immediately connect to the object storage server.
+            await self._connector_storage.connect(self._object_storage_address.host, self._object_storage_address.port)
+
+        try:
+            await asyncio.gather(
+                create_async_loop_routine(self._connector_external.routine, 0),
+                create_async_loop_routine(self._connector_storage.routine, 0),
+                create_async_loop_routine(self._heartbeat_manager.routine, self._heartbeat_interval_seconds),
+                create_async_loop_routine(self._timeout_manager.routine, 1),
+                create_async_loop_routine(self._task_manager.process_task, 0),
+                create_async_loop_routine(self._task_manager.resolve_tasks, 0),
+            )
+        except asyncio.CancelledError:
+            pass
+        except (ClientShutdownException, TimeoutError) as e:
+            logging.info(f"{self.identity!r}: {str(e)}")
+        except Exception as e:
+            logging.exception(f"{self.identity!r}: failed with unhandled exception:\n{e}")
+
+        await self._connector_external.send(DisconnectRequest.new_msg(self.identity))
+
+        self._connector_external.destroy()
+        logging.info(f"{self.identity!r}: quit")
+
+    def __run_forever(self):
+        self._loop.run_until_complete(self._task)
+
+    def __register_signal(self):
+        self._loop.add_signal_handler(signal.SIGINT, self.__destroy)
+
+    def __destroy(self):
+        self._task.cancel()

scaler/worker_adapter/symphony/worker_adapter.py

@@ -0,0 +1,123 @@
+import os
+import signal
+import uuid
+from typing import Dict
+
+from aiohttp import web
+from aiohttp.web_request import Request
+
+from scaler.config.section.symphony_worker_adapter import SymphonyWorkerConfig
+from scaler.utility.identifiers import WorkerID
+from scaler.worker_adapter.common import CapacityExceededError, WorkerGroupID, WorkerGroupNotFoundError
+from scaler.worker_adapter.symphony.worker import SymphonyWorker
+
+
+class SymphonyWorkerAdapter:
+    def __init__(self, config: SymphonyWorkerConfig):
+        self._address = config.worker_adapter_config.scheduler_address
+        self._object_storage_address = config.worker_adapter_config.object_storage_address
+        self._service_name = config.service_name
+        self._base_concurrency = config.worker_adapter_config.max_workers
+        self._capabilities = config.worker_config.per_worker_capabilities.capabilities
+        self._io_threads = config.worker_io_threads
+        self._task_queue_size = config.worker_config.per_worker_task_queue_size
+        self._heartbeat_interval_seconds = config.worker_config.heartbeat_interval_seconds
+        self._death_timeout_seconds = config.worker_config.death_timeout_seconds
+        self._event_loop = config.event_loop
+        self._logging_paths = config.logging_config.paths
+        self._logging_level = config.logging_config.level
+        self._logging_config_file = config.logging_config.config_file
+
+        """
+        Although a worker group can contain multiple workers, in this Symphony adapter implementation,
+        there will be only one worker group which contains one Symphony worker.
+        """
+        self._worker_groups: Dict[WorkerGroupID, Dict[WorkerID, SymphonyWorker]] = {}
+
+    async def start_worker_group(self) -> WorkerGroupID:
+        if self._worker_groups:
+            raise CapacityExceededError("Symphony worker already started")
+
+        worker = SymphonyWorker(
+            name=f"SYM|{uuid.uuid4().hex}",
+            address=self._address,
+            object_storage_address=self._object_storage_address,
+            service_name=self._service_name,
+            base_concurrency=self._base_concurrency,
+            capabilities=self._capabilities,
+            io_threads=self._io_threads,
+            task_queue_size=self._task_queue_size,
+            heartbeat_interval_seconds=self._heartbeat_interval_seconds,
+            death_timeout_seconds=self._death_timeout_seconds,
+            event_loop=self._event_loop,
+        )
+
+        worker.start()
+        worker_group_id = f"symphony-{uuid.uuid4().hex}".encode()
+        self._worker_groups[worker_group_id] = {worker.identity: worker}
+        return worker_group_id
+
+    async def shutdown_worker_group(self, worker_group_id: WorkerGroupID):
+        if worker_group_id not in self._worker_groups:
+            raise WorkerGroupNotFoundError(f"Worker group with ID {worker_group_id.decode()} does not exist.")
+
+        for worker in self._worker_groups[worker_group_id].values():
+            os.kill(worker.pid, signal.SIGINT)
+            worker.join()
+
+        self._worker_groups.pop(worker_group_id)
+
+    async def webhook_handler(self, request: Request):
+        request_json = await request.json()
+
+        if "action" not in request_json:
+            return web.json_response({"error": "No action specified"}, status=web.HTTPBadRequest.status_code)
+
+        action = request_json["action"]
+
+        if action == "get_worker_adapter_info":
+            return web.json_response(
+                {"max_worker_groups": 1, "workers_per_group": 1, "base_capabilities": self._capabilities},
+                status=web.HTTPOk.status_code,
+            )
+
+        elif action == "start_worker_group":
+            try:
+                worker_group_id = await self.start_worker_group()
+            except CapacityExceededError as e:
+                return web.json_response({"error": str(e)}, status=web.HTTPTooManyRequests.status_code)
+            except Exception as e:
+                return web.json_response({"error": str(e)}, status=web.HTTPInternalServerError.status_code)
+
+            return web.json_response(
+                {
+                    "status": "Worker group started",
+                    "worker_group_id": worker_group_id.decode(),
+                    "worker_ids": [worker_id.decode() for worker_id in self._worker_groups[worker_group_id].keys()],
+                },
+                status=web.HTTPOk.status_code,
+            )
+
+        elif action == "shutdown_worker_group":
+            if "worker_group_id" not in request_json:
+                return web.json_response(
+                    {"error": "No worker_group_id specified"}, status=web.HTTPBadRequest.status_code
+                )
+
+            worker_group_id = request_json["worker_group_id"].encode()
+            try:
+                await self.shutdown_worker_group(worker_group_id)
+            except WorkerGroupNotFoundError as e:
+                return web.json_response({"error": str(e)}, status=web.HTTPNotFound.status_code)
+            except Exception as e:
+                return web.json_response({"error": str(e)}, status=web.HTTPInternalServerError.status_code)
+
+            return web.json_response({"status": "Worker group shutdown"}, status=web.HTTPOk.status_code)
+
+        else:
+            return web.json_response({"error": "Unknown action"}, status=web.HTTPBadRequest.status_code)
+
+    def create_app(self):
+        app = web.Application()
+        app.router.add_post("/", self.webhook_handler)
+        return app