opengris-scaler 1.12.37__cp38-cp38-musllinux_1_2_x86_64.whl

scaler/utility/queues/indexed_queue.py

@@ -0,0 +1,114 @@
+import dataclasses
+from typing import Any, Dict, Hashable, Optional
+
+
+@dataclasses.dataclass
+class _Node:
+    value: Any
+    prev: Optional["_Node"] = None
+    next: Optional["_Node"] = None
+
+
+class _DoubleLinkedList:
+    def __init__(self):
+        self._head: Optional[_Node] = None
+        self._tail: Optional[_Node] = None
+        self._size = 0
+
+    def __len__(self):
+        return self._size
+
+    def add_to_head(self, node: _Node):
+        if self._head is None:
+            self._head = node
+            self._tail = node
+        else:
+            node.next = self._head
+            self._head.prev = node
+            self._head = node
+
+        self._size += 1
+
+    def remove_tail(self):
+        if self._tail is None:
+            raise IndexError(f"{self.__class__.__name__} queue empty")
+
+        node = self._tail
+        if self._tail.prev is None:
+            self._head = None
+            self._tail = None
+        else:
+            self._tail = self._tail.prev
+            self._tail.next = None
+
+        self._size -= 1
+        return node
+
+    def remove(self, node: _Node):
+        prev_node = node.prev
+        next_node = node.next
+        if prev_node and next_node:
+            prev_node.next = next_node
+            next_node.prev = prev_node
+
+        elif not prev_node and not next_node:
+            assert self._head is node
+            assert self._tail is node
+            self._head = None
+            self._tail = None
+
+        elif prev_node and not next_node:
+            assert self._tail is node
+            prev_node.next = None
+            self._tail = prev_node
+
+        elif not prev_node and next_node:
+            assert self._head is node
+            next_node.prev = None
+            self._head = next_node
+
+        self._size -= 1
+        del node
+
+
+class IndexedQueue:
+    """A queue that provides O(1) operations for adding and removing any item."""
+
+    def __init__(self):
+        self._double_linked_list = _DoubleLinkedList()
+        self._hash_map: Dict[int, _Node] = {}
+
+    def __contains__(self, item: Hashable):
+        key = hash(item)
+        return key in self._hash_map
+
+    def __len__(self):
+        return self._double_linked_list.__len__()
+
+    def __iter__(self):
+        node = self._double_linked_list._tail
+        while node is not None:
+            yield node.value
+            node = node.prev
+
+    def put(self, item: Hashable):
+        key = hash(item)
+        if key in self._hash_map:
+            raise KeyError(f"{self.__class__.__name__} already have item: {item}")
+
+        node = _Node(item)
+        self._double_linked_list.add_to_head(node)
+        self._hash_map[key] = node
+
+    def get(self):
+        node = self._double_linked_list.remove_tail()
+        del self._hash_map[hash(node.value)]
+        return node.value
+
+    def remove(self, item: Hashable):
+        key = hash(item)
+        if key not in self._hash_map:
+            raise ValueError(f"{self.__class__.__name__} doesn't have item: {item}")
+
+        node = self._hash_map.pop(key)
+        self._double_linked_list.remove(node)

scaler/utility/serialization.py

@@ -0,0 +1,9 @@
+import pickle
+
+
+def serialize_failure(exp: Exception) -> bytes:
+    return pickle.dumps(exp, protocol=pickle.HIGHEST_PROTOCOL)
+
+
+def deserialize_failure(result: bytes) -> Exception:
+    return pickle.loads(result)

scaler/version.txt

@@ -0,0 +1 @@
+1.12.37

scaler/worker/agent/heartbeat_manager.py

@@ -0,0 +1,110 @@
+import time
+from typing import Dict, Optional
+
+import psutil
+
+from scaler.config.types.object_storage_server import ObjectStorageAddressConfig
+from scaler.io.mixins import AsyncConnector, AsyncObjectStorageConnector
+from scaler.protocol.python.message import Resource, WorkerHeartbeat, WorkerHeartbeatEcho
+from scaler.protocol.python.status import ProcessorStatus
+from scaler.utility.mixins import Looper
+from scaler.worker.agent.mixins import HeartbeatManager, ProcessorManager, TaskManager, TimeoutManager
+from scaler.worker.agent.processor_holder import ProcessorHolder
+
+
+class VanillaHeartbeatManager(Looper, HeartbeatManager):
+    def __init__(
+        self,
+        object_storage_address: Optional[ObjectStorageAddressConfig],
+        capabilities: Dict[str, int],
+        task_queue_size: int,
+    ):
+        self._agent_process = psutil.Process()
+        self._capabilities = capabilities
+        self._task_queue_size = task_queue_size
+
+        self._connector_external: Optional[AsyncConnector] = None
+        self._connector_storage: Optional[AsyncObjectStorageConnector] = None
+        self._worker_task_manager: Optional[TaskManager] = None
+        self._timeout_manager: Optional[TimeoutManager] = None
+        self._processor_manager: Optional[ProcessorManager] = None
+
+        self._start_timestamp_ns = 0
+        self._latency_us = 0
+
+        self._object_storage_address: Optional[ObjectStorageAddressConfig] = object_storage_address
+
+    def register(
+        self,
+        connector_external: AsyncConnector,
+        connector_storage: AsyncObjectStorageConnector,
+        worker_task_manager: TaskManager,
+        timeout_manager: TimeoutManager,
+        processor_manager: ProcessorManager,
+    ):
+        self._connector_external = connector_external
+        self._connector_storage = connector_storage
+        self._worker_task_manager = worker_task_manager
+        self._timeout_manager = timeout_manager
+        self._processor_manager = processor_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 = ObjectStorageAddressConfig(address_message.host, address_message.port)
+            await self._connector_storage.connect(self._object_storage_address.host, self._object_storage_address.port)
+
+    async def routine(self):
+        processors = self._processor_manager.processors()
+
+        if self._start_timestamp_ns != 0:
+            # already sent heartbeat, expecting heartbeat echo, so not sending
+            return
+
+        for processor_holder in processors:
+            status = processor_holder.process().status()
+            if status in {psutil.STATUS_ZOMBIE, psutil.STATUS_DEAD}:
+                await self._processor_manager.on_failing_processor(processor_holder.processor_id(), status)
+
+        processors = self._processor_manager.processors()  # refreshes for removed dead and zombie processors
+        num_suspended_processors = self._processor_manager.num_suspended_processors()
+
+        # TODO: add task queue size to WorkerHeartbeat
+        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() - num_suspended_processors,
+                self._latency_us,
+                self._processor_manager.can_accept_task(),
+                [self.__get_processor_status_from_holder(processor) for processor in processors],
+                self._capabilities,
+            )
+        )
+        self._start_timestamp_ns = time.time_ns()
+
+    def get_object_storage_address(self) -> Optional[ObjectStorageAddressConfig]:
+        return self._object_storage_address
+
+    @staticmethod
+    def __get_processor_status_from_holder(processor: ProcessorHolder) -> ProcessorStatus:
+        process = processor.process()
+
+        try:
+            resource = Resource.new_msg(int(process.cpu_percent() * 10), process.memory_info().rss)
+        except psutil.ZombieProcess:
+            # Assumes dead processes do not use any resources
+            resource = Resource.new_msg(0, 0)
+
+        return ProcessorStatus.new_msg(
+            processor.pid(), processor.initialized(), processor.task() is not None, processor.suspended(), resource
+        )

scaler/worker/agent/mixins.py

@@ -0,0 +1,137 @@
+import abc
+from typing import List, Optional
+
+from scaler.config.types.object_storage_server import ObjectStorageAddressConfig
+from scaler.protocol.python.message import (
+    ObjectInstruction,
+    ProcessorInitialized,
+    Task,
+    TaskCancel,
+    TaskResult,
+    WorkerHeartbeatEcho,
+)
+from scaler.utility.identifiers import ProcessorID, TaskID
+from scaler.utility.metadata.profile_result import ProfileResult
+from scaler.worker.agent.processor_holder import ProcessorHolder
+
+
+class HeartbeatManager(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    async def on_heartbeat_echo(self, heartbeat: WorkerHeartbeatEcho):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_object_storage_address(self) -> ObjectStorageAddressConfig:
+        raise NotImplementedError()
+
+
+class TimeoutManager(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def update_last_seen_time(self):
+        raise NotImplementedError()
+
+
+class TaskManager(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    async def on_task_new(self, task: Task):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_task_result(self, result: TaskResult):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def on_cancel_task(self, task_cancel: TaskCancel):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_queued_size(self):
+        raise NotImplementedError()
+
+
+class ProcessorManager(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def can_accept_task(self) -> bool:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def wait_until_can_accept_task(self):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_processor_initialized(self, processor_id: ProcessorID, processor_initialized: ProcessorInitialized):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_task(self, task: Task) -> bool:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_cancel_task(self, task_id: TaskID) -> Optional[Task]:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_failing_processor(self, processor_id: ProcessorID, process_status: str):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_suspend_task(self, task_id: TaskID) -> bool:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def on_resume_task(self, task_id: TaskID) -> bool:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_task_result(self, processor_id: ProcessorID, task_result: TaskResult):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_external_object_instruction(self, instruction: ObjectInstruction):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    async def on_internal_object_instruction(self, processor_id: ProcessorID, instruction: ObjectInstruction):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def destroy(self, reason: str):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def current_processor_is_initialized(self) -> bool:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def current_task(self) -> Optional[Task]:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def current_task_id(self) -> TaskID:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def processors(self) -> List[ProcessorHolder]:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def num_suspended_processors(self) -> int:
+        raise NotImplementedError()
+
+
+class ProfilingManager(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def on_process_start(self, pid: int):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def on_process_end(self, pid: int):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def on_task_start(self, pid: int, task_id: TaskID):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def on_task_end(self, pid: int, task_id: TaskID) -> ProfileResult:
+        raise NotImplementedError()

scaler/worker/agent/processor/object_cache.py

@@ -0,0 +1,107 @@
+import ctypes
+import gc
+import logging
+import multiprocessing
+import platform
+import threading
+import time
+from typing import Any, Dict, Optional
+
+import cloudpickle
+import psutil
+
+from scaler.client.serializer.mixins import Serializer
+from scaler.config.defaults import CLEANUP_INTERVAL_SECONDS
+from scaler.utility.exceptions import DeserializeObjectError
+from scaler.utility.identifiers import ClientID, ObjectID
+
+
+class ObjectCache(threading.Thread):
+    def __init__(self, garbage_collect_interval_seconds: int, trim_memory_threshold_bytes: int):
+        threading.Thread.__init__(self)
+
+        self._serializers: Dict[ClientID, Serializer] = dict()
+
+        self._garbage_collect_interval_seconds = garbage_collect_interval_seconds
+        self._previous_garbage_collect_time = time.time()
+        self._trim_memory_threshold_bytes = trim_memory_threshold_bytes
+
+        self._cached_objects: Dict[ObjectID, Any] = {}
+        self._cached_objects_alive_since: Dict[ObjectID, float] = dict()
+        self._process = psutil.Process(multiprocessing.current_process().pid)
+        self._libc = ctypes.cdll.LoadLibrary("libc.{}".format("so.6" if platform.uname()[0] != "Darwin" else "dylib"))
+
+        self._stop_event = threading.Event()
+
+    def run(self) -> None:
+        try:
+            while not self._stop_event.wait(timeout=CLEANUP_INTERVAL_SECONDS):
+                self.__clean_memory()
+        finally:
+            self.__clear()  # gracefully destroy all cached objects
+
+    def destroy(self) -> None:
+        self._stop_event.set()
+
+    def add_serializer(self, client: ClientID, serializer: Serializer):
+        self._serializers[client] = serializer
+
+    def serialize(self, client: ClientID, obj: Any) -> bytes:
+        return self.get_serializer(client).serialize(obj)
+
+    def deserialize(self, client: ClientID, payload: bytes) -> Any:
+        return self.get_serializer(client).deserialize(payload)
+
+    def add_object(self, client: ClientID, object_id: ObjectID, object_bytes: bytes) -> None:
+
+        if object_id.is_serializer():
+            self.add_serializer(client, cloudpickle.loads(object_bytes))
+        else:
+            try:
+                deserialized = self.deserialize(client, object_bytes)
+            except Exception:  # noqa
+                logging.exception(f"failed to deserialize received {object_id!r}, length={len(object_bytes)}")
+
+            self._cached_objects[object_id] = deserialized
+            self._cached_objects_alive_since[object_id] = time.time()
+
+    def del_object(self, object_id: ObjectID):
+        self._cached_objects_alive_since.pop(object_id, None)
+        self._cached_objects.pop(object_id, None)
+
+    def has_object(self, object_id: ObjectID):
+        return object_id in self._cached_objects or object_id in self._serializers
+
+    def get_object(self, object_id: ObjectID) -> Optional[Any]:
+        if object_id not in self._cached_objects:
+            raise ValueError(f"cannot get object for {object_id!r}")
+
+        obj = self._cached_objects[object_id]
+
+        self._cached_objects_alive_since[object_id] = time.time()
+        return obj
+
+    def get_serializer(self, client: ClientID) -> Serializer:
+        serializer = self._serializers.get(client)
+
+        if serializer is None:
+            raise DeserializeObjectError(f"cannot get serializer for {client!r}")
+
+        return serializer
+
+    def __clean_memory(self):
+        if time.time() - self._previous_garbage_collect_time < self._garbage_collect_interval_seconds:
+            return
+
+        self._previous_garbage_collect_time = time.time()
+
+        gc.collect()
+
+        if self._process.memory_info().rss < self._trim_memory_threshold_bytes:
+            return
+
+        self._libc.malloc_trim(0)
+
+    def __clear(self) -> None:
+        self._cached_objects.clear()
+        self._cached_objects_alive_since.clear()