opengris-scaler 1.12.28__cp313-cp313-musllinux_1_2_x86_64.whl

scaler/utility/metadata/task_flags.py

@@ -0,0 +1,30 @@
+import dataclasses
+import struct
+
+from scaler.protocol.python.message import Task
+
+
+@dataclasses.dataclass
+class TaskFlags:
+    profiling: bool = dataclasses.field(default=True)
+    priority: int = dataclasses.field(default=0)
+    stream_output: bool = dataclasses.field(default=False)
+
+    FORMAT = "!?i?"
+
+    def serialize(self) -> bytes:
+        return struct.pack(TaskFlags.FORMAT, self.profiling, self.priority, self.stream_output)
+
+    @staticmethod
+    def deserialize(data: bytes) -> "TaskFlags":
+        return TaskFlags(*struct.unpack(TaskFlags.FORMAT, data))
+
+
+def retrieve_task_flags_from_task(task: Task) -> TaskFlags:
+    if task.metadata == b"":
+        return TaskFlags()
+
+    try:
+        return TaskFlags.deserialize(task.metadata)
+    except struct.error:
+        raise ValueError(f"unexpected metadata value (expected {TaskFlags.__name__}).")

scaler/utility/mixins.py

@@ -0,0 +1,13 @@
+import abc
+
+
+class Looper(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    async def routine(self):
+        raise NotImplementedError()
+
+
+class Reporter(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def get_status(self):
+        raise NotImplementedError()

scaler/utility/network_util.py

@@ -0,0 +1,7 @@
+import socket
+
+
+def get_available_tcp_port(hostname: str = "127.0.0.1") -> int:
+    with socket.socket(socket.AddressFamily.AF_INET, socket.SocketKind.SOCK_STREAM) as sock:
+        sock.bind((hostname, 0))
+        return sock.getsockname()[1]

scaler/utility/one_to_many_dict.py

@@ -0,0 +1,72 @@
+from typing import Dict, Generic, Set, TypeVar
+
+KeyT = TypeVar("KeyT")
+ValueT = TypeVar("ValueT")
+
+
+class OneToManyDict(Generic[KeyT, ValueT]):
+    def __init__(self):
+        self._key_to_value_set: Dict[KeyT, Set[ValueT]] = dict()
+        self._value_to_key: Dict[ValueT, KeyT] = dict()
+
+    def __contains__(self, key) -> bool:
+        return self.has_key(key)
+
+    def keys(self):
+        return self._key_to_value_set.keys()
+
+    def values(self):
+        return self._key_to_value_set.values()
+
+    def items(self):
+        return self._key_to_value_set.items()
+
+    def add(self, key: KeyT, value: ValueT):
+        if value in self._value_to_key and self._value_to_key[value] != key:
+            raise ValueError("value has to be unique in OneToManyDict")
+
+        self._value_to_key[value] = key
+
+        if key not in self._key_to_value_set:
+            self._key_to_value_set[key] = set()
+
+        self._key_to_value_set[key].add(value)
+
+    def has_key(self, key: KeyT) -> bool:
+        return key in self._key_to_value_set
+
+    def has_value(self, value: ValueT) -> bool:
+        return value in self._value_to_key
+
+    def get_key(self, value: ValueT) -> KeyT:
+        if value not in self._value_to_key:
+            raise ValueError(f"cannot find {value=} in OneToManyDict")
+
+        return self._value_to_key[value]
+
+    def get_values(self, key: KeyT) -> Set[ValueT]:
+        if key not in self._key_to_value_set:
+            raise ValueError(f"cannot find {key=} in OneToManyDict")
+
+        return self._key_to_value_set[key]
+
+    def remove_key(self, key: KeyT) -> Set[ValueT]:
+        if key not in self._key_to_value_set:
+            raise KeyError(f"cannot find {key=} in OneToManyDict")
+
+        values = self._key_to_value_set.pop(key)
+        for value in values:
+            self._value_to_key.pop(value)
+
+        return values
+
+    def remove_value(self, value: ValueT) -> KeyT:
+        if value not in self._value_to_key:
+            raise ValueError(f"cannot find {value=} in OneToManyDict")
+
+        key = self._value_to_key.pop(value)
+        self._key_to_value_set[key].remove(value)
+        if not self._key_to_value_set[key]:
+            self._key_to_value_set.pop(key)
+
+        return key

scaler/utility/queues/async_indexed_queue.py

@@ -0,0 +1,37 @@
+from asyncio import Queue, QueueEmpty
+from typing import Generic, TypeVar
+
+from scaler.utility.queues.indexed_queue import IndexedQueue
+
+ItemType = TypeVar("ItemType")
+
+
+class AsyncIndexedQueue(Queue, Generic[ItemType]):
+    """This should have same set of features as asyncio.Queue, with additional methods like remove
+    - it behaves like regular async queue, except:
+      - all the items pushed to queue should be hashable
+      - those items should be unique in queue
+    - IndexedQueue.put(), IndexedQueue.get(), IndexedQueue.remove() should all take O(1) time complexity
+    """
+
+    def __contains__(self, item: ItemType):
+        return item in self._queue
+
+    def __len__(self):
+        return self._queue.__len__()
+
+    def _init(self, maxsize):
+        self._queue = IndexedQueue()
+
+    def _put(self, item: ItemType):
+        self._queue.put(item)
+
+    def _get(self):
+        try:
+            return self._queue.get()
+        except IndexError:
+            raise QueueEmpty(f"{self.__class__.__name__} queue empty")
+
+    def remove(self, item: ItemType):
+        """remove the item in the queue in O(1) time complexity"""
+        self._queue.remove(item)

scaler/utility/queues/async_priority_queue.py

@@ -0,0 +1,70 @@
+import heapq
+import sys
+from asyncio import Queue
+from typing import Any, Dict, List, Tuple, Union
+
+PriorityType = Union[int, Tuple["PriorityType", ...]]
+
+
+class AsyncPriorityQueue(Queue):
+    """A subclass of Queue; retrieves entries in priority order (lowest first).
+
+    Entries are typically list of the form: [priority, data].
+    """
+
+    def __len__(self):
+        return len(self._queue)
+
+    def _init(self, maxsize):
+        self._queue: List[List] = []
+        self._locator: Dict[bytes, List] = {}
+
+    def _put(self, item):
+        if not isinstance(item, list):
+            item = list(item)
+
+        heapq.heappush(self._queue, item)
+        self._locator[item[1]] = item
+
+    def _get(self):
+        priority, data = heapq.heappop(self._queue)
+        self._locator.pop(data)
+        return priority, data
+
+    def remove(self, data):
+        # this operation is O(n), first change priority to -1 and pop from top of the heap, mark it as invalid
+        # entry in the heap is not good idea as those invalid, entry will never get removed, so we used heapq internal
+        # function _siftdown to maintain min heap invariant
+        item = self._locator.pop(data)
+        i = self._queue.index(item)  # O(n)
+        item[0] = self.__to_lowest_priority(item[0])
+        heapq._siftdown(self._queue, 0, i)  # type: ignore[attr-defined]
+        assert heapq.heappop(self._queue) == item
+
+    def decrease_priority(self, data):
+        # this operation should be O(n), mark it as invalid entry in the heap is not good idea as those invalid
+        # entry will never get removed, so we used heapq internal function _siftdown to maintain min heap invariant
+        item = self._locator[data]
+        i = self._queue.index(item)  # O(n)
+        item[0] = self.__to_lower_priority(item[0])
+        heapq._siftdown(self._queue, 0, i)  # type: ignore[attr-defined]
+
+    def max_priority_item(self) -> Tuple[PriorityType, Any]:
+        """output the Tuple of top priority number and top priority item"""
+        item = heapq.heappop(self._queue)
+        heapq.heappush(self._queue, item)
+        return item[0], item[1]
+
+    @classmethod
+    def __to_lowest_priority(cls, original_priority: PriorityType) -> PriorityType:
+        if isinstance(original_priority, tuple):
+            return tuple(cls.__to_lowest_priority(value) for value in original_priority)
+        else:
+            return -sys.maxsize - 1
+
+    @classmethod
+    def __to_lower_priority(cls, original_priority: PriorityType) -> PriorityType:
+        if isinstance(original_priority, tuple):
+            return tuple(cls.__to_lower_priority(value) for value in original_priority)
+        else:
+            return original_priority - 1

scaler/utility/queues/async_sorted_priority_queue.py

@@ -0,0 +1,45 @@
+from asyncio import Queue
+from typing import Any, Dict
+
+from scaler.utility.queues.async_priority_queue import AsyncPriorityQueue
+
+
+class AsyncSortedPriorityQueue(Queue):
+    """A subclass of Queue; retrieves entries in priority order (lowest first), and then by adding order.
+
+    Entries are typically list of the form: [priority number, data].
+    """
+
+    def __len__(self):
+        return len(self._queue)
+
+    def _init(self, maxsize: int):
+        self._queue = AsyncPriorityQueue()
+
+        # Keeps an item count to assign monotonic integer to queued items, so to also keep the priority queue sorted by
+        # adding order.
+        # See https://docs.python.org/3/library/heapq.html#priority-queue-implementation-notes.
+        self._item_counter: int = 0
+        self._data_to_item_id: Dict[Any, int] = dict()
+
+    def _put(self, item) -> None:
+        priority, data = item
+
+        if data in self._data_to_item_id:
+            raise ValueError(f"item `{data}` already in the queue")
+
+        item_id = self._item_counter
+        self._item_counter += 1
+
+        self._queue._put([priority, (item_id, data)])
+        self._data_to_item_id[data] = item_id
+
+    def _get(self):
+        priority, (_, data) = self._queue._get()
+        self._data_to_item_id.pop(data)
+
+        return [priority, data]
+
+    def remove(self, data: Any) -> None:
+        item_id = self._data_to_item_id.pop(data)
+        self._queue.remove((item_id, data))

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.28

scaler/worker/agent/heartbeat_manager.py

@@ -0,0 +1,107 @@
+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 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[ObjectStorageConfig], 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[ObjectStorageConfig] = 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 = ObjectStorageConfig(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[ObjectStorageConfig]:
+        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 ObjectStorageConfig
+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) -> ObjectStorageConfig:
+        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()