opengris-scaler 1.12.7__cp311-cp311-musllinux_1_2_x86_64.whl

scaler/worker/agent/processor/processor.py

@@ -0,0 +1,279 @@
+import contextlib
+import logging
+import multiprocessing
+import os
+import signal
+from contextlib import redirect_stderr, redirect_stdout
+from contextvars import ContextVar, Token
+from multiprocessing.synchronize import Event as EventType
+from typing import IO, Callable, List, Optional, Tuple, cast
+
+import tblib.pickling_support
+import zmq
+
+from scaler.config.types.object_storage_server import ObjectStorageConfig
+from scaler.config.types.zmq import ZMQConfig
+from scaler.io.mixins import SyncConnector, SyncObjectStorageConnector
+from scaler.io.sync_connector import ZMQSyncConnector
+from scaler.io.sync_object_storage_connector import PySyncObjectStorageConnector
+from scaler.protocol.python.common import ObjectMetadata, TaskResultType
+from scaler.protocol.python.message import ObjectInstruction, ProcessorInitialized, Task, TaskLog, TaskResult
+from scaler.protocol.python.mixins import Message
+from scaler.utility.identifiers import ClientID, ObjectID, TaskID
+from scaler.utility.logging.utility import setup_logger
+from scaler.utility.metadata.task_flags import retrieve_task_flags_from_task
+from scaler.utility.serialization import serialize_failure
+from scaler.worker.agent.processor.object_cache import ObjectCache
+from scaler.worker.agent.processor.streaming_buffer import StreamingBuffer
+from scaler.worker.preload import execute_preload
+
+SUSPEND_SIGNAL = "SIGUSR1"  # use str instead of a signal.Signal to not trigger an import error on unsupported systems.
+
+_current_processor: ContextVar[Optional["Processor"]] = ContextVar("_current_processor", default=None)
+
+
+class Processor(multiprocessing.get_context("spawn").Process):  # type: ignore
+    def __init__(
+        self,
+        event_loop: str,
+        agent_address: ZMQConfig,
+        storage_address: ObjectStorageConfig,
+        preload: Optional[str],
+        resume_event: Optional[EventType],
+        resumed_event: Optional[EventType],
+        garbage_collect_interval_seconds: int,
+        trim_memory_threshold_bytes: int,
+        logging_paths: Tuple[str, ...],
+        logging_level: str,
+    ):
+        multiprocessing.Process.__init__(self, name="Processor")
+
+        self._event_loop = event_loop
+        self._agent_address = agent_address
+        self._storage_address = storage_address
+        self._preload = preload
+
+        self._resume_event = resume_event
+        self._resumed_event = resumed_event
+
+        self._garbage_collect_interval_seconds = garbage_collect_interval_seconds
+        self._trim_memory_threshold_bytes = trim_memory_threshold_bytes
+        self._logging_paths = logging_paths
+        self._logging_level = logging_level
+
+        self._object_cache: Optional[ObjectCache] = None
+
+        self._current_task: Optional[Task] = None
+
+    def run(self) -> None:
+        self.__initialize()
+        self.__run_forever()
+
+    @staticmethod
+    def get_current_processor() -> Optional["Processor"]:
+        """Returns the current Processor instance controlling the current process, if any."""
+        return _current_processor.get()
+
+    def current_task(self) -> Optional[Task]:
+        return self._current_task
+
+    def __initialize(self):
+        # modify the logging path and add process id to the path
+        logging_paths = [f"{path}-{os.getpid()}" for path in self._logging_paths if path != "/dev/stdout"]
+        if "/dev/stdout" in self._logging_paths:
+            logging_paths.append("/dev/stdout")
+
+        setup_logger(log_paths=tuple(logging_paths), logging_level=self._logging_level)
+        tblib.pickling_support.install()
+
+        self._connector_agent: SyncConnector = ZMQSyncConnector(
+            context=zmq.Context(), socket_type=zmq.DEALER, address=self._agent_address, identity=None
+        )
+        self._connector_storage: SyncObjectStorageConnector = PySyncObjectStorageConnector(
+            self._storage_address.host, self._storage_address.port
+        )
+
+        self._object_cache = ObjectCache(
+            garbage_collect_interval_seconds=self._garbage_collect_interval_seconds,
+            trim_memory_threshold_bytes=self._trim_memory_threshold_bytes,
+        )
+        self._object_cache.start()
+
+        self.__register_signals()
+
+        # Execute optional preload hook if provided
+        if self._preload is not None:
+            try:
+                execute_preload(self._preload)
+            except Exception as e:
+                raise RuntimeError(
+                    f"Processor[{self.pid}] initialization failed due to preload error: {self._preload}"
+                ) from e
+
+    def __register_signals(self):
+        self.__register_signal("SIGTERM", self.__interrupt)
+
+        if self._resume_event is not None:
+            self.__register_signal(SUSPEND_SIGNAL, self.__suspend)
+
+    def __interrupt(self, *args):
+        self._connector_agent.destroy()  # interrupts any blocking socket.
+
+    def __suspend(self, *args):
+        assert self._resume_event is not None
+        assert self._resumed_event is not None
+
+        self._resume_event.wait()  # stops any computation in the main thread until the event is triggered
+
+        # Ensures the processor agent knows we stopped waiting on `_resume_event`, as to avoid re-entrant wait on the
+        # event.
+        self._resumed_event.set()
+
+    def __run_forever(self):
+        try:
+            self._connector_agent.send(ProcessorInitialized.new_msg())
+            while True:
+                message = self._connector_agent.receive()
+                if message is None:
+                    continue
+
+                self.__on_connector_receive(message)
+
+        except zmq.error.ZMQError as e:
+            if e.errno != zmq.ENOTSOCK:  # ignore if socket got closed
+                raise
+
+        except (KeyboardInterrupt, InterruptedError):
+            pass
+
+        except Exception as e:
+            logging.exception(f"Processor[{self.pid}]: failed with unhandled exception:\n{e}")
+
+        finally:
+            self._object_cache.destroy()
+            self._connector_agent.destroy()
+
+            self._object_cache.join()
+
+    def __on_connector_receive(self, message: Message):
+        if isinstance(message, ObjectInstruction):
+            self.__on_receive_object_instruction(message)
+            return
+
+        if isinstance(message, Task):
+            self.__on_received_task(message)
+            return
+
+        logging.error(f"unknown {message=}")
+
+    def __on_receive_object_instruction(self, instruction: ObjectInstruction):
+        if instruction.instruction_type == ObjectInstruction.ObjectInstructionType.Delete:
+            for object_id in instruction.object_metadata.object_ids:
+                self._object_cache.del_object(object_id)
+            return
+
+        logging.error(f"worker received unknown object instruction type {instruction=}")
+
+    def __on_received_task(self, task: Task):
+        self._current_task = task
+
+        self.__cache_required_object_ids(task)
+
+        self.__process_task(task)
+
+    def __cache_required_object_ids(self, task: Task) -> None:
+        required_object_ids = self.__get_required_object_ids_for_task(task)
+
+        for object_id in required_object_ids:
+            if self._object_cache.has_object(object_id):
+                continue
+
+            object_content = self._connector_storage.get_object(object_id)
+            self._object_cache.add_object(task.source, object_id, object_content)
+
+    @staticmethod
+    def __get_required_object_ids_for_task(task: Task) -> List[ObjectID]:
+        serializer_id = ObjectID.generate_serializer_object_id(task.source)
+        object_ids = [
+            serializer_id,
+            task.func_object_id,
+            *(cast(ObjectID, argument) for argument in task.function_args),
+        ]
+        return object_ids
+
+    def __process_task(self, task: Task):
+        task_flags = retrieve_task_flags_from_task(task)
+
+        try:
+            function = self._object_cache.get_object(task.func_object_id)
+
+            args = [self._object_cache.get_object(cast(ObjectID, arg)) for arg in task.function_args]
+
+            if task_flags.stream_output:
+                with StreamingBuffer(
+                    task.task_id, TaskLog.LogType.Stdout, self._connector_agent
+                ) as stdout_buf, StreamingBuffer(
+                    task.task_id, TaskLog.LogType.Stderr, self._connector_agent
+                ) as stderr_buf, self.__processor_context(), redirect_stdout(
+                    cast(IO[str], stdout_buf)
+                ), redirect_stderr(
+                    cast(IO[str], stderr_buf)
+                ):
+                    result = function(*args)
+            else:
+                with self.__processor_context():
+                    result = function(*args)
+
+            result_bytes = self._object_cache.serialize(task.source, result)
+            task_result_type = TaskResultType.Success
+
+        except Exception as e:
+            logging.exception(f"exception when processing task_id={task.task_id.hex()}:")
+            task_result_type = TaskResultType.Failed
+            result_bytes = serialize_failure(e)
+
+        self.__send_result(task.source, task.task_id, task_result_type, result_bytes)
+
+    def __send_result(self, source: ClientID, task_id: TaskID, task_result_type: TaskResultType, result_bytes: bytes):
+        self._current_task = None
+
+        result_object_id = ObjectID.generate_object_id(source)
+
+        self._connector_storage.set_object(result_object_id, result_bytes)
+        self._connector_agent.send(
+            ObjectInstruction.new_msg(
+                ObjectInstruction.ObjectInstructionType.Create,
+                source,
+                ObjectMetadata.new_msg(
+                    (result_object_id,),
+                    (ObjectMetadata.ObjectContentType.Object,),
+                    (f"<res {repr(result_object_id)}>".encode(),),
+                ),
+            )
+        )
+        self._connector_agent.send(
+            TaskResult.new_msg(task_id, task_result_type, metadata=b"", results=[bytes(result_object_id)])
+        )
+
+    @staticmethod
+    def __set_current_processor(context: Optional["Processor"]) -> Token:
+        if context is not None and _current_processor.get() is not None:
+            raise ValueError("cannot override a previously set processor context.")
+
+        return _current_processor.set(context)
+
+    @contextlib.contextmanager
+    def __processor_context(self):
+        self.__set_current_processor(self)
+        try:
+            yield
+        finally:
+            self.__set_current_processor(None)
+
+    @staticmethod
+    def __register_signal(signal_name: str, handler: Callable) -> None:
+        signal_instance = getattr(signal, signal_name, None)
+        if signal_instance is None:
+            raise RuntimeError(f"unsupported platform, signal not available: {signal_name}.")
+
+        signal.signal(signal_instance, handler)

scaler/worker/agent/processor/streaming_buffer.py

@@ -0,0 +1,28 @@
+import io
+import logging
+
+from scaler.io.mixins import SyncConnector
+from scaler.protocol.python.message import TaskLog
+from scaler.utility.identifiers import TaskID
+
+
+class StreamingBuffer(io.TextIOBase):
+    """A custom IO buffer that sends content as it's written."""
+
+    def __init__(self, task_id: TaskID, log_type: TaskLog.LogType, connector_agent: SyncConnector):
+        super().__init__()
+        self._task_id = task_id
+        self._log_type = log_type
+        self._connector_agent = connector_agent
+
+    def write(self, content: str) -> int:
+        if self.closed:
+            return 0
+
+        if content:
+            try:
+                self._connector_agent.send(TaskLog.new_msg(self._task_id, self._log_type, content))
+            except Exception as e:
+                logging.warning(f"Failed to send stream content: {e}")
+
+        return 0

scaler/worker/agent/processor_holder.py

@@ -0,0 +1,145 @@
+import logging
+import multiprocessing
+import os
+import signal
+from typing import Optional, Tuple
+
+import psutil
+
+from scaler.config.defaults import DEFAULT_PROCESSOR_KILL_DELAY_SECONDS
+from scaler.protocol.python.message import Task
+from scaler.utility.identifiers import ProcessorID
+from scaler.config.types.object_storage_server import ObjectStorageConfig
+from scaler.config.types.zmq import ZMQConfig
+from scaler.worker.agent.processor.processor import SUSPEND_SIGNAL, Processor
+
+
+class ProcessorHolder:
+    def __init__(
+        self,
+        event_loop: str,
+        agent_address: ZMQConfig,
+        storage_address: ObjectStorageConfig,
+        preload: Optional[str],
+        garbage_collect_interval_seconds: int,
+        trim_memory_threshold_bytes: int,
+        hard_suspend: bool,
+        logging_paths: Tuple[str, ...],
+        logging_level: str,
+    ):
+        self._processor_id: Optional[ProcessorID] = None
+        self._task: Optional[Task] = None
+        self._suspended = False
+
+        self._hard_suspend = hard_suspend
+        if hard_suspend:
+            self._resume_event = None
+            self._resumed_event = None
+        else:
+            context = multiprocessing.get_context("spawn")
+            self._resume_event = context.Event()
+            self._resumed_event = context.Event()
+
+        self._processor = Processor(
+            event_loop=event_loop,
+            agent_address=agent_address,
+            storage_address=storage_address,
+            preload=preload,
+            resume_event=self._resume_event,
+            resumed_event=self._resumed_event,
+            garbage_collect_interval_seconds=garbage_collect_interval_seconds,
+            trim_memory_threshold_bytes=trim_memory_threshold_bytes,
+            logging_paths=logging_paths,
+            logging_level=logging_level,
+        )
+        self._processor.start()
+        self._process = psutil.Process(self._processor.pid)
+
+    def pid(self) -> int:
+        assert self._processor.pid is not None
+        return self._processor.pid
+
+    def process(self) -> psutil.Process:
+        return self._process
+
+    def processor_id(self) -> ProcessorID:
+        assert self._processor_id is not None
+        return self._processor_id
+
+    def initialized(self) -> bool:
+        return self._processor_id is not None
+
+    def initialize(self, processor_id: ProcessorID):
+        self._processor_id = processor_id
+
+    def task(self) -> Optional[Task]:
+        return self._task
+
+    def set_task(self, task: Optional[Task]):
+        self._task = task
+
+    def suspended(self) -> bool:
+        return self._suspended
+
+    def suspend(self):
+        assert self._processor is not None
+        assert self._task is not None
+        assert self._suspended is False
+        assert self.initialized()
+
+        if self._hard_suspend:
+            self.__send_signal("SIGSTOP")
+        else:
+            # If we do not want to hardly suspend the processor's process (e.g. to keep network links alive), we request
+            # the process to wait on a synchronization event. That will stop the main thread while allowing the helper
+            # threads to continue running.
+            #
+            # See https://github.com/Citi/scaler/issues/14
+
+            assert self._resume_event is not None
+            assert self._resumed_event is not None
+            self._resume_event.clear()
+            self._resumed_event.clear()
+
+            self.__send_signal(SUSPEND_SIGNAL)
+
+        self._suspended = True
+
+    def resume(self):
+        assert self._task is not None
+        assert self._suspended is True
+
+        if self._hard_suspend:
+            self.__send_signal("SIGCONT")
+        else:
+            assert self._resume_event is not None
+            assert self._resumed_event is not None
+
+            self._resume_event.set()
+
+            # Waits until the processor resumes processing. This avoids any future call to `suspend()` while the
+            # processor hasn't returned from the `_resumed_event.wait()` call yet (causes a re-entrant error on Linux).
+            self._resumed_event.wait()
+
+        self._suspended = False
+
+    def kill(self):
+        self.__send_signal("SIGTERM")
+        self._processor.join(DEFAULT_PROCESSOR_KILL_DELAY_SECONDS)
+
+        if self._processor.exitcode is None:
+            # TODO: some processors fail to interrupt because of a blocking 0mq call. Ideally we should interrupt
+            # these blocking calls instead of sending a SIGKILL signal.
+
+            logging.warning(f"Processor[{self.pid()}] does not terminate in time, send SIGKILL.")
+            self.__send_signal("SIGKILL")
+            self._processor.join()
+
+        self.set_task(None)
+
+    def __send_signal(self, signal_name: str):
+        signal_instance = getattr(signal, signal_name, None)
+        if signal_instance is None:
+            raise RuntimeError(f"unsupported platform, signal not available: {signal_name}.")
+
+        os.kill(self.pid(), signal_instance)