opengris-scaler 1.12.28__cp313-cp313-musllinux_1_2_x86_64.whl

scaler/worker/agent/processor_manager.py

@@ -0,0 +1,369 @@
+import asyncio
+import logging
+from typing import Dict, List, Optional, Tuple
+
+import tblib.pickling_support
+
+from scaler.config.types.zmq import ZMQConfig
+
+# from scaler.utility.logging.utility import setup_logger
+from scaler.io.mixins import AsyncBinder, AsyncConnector, AsyncObjectStorageConnector
+from scaler.protocol.python.common import ObjectMetadata, TaskResultType
+from scaler.protocol.python.message import ObjectInstruction, ProcessorInitialized, Task, TaskResult
+from scaler.utility.exceptions import ProcessorDiedError
+from scaler.utility.identifiers import ObjectID, ProcessorID, TaskID, WorkerID
+from scaler.utility.metadata.profile_result import ProfileResult
+from scaler.utility.serialization import serialize_failure
+from scaler.worker.agent.mixins import HeartbeatManager, ProcessorManager, ProfilingManager, TaskManager
+from scaler.worker.agent.processor_holder import ProcessorHolder
+
+
+class VanillaProcessorManager(ProcessorManager):
+    def __init__(
+        self,
+        identity: WorkerID,
+        event_loop: str,
+        address_internal: ZMQConfig,
+        scheduler_address: ZMQConfig,
+        preload: Optional[str],
+        garbage_collect_interval_seconds: int,
+        trim_memory_threshold_bytes: int,
+        hard_processor_suspend: bool,
+        logging_paths: Tuple[str, ...],
+        logging_level: str,
+    ):
+        tblib.pickling_support.install()
+
+        self._identity = identity
+        self._event_loop = event_loop
+        self._scheduler_address = scheduler_address
+        self._preload = preload
+
+        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._logging_paths = logging_paths
+        self._logging_level = logging_level
+
+        self._heartbeat_manager: Optional[HeartbeatManager] = None
+        self._task_manager: Optional[TaskManager] = None
+        self._profiling_manager: Optional[ProfilingManager] = None
+        self._connector_external: Optional[AsyncConnector] = None
+        self._connector_storage: Optional[AsyncObjectStorageConnector] = None
+
+        self._address_internal: ZMQConfig = address_internal
+
+        self._current_holder: Optional[ProcessorHolder] = None
+        self._suspended_holders_by_task_id: Dict[bytes, ProcessorHolder] = {}
+        self._holders_by_processor_id: Dict[ProcessorID, ProcessorHolder] = {}
+
+        self._can_accept_task_lock: asyncio.Lock = asyncio.Lock()
+
+        self._binder_internal: Optional[AsyncBinder] = None
+
+    def register(
+        self,
+        heartbeat_manager: HeartbeatManager,
+        task_manager: TaskManager,
+        profiling_manager: ProfilingManager,
+        connector_external: AsyncConnector,
+        binder_internal: AsyncBinder,
+        connector_storage: AsyncObjectStorageConnector,
+    ):
+        self._heartbeat_manager = heartbeat_manager
+        self._task_manager = task_manager
+        self._profiling_manager = profiling_manager
+        self._connector_external = connector_external
+        self._binder_internal = binder_internal
+        self._connector_storage = connector_storage
+
+    async def initialize(self):
+        await self._can_accept_task_lock.acquire()  # prevents any processor to accept task until initialized
+
+        await self._connector_storage.wait_until_connected()
+
+        self.__start_new_processor()  # we can start the processor now that we know the storage address.
+
+    def can_accept_task(self) -> bool:
+        return not self._can_accept_task_lock.locked()
+
+    async def wait_until_can_accept_task(self):
+        """
+        Makes sure a processor is ready to start processing a new or suspended task.
+
+        Must be called before any call to `on_task()` or `on_task_resume()`.
+        """
+
+        await self._can_accept_task_lock.acquire()
+
+    async def on_processor_initialized(self, processor_id: ProcessorID, processor_initialized: ProcessorInitialized):
+        assert self._current_holder is not None
+
+        if self._current_holder.initialized():
+            return
+
+        self._holders_by_processor_id[processor_id] = self._current_holder
+        self._current_holder.initialize(processor_id)
+
+        self._can_accept_task_lock.release()
+
+    async def on_task(self, task: Task) -> bool:
+        assert self._can_accept_task_lock.locked()
+        assert self.current_processor_is_initialized()
+
+        holder = self._current_holder
+
+        assert holder.task() is None
+        holder.set_task(task)
+
+        self._profiling_manager.on_task_start(holder.pid(), task.task_id)
+
+        await self._binder_internal.send(holder.processor_id(), task)
+
+        return True
+
+    async def on_cancel_task(self, task_id: TaskID) -> Optional[Task]:
+        assert self._current_holder is not None
+
+        if self.current_task_id() == task_id:
+            current_task = self.current_task()
+            self.__restart_current_processor(f"cancel task_id={task_id.hex()}")
+            return current_task
+
+        if task_id in self._suspended_holders_by_task_id:
+            suspended_holder = self._suspended_holders_by_task_id.pop(task_id)
+            task = suspended_holder.task()
+            self.__kill_processor(f"cancel suspended task_id={task_id.hex()}", suspended_holder)
+            return task
+
+        return None
+
+    async def on_failing_processor(self, processor_id: ProcessorID, process_status: str):
+        assert self._current_holder is not None
+
+        holder = self._holders_by_processor_id.get(processor_id)
+
+        if holder is None:
+            return
+
+        task = holder.task()
+        if task is not None:
+            profile_result = self.__end_task(holder)  # profiling the task should happen before killing the processor
+        else:
+            profile_result = None
+
+        reason = f"process died {process_status=}"
+        if holder == self._current_holder:
+            self.__restart_current_processor(reason)
+        else:
+            self.__kill_processor(reason, holder)
+
+        if task is not None:
+            source = task.source
+            task_id = task.task_id
+
+            result_object_id = ObjectID.generate_object_id(source)
+            result_object_bytes = serialize_failure(ProcessorDiedError(f"{process_status=}"))
+
+            await self._connector_storage.set_object(result_object_id, result_object_bytes)
+            await self._connector_external.send(
+                ObjectInstruction.new_msg(
+                    ObjectInstruction.ObjectInstructionType.Create,
+                    source,
+                    ObjectMetadata.new_msg((result_object_id,), (ObjectMetadata.ObjectContentType.Object,), (b"",)),
+                )
+            )
+
+            await self._task_manager.on_task_result(
+                TaskResult.new_msg(
+                    task_id, TaskResultType.Failed, profile_result.serialize(), [bytes(result_object_id)]
+                )
+            )
+
+    async def on_suspend_task(self, task_id: TaskID) -> bool:
+        assert self._current_holder is not None
+        holder = self._current_holder
+
+        current_task = holder.task()
+
+        if current_task is None or current_task.task_id != task_id:
+            return False
+
+        holder.suspend()
+        self._suspended_holders_by_task_id[task_id] = holder
+
+        logging.info(f"{self._identity!r}: suspend Processor[{holder.pid()}]")
+
+        self.__start_new_processor()
+
+        return True
+
+    def on_resume_task(self, task_id: TaskID) -> bool:
+        assert self._can_accept_task_lock.locked()
+        assert self.current_processor_is_initialized()
+
+        if self.current_task() is not None:
+            return False
+
+        suspended_holder = self._suspended_holders_by_task_id.pop(task_id, None)
+
+        if suspended_holder is None:
+            return False
+
+        self.__kill_processor("replaced by suspended processor", self._current_holder)
+
+        self._current_holder = suspended_holder
+        suspended_holder.resume()
+
+        logging.info(f"{self._identity!r}: resume Processor[{self._current_holder.pid()}]")
+
+        return True
+
+    async def on_task_result(self, processor_id: ProcessorID, task_result: TaskResult):
+        assert self._current_holder is not None
+        task_id = task_result.task_id
+
+        if task_id == self.current_task_id():
+            assert self._current_holder.processor_id() == processor_id
+
+            profile_result = self.__end_task(self._current_holder)
+
+            release_task_lock = True
+        elif task_id in self._suspended_holders_by_task_id:
+            # Receiving a task result from a suspended processor is possible as the message might have been queued while
+            # we were suspending the process.
+
+            holder = self._suspended_holders_by_task_id.pop(task_id)
+            assert holder.processor_id() == processor_id
+
+            profile_result = self.__end_task(holder)
+
+            self.__kill_processor("task finished in suspended processor", holder)
+
+            release_task_lock = False
+        else:
+            return
+
+        await self._task_manager.on_task_result(
+            TaskResult.new_msg(
+                task_id=task_id,
+                result_type=task_result.result_type,
+                metadata=profile_result.serialize(),
+                results=task_result.results,
+            )
+        )
+
+        # task lock must be released after calling `TaskManager.on_task_result()`
+        if release_task_lock:
+            self._can_accept_task_lock.release()
+
+    async def on_external_object_instruction(self, instruction: ObjectInstruction):
+        for processor_id in self._holders_by_processor_id.keys():
+            await self._binder_internal.send(processor_id, instruction)
+
+    async def on_internal_object_instruction(self, processor_id: ProcessorID, instruction: ObjectInstruction):
+        if not self.__processor_ready_to_process_object(processor_id):
+            return
+
+        await self._connector_external.send(instruction)
+
+    def destroy(self, reason: str):
+        if self._connector_storage is not None:
+            self._connector_external.destroy()
+
+        self.__kill_all_processors(reason)
+
+    def current_processor_is_initialized(self) -> bool:
+        return self._current_holder is not None and self._current_holder.initialized()
+
+    def current_task(self) -> Optional[Task]:
+        if self._current_holder is None:  # worker is not yet initialized
+            return None
+
+        return self._current_holder.task()
+
+    def current_task_id(self) -> Optional[TaskID]:
+        task = self.current_task()
+
+        if task is None:
+            return None
+        else:
+            return task.task_id
+
+    def processors(self) -> List[ProcessorHolder]:
+        return list(self._holders_by_processor_id.values())
+
+    def num_suspended_processors(self) -> int:
+        return len(self._suspended_holders_by_task_id)
+
+    def __start_new_processor(self):
+        object_storage_address = self._heartbeat_manager.get_object_storage_address()
+
+        self._current_holder = ProcessorHolder(
+            self._event_loop,
+            self._address_internal,
+            self._scheduler_address,
+            object_storage_address,
+            self._preload,
+            self._garbage_collect_interval_seconds,
+            self._trim_memory_threshold_bytes,
+            self._hard_processor_suspend,
+            self._logging_paths,
+            self._logging_level,
+        )
+
+        processor_pid = self._current_holder.pid()
+
+        self._profiling_manager.on_process_start(processor_pid)
+
+        logging.info(f"{self._identity!r}: start Processor[{processor_pid}]")
+
+    def __kill_processor(self, reason: str, holder: ProcessorHolder):
+        processor_pid = holder.pid()
+
+        self._profiling_manager.on_process_end(processor_pid)
+
+        if holder.initialized():
+            self._holders_by_processor_id.pop(holder.processor_id(), None)
+
+        holder.kill()
+
+        logging.info(f"{self._identity!r}: stop Processor[{processor_pid}], reason: {reason}")
+
+    def __restart_current_processor(self, reason: str):
+        assert self._current_holder is not None
+
+        self.__kill_processor(reason, self._current_holder)
+        self.__start_new_processor()
+
+    def __kill_all_processors(self, reason: str):
+        if self._current_holder is not None:
+            self.__kill_processor(reason, self._current_holder)
+            self._current_holder = None
+
+        for processor_holder in self._suspended_holders_by_task_id.values():
+            self.__kill_processor(reason, processor_holder)
+
+        self._suspended_holders_by_task_id = {}
+        self._holders_by_processor_id = {}
+
+    def __end_task(self, processor_holder: ProcessorHolder) -> ProfileResult:
+        profile_result = self._profiling_manager.on_task_end(processor_holder.pid(), processor_holder.task().task_id)
+        processor_holder.set_task(None)
+
+        return profile_result
+
+    def __processor_ready_to_process_object(self, processor_id: ProcessorID) -> bool:
+        holder = self._holders_by_processor_id.get(processor_id)
+
+        if holder is None:
+            return False
+
+        assert holder.initialized()
+
+        if holder.task() is None:
+            return False
+
+        # TODO: check if the objects belong to the task
+
+        return True

scaler/worker/agent/profiling_manager.py

@@ -0,0 +1,109 @@
+import dataclasses
+import logging
+import time
+from typing import Dict, Optional
+
+import psutil
+
+from scaler.utility.identifiers import TaskID
+from scaler.utility.metadata.profile_result import ProfileResult
+from scaler.utility.mixins import Looper
+from scaler.worker.agent.mixins import ProfilingManager
+
+
+@dataclasses.dataclass
+class _ProcessProfiler:
+    process: psutil.Process
+
+    current_task_id: Optional[TaskID] = None
+
+    start_time: Optional[float] = None
+    start_cpu_time: Optional[float] = None
+    init_memory_rss: Optional[int] = None
+    peak_memory_rss: Optional[int] = None
+
+
+class VanillaProfilingManager(ProfilingManager, Looper):
+    def __init__(self):
+        self._process_profiler_by_pid: Dict[int, _ProcessProfiler] = {}
+
+    def on_process_start(self, pid: int):
+        if pid in self._process_profiler_by_pid:
+            raise ValueError(f"process {pid=} is already registered.")
+
+        self._process_profiler_by_pid[pid] = _ProcessProfiler(psutil.Process(pid))
+
+    def on_process_end(self, pid: int):
+        if pid not in self._process_profiler_by_pid:
+            raise ValueError(f"process {pid=} is not registered.")
+
+        self._process_profiler_by_pid.pop(pid)
+
+    def on_task_start(self, pid: int, task_id: TaskID):
+        process_profiler = self._process_profiler_by_pid.get(pid)
+
+        if process_profiler is None:
+            raise ValueError(f"process {pid=} is not registered.")
+
+        process_profiler.current_task_id = task_id
+
+        process = process_profiler.process
+
+        process_profiler.start_time = self.__process_time()
+        process_profiler.start_cpu_time = self.__process_cpu_time(process)
+        process_profiler.init_memory_rss = self.__process_memory_rss(process)
+        process_profiler.peak_memory_rss = process_profiler.init_memory_rss
+
+    def on_task_end(self, pid: int, task_id: TaskID) -> ProfileResult:
+        process_profiler = self._process_profiler_by_pid.get(pid)
+
+        if process_profiler is None:
+            raise ValueError(f"process {pid=} is not registered.")
+
+        if task_id != process_profiler.current_task_id:
+            raise ValueError(f"task {task_id=!r} is not the current task task_id={process_profiler.current_task_id!r}.")
+
+        assert process_profiler.start_time is not None
+        assert process_profiler.init_memory_rss is not None
+        assert process_profiler.peak_memory_rss is not None
+
+        process = process_profiler.process
+
+        time_delta = self.__process_time() - process_profiler.start_time
+
+        try:
+            cpu_time_delta = self.__process_cpu_time(process) - process_profiler.start_cpu_time
+        except psutil.ZombieProcess:
+            logging.warning(f"profiling zombie process: {pid=}")
+            cpu_time_delta = 0
+
+        memory_delta = process_profiler.peak_memory_rss - process_profiler.init_memory_rss
+
+        process_profiler.current_task_id = None
+        process_profiler.init_memory_rss = None
+        process_profiler.peak_memory_rss = None
+
+        return ProfileResult(time_delta, memory_delta, cpu_time_delta)
+
+    async def routine(self):
+        for process_profiler in self._process_profiler_by_pid.values():
+            if process_profiler.current_task_id is not None:
+                try:
+                    process_profiler.peak_memory_rss = max(
+                        process_profiler.peak_memory_rss, self.__process_memory_rss(process_profiler.process)
+                    )
+                except psutil.ZombieProcess:
+                    logging.warning(f"profiling zombie process: pid={process_profiler.process.pid}")
+
+    @staticmethod
+    def __process_time():
+        return time.monotonic()
+
+    @staticmethod
+    def __process_cpu_time(process: psutil.Process) -> float:
+        cpu_times = process.cpu_times()
+        return cpu_times.user + cpu_times.system
+
+    @staticmethod
+    def __process_memory_rss(process: psutil.Process) -> int:
+        return process.memory_info().rss

scaler/worker/agent/task_manager.py

@@ -0,0 +1,150 @@
+from typing import Dict, Optional, Set
+
+from scaler.io.mixins import AsyncConnector
+from scaler.protocol.python.common import TaskCancelConfirmType
+from scaler.protocol.python.message import Task, TaskCancel, TaskCancelConfirm, TaskResult
+from scaler.utility.identifiers import 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.worker.agent.mixins import ProcessorManager, TaskManager
+
+_SUSPENDED_TASKS_PRIORITY = 1
+_QUEUED_TASKS_PRIORITY = 2
+
+
+class VanillaTaskManager(Looper, TaskManager):
+    def __init__(self, task_timeout_seconds: int):
+        self._task_timeout_seconds = task_timeout_seconds
+
+        self._queued_task_id_to_task: Dict[TaskID, Task] = dict()
+
+        # Queued tasks are sorted first by task's priorities, then suspended tasks are prioritized over non yet started
+        # tasks, finally the sorted queue ensure we execute the oldest tasks first.
+        #
+        # For example, if we receive these tasks in this order:
+        #   1. Task(priority=0) [suspended]
+        #   2. Task(priority=3) [suspended]
+        #   3. Task(priority=3)
+        #   4. Task(priority=0)
+        #
+        # We want to execute the tasks in this order: 2-3-1-4.
+        self._queued_task_ids = AsyncSortedPriorityQueue()
+
+        self._processing_task_ids: Set[TaskID] = set()  # Tasks associated with a processor, including suspended tasks
+
+        self._connector_external: Optional[AsyncConnector] = None
+        self._processor_manager: Optional[ProcessorManager] = None
+
+    def register(self, connector: AsyncConnector, processor_manager: ProcessorManager):
+        self._connector_external = connector
+        self._processor_manager = processor_manager
+
+    async def on_task_new(self, task: Task):
+        self.__enqueue_task(task, is_suspended=False)
+
+        await self.__suspend_if_priority_is_higher(task)
+
+    async def on_cancel_task(self, task_cancel: TaskCancel):
+        task_not_found = (
+            task_cancel.task_id not in self._processing_task_ids
+            and task_cancel.task_id not in self._queued_task_id_to_task
+        )
+        if task_not_found:
+            await self._connector_external.send(
+                TaskCancelConfirm.new_msg(
+                    task_id=task_cancel.task_id, cancel_confirm_type=TaskCancelConfirmType.CancelNotFound
+                )
+            )
+            return
+
+        if task_cancel.task_id in self._processing_task_ids and not task_cancel.flags.force:
+            # ignore cancel task while in processing if is not force cancel
+            await self._connector_external.send(
+                TaskCancelConfirm.new_msg(
+                    task_id=task_cancel.task_id, cancel_confirm_type=TaskCancelConfirmType.CancelFailed
+                )
+            )
+            return
+
+        # A suspended task will be both processing AND queued
+
+        if task_cancel.task_id in self._processing_task_ids:
+            # if task is in processing
+            self._processing_task_ids.remove(task_cancel.task_id)
+            _ = await self._processor_manager.on_cancel_task(task_cancel.task_id)
+        else:
+            # if task is queued
+            assert task_cancel.task_id in self._queued_task_id_to_task
+            self._queued_task_ids.remove(task_cancel.task_id)
+            _ = self._queued_task_id_to_task.pop(task_cancel.task_id)
+
+        await self._connector_external.send(
+            TaskCancelConfirm.new_msg(task_id=task_cancel.task_id, cancel_confirm_type=TaskCancelConfirmType.Canceled)
+        )
+
+    async def on_task_result(self, result: TaskResult):
+        if result.task_id in self._queued_task_id_to_task:
+            # Finishing a queued task might happen if a task ended during the suspension process.
+            self._queued_task_id_to_task.pop(result.task_id)
+            self._queued_task_ids.remove(result.task_id)
+
+        self._processing_task_ids.remove(result.task_id)
+
+        await self._connector_external.send(result)
+
+    async def routine(self):
+        await self.__processing_task()
+
+    def get_queued_size(self):
+        return self._queued_task_ids.qsize()
+
+    async def __processing_task(self):
+        await self._processor_manager.wait_until_can_accept_task()
+
+        _, task_id = await self._queued_task_ids.get()
+        task = self._queued_task_id_to_task.pop(task_id)
+
+        if task_id not in self._processing_task_ids:
+            self._processing_task_ids.add(task_id)
+            await self._processor_manager.on_task(task)
+        else:
+            self._processor_manager.on_resume_task(task_id)
+
+    async def __suspend_if_priority_is_higher(self, new_task: Task):
+        current_task = self._processor_manager.current_task()
+
+        if current_task is None:
+            return
+
+        new_task_priority = self.__get_task_priority(new_task)
+        current_task_priority = self.__get_task_priority(current_task)
+
+        if new_task_priority <= current_task_priority:
+            return
+
+        self.__enqueue_task(current_task, is_suspended=True)
+
+        await self._processor_manager.on_suspend_task(current_task.task_id)
+
+    def __enqueue_task(self, task: Task, is_suspended: bool):
+        task_priority = self.__get_task_priority(task)
+
+        # Higher-priority tasks have a higher priority value. But as the queue is sorted by increasing order, we negate
+        # the inserted value that it will be at the head of the queue.
+        if is_suspended:
+            queue_priority = (-task_priority, _SUSPENDED_TASKS_PRIORITY)
+        else:
+            queue_priority = (-task_priority, _QUEUED_TASKS_PRIORITY)
+
+        self._queued_task_ids.put_nowait((queue_priority, task.task_id))
+        self._queued_task_id_to_task[task.task_id] = task
+
+    @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

scaler/worker/agent/timeout_manager.py

@@ -0,0 +1,19 @@
+import time
+
+from scaler.utility.mixins import Looper
+from scaler.worker.agent.mixins import TimeoutManager
+
+
+class VanillaTimeoutManager(Looper, TimeoutManager):
+    def __init__(self, death_timeout_seconds: int):
+        self._death_timeout_seconds = death_timeout_seconds
+        self._last_seen_time = time.time()
+
+    def update_last_seen_time(self):
+        self._last_seen_time = time.time()
+
+    async def routine(self):
+        if (time.time() - self._last_seen_time) < self._death_timeout_seconds:
+            return
+
+        raise TimeoutError("timeout when connect to scheduler, quitting")