opengris-scaler 1.12.28__cp313-cp313-musllinux_1_2_x86_64.whl

scaler/protocol/python/mixins.py

@@ -0,0 +1,13 @@
+import abc
+from typing import TypeVar
+
+
+class Message(metaclass=abc.ABCMeta):
+    def __init__(self, msg):
+        self._msg = msg
+
+    def get_message(self):
+        return self._msg
+
+
+MessageType = TypeVar("MessageType", bound=Message)

scaler/protocol/python/object_storage.py

@@ -0,0 +1,118 @@
+import dataclasses
+import enum
+import struct
+
+from scaler.protocol.capnp._python import _object_storage  # noqa
+from scaler.protocol.python.mixins import Message
+from scaler.utility.identifiers import ObjectID
+
+OBJECT_ID_FORMAT = "!QQQQ"
+
+
+@dataclasses.dataclass
+class ObjectRequestHeader(Message):
+    class ObjectRequestType(enum.Enum):
+        SetObject = _object_storage.ObjectRequestHeader.ObjectRequestType.setObject
+        GetObject = _object_storage.ObjectRequestHeader.ObjectRequestType.getObject
+        DeleteObject = _object_storage.ObjectRequestHeader.ObjectRequestType.deleteObject
+        DuplicateObjectID = _object_storage.ObjectRequestHeader.ObjectRequestType.duplicateObjectID
+
+    def __init__(self, msg):
+        super().__init__(msg)
+
+    @property
+    def object_id(self) -> ObjectID:
+        return from_capnp_object_id(self._msg.objectID)
+
+    @property
+    def payload_length(self) -> int:
+        return self._msg.payloadLength
+
+    @property
+    def request_id(self) -> int:
+        return self._msg.requestID
+
+    @property
+    def request_type(self) -> ObjectRequestType:
+        return ObjectRequestHeader.ObjectRequestType(self._msg.requestType.raw)
+
+    @staticmethod
+    def new_msg(
+        object_id: ObjectID, payload_length: int, request_id: int, request_type: ObjectRequestType
+    ) -> "ObjectRequestHeader":
+        return ObjectRequestHeader(
+            _object_storage.ObjectRequestHeader(
+                objectID=to_capnp_object_id(object_id),
+                payloadLength=payload_length,
+                requestID=request_id,
+                requestType=request_type.value,
+            )
+        )
+
+    def get_message(self):
+        return self._msg
+
+
+@dataclasses.dataclass
+class ObjectResponseHeader(Message):
+    MESSAGE_LENGTH = 80  # there does not seem to be a way to statically know the size of a pycapnp message
+
+    class ObjectResponseType(enum.Enum):
+        SetOK = _object_storage.ObjectResponseHeader.ObjectResponseType.setOK
+        GetOK = _object_storage.ObjectResponseHeader.ObjectResponseType.getOK
+        DelOK = _object_storage.ObjectResponseHeader.ObjectResponseType.delOK
+        DelNotExists = _object_storage.ObjectResponseHeader.ObjectResponseType.delNotExists
+        DuplicateOK = _object_storage.ObjectResponseHeader.ObjectResponseType.duplicateOK
+
+    def __init__(self, msg):
+        super().__init__(msg)
+
+    @property
+    def object_id(self) -> ObjectID:
+        return from_capnp_object_id(self._msg.objectID)
+
+    @property
+    def payload_length(self) -> int:
+        return self._msg.payloadLength
+
+    @property
+    def response_id(self) -> int:
+        return self._msg.responseID
+
+    @property
+    def response_type(self) -> ObjectResponseType:
+        return ObjectResponseHeader.ObjectResponseType(self._msg.responseType.raw)
+
+    @staticmethod
+    def new_msg(
+        object_id: ObjectID, payload_length: int, response_id: int, response_type: ObjectResponseType
+    ) -> "ObjectResponseHeader":
+        return ObjectResponseHeader(
+            _object_storage.ObjectResponseHeader(
+                objectID=to_capnp_object_id(object_id),
+                payloadLength=payload_length,
+                responseID=response_id,
+                responseType=response_type.value,
+            )
+        )
+
+    def get_message(self):
+        return self._msg
+
+
+def to_capnp_object_id(object_id: ObjectID) -> _object_storage.ObjectID:
+    field0, field1, field2, field3 = struct.unpack(OBJECT_ID_FORMAT, object_id)
+
+    return _object_storage.ObjectID(field0=field0, field1=field1, field2=field2, field3=field3)
+
+
+def from_capnp_object_id(capnp_object_id: _object_storage.ObjectID) -> ObjectID:
+    return ObjectID(
+        struct.pack(
+            OBJECT_ID_FORMAT,
+            capnp_object_id.field0,
+            capnp_object_id.field1,
+            capnp_object_id.field2,
+            capnp_object_id.field3,
+        )
+    )

scaler/protocol/python/status.py

@@ -0,0 +1,279 @@
+from typing import Dict, List
+
+from scaler.protocol.capnp._python import _status  # noqa
+from scaler.protocol.python.common import TaskState
+from scaler.protocol.python.mixins import Message
+from scaler.utility.identifiers import ClientID, WorkerID
+
+CPU_MAXIMUM = 1000
+WorkerGroupID = bytes
+
+
+class Resource(Message):
+    def __init__(self, msg):
+        super().__init__(msg)
+
+    @property
+    def cpu(self) -> int:
+        return self._msg.cpu
+
+    @property
+    def rss(self) -> int:
+        return self._msg.rss
+
+    @staticmethod
+    def new_msg(cpu: int, rss: int) -> "Resource":  # type: ignore[override]
+        return Resource(_status.Resource(cpu=min(cpu, CPU_MAXIMUM), rss=rss))
+
+    def get_message(self):
+        return self._msg
+
+
+class ObjectManagerStatus(Message):
+    def __init__(self, msg):
+        super().__init__(msg)
+
+    @property
+    def number_of_objects(self) -> int:
+        return self._msg.numberOfObjects
+
+    @staticmethod
+    def new_msg(number_of_objects: int) -> "ObjectManagerStatus":  # type: ignore[override]
+        return ObjectManagerStatus(_status.ObjectManagerStatus(numberOfObjects=number_of_objects))
+
+    def get_message(self):
+        return self._msg
+
+
+class ClientManagerStatus(Message):
+    def __init__(self, msg):
+        super().__init__(msg)
+
+    @property
+    def client_to_num_of_tasks(self) -> Dict[ClientID, int]:
+        return {p.client: p.numTask for p in self._msg.clientToNumOfTask}
+
+    @staticmethod
+    def new_msg(client_to_num_of_tasks: Dict[ClientID, int]) -> "ClientManagerStatus":  # type: ignore[override]
+        return ClientManagerStatus(
+            _status.ClientManagerStatus(
+                clientToNumOfTask=[
+                    _status.ClientManagerStatus.Pair(client=client_id.decode(), numTask=num_tasks)
+                    for client_id, num_tasks in client_to_num_of_tasks.items()
+                ]
+            )
+        )
+
+    def get_message(self):
+        return self._msg
+
+
+class TaskManagerStatus(Message):
+    VALUE_SIZE_LIMIT = 2**32
+
+    def __init__(self, msg):
+        super().__init__(msg)
+
+    @property
+    def state_to_count(self) -> Dict[TaskState, int]:
+        return {TaskState(p.state): p.count for p in self._msg.stateToCount}
+
+    @staticmethod
+    def new_msg(state_to_count: Dict[TaskState, int]) -> "TaskManagerStatus":  # type: ignore[override]
+        return TaskManagerStatus(
+            _status.TaskManagerStatus(
+                stateToCount=[
+                    _status.TaskManagerStatus.Pair(state=p[0].value, count=p[1] % TaskManagerStatus.VALUE_SIZE_LIMIT)
+                    for p in state_to_count.items()
+                ]
+            )
+        )
+
+    def get_message(self):
+        return self._msg
+
+
+class ProcessorStatus(Message):
+    def __init__(self, msg):
+        super().__init__(msg)
+
+    @property
+    def pid(self) -> int:
+        return self._msg.pid
+
+    @property
+    def initialized(self) -> int:
+        return self._msg.initialized
+
+    @property
+    def has_task(self) -> bool:
+        return self._msg.hasTask
+
+    @property
+    def suspended(self) -> bool:
+        return self._msg.suspended
+
+    @property
+    def resource(self) -> Resource:
+        return Resource(self._msg.resource)
+
+    @staticmethod
+    def new_msg(
+        pid: int, initialized: int, has_task: bool, suspended: bool, resource: Resource  # type: ignore[override]
+    ) -> "ProcessorStatus":
+        return ProcessorStatus(
+            _status.ProcessorStatus(
+                pid=pid, initialized=initialized, hasTask=has_task, suspended=suspended, resource=resource.get_message()
+            )
+        )
+
+    def get_message(self):
+        return self._msg
+
+
+class WorkerStatus(Message):
+    def __init__(self, msg):
+        super().__init__(msg)
+
+    @property
+    def worker_id(self) -> WorkerID:
+        return WorkerID(self._msg.workerId)
+
+    @property
+    def agent(self) -> Resource:
+        return Resource(self._msg.agent)
+
+    @property
+    def rss_free(self) -> int:
+        return self._msg.rssFree
+
+    @property
+    def free(self) -> int:
+        return self._msg.free
+
+    @property
+    def sent(self) -> int:
+        return self._msg.sent
+
+    @property
+    def queued(self) -> int:
+        return self._msg.queued
+
+    @property
+    def suspended(self) -> bool:
+        return self._msg.suspended
+
+    @property
+    def lag_us(self) -> int:
+        return self._msg.lagUS
+
+    @property
+    def last_s(self) -> int:
+        return self._msg.lastS
+
+    @property
+    def itl(self) -> str:
+        return self._msg.itl
+
+    @property
+    def processor_statuses(self) -> List[ProcessorStatus]:
+        return [ProcessorStatus(ps) for ps in self._msg.processorStatuses]
+
+    @staticmethod
+    def new_msg(  # type: ignore[override]
+        worker_id: WorkerID,
+        agent: Resource,
+        rss_free: int,
+        free: int,
+        sent: int,
+        queued: int,
+        suspended: int,
+        lag_us: int,
+        last_s: int,
+        itl: str,
+        processor_statuses: List[ProcessorStatus],
+    ) -> "WorkerStatus":
+        return WorkerStatus(
+            _status.WorkerStatus(
+                workerId=bytes(worker_id),
+                agent=agent.get_message(),
+                rssFree=rss_free,
+                free=free,
+                sent=sent,
+                queued=queued,
+                suspended=suspended,
+                lagUS=lag_us,
+                lastS=last_s,
+                itl=itl,
+                processorStatuses=[ps.get_message() for ps in processor_statuses],
+            )
+        )
+
+    def get_message(self):
+        return self._msg
+
+
+class WorkerManagerStatus(Message):
+    def __init__(self, msg):
+        super().__init__(msg)
+
+    @property
+    def workers(self) -> List[WorkerStatus]:
+        return [WorkerStatus(ws) for ws in self._msg.workers]
+
+    @staticmethod
+    def new_msg(workers: List[WorkerStatus]) -> "WorkerManagerStatus":  # type: ignore[override]
+        return WorkerManagerStatus(_status.WorkerManagerStatus(workers=[ws.get_message() for ws in workers]))
+
+    def get_message(self):
+        return self._msg
+
+
+class ScalingManagerStatus(Message):
+    def __init__(self, msg):
+        super().__init__(msg)
+
+    @property
+    def worker_groups(self) -> Dict[WorkerGroupID, List[WorkerID]]:
+        return {wg.workerGroupID: [WorkerID(wid) for wid in wg.workerIDs] for wg in self._msg.workerGroups}
+
+    @staticmethod
+    def new_msg(worker_groups: Dict[WorkerGroupID, List[WorkerID]]) -> "ScalingManagerStatus":  # type: ignore[override]
+        return ScalingManagerStatus(
+            _status.ScalingManagerStatus(
+                workerGroups=[
+                    _status.ScalingManagerStatus.Pair(
+                        workerGroupID=worker_group_id, workerIDs=[bytes(worker_id) for worker_id in worker_ids]
+                    )
+                    for worker_group_id, worker_ids in worker_groups.items()
+                ]
+            )
+        )
+
+    def get_message(self):
+        return self._msg
+
+
+class BinderStatus(Message):
+    def __init__(self, msg):
+        super().__init__(msg)
+
+    @property
+    def received(self) -> Dict[str, int]:
+        return {p.client: p.number for p in self._msg.received}
+
+    @property
+    def sent(self) -> Dict[str, int]:
+        return {p.client: p.number for p in self._msg.sent}
+
+    @staticmethod
+    def new_msg(received: Dict[str, int], sent: Dict[str, int]) -> "BinderStatus":  # type: ignore[override]
+        return BinderStatus(
+            _status.BinderStatus(
+                received=[_status.BinderStatus.Pair(client=p[0], number=p[1]) for p in received.items()],
+                sent=[_status.BinderStatus.Pair(client=p[0], number=p[1]) for p in sent.items()],
+            )
+        )
+
+    def get_message(self):
+        return self._msg

scaler/protocol/worker.md

@@ -0,0 +1,228 @@
+# Custom Worker Implementation
+
+## Overview
+
+The existed python worker implementation can work with scheduler and scaler protocol is under `scaler.worker`
+
+- workers are connecting to 1 single TCP port of scheduler by using zmq
+- each worker has a fixed length task queue to store the incoming tasks
+- in Task message or `TaskResult` message, it will contain object ids instead of actual data, once worker received one
+  task, it should ask scheduler for the actual data by sending `ObjectRequest` message
+- once worker finished the task, it should send `TaskResult` message back to scheduler, the `TaskResult` message should
+  have the object id of the result, and worker should send `ObjectInstruction` message to scheduler to store the result
+- each worker have a unique id and can be identified by the scheduler that can be set by zmq (see below setup)
+- worker is responsible for sending the heartbeat to scheduler, if scheduler didn't receive heartbeat from worker for
+  a period of time, then scheduler will consider the worker is dead and will reallocate all tasks to other workers, also
+  heartbeat message will contain the worker's resource usage information and queue capacity and usage status
+- each Task is tied to a source, each source has a dedicated serializer, the source is used to choose the proper
+  serializer for deserializing the function and arguments, worker will need retrieve the serializer bytes using an
+  `ObjectRequest` message, and then deserialize the serializer bytes using cloudpickle, deserialized serializer is
+  implements the interface `scaler.client.serializer.mixins.Serializer`
+- When worker received TaskCancel message, it should cancel the task regardless of tasks status with the given task ID
+  and send a `TaskResult` back to the scheduler
+
+- worker can request to scheduler to balance the tasks, scheduler will send the tasks to be given up by the worker (this
+will be replaced by TaskCancel message in the future)
+
+## Setup
+
+The zmq `IDENTITY` must be unique. The zmq `SNDHWM` and `RCVHWM` should both be set to 0 to prevent messages from being
+dropped unexpectedly, and it should be DEALER socket type.
+
+All message will be list of bytes represented by frames, the first frame is the message type, and the rest of the frames
+are the message data. One message is the whole list of frames.
+
+To send a single message using pyzmq, use the `send_multipart` method:
+
+```python
+# Synchronous API
+socket.send_multipart(frames)
+
+# Asyncronous API
+await socket.send_multipart(frames)
+```
+
+Each message may contain multiple zmq frames. To read a single message using pyzmq, use the `recv_multipart` method:
+
+```python
+# Synchronous API
+frames = socket.recv_multipart()
+
+# Asyncronous API
+frames = await socket.recv_multipart()
+```
+
+All below messages, please see `scaler.protocol.python.message.py` for actual message structure to help you understand
+
+## Recv messages
+
+### Task `TK`
+
+| message_type | task_id | source  | metadata | func_object_id | arg 1 type | arg 1 data | (...) | arg N type | arg N data |
+|:------------:|:-------:|:-------:|:--------:|:--------------:|:----------:|:----------:|:-----:|:----------:|:----------:|
+|    b"TK"     | X bytes | X bytes | X bytes  |    X bytes     |    b"R"    |  X bytes   |       |    b"R"    |  X bytes   |
+
+* task_id: Task ID
+* source: Source ID. This is used to choose the proper serializer for deserializing the
+  function and arguments. The serializer object ID is md5 hash of `source + b"serializer"` and the serializer object
+  bytes data must first be deserialized using cloudpickle. source is essentially the client id, indicate which client
+  this task it belongs to, please refer to `scaler.utility.object_utility` for `generate_serializer_object_id`
+* metadata: The metadata of the task, can be empty bytes like `b""`
+* func_object_id: Function object ID
+* arg type: Must be type `b"R"` for ObjectID
+* arg data: Object ID
+
+After executing `Task`, the worker should send a `TaskResult` message with the final task result.
+
+### TaskCancel `TC`
+
+| message_type | task_id |
+|:------------:|:-------:|
+|    b"TC"     | X bytes |
+
+* task_id: Task ID to cancel
+
+When a `TaskCancel` message is received, the worker should cancel the task with the given task ID and send a
+`TaskResult`
+
+### ObjectInstruction `OI`
+
+| message_type | source  | type | num_object_ids | num_object_names | num_object_bytes | object 1 id | (...) | object N id |
+|:------------:|:-------:|:----:|:--------------:|:----------------:|:----------------:|:-----------:|:-----:|:-----------:|
+|    b"OI"     | X bytes | b"D" |  unsigned int  |   unsigned int   |   unsigned int   |   X bytes   |       |   X bytes   |
+
+* source: Source ID
+* type: Must be `b"D` for Delete
+* num_object_ids: Number of object IDs
+* num_object_names: Number of object names, value must be zero
+* num_object_bytes: Number of object bytes, value must be zero
+* object id: Object ID
+
+When a Delete `ObjectInstruction` message is received, the worker should delete the objects with the given object IDs on
+worker side
+
+### ObjectResponse `OA`
+
+| message_type |  type  | num_object_ids | num_object_names | num_object_bytes | object 1 id | (...) | object N id | object 1 name | (...) | object N name | object 1 bytes | (...) | object N bytes |
+|:------------:|:------:|:--------------:|:----------------:|:----------------:|:-----------:|:-----:|:-----------:|:-------------:|:-----:|:-------------:|:--------------:|:-----:|:--------------:|
+|    b"OA"     | 1 byte |  unsigned int  |   unsigned int   |   unsigned int   |   X bytes   |       |   X bytes   |    X bytes    |       |    X bytes    |    X bytes     |       |    X bytes     |
+
+* type: `b"C"` for object found, `b"N"` for object not found
+* num_object_ids: Number of object IDs
+* num_object_names: Number of object names
+* num_object_bytes: Number of object bytes
+* object id: Object ID
+* object name: Object name
+* object bytes: Object bytes
+
+### BalanceRequest `BQ` (will be replaced by TaskCancel in the future, low priority to implement this)
+
+| message_type |  num_tasks   |
+|:------------:|:------------:|
+|    b"BQ"     | unsigned int |
+
+* num_tasks: Number of tasks to give up
+
+When a `BalanceRequest` message is received, the worker should send a `BalanceResponse` message with num_tasks number of
+task IDs.
+
+### WorkerHeartbeatEcho `HE` (optional)
+
+| message_type |  empty  |
+|:------------:|:-------:|
+|    b"HE"     | 0 bytes |
+
+A `WorkerHeartbeatEcho` message indicates that the scheduler has received the worker's `WorkerHeartbeat` message.
+
+### ClientDisconnect `CS` (optional)
+
+| message_type | type |
+|:------------:|:----:|
+|    b"CS"     | b"S" |
+
+* type: Must be type `b"S"` for Shutdown.
+
+When a Shutdown `ClientDisconnect` message is received, the worker should shutdown.
+
+## Send messages
+
+### TaskResult `TR`
+
+| message_type | task_id | status | result  | metadata |
+|:------------:|:-------:|:------:|:-------:|:--------:|
+|    b"TR"     | X bytes | 1 byte | X bytes | X bytes  |
+
+* task_id: Task ID
+* status: `b"S"` for Success, `b"F"` for Failed, `b"C"` for Canceled, `b"K"` for WorkerDied, `b"W"` for NoWorker, `b"I"`
+  for Inactive, `b"R"` for Running, and `b"X"` for Canceling
+* result: Task result object ID
+* metadata: Task metadata
+
+Worker must submit a Create `ObjectInstruction` message BEFORE returning the task result containing a task result object
+ID.
+
+### ObjectInstruction `OI`
+
+| message_type | source  | type | num_object_ids | num_object_names | num_object_bytes | object 1 id | (...) | object N id | object 1 name | (...) | object N name | object 1 bytes | (...) | object N bytes |
+|:------------:|:-------:|:----:|:--------------:|:----------------:|:----------------:|:-----------:|:-----:|:-----------:|:-------------:|:-----:|:-------------:|:--------------:|:-----:|:--------------:|
+|    b"OI"     | X bytes | b"C" |  unsigned int  |   unsigned int   |   unsigned int   |   X bytes   |       |   X bytes   |    X bytes    |       |    X bytes    |    X bytes     |       |    X bytes     |
+
+* source: Source ID (this should be the same as the source in the corresponding `Task` message)
+* type: Must be `b"C` for Create
+* num_object_ids: Number of object IDs
+* num_object_names: Number of object names
+* num_object_bytes: Number of object bytes
+* object id: Object ID (please use uuid.uuid4.bytes for object id)
+* object name: Object name
+* object bytes: Object bytes
+
+### ObjectRequest `OR`
+
+| message_type | type | object 1 id | (...) | object N id |
+|:------------:|:----:|:-----------:|:-----:|:-----------:|
+|    b"OR"     | b"A" |   X bytes   |       |   X bytes   |
+
+* type: Must be `b"A"` for Get
+* object id: Object ID
+
+When received a `Task` message, all the functions and arguments are just object IDs, the worker will need get object
+content by sending `ObjectRequest` message to the scheduler, the scheduler will reply with `ObjectResponse` message,
+then worker can deserialize the object content and start executing the task
+
+### BalanceResponse `BR`
+
+| message_type | task 1 id | (...) | task N id |
+|:------------:|:---------:|:-----:|:---------:|
+|    b"BR"     |  X bytes  |       |  X bytes  |
+
+* task id: Task ID
+
+### WorkerHeartbeat `HB`
+
+| message_type |   agent_cpu    |     agent_rss      |   worker_cpu   |     worker_rss     |      rss_free      |  queued_tasks  |  latency_us  | initialized | has_task | task_lock |
+|:------------:|:--------------:|:------------------:|:--------------:|:------------------:|:------------------:|:--------------:|:------------:|:-----------:|:--------:|:---------:|
+|    b"HB"     | unsigned short | unsigned long long | unsigned short | unsigned long long | unsigned long long | unsigned short | unsigned int |    _Bool    |  _Bool   |   _Bool   |
+
+* agent_cpu: Agent CPU usage
+* agent_rss: Agent resident set size in bytes
+* worker_cpu: Worker CPU usage
+* worker_rss: Worker resident set size in bytes
+* rss_free: Free memory in bytes
+* queued_tasks: Number of queued tasks
+* latency_us: Latency in microseconds
+* initialized: Worker initialized
+* has_task: Worker has task
+* task_lock: Worker task lock
+
+Worker must send a `WorkerHeartbeat` message every heartbeat interval (at least 1 second) or else the scheduler will
+consider the worker dead.
+
+### DisconnectRequest `DR` (optional)
+
+| message_type | worker  |
+|:------------:|:-------:|
+|    b"DR"     | X bytes |
+
+* worker: Worker ID
+
+When a `DisconnectRequest` message is sent, the worker should disconnect from the scheduler.

scaler/scheduler/allocate_policy/allocate_policy.py

@@ -0,0 +1,9 @@
+import enum
+
+from scaler.scheduler.allocate_policy.capability_allocate_policy import CapabilityAllocatePolicy
+from scaler.scheduler.allocate_policy.even_load_allocate_policy import EvenLoadAllocatePolicy
+
+
+class AllocatePolicy(enum.Enum):
+    even = EvenLoadAllocatePolicy
+    capability = CapabilityAllocatePolicy