opengris-scaler 1.12.28__cp313-cp313-musllinux_1_2_x86_64.whl

scaler/entry_points/worker_adapter_native.py

@@ -0,0 +1,137 @@
+import argparse
+
+from aiohttp import web
+
+from scaler.config.loader import load_config
+from scaler.config.section.native_worker_adapter import NativeWorkerAdapterConfig
+from scaler.utility.event_loop import EventLoopType, register_event_loop
+from scaler.utility.logging.utility import setup_logger
+from scaler.worker_adapter.native import NativeWorkerAdapter
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        "scaler_native_worker_adapter", formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+
+    parser.add_argument("--config", "-c", type=str, default=None, help="Path to the TOML configuration file.")
+
+    # Server configuration
+    parser.add_argument("--adapter-web-host", type=str, help="Host for the native worker adapter HTTP server.")
+    parser.add_argument("--adapter-web-port", "-p", type=int, help="Port for the native worker adapter HTTP server.")
+
+    # Worker configuration
+    parser.add_argument("--io-threads", type=int, help="number of io threads for zmq")
+    parser.add_argument(
+        "--per-worker-capabilities",
+        "-pwc",
+        type=str,
+        help='comma-separated capabilities provided by the workers (e.g. "-pwc linux,cpu=4")',
+    )
+    parser.add_argument("--worker-task-queue-size", "-wtqs", type=int, default=10, help="specify worker queue size")
+    parser.add_argument(
+        "--max-workers", "-mw", type=int, help="maximum number of workers that can be started, -1 means no limit"
+    )
+    parser.add_argument(
+        "--heartbeat-interval", "-hi", type=int, help="number of seconds that worker agent send heartbeat to scheduler"
+    )
+    parser.add_argument(
+        "--task-timeout-seconds", "-tt", type=int, help="default task timeout seconds, 0 means never timeout"
+    )
+    parser.add_argument(
+        "--death-timeout-seconds",
+        "-dt",
+        type=int,
+        help="number of seconds without scheduler contact before worker shuts down",
+    )
+    parser.add_argument(
+        "--garbage-collect-interval-seconds", "-gc", type=int, help="number of seconds worker doing garbage collection"
+    )
+    parser.add_argument(
+        "--trim-memory-threshold-bytes",
+        "-tm",
+        type=int,
+        help="number of bytes threshold for worker process that trigger deep garbage collection",
+    )
+    parser.add_argument(
+        "--hard-processor-suspend",
+        "-hps",
+        action="store_true",
+        help="if true, suspended worker's processors will be actively suspended with a SIGTSTP signal",
+    )
+    parser.add_argument("--event-loop", "-e", choices=EventLoopType.allowed_types(), help="select event loop type")
+    parser.add_argument(
+        "--logging-paths",
+        "-lp",
+        nargs="*",
+        type=str,
+        help="specify where worker logs should be logged to, it can accept multiple files, default is /dev/stdout",
+    )
+    parser.add_argument(
+        "--logging-level",
+        "-ll",
+        type=str,
+        choices=("DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"),
+        help="specify the logging level",
+    )
+    parser.add_argument(
+        "--logging-config-file",
+        "-lc",
+        type=str,
+        help="use standard python .conf file to specify python logging file configuration format",
+    )
+    parser.add_argument(
+        "--object-storage-address",
+        "-osa",
+        type=str,
+        help="specify the object storage server address, e.g.: tcp://localhost:2346",
+    )
+    parser.add_argument(
+        "scheduler_address",
+        nargs="?",
+        type=str,
+        help="scheduler address to connect workers to, e.g.: `tcp://localhost:6378",
+    )
+
+    return parser.parse_args()
+
+
+def main():
+    args = get_args()
+    native_adapter_config = load_config(
+        NativeWorkerAdapterConfig, args.config, args, section_name="native_worker_adapter"
+    )
+
+    register_event_loop(native_adapter_config.event_loop)
+
+    setup_logger(
+        native_adapter_config.logging_paths,
+        native_adapter_config.logging_config_file,
+        native_adapter_config.logging_level,
+    )
+
+    native_worker_adapter = NativeWorkerAdapter(
+        address=native_adapter_config.scheduler_address,
+        object_storage_address=native_adapter_config.object_storage_address,
+        capabilities=native_adapter_config.per_worker_capabilities.capabilities,
+        io_threads=native_adapter_config.io_threads,
+        task_queue_size=native_adapter_config.worker_task_queue_size,
+        max_workers=native_adapter_config.max_workers,
+        heartbeat_interval_seconds=native_adapter_config.heartbeat_interval_seconds,
+        task_timeout_seconds=native_adapter_config.task_timeout_seconds,
+        death_timeout_seconds=native_adapter_config.death_timeout_seconds,
+        garbage_collect_interval_seconds=native_adapter_config.garbage_collect_interval_seconds,
+        trim_memory_threshold_bytes=native_adapter_config.trim_memory_threshold_bytes,
+        hard_processor_suspend=native_adapter_config.hard_processor_suspend,
+        event_loop=native_adapter_config.event_loop,
+        logging_paths=native_adapter_config.logging_paths,
+        logging_level=native_adapter_config.logging_level,
+        logging_config_file=native_adapter_config.logging_config_file,
+    )
+
+    app = native_worker_adapter.create_app()
+    web.run_app(app, host=native_adapter_config.adapter_web_host, port=native_adapter_config.adapter_web_port)
+
+
+if __name__ == "__main__":
+    main()

scaler/entry_points/worker_adapter_symphony.py

@@ -0,0 +1,98 @@
+import argparse
+
+from aiohttp import web
+
+from scaler.config.loader import load_config
+from scaler.config.section.symphony_worker_adapter import SymphonyWorkerConfig
+from scaler.utility.event_loop import EventLoopType, register_event_loop
+from scaler.utility.logging.utility import setup_logger
+from scaler.worker_adapter.symphony.worker_adapter import SymphonyWorkerAdapter
+
+
+def get_args():
+    parser = argparse.ArgumentParser(
+        "scaler Symphony worker adapter", formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument("--config", "-c", type=str, default=None, help="Path to the TOML configuration file.")
+
+    # Server configuration
+    parser.add_argument("--adapter-web-host", type=str, help="host address for symphony worker adapter HTTP server")
+    parser.add_argument("--adapter-web-port", "-p", type=int, help="port for symphony worker adapter HTTP server")
+
+    # Symphony configuration
+    parser.add_argument("--service-name", "-sn", type=str, help="symphony service name")
+    parser.add_argument("--base-concurrency", "-n", type=int, help="base task concurrency")
+
+    # Worker configuration
+    parser.add_argument("--io-threads", "-it", help="specify number of io threads per worker")
+    parser.add_argument(
+        "--worker-capabilities",
+        "-wc",
+        type=str,
+        help='comma-separated capabilities provided by the worker (e.g. "-wr linux,cpu=4")',
+    )
+    parser.add_argument("--worker-task-queue-size", "-wtqs", type=int, help="specify symphony worker queue size")
+    parser.add_argument("--heartbeat-interval", "-hi", type=int, help="number of seconds to send heartbeat interval")
+    parser.add_argument("--death-timeout-seconds", "-ds", type=int, help="death timeout seconds")
+    parser.add_argument("--event-loop", "-el", choices=EventLoopType.allowed_types(), help="select event loop type")
+    parser.add_argument(
+        "--logging-paths",
+        "-lp",
+        nargs="*",
+        type=str,
+        help='specify where cluster log should logged to, it can be multiple paths, "/dev/stdout" is default for '
+        "standard output, each worker will have its own log file with process id appended to the path",
+    )
+    parser.add_argument(
+        "--logging-level",
+        "-ll",
+        type=str,
+        choices=("DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"),
+        help="specify the logging level",
+    )
+    parser.add_argument(
+        "--logging-config-file",
+        type=str,
+        default=None,
+        help="use standard python the .conf file the specify python logging file configuration format, this will "
+        "bypass --logging-paths and --logging-level at the same time, and this will not work on per worker logging",
+    )
+    parser.add_argument(
+        "--object-storage-address",
+        "-osa",
+        default=None,
+        help="specify the object storage server address, e.g.: tcp://localhost:2346",
+    )
+    parser.add_argument("scheduler_address", nargs="?", type=str, help="scheduler address to connect to")
+    return parser.parse_args()
+
+
+def main():
+    args = get_args()
+    symphony_config = load_config(SymphonyWorkerConfig, args.config, args, section_name="symphony_worker_adapter")
+    register_event_loop(symphony_config.event_loop)
+
+    setup_logger(symphony_config.logging_paths, symphony_config.logging_config_file, symphony_config.logging_level)
+
+    symphony_worker_adapter = SymphonyWorkerAdapter(
+        address=symphony_config.scheduler_address,
+        object_storage_address=symphony_config.object_storage_address,
+        capabilities=symphony_config.worker_capabilities.capabilities,
+        task_queue_size=symphony_config.worker_task_queue_size,
+        service_name=symphony_config.service_name,
+        base_concurrency=symphony_config.base_concurrency,
+        heartbeat_interval_seconds=symphony_config.heartbeat_interval,
+        death_timeout_seconds=symphony_config.death_timeout_seconds,
+        event_loop=symphony_config.event_loop,
+        io_threads=symphony_config.io_threads,
+        logging_paths=symphony_config.logging_paths,
+        logging_level=symphony_config.logging_level,
+        logging_config_file=symphony_config.logging_config_file,
+    )
+
+    app = symphony_worker_adapter.create_app()
+    web.run_app(app, host=symphony_config.adapter_web_host, port=symphony_config.adapter_web_port)
+
+
+if __name__ == "__main__":
+    main()

scaler/io/async_binder.py

@@ -0,0 +1,89 @@
+import logging
+import os
+import uuid
+from collections import defaultdict
+from typing import Awaitable, Callable, Dict, List, Optional
+
+import zmq.asyncio
+from zmq import Frame
+
+from scaler.config.types.zmq import ZMQConfig
+from scaler.io.mixins import AsyncBinder
+from scaler.io.utility import deserialize, serialize
+from scaler.protocol.python.mixins import Message
+from scaler.protocol.python.status import BinderStatus
+
+
+class ZMQAsyncBinder(AsyncBinder):
+    def __init__(self, context: zmq.asyncio.Context, name: str, address: ZMQConfig, identity: Optional[bytes] = None):
+        self._address = address
+
+        if identity is None:
+            identity = f"{os.getpid()}|{name}|{uuid.uuid4()}".encode()
+        self._identity = identity
+
+        self._context = context
+        self._socket = self._context.socket(zmq.ROUTER)
+        self.__set_socket_options()
+        self._socket.bind(self._address.to_address())
+
+        self._callback: Optional[Callable[[bytes, Message], Awaitable[None]]] = None
+
+        self._received: Dict[str, int] = defaultdict(lambda: 0)
+        self._sent: Dict[str, int] = defaultdict(lambda: 0)
+
+    @property
+    def identity(self):
+        return self._identity
+
+    def destroy(self):
+        self._context.destroy(linger=0)
+
+    def register(self, callback: Callable[[bytes, Message], Awaitable[None]]):
+        self._callback = callback
+
+    async def routine(self):
+        frames: List[Frame] = await self._socket.recv_multipart(copy=False)
+        if not self.__is_valid_message(frames):
+            return
+
+        source, payload = frames
+        try:
+            message: Optional[Message] = deserialize(payload.bytes)
+            if message is None:
+                logging.error(f"received unknown message from {source.bytes!r}: {payload!r}")
+                return
+        except Exception as e:
+            logging.error(f"{self.__get_prefix()} failed to deserialize message from {source.bytes!r}: {e}")
+            return
+
+        self.__count_received(message.__class__.__name__)
+        await self._callback(source.bytes, message)
+
+    async def send(self, to: bytes, message: Message):
+        self.__count_sent(message.__class__.__name__)
+        await self._socket.send_multipart([to, serialize(message)], copy=False)
+
+    def get_status(self) -> BinderStatus:
+        return BinderStatus.new_msg(received=self._received, sent=self._sent)
+
+    def __set_socket_options(self):
+        self._socket.setsockopt(zmq.IDENTITY, self._identity)
+        self._socket.setsockopt(zmq.SNDHWM, 0)
+        self._socket.setsockopt(zmq.RCVHWM, 0)
+
+    def __is_valid_message(self, frames: List[Frame]) -> bool:
+        if len(frames) != 2:
+            logging.error(f"{self.__get_prefix()} received unexpected frames {frames}")
+            return False
+
+        return True
+
+    def __count_received(self, message_type: str):
+        self._received[message_type] += 1
+
+    def __count_sent(self, message_type: str):
+        self._sent[message_type] += 1
+
+    def __get_prefix(self):
+        return f"{self.__class__.__name__}[{self._identity.decode()}]:"

scaler/io/async_connector.py

@@ -0,0 +1,95 @@
+import logging
+import os
+import uuid
+from typing import Awaitable, Callable, Literal, Optional
+
+import zmq.asyncio
+
+from scaler.config.types.zmq import ZMQConfig
+from scaler.io.mixins import AsyncConnector
+from scaler.io.utility import deserialize, serialize
+from scaler.protocol.python.mixins import Message
+
+
+class ZMQAsyncConnector(AsyncConnector):
+    def __init__(
+        self,
+        context: zmq.asyncio.Context,
+        name: str,
+        socket_type: int,
+        address: ZMQConfig,
+        bind_or_connect: Literal["bind", "connect"],
+        callback: Optional[Callable[[Message], Awaitable[None]]],
+        identity: Optional[bytes],
+    ):
+        self._address = address
+
+        self._context = context
+        self._socket = self._context.socket(socket_type)
+
+        if identity is None:
+            identity = f"{os.getpid()}|{name}|{uuid.uuid4().bytes.hex()}".encode()
+        self._identity = identity
+
+        # set socket option
+        self._socket.setsockopt(zmq.IDENTITY, self._identity)
+        self._socket.setsockopt(zmq.SNDHWM, 0)
+        self._socket.setsockopt(zmq.RCVHWM, 0)
+
+        if bind_or_connect == "bind":
+            self._socket.bind(self.address)
+        elif bind_or_connect == "connect":
+            self._socket.connect(self.address)
+        else:
+            raise TypeError("bind_or_connect has to be 'bind' or 'connect'")
+
+        self._callback: Optional[Callable[[Message], Awaitable[None]]] = callback
+
+    def __del__(self):
+        self.destroy()
+
+    def destroy(self):
+        if self._socket.closed:
+            return
+
+        self._socket.close(linger=1)
+
+    @property
+    def identity(self) -> bytes:
+        return self._identity
+
+    @property
+    def socket(self) -> zmq.asyncio.Socket:
+        return self._socket
+
+    @property
+    def address(self) -> str:
+        return self._address.to_address()
+
+    async def routine(self):
+        if self._callback is None:
+            return
+
+        message: Optional[Message] = await self.receive()
+        if message is None:
+            return
+
+        await self._callback(message)
+
+    async def receive(self) -> Optional[Message]:
+        if self._context.closed:
+            return None
+
+        if self._socket.closed:
+            return None
+
+        payload = await self._socket.recv(copy=False)
+        result: Optional[Message] = deserialize(payload.bytes)
+        if result is None:
+            logging.error(f"received unknown message: {payload.bytes!r}")
+            return None
+
+        return result
+
+    async def send(self, message: Message):
+        await self._socket.send(serialize(message), copy=False)

scaler/io/async_object_storage_connector.py

@@ -0,0 +1,225 @@
+import asyncio
+import logging
+import socket
+import struct
+import uuid
+from typing import Dict, Optional, Tuple
+
+from scaler.io.mixins import AsyncObjectStorageConnector
+from scaler.protocol.capnp._python import _object_storage  # noqa
+from scaler.protocol.python.object_storage import ObjectRequestHeader, ObjectResponseHeader, to_capnp_object_id
+from scaler.utility.exceptions import ObjectStorageException
+from scaler.utility.identifiers import ObjectID
+
+
+class PyAsyncObjectStorageConnector(AsyncObjectStorageConnector):
+    """An asyncio connector that uses an raw TCP socket to connect to a Scaler's object storage instance."""
+
+    def __init__(self):
+        self._host: Optional[str] = None
+        self._port: Optional[int] = None
+
+        self._connected_event = asyncio.Event()
+
+        self._reader: Optional[asyncio.StreamReader] = None
+        self._writer: Optional[asyncio.StreamWriter] = None
+
+        self._next_request_id = 0
+        self._pending_get_requests: Dict[ObjectID, asyncio.Future] = {}
+
+        self._identity: bytes = (
+            f"{self.__class__.__name__}|{socket.gethostname().split('.')[0]}|{uuid.uuid4()}".encode()
+        )
+
+    def __del__(self):
+        if not self.is_connected():
+            return
+
+        self._writer.close()
+
+    async def connect(self, host: str, port: int):
+        self._host = host
+        self._port = port
+
+        if self.is_connected():
+            raise ObjectStorageException("connector is already connected.")
+
+        self._reader, self._writer = await asyncio.open_connection(self._host, self._port)
+        await self.__read_framed_message()
+        self.__write_framed(self._identity)
+
+        try:
+            await self._writer.drain()
+        except ConnectionResetError:
+            self.__raise_connection_failure()
+
+        # Makes sure the socket is TCP_NODELAY. It seems to be the case by default, but that's not specified in the
+        # asyncio's documentation and might change in the future.
+        self._writer.get_extra_info("socket").setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+
+        self._connected_event.set()
+
+    async def wait_until_connected(self):
+        await self._connected_event.wait()
+
+    def is_connected(self) -> bool:
+        return self._connected_event.is_set()
+
+    async def destroy(self):
+        if not self.is_connected():
+            return
+
+        if not self._writer.is_closing:
+            self._writer.close()
+
+        await self._writer.wait_closed()
+
+    @property
+    def reader(self) -> Optional[asyncio.StreamReader]:
+        return self._reader
+
+    @property
+    def writer(self) -> Optional[asyncio.StreamWriter]:
+        return self._writer
+
+    @property
+    def address(self) -> str:
+        self.__ensure_is_connected()
+        return f"tcp://{self._host}:{self._port}"
+
+    async def routine(self):
+        await self.wait_until_connected()
+
+        response = await self.__receive_response()
+        if response is None:
+            return
+
+        header, payload = response
+
+        if header.response_type != ObjectResponseHeader.ObjectResponseType.GetOK:
+            return
+
+        pending_get_future = self._pending_get_requests.pop(header.object_id, None)
+
+        if pending_get_future is None:
+            logging.warning(f"unknown get-ok response for unrequested object_id={repr(header.object_id)}.")
+            return
+
+        pending_get_future.set_result(payload)
+
+    async def set_object(self, object_id: ObjectID, payload: bytes) -> None:
+        await self.__send_request(object_id, len(payload), ObjectRequestHeader.ObjectRequestType.SetObject, payload)
+
+    async def get_object(self, object_id: ObjectID, max_payload_length: int = 2**64 - 1) -> bytes:
+        pending_get_future = self._pending_get_requests.get(object_id)
+
+        if pending_get_future is None:
+            pending_get_future = asyncio.Future()
+            self._pending_get_requests[object_id] = pending_get_future
+
+            await self.__send_request(
+                object_id, max_payload_length, ObjectRequestHeader.ObjectRequestType.GetObject, None
+            )
+
+        return await pending_get_future
+
+    async def delete_object(self, object_id: ObjectID) -> None:
+        await self.__send_request(object_id, 0, ObjectRequestHeader.ObjectRequestType.DeleteObject, None)
+
+    async def duplicate_object_id(self, object_id: ObjectID, new_object_id: ObjectID) -> None:
+        object_id_payload = to_capnp_object_id(object_id).to_bytes()
+
+        await self.__send_request(
+            new_object_id,
+            len(object_id_payload),
+            ObjectRequestHeader.ObjectRequestType.DuplicateObjectID,
+            object_id_payload,
+        )
+
+    def __ensure_is_connected(self):
+        if self._writer is None:
+            raise ObjectStorageException("connector is not connected.")
+
+        if self._writer.is_closing():
+            raise ObjectStorageException("connection is closed.")
+
+    async def __send_request(
+        self,
+        object_id: ObjectID,
+        payload_length: int,
+        request_type: ObjectRequestHeader.ObjectRequestType,
+        payload: Optional[bytes],
+    ):
+        self.__ensure_is_connected()
+        assert self._writer is not None
+
+        request_id = self._next_request_id
+        self._next_request_id += 1
+        self._next_request_id %= 2**64 - 1  # UINT64_MAX
+
+        header = ObjectRequestHeader.new_msg(object_id, payload_length, request_id, request_type)
+
+        self.__write_request_header(header)
+
+        if payload is not None:
+            self.__write_request_payload(payload)
+
+        try:
+            await self._writer.drain()
+        except ConnectionResetError:
+            self.__raise_connection_failure()
+
+    def __write_request_header(self, header: ObjectRequestHeader):
+        assert self._writer is not None
+        self.__write_framed(header.get_message().to_bytes())
+
+    def __write_request_payload(self, payload: bytes):
+        assert self._writer is not None
+        self.__write_framed(payload)
+
+    async def __receive_response(self) -> Optional[Tuple[ObjectResponseHeader, bytes]]:
+        assert self._reader is not None
+
+        if self._writer.is_closing():
+            return None
+
+        try:
+            header = await self.__read_response_header()
+            payload = await self.__read_response_payload(header)
+        except asyncio.IncompleteReadError:
+            self.__raise_connection_failure()
+
+        return header, payload
+
+    async def __read_response_header(self) -> ObjectResponseHeader:
+        assert self._reader is not None
+
+        header_data = await self.__read_framed_message()
+        assert len(header_data) == ObjectResponseHeader.MESSAGE_LENGTH
+
+        with _object_storage.ObjectResponseHeader.from_bytes(header_data) as header_message:
+            return ObjectResponseHeader(header_message)
+
+    async def __read_response_payload(self, header: ObjectResponseHeader) -> bytes:
+        assert self._reader is not None
+
+        if header.payload_length > 0:
+            res = await self.__read_framed_message()
+            assert len(res) == header.payload_length
+            return res
+        else:
+            return b""
+
+    async def __read_framed_message(self) -> bytes:
+        length_bytes = await self._reader.readexactly(8)
+        (payload_length,) = struct.unpack("<Q", length_bytes)
+        return await self._reader.readexactly(payload_length) if payload_length > 0 else bytes()
+
+    def __write_framed(self, payload: bytes):
+        self._writer.write(struct.pack("<Q", len(payload)))
+        self._writer.write(payload)
+        return
+
+    @staticmethod
+    def __raise_connection_failure():
+        raise ObjectStorageException("connection failure to object storage server.")