opengris-scaler 1.12.28__cp313-cp313-musllinux_1_2_x86_64.whl

scaler/worker_adapter/native.py

@@ -0,0 +1,155 @@
+import os
+import signal
+import uuid
+from typing import Dict, Optional, Tuple
+
+from aiohttp import web
+from aiohttp.web_request import Request
+
+from scaler.config.types.object_storage_server import ObjectStorageConfig
+from scaler.config.types.zmq import ZMQConfig
+from scaler.utility.identifiers import WorkerID
+from scaler.worker.worker import Worker
+from scaler.worker_adapter.common import CapacityExceededError, WorkerGroupID, WorkerGroupNotFoundError
+
+
+class NativeWorkerAdapter:
+    def __init__(
+        self,
+        address: ZMQConfig,
+        object_storage_address: Optional[ObjectStorageConfig],
+        capabilities: Dict[str, int],
+        io_threads: int,
+        task_queue_size: int,
+        max_workers: int,
+        heartbeat_interval_seconds: int,
+        task_timeout_seconds: int,
+        death_timeout_seconds: int,
+        garbage_collect_interval_seconds: int,
+        trim_memory_threshold_bytes: int,
+        hard_processor_suspend: bool,
+        event_loop: str,
+        logging_paths: Tuple[str, ...],
+        logging_level: str,
+        logging_config_file: Optional[str],
+    ):
+        self._address = address
+        self._object_storage_address = object_storage_address
+        self._capabilities = capabilities
+        self._io_threads = io_threads
+        self._task_queue_size = task_queue_size
+        self._max_workers = max_workers
+        self._heartbeat_interval_seconds = heartbeat_interval_seconds
+        self._task_timeout_seconds = task_timeout_seconds
+        self._death_timeout_seconds = death_timeout_seconds
+        self._garbage_collect_interval_seconds = garbage_collect_interval_seconds
+        self._trim_memory_threshold_bytes = trim_memory_threshold_bytes
+        self._hard_processor_suspend = hard_processor_suspend
+        self._event_loop = event_loop
+        self._logging_paths = logging_paths
+        self._logging_level = logging_level
+        self._logging_config_file = logging_config_file
+
+        """
+        Although a worker group can contain multiple workers, in this native adapter implementation,
+        each worker group will only contain one worker.
+        """
+        self._worker_groups: Dict[WorkerGroupID, Dict[WorkerID, Worker]] = {}
+
+    async def start_worker_group(self) -> WorkerGroupID:
+        num_of_workers = sum(len(workers) for workers in self._worker_groups.values())
+        if num_of_workers >= self._max_workers != -1:
+            raise CapacityExceededError(f"Maximum number of workers ({self._max_workers}) reached.")
+
+        worker = Worker(
+            name=f"NAT|{uuid.uuid4().hex}",
+            address=self._address,
+            object_storage_address=self._object_storage_address,
+            preload=None,
+            capabilities=self._capabilities,
+            io_threads=self._io_threads,
+            task_queue_size=self._task_queue_size,
+            heartbeat_interval_seconds=self._heartbeat_interval_seconds,
+            task_timeout_seconds=self._task_timeout_seconds,
+            death_timeout_seconds=self._death_timeout_seconds,
+            garbage_collect_interval_seconds=self._garbage_collect_interval_seconds,
+            trim_memory_threshold_bytes=self._trim_memory_threshold_bytes,
+            hard_processor_suspend=self._hard_processor_suspend,
+            event_loop=self._event_loop,
+            logging_paths=self._logging_paths,
+            logging_level=self._logging_level,
+        )
+
+        worker.start()
+        worker_group_id = f"native-{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": self._max_workers,
+                    "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

scaler/worker_adapter/symphony/callback.py

@@ -0,0 +1,45 @@
+import concurrent.futures
+import threading
+from typing import Dict
+
+import cloudpickle
+
+from scaler.worker_adapter.symphony.message import SoamMessage
+
+try:
+    import soamapi
+except ImportError:
+    raise ImportError("IBM Spectrum Symphony API not found, please install it with 'pip install soamapi'.")
+
+
+class SessionCallback(soamapi.SessionCallback):
+    def __init__(self):
+        self._callback_lock = threading.Lock()
+        self._task_id_to_future: Dict[str, concurrent.futures.Future] = {}
+
+    def on_response(self, task_output_handle):
+        with self._callback_lock:
+            task_id = task_output_handle.get_id()
+
+            future = self._task_id_to_future.pop(task_id)
+
+            if task_output_handle.is_successful():
+                output_message = SoamMessage()
+                task_output_handle.populate_task_output(output_message)
+                result = cloudpickle.loads(output_message.get_payload())
+                future.set_result(result)
+            else:
+                future.set_exception(task_output_handle.get_exception().get_embedded_exception())
+
+    def on_exception(self, exception):
+        with self._callback_lock:
+            for future in self._task_id_to_future.values():
+                future.set_exception(exception)
+
+            self._task_id_to_future.clear()
+
+    def submit_task(self, task_id: str, future: concurrent.futures.Future):
+        self._task_id_to_future[task_id] = future
+
+    def get_callback_lock(self) -> threading.Lock:
+        return self._callback_lock

scaler/worker_adapter/symphony/heartbeat_manager.py

@@ -0,0 +1,79 @@
+import time
+from typing import Dict, Optional
+
+import psutil
+
+from scaler.config.types.object_storage_server import ObjectStorageConfig
+from scaler.io.mixins import AsyncConnector, AsyncObjectStorageConnector
+from scaler.protocol.python.message import WorkerHeartbeat, WorkerHeartbeatEcho
+from scaler.protocol.python.status import Resource
+from scaler.utility.mixins import Looper
+from scaler.worker.agent.mixins import HeartbeatManager, TimeoutManager
+from scaler.worker_adapter.symphony.task_manager import SymphonyTaskManager
+
+
+class SymphonyHeartbeatManager(Looper, HeartbeatManager):
+    def __init__(
+        self, object_storage_address: Optional[ObjectStorageConfig], capabilities: Dict[str, int], task_queue_size: int
+    ):
+        self._capabilities = capabilities
+        self._task_queue_size = task_queue_size
+
+        self._agent_process = psutil.Process()
+
+        self._connector_external: Optional[AsyncConnector] = None
+        self._connector_storage: Optional[AsyncObjectStorageConnector] = None
+        self._worker_task_manager: Optional[SymphonyTaskManager] = None
+        self._timeout_manager: Optional[TimeoutManager] = None
+
+        self._start_timestamp_ns = 0
+        self._latency_us = 0
+
+        self._object_storage_address: Optional[ObjectStorageConfig] = object_storage_address
+
+    def register(
+        self,
+        connector_external: AsyncConnector,
+        connector_storage: AsyncObjectStorageConnector,
+        worker_task_manager: SymphonyTaskManager,
+        timeout_manager: TimeoutManager,
+    ):
+        self._connector_external = connector_external
+        self._connector_storage = connector_storage
+        self._worker_task_manager = worker_task_manager
+        self._timeout_manager = timeout_manager
+
+    async def on_heartbeat_echo(self, heartbeat: WorkerHeartbeatEcho):
+        if self._start_timestamp_ns == 0:
+            # not handling echo if we didn't send out heartbeat
+            return
+
+        self._latency_us = int(((time.time_ns() - self._start_timestamp_ns) / 2) // 1_000)
+        self._start_timestamp_ns = 0
+        self._timeout_manager.update_last_seen_time()
+
+        if self._object_storage_address is None:
+            address_message = heartbeat.object_storage_address()
+            self._object_storage_address = ObjectStorageConfig(address_message.host, address_message.port)
+            await self._connector_storage.connect(self._object_storage_address.host, self._object_storage_address.port)
+
+    def get_object_storage_address(self) -> Optional[ObjectStorageConfig]:
+        return self._object_storage_address
+
+    async def routine(self):
+        if self._start_timestamp_ns != 0:
+            return
+
+        await self._connector_external.send(
+            WorkerHeartbeat.new_msg(
+                Resource.new_msg(int(self._agent_process.cpu_percent() * 10), self._agent_process.memory_info().rss),
+                psutil.virtual_memory().available,
+                self._task_queue_size,
+                self._worker_task_manager.get_queued_size(),
+                self._latency_us,
+                self._worker_task_manager.can_accept_task(),
+                [],
+                self._capabilities,
+            )
+        )
+        self._start_timestamp_ns = time.time_ns()

scaler/worker_adapter/symphony/message.py

@@ -0,0 +1,24 @@
+import array
+
+try:
+    import soamapi
+except ImportError:
+    raise ImportError("IBM Spectrum Symphony API not found, please install it with 'pip install soamapi'.")
+
+
+class SoamMessage(soamapi.Message):
+    def __init__(self, payload: bytes = b""):
+        self.__payload = payload
+
+    def set_payload(self, payload: bytes):
+        self.__payload = payload
+
+    def get_payload(self) -> bytes:
+        return self.__payload
+
+    def on_serialize(self, stream):
+        payload_array = array.array("b", self.get_payload())
+        stream.write_byte_array(payload_array, 0, len(payload_array))
+
+    def on_deserialize(self, stream):
+        self.set_payload(stream.read_byte_array("b"))

scaler/worker_adapter/symphony/task_manager.py

@@ -0,0 +1,289 @@
+import asyncio
+import logging
+from concurrent.futures import Future
+from typing import Dict, Optional, Set, cast
+
+import cloudpickle
+from bidict import bidict
+
+from scaler import Serializer
+from scaler.io.mixins import AsyncConnector, AsyncObjectStorageConnector
+from scaler.protocol.python.common import ObjectMetadata, ObjectStorageAddress, TaskCancelConfirmType, TaskResultType
+from scaler.protocol.python.message import ObjectInstruction, Task, TaskCancel, TaskCancelConfirm, TaskResult
+from scaler.utility.identifiers import ObjectID, TaskID
+from scaler.utility.metadata.task_flags import retrieve_task_flags_from_task
+from scaler.utility.mixins import Looper
+from scaler.utility.queues.async_sorted_priority_queue import AsyncSortedPriorityQueue
+from scaler.utility.serialization import serialize_failure
+from scaler.worker.agent.mixins import HeartbeatManager, TaskManager
+from scaler.worker_adapter.symphony.callback import SessionCallback
+from scaler.worker_adapter.symphony.message import SoamMessage
+
+try:
+    import soamapi
+except ImportError:
+    raise ImportError("IBM Spectrum Symphony API not found, please install it with 'pip install soamapi'.")
+
+
+class SymphonyTaskManager(Looper, TaskManager):
+    def __init__(self, base_concurrency: int, service_name: str):
+        if isinstance(base_concurrency, int) and base_concurrency <= 0:
+            raise ValueError(f"base_concurrency must be a possible integer, got {base_concurrency}")
+
+        self._base_concurrency = base_concurrency
+        self._service_name = service_name
+
+        self._executor_semaphore = asyncio.Semaphore(value=self._base_concurrency)
+
+        self._task_id_to_task: Dict[TaskID, Task] = dict()
+        self._task_id_to_future: bidict[TaskID, asyncio.Future] = bidict()
+
+        self._serializers: Dict[bytes, Serializer] = dict()
+
+        self._queued_task_id_queue = AsyncSortedPriorityQueue()
+        self._queued_task_ids: Set[bytes] = set()
+
+        self._acquiring_task_ids: Set[TaskID] = set()  # tasks contesting the semaphore
+        self._processing_task_ids: Set[TaskID] = set()
+        self._canceled_task_ids: Set[TaskID] = set()
+
+        self._object_storage_address: Optional[ObjectStorageAddress] = None
+
+        self._connector_external: Optional[AsyncConnector] = None
+        self._connector_storage: Optional[AsyncObjectStorageConnector] = None
+
+        """
+        SOAM specific code
+        """
+        soamapi.initialize()
+
+        self._session_callback = SessionCallback()
+
+        self._ibm_soam_connection = soamapi.connect(
+            self._service_name, soamapi.DefaultSecurityCallback("Guest", "Guest")
+        )
+        logging.info(f"established IBM Spectrum Symphony connection {self._ibm_soam_connection.get_id()}")
+
+        ibm_soam_session_attr = soamapi.SessionCreationAttributes()
+        ibm_soam_session_attr.set_session_type("RecoverableAllHistoricalData")
+        ibm_soam_session_attr.set_session_name("ScalerSession")
+        ibm_soam_session_attr.set_session_flags(soamapi.SessionFlags.PARTIAL_ASYNC)
+        ibm_soam_session_attr.set_session_callback(self._session_callback)
+        self._ibm_soam_session = self._ibm_soam_connection.create_session(ibm_soam_session_attr)
+        logging.info(f"established IBM Spectrum Symphony session {self._ibm_soam_session.get_id()}")
+
+    def register(
+        self,
+        connector_external: AsyncConnector,
+        connector_storage: AsyncObjectStorageConnector,
+        heartbeat_manager: HeartbeatManager,
+    ):
+        self._connector_external = connector_external
+        self._connector_storage = connector_storage
+        self._heartbeat_manager = heartbeat_manager
+
+    async def routine(self):  # SymphonyTaskManager has two loops
+        pass
+
+    async def on_object_instruction(self, instruction: ObjectInstruction):
+        if instruction.instruction_type == ObjectInstruction.ObjectInstructionType.Delete:
+            for object_id in instruction.object_metadata.object_ids:
+                self._serializers.pop(object_id, None)  # we only cache serializers
+
+            return
+
+        logging.error(f"worker received unknown object instruction type {instruction=}")
+
+    async def on_task_new(self, task: Task):
+        task_priority = self.__get_task_priority(task)
+
+        # if semaphore is locked, check if task is higher priority than all acquired tasks
+        # if so, bypass acquiring and execute the task immediately
+        if self._executor_semaphore.locked():
+            for acquired_task_id in self._acquiring_task_ids:
+                acquired_task = self._task_id_to_task[acquired_task_id]
+                acquired_task_priority = self.__get_task_priority(acquired_task)
+                if task_priority <= acquired_task_priority:
+                    break
+            else:
+                self._task_id_to_task[task.task_id] = task
+                self._processing_task_ids.add(task.task_id)
+                self._task_id_to_future[task.task_id] = await self.__execute_task(task)
+                return
+
+        self._task_id_to_task[task.task_id] = task
+        self._queued_task_id_queue.put_nowait((-task_priority, task.task_id))
+        self._queued_task_ids.add(task.task_id)
+
+    async def on_cancel_task(self, task_cancel: TaskCancel):
+        task_queued = task_cancel.task_id in self._queued_task_ids
+        task_processing = task_cancel.task_id in self._processing_task_ids
+
+        if not task_queued and not task_processing:
+            await self._connector_external.send(
+                TaskCancelConfirm.new_msg(
+                    task_id=task_cancel.task_id, cancel_confirm_type=TaskCancelConfirmType.CancelNotFound
+                )
+            )
+            return
+
+        if task_processing and not task_cancel.flags.force:
+            await self._connector_external.send(
+                TaskCancelConfirm.new_msg(
+                    task_id=task_cancel.task_id, cancel_confirm_type=TaskCancelConfirmType.CancelFailed
+                )
+            )
+            return
+
+        if task_queued:
+            self._queued_task_ids.remove(task_cancel.task_id)
+            self._queued_task_id_queue.remove(task_cancel.task_id)
+
+            # task can be discarded because task was never submitted
+            self._task_id_to_task.pop(task_cancel.task_id)
+
+        if task_processing:
+            future = self._task_id_to_future[task_cancel.task_id]
+            future.cancel()
+
+            # regardless of the future being canceled, the task is considered canceled and cleanup will occur later
+            self._processing_task_ids.remove(task_cancel.task_id)
+            self._canceled_task_ids.add(task_cancel.task_id)
+
+        result = TaskCancelConfirm.new_msg(
+            task_id=task_cancel.task_id, cancel_confirm_type=TaskCancelConfirmType.Canceled
+        )
+        await self._connector_external.send(result)
+
+    async def on_task_result(self, result: TaskResult):
+        if result.task_id in self._queued_task_ids:
+            self._queued_task_ids.remove(result.task_id)
+            self._queued_task_id_queue.remove(result.task_id)
+
+        self._processing_task_ids.remove(result.task_id)
+        self._task_id_to_task.pop(result.task_id)
+
+        await self._connector_external.send(result)
+
+    def get_queued_size(self):
+        return self._queued_task_id_queue.qsize()
+
+    def can_accept_task(self):
+        return not self._executor_semaphore.locked()
+
+    async def resolve_tasks(self):
+        if not self._task_id_to_future:
+            await asyncio.sleep(0)
+            return
+
+        done, _ = await asyncio.wait(self._task_id_to_future.values(), return_when=asyncio.FIRST_COMPLETED)
+        for future in done:
+            task_id = self._task_id_to_future.inv.pop(future)
+            task = self._task_id_to_task[task_id]
+
+            if task_id in self._processing_task_ids:
+                self._processing_task_ids.remove(task_id)
+
+                if future.exception() is None:
+                    serializer_id = ObjectID.generate_serializer_object_id(task.source)
+                    serializer = self._serializers[serializer_id]
+                    result_bytes = serializer.serialize(future.result())
+                    result_type = TaskResultType.Success
+                else:
+                    result_bytes = serialize_failure(cast(Exception, future.exception()))
+                    result_type = TaskResultType.Failed
+
+                result_object_id = ObjectID.generate_object_id(task.source)
+
+                await self._connector_storage.set_object(result_object_id, result_bytes)
+                await self._connector_external.send(
+                    ObjectInstruction.new_msg(
+                        ObjectInstruction.ObjectInstructionType.Create,
+                        task.source,
+                        ObjectMetadata.new_msg(
+                            object_ids=(result_object_id,),
+                            object_types=(ObjectMetadata.ObjectContentType.Object,),
+                            object_names=(f"<res {result_object_id.hex()[:6]}>".encode(),),
+                        ),
+                    )
+                )
+
+                await self._connector_external.send(
+                    TaskResult.new_msg(task_id, result_type, metadata=b"", results=[bytes(result_object_id)])
+                )
+
+            elif task_id in self._canceled_task_ids:
+                self._canceled_task_ids.remove(task_id)
+
+            else:
+                raise ValueError(f"task_id {task_id.hex()} not found in processing or canceled tasks")
+
+            if task_id in self._acquiring_task_ids:
+                self._acquiring_task_ids.remove(task_id)
+                self._executor_semaphore.release()
+
+            self._task_id_to_task.pop(task_id)
+
+    async def process_task(self):
+        await self._executor_semaphore.acquire()
+
+        _, task_id = await self._queued_task_id_queue.get()
+        task = self._task_id_to_task[task_id]
+
+        self._acquiring_task_ids.add(task_id)
+        self._processing_task_ids.add(task_id)
+        self._task_id_to_future[task.task_id] = await self.__execute_task(task)
+
+    async def __execute_task(self, task: Task) -> asyncio.Future:
+        """
+        This method is not very efficient because it does let objects linger in the cache. Each time inputs are
+        requested, all object data are requested.
+        """
+        serializer_id = ObjectID.generate_serializer_object_id(task.source)
+
+        if serializer_id not in self._serializers:
+            serializer_bytes = await self._connector_storage.get_object(serializer_id)
+            serializer = cloudpickle.loads(serializer_bytes)
+            self._serializers[serializer_id] = serializer
+        else:
+            serializer = self._serializers[serializer_id]
+
+        # Fetches the function object and the argument objects concurrently
+
+        get_tasks = [
+            self._connector_storage.get_object(object_id)
+            for object_id in [task.func_object_id, *(cast(ObjectID, arg) for arg in task.function_args)]
+        ]
+
+        function_bytes, *arg_bytes = await asyncio.gather(*get_tasks)
+
+        function = serializer.deserialize(function_bytes)
+        arg_objects = [serializer.deserialize(object_bytes) for object_bytes in arg_bytes]
+
+        """
+        SOAM specific code
+        """
+        input_message = SoamMessage()
+        input_message.set_payload(cloudpickle.dumps((function, *arg_objects)))
+
+        task_attr = soamapi.TaskSubmissionAttributes()
+        task_attr.set_task_input(input_message)
+
+        with self._session_callback.get_callback_lock():
+            symphony_task = self._ibm_soam_session.send_task_input(task_attr)
+
+            future: Future = Future()
+            future.set_running_or_notify_cancel()
+
+            self._session_callback.submit_task(symphony_task.get_id(), future)
+
+        return asyncio.wrap_future(future)
+
+    @staticmethod
+    def __get_task_priority(task: Task) -> int:
+        priority = retrieve_task_flags_from_task(task).priority
+
+        if priority < 0:
+            raise ValueError(f"invalid task priority, must be positive or zero, got {priority}")
+
+        return priority