opengris-scaler 1.12.28__cp313-cp313-musllinux_1_2_x86_64.whl

scaler/client/future.py

@@ -0,0 +1,252 @@
+import concurrent.futures
+from typing import Any, Callable, Optional
+
+from scaler.client.serializer.mixins import Serializer
+from scaler.io.mixins import SyncConnector, SyncObjectStorageConnector
+from scaler.protocol.python.common import TaskState
+from scaler.protocol.python.message import Task, TaskCancel
+from scaler.utility.event_list import EventList
+from scaler.utility.identifiers import ObjectID, TaskID
+from scaler.utility.metadata.profile_result import ProfileResult
+from scaler.utility.serialization import deserialize_failure
+
+
+class ScalerFuture(concurrent.futures.Future):
+    """
+    A drop-in replacement for Python's `concurrent.futures.Future`.
+
+    This class is designed to be compatible with Python's Future API, but with some key differences:
+
+    - Delayed futures (`is_delayed` set to `True`) might not fetch the result data when the future is done.
+      Instead, the result is lazily fetched when `result()` or `exception()` is called, or when a callback or waiter is
+      added. That is, `result()` might temporarily be blocking even if `done()` is `True`.
+
+    - `cancel()` may block until a cancellation confirmation is received from Scaler's scheduler.
+    """
+
+    def __init__(
+        self,
+        task: Task,
+        is_delayed: bool,
+        group_task_id: Optional[TaskID],
+        serializer: Serializer,
+        connector_agent: SyncConnector,
+        connector_storage: SyncObjectStorageConnector,
+    ):
+        super().__init__()
+
+        self._waiters = EventList(self._waiters)  # type: ignore[assignment]
+        self._waiters.add_update_callback(self._on_waiters_updated)  # type: ignore[attr-defined]
+
+        self._task_id: TaskID = task.task_id
+        self._is_delayed: bool = is_delayed
+        self._group_task_id: Optional[TaskID] = group_task_id
+        self._serializer: Serializer = serializer
+        self._connector_agent: SyncConnector = connector_agent
+        self._connector_storage: SyncObjectStorageConnector = connector_storage
+
+        self._result_object_id: Optional[ObjectID] = None
+        self._result_received = False
+        self._task_state: Optional[TaskState] = None
+        self._cancel_requested: bool = False
+
+        self._profiling_info: Optional[ProfileResult] = None
+
+    @property
+    def task_id(self) -> TaskID:
+        return self._task_id
+
+    def profiling_info(self) -> ProfileResult:
+        with self._condition:  # type: ignore[attr-defined]
+            if self._profiling_info is None:
+                raise ValueError(f"didn't receive profiling info for {self} yet")
+
+            return self._profiling_info
+
+    def set_result_ready(
+        self, object_id: Optional[ObjectID], task_state: TaskState, profile_result: Optional[ProfileResult] = None
+    ) -> None:
+        with self._condition:  # type: ignore[attr-defined]
+            if self.done():
+                raise concurrent.futures.InvalidStateError(f"invalid future state: {self._state}")
+
+            self._state = "FINISHED"
+
+            self._result_object_id = object_id
+
+            self._task_state = task_state
+
+            if profile_result is not None:
+                self._profiling_info = profile_result
+
+            # if it's not delayed future, or if there is any listener (waiter or callback), get the result immediately
+            if not self._is_delayed or self._has_result_listeners():
+                self._get_result_object()
+
+            self._condition.notify_all()  # type: ignore[attr-defined]
+
+    def set_canceled(self):
+        with self._condition:
+            if self.done():
+                return
+
+            self._state = "CANCELLED_AND_NOTIFIED"
+            self._result_received = True
+            self._cancel_requested = True
+
+            for waiter in self._waiters:
+                waiter.add_cancelled(self)
+
+            self._condition.notify_all()  # type: ignore[attr-defined]
+
+        self._invoke_callbacks()  # type: ignore[attr-defined]
+
+    def _set_result_or_exception(
+        self,
+        result: Optional[Any] = None,
+        exception: Optional[BaseException] = None,
+        profiling_info: Optional[ProfileResult] = None,
+    ) -> None:
+        with self._condition:  # type: ignore[attr-defined]
+            if self.cancelled():
+                raise concurrent.futures.InvalidStateError(f"invalid future state: {self._state}")
+
+            if self._result_received:
+                raise concurrent.futures.InvalidStateError("future already received object data.")
+
+            if profiling_info is not None:
+                if self._profiling_info is not None:
+                    raise concurrent.futures.InvalidStateError("cannot set profiling info twice.")
+
+                self._profiling_info = profiling_info
+
+            self._state = "FINISHED"
+            self._result_received = True
+
+            if exception is not None:
+                assert result is None
+                self._exception = exception
+                for waiter in self._waiters:
+                    waiter.add_exception(self)
+            else:
+                self._result = result
+                for waiter in self._waiters:
+                    waiter.add_result(self)
+
+            self._condition.notify_all()  # type: ignore[attr-defined]
+
+        self._invoke_callbacks()  # type: ignore[attr-defined]
+
+    def set_result(self, result: Any, profiling_info: Optional[ProfileResult] = None) -> None:
+        self._set_result_or_exception(result=result, profiling_info=profiling_info)
+
+    def set_exception(self, exception: Optional[BaseException], profiling_info: Optional[ProfileResult] = None) -> None:
+        self._set_result_or_exception(exception=exception, profiling_info=profiling_info)
+
+    def result(self, timeout: Optional[float] = None) -> Any:
+        with self._condition:  # type: ignore[attr-defined]
+            self._wait_result_ready(timeout)
+
+            # if it's delayed future, get the result when future.result() gets called
+            if self._is_delayed:
+                self._get_result_object()
+
+            return super().result()
+
+    def exception(self, timeout: Optional[float] = None) -> Optional[BaseException]:
+        with self._condition:  # type: ignore[attr-defined]
+            self._wait_result_ready(timeout)
+
+            # if it's delayed future, get the result when future.exception() gets called
+            if self._is_delayed:
+                self._get_result_object()
+
+            return super().exception()
+
+    def cancel(self, timeout: Optional[float] = None) -> bool:
+        with self._condition:  # type: ignore[attr-defined]
+            if self.cancelled():
+                return True
+
+            if self.done():
+                return False
+
+            if not self._cancel_requested:
+                # Send cancellation request to the server
+                cancel_flags = TaskCancel.TaskCancelFlags(force=True)
+
+                if self._group_task_id is not None:
+                    self._connector_agent.send(TaskCancel.new_msg(self._group_task_id, flags=cancel_flags))
+                else:
+                    self._connector_agent.send(TaskCancel.new_msg(self._task_id, flags=cancel_flags))
+
+                self._cancel_requested = True
+
+            # Wait for the answer from the server, can either be a cancel confirmation, or the results if the task
+            # finished while being canceled.
+            self._wait_result_ready(timeout)
+
+        return self.cancelled()
+
+    def add_done_callback(self, fn: Callable[["ScalerFuture"], Any]) -> None:
+        with self._condition:
+            if self.done():
+                self._get_result_object()
+            else:
+                self._done_callbacks.append(fn)  # type: ignore[attr-defined]
+                return
+
+        try:
+            fn(self)
+        except Exception:
+            concurrent.futures._base.LOGGER.exception(f"exception calling callback for {self!r}")
+            raise
+
+    def _on_waiters_updated(self, waiters: EventList):
+        with self._condition:  # type: ignore[attr-defined]
+            # if it's delayed future, get the result when waiter gets added
+            if self._is_delayed and len(self._waiters) > 0:
+                self._get_result_object()
+
+    def _has_result_listeners(self) -> bool:
+        return len(self._done_callbacks) > 0 or len(self._waiters) > 0  # type: ignore[attr-defined]
+
+    def _get_result_object(self):
+        with self._condition:  # type: ignore[attr-defined]
+            if self._result_object_id is None or self.cancelled() or self._result_received:
+                return
+
+            object_bytes = self._connector_storage.get_object(self._result_object_id)
+
+            if self._is_simple_task():
+                # immediately delete non graph result objects
+                # TODO: graph task results could also be deleted if these are not required by another task of the graph.
+                self._connector_storage.delete_object(self._result_object_id)
+
+            if self._task_state == TaskState.Success:
+                self.set_result(self._serializer.deserialize(object_bytes))
+            elif self._task_state == TaskState.Failed:
+                self.set_exception(deserialize_failure(object_bytes))
+            else:
+                raise ValueError(f"unexpected task status: {self._task_state}")
+
+    def _wait_result_ready(self, timeout: Optional[float] = None):
+        """
+        Blocks until the future is done (either successfully, or on failure/cancellation).
+
+        Raises a `TimeoutError` if it blocks more than `timeout` seconds.
+        """
+        if not self.done() and not self._condition.wait(timeout):
+            raise concurrent.futures.TimeoutError()
+
+    def _is_simple_task(self):
+        return self._group_task_id is None and self._task_id is not None
+
+    def __task_type(self) -> str:
+        if self._group_task_id is None:
+            return "SimpleTask"
+
+        if self._group_task_id == self._task_id:
+            return "GraphUmbrellaTask"
+        else:
+            return "GraphSubTask"

scaler/client/object_buffer.py

@@ -0,0 +1,129 @@
+import dataclasses
+import pickle
+from typing import Any, Callable, List, Optional, Set
+
+import cloudpickle
+
+from scaler.client.serializer.mixins import Serializer
+from scaler.io.mixins import SyncConnector, SyncObjectStorageConnector
+from scaler.protocol.python.common import ObjectMetadata
+from scaler.protocol.python.message import ObjectInstruction
+from scaler.utility.identifiers import ClientID, ObjectID
+
+
+@dataclasses.dataclass
+class ObjectCache:
+    object_id: ObjectID
+    object_type: ObjectMetadata.ObjectContentType
+    object_name: bytes
+    object_payload: bytes
+
+
+class ObjectBuffer:
+    def __init__(
+        self,
+        identity: ClientID,
+        serializer: Serializer,
+        connector_agent: SyncConnector,
+        connector_storage: SyncObjectStorageConnector,
+    ):
+        self._identity = identity
+        self._serializer = serializer
+
+        self._connector_agent = connector_agent
+        self._connector_storage = connector_storage
+
+        self._valid_object_ids: Set[ObjectID] = set()
+        self._pending_objects: List[ObjectCache] = list()
+
+        self._serializer_object_id = self.__send_serializer()
+
+    def buffer_send_function(self, fn: Callable) -> ObjectCache:
+        return self.__buffer_send_serialized_object(self.__construct_function(fn))
+
+    def buffer_send_object(self, obj: Any, name: Optional[str] = None) -> ObjectCache:
+        return self.__buffer_send_serialized_object(self.__construct_object(obj, name))
+
+    def commit_send_objects(self):
+        if not self._pending_objects:
+            return
+
+        object_instructions_to_send = [
+            (obj_cache.object_id, obj_cache.object_type, obj_cache.object_name) for obj_cache in self._pending_objects
+        ]
+
+        self._connector_agent.send(
+            ObjectInstruction.new_msg(
+                ObjectInstruction.ObjectInstructionType.Create,
+                self._identity,
+                ObjectMetadata.new_msg(*zip(*object_instructions_to_send)),
+            )
+        )
+
+        for obj_cache in self._pending_objects:
+            self._connector_storage.set_object(obj_cache.object_id, obj_cache.object_payload)
+
+        self._pending_objects.clear()
+
+    def clear(self):
+        """
+        remove all committed and pending objects.
+        """
+
+        self._pending_objects.clear()
+
+        # the Clear instruction does not clear the serializer.
+        self._valid_object_ids.clear()
+        self._valid_object_ids.add(self._serializer_object_id)
+
+        self._connector_agent.send(
+            ObjectInstruction.new_msg(
+                ObjectInstruction.ObjectInstructionType.Clear, self._identity, ObjectMetadata.new_msg(tuple())
+            )
+        )
+
+    def is_valid_object_id(self, object_id: ObjectID) -> bool:
+        return object_id in self._valid_object_ids
+
+    def __construct_serializer(self) -> ObjectCache:
+        serializer_payload = cloudpickle.dumps(self._serializer, protocol=pickle.HIGHEST_PROTOCOL)
+        object_id = ObjectID.generate_serializer_object_id(self._identity)
+        serializer_cache = ObjectCache(
+            object_id, ObjectMetadata.ObjectContentType.Serializer, b"serializer", serializer_payload
+        )
+
+        return serializer_cache
+
+    def __construct_function(self, fn: Callable) -> ObjectCache:
+        function_payload = self._serializer.serialize(fn)
+        object_id = ObjectID.generate_object_id(self._identity)
+        function_cache = ObjectCache(
+            object_id,
+            ObjectMetadata.ObjectContentType.Object,
+            getattr(fn, "__name__", f"<func {repr(object_id)}>").encode(),
+            function_payload,
+        )
+
+        return function_cache
+
+    def __construct_object(self, obj: Any, name: Optional[str] = None) -> ObjectCache:
+        object_payload = self._serializer.serialize(obj)
+        object_id = ObjectID.generate_object_id(self._identity)
+        name_bytes = name.encode() if name else f"<obj {repr(object_id)}>".encode()
+        object_cache = ObjectCache(object_id, ObjectMetadata.ObjectContentType.Object, name_bytes, object_payload)
+
+        return object_cache
+
+    def __buffer_send_serialized_object(self, object_cache: ObjectCache) -> ObjectCache:
+        if object_cache.object_id not in self._valid_object_ids:
+            self._pending_objects.append(object_cache)
+            self._valid_object_ids.add(object_cache.object_id)
+
+        return object_cache
+
+    def __send_serializer(self) -> ObjectID:
+        serialized_serializer = self.__construct_serializer()
+        self.__buffer_send_serialized_object(serialized_serializer)
+        self.commit_send_objects()
+
+        return serialized_serializer.object_id

scaler/client/object_reference.py

@@ -0,0 +1,25 @@
+import dataclasses
+
+from scaler.utility.identifiers import ObjectID
+
+
+@dataclasses.dataclass
+class ObjectReference:
+    name: bytes
+    size: int
+    object_id: ObjectID
+
+    def __repr__(self):
+        return f"ObjectReference(name={self.name!r}, size={self.size} bytes, id={self.object_id!r})"
+
+    def __hash__(self):
+        return hash(self.object_id)
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, ObjectReference):
+            return NotImplemented
+
+        return self.object_id == other.object_id
+
+    def __ne__(self, other):
+        return not self.__eq__(other)

scaler/client/serializer/default.py

@@ -0,0 +1,16 @@
+import pickle
+from typing import Any
+
+import cloudpickle
+
+from scaler.client.serializer.mixins import Serializer
+
+
+class DefaultSerializer(Serializer):
+    @staticmethod
+    def serialize(obj: Any) -> bytes:
+        return cloudpickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
+
+    @staticmethod
+    def deserialize(payload: bytes) -> Any:
+        return cloudpickle.loads(payload)

scaler/client/serializer/mixins.py

@@ -0,0 +1,38 @@
+import abc
+from typing import Any
+
+
+class Serializer(metaclass=abc.ABCMeta):
+    @staticmethod
+    @abc.abstractmethod
+    def serialize(obj: Any) -> bytes:
+        """
+        Serialize the object to bytes, this serialization method is used to call for function object and EACH argument
+        object and function result object, for example:
+
+        def add(a, b):
+            return a + b
+
+        client.submit(add, 1, 2)
+
+        The add function and the arguments 1 and 2 will be serialized and sent to the worker, and the result of the a+b
+        will be serialized and sent back to the client, client will use deserialize function below to deserialize
+
+        :param obj: the object to be serialized, can be function object, argument object, or function result object
+        :return: serialized bytes of the object
+        """
+
+        raise NotImplementedError()
+
+    @staticmethod
+    @abc.abstractmethod
+    def deserialize(payload: bytes) -> Any:
+        """
+        Deserialize the bytes to the original object, this de-serialize method is used to deserialize the function
+        object bytes and EACH serialized argument and serialized function result.
+
+        :param payload: the serialized bytes of the object, can be function object, argument object, or function result
+            object
+        :return: any deserialized object
+        """
+        raise NotImplementedError()

scaler/cluster/cluster.py

@@ -0,0 +1,115 @@
+import logging
+import multiprocessing
+import os
+import signal
+from typing import Dict, List, Optional, Tuple
+
+from scaler.config.types.object_storage_server import ObjectStorageConfig
+from scaler.config.types.zmq import ZMQConfig
+from scaler.utility.logging.utility import setup_logger
+from scaler.worker.worker import Worker
+
+
+class Cluster(multiprocessing.get_context("spawn").Process):  # type: ignore[misc]
+
+    def __init__(
+        self,
+        address: ZMQConfig,
+        object_storage_address: Optional[ObjectStorageConfig],
+        preload: Optional[str],
+        worker_io_threads: int,
+        worker_names: List[str],
+        per_worker_capabilities: Dict[str, int],
+        per_worker_task_queue_size: int,
+        heartbeat_interval_seconds: int,
+        task_timeout_seconds: int,
+        death_timeout_seconds: int,
+        garbage_collect_interval_seconds: int,
+        trim_memory_threshold_bytes: int,
+        hard_processor_suspend: bool,
+        event_loop: str,
+        logging_paths: Tuple[str, ...],
+        logging_config_file: Optional[str],
+        logging_level: str,
+    ):
+        multiprocessing.Process.__init__(self, name="WorkerMaster")
+
+        self._address = address
+        self._object_storage_address = object_storage_address
+        self._preload = preload
+        self._worker_io_threads = worker_io_threads
+        self._worker_names = worker_names
+        self._per_worker_capabilities = per_worker_capabilities
+
+        self._per_worker_task_queue_size = per_worker_task_queue_size
+        self._heartbeat_interval_seconds = heartbeat_interval_seconds
+        self._task_timeout_seconds = task_timeout_seconds
+        self._death_timeout_seconds = death_timeout_seconds
+        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._event_loop = event_loop
+
+        self._logging_paths = logging_paths
+        self._logging_config_file = logging_config_file
+        self._logging_level = logging_level
+
+        self._workers: List[Worker] = []
+
+    def run(self):
+        setup_logger(self._logging_paths, self._logging_config_file, self._logging_level)
+        self.__register_signal()
+        self.__start_workers_and_run_forever()
+
+    def __destroy(self, *args):
+        assert args is not None
+        logging.info(f"{self.__get_prefix()} received signal, shutting down")
+        for worker in self._workers:
+            logging.info(f"{self.__get_prefix()} shutting down {worker.identity!r}")
+            os.kill(worker.pid, signal.SIGINT)
+
+    def __register_signal(self):
+        signal.signal(signal.SIGINT, self.__destroy)
+        signal.signal(signal.SIGTERM, self.__destroy)
+
+    def __start_workers_and_run_forever(self):
+        logging.info(
+            f"{self.__get_prefix()} starting {len(self._worker_names)} workers, heartbeat_interval_seconds="
+            f"{self._heartbeat_interval_seconds}, task_timeout_seconds={self._task_timeout_seconds}"
+        )
+
+        self._workers = [
+            Worker(
+                event_loop=self._event_loop,
+                name=name,
+                address=self._address,
+                object_storage_address=self._object_storage_address,
+                capabilities=self._per_worker_capabilities,
+                preload=self._preload,
+                io_threads=self._worker_io_threads,
+                task_queue_size=self._per_worker_task_queue_size,
+                heartbeat_interval_seconds=self._heartbeat_interval_seconds,
+                garbage_collect_interval_seconds=self._garbage_collect_interval_seconds,
+                trim_memory_threshold_bytes=self._trim_memory_threshold_bytes,
+                task_timeout_seconds=self._task_timeout_seconds,
+                death_timeout_seconds=self._death_timeout_seconds,
+                hard_processor_suspend=self._hard_processor_suspend,
+                logging_paths=self._logging_paths,
+                logging_level=self._logging_level,
+            )
+            for name in self._worker_names
+        ]
+
+        for worker in self._workers:
+            worker.start()
+
+        for worker in self._workers:
+            logging.info(f"{worker.identity!r} started")
+
+        for worker in self._workers:
+            worker.join()
+
+        logging.info(f"{self.__get_prefix()} shutdown")
+
+    def __get_prefix(self):
+        return f"{self.__class__.__name__}:"