opengris-scaler 1.12.37__cp38-cp38-musllinux_1_2_x86_64.whl

scaler/io/ymq_sync_object_storage_connector.py

@@ -0,0 +1,184 @@
+import socket
+import uuid
+from threading import Lock
+from typing import Iterable, Optional
+
+from scaler.io.mixins import SyncObjectStorageConnector
+from scaler.io.ymq.ymq import IOContext, IOSocket, IOSocketType, Message, YMQException
+from scaler.protocol.capnp._python import _object_storage  # noqa
+from scaler.protocol.python.object_storage import ObjectRequestHeader, ObjectResponseHeader, to_capnp_object_id
+from scaler.utility.exceptions import ObjectStorageException
+from scaler.utility.identifiers import ObjectID
+
+# Some OSes raise an OSError when sending buffers too large with send() or sendmsg().
+MAX_CHUNK_SIZE = 128 * 1024 * 1024
+
+
+class PyYMQSyncObjectStorageConnector(SyncObjectStorageConnector):
+    """An synchronous connector that uses an raw TCP socket to connect to a Scaler's object storage instance."""
+
+    def __init__(self, host: str, port: int):
+        self._host = host
+        self._port = port
+
+        self._next_request_id = 0
+
+        self._socket_lock = Lock()
+
+        self._identity: str = f"{self.__class__.__name__}|{socket.gethostname().split('.')[0]}|{uuid.uuid4()}"
+        self._io_context = IOContext()
+        self._io_socket: IOSocket = self._io_context.createIOSocket_sync(self._identity, IOSocketType.Connector)
+        self._io_socket.connect_sync(self.address)
+
+    def __del__(self):
+        self.destroy()
+
+    def destroy(self):
+        with self._socket_lock:
+            if self._io_socket is not None:
+                self._io_socket = None
+
+    @property
+    def address(self) -> str:
+        return f"tcp://{self._host}:{self._port}"
+
+    def set_object(self, object_id: ObjectID, payload: bytes):
+        """
+        Sets the object's payload on the object storage server.
+        """
+
+        with self._socket_lock:
+            self.__send_request(object_id, len(payload), ObjectRequestHeader.ObjectRequestType.SetObject, payload)
+            response_header, response_payload = self.__receive_response()
+
+        self.__ensure_response_type(response_header, [ObjectResponseHeader.ObjectResponseType.SetOK])
+        self.__ensure_empty_payload(response_payload)
+
+    def get_object(self, object_id: ObjectID, max_payload_length: int = 2**64 - 1) -> bytearray:
+        """
+        Returns the object's payload from the object storage server.
+
+        Will block until the object is available.
+        """
+
+        with self._socket_lock:
+            self.__send_request(object_id, max_payload_length, ObjectRequestHeader.ObjectRequestType.GetObject)
+            response_header, response_payload = self.__receive_response()
+
+        self.__ensure_response_type(response_header, [ObjectResponseHeader.ObjectResponseType.GetOK])
+
+        return response_payload
+
+    def delete_object(self, object_id: ObjectID) -> bool:
+        """
+        Removes the object from the object storage server.
+
+        Returns `False` if the object wasn't found in the server. Otherwise returns `True`.
+        """
+
+        with self._socket_lock:
+            self.__send_request(object_id, 0, ObjectRequestHeader.ObjectRequestType.DeleteObject)
+            response_header, response_payload = self.__receive_response()
+
+        self.__ensure_response_type(
+            response_header,
+            [ObjectResponseHeader.ObjectResponseType.DelOK, ObjectResponseHeader.ObjectResponseType.DelNotExists],
+        )
+        self.__ensure_empty_payload(response_payload)
+
+        return response_header.response_type == ObjectResponseHeader.ObjectResponseType.DelOK
+
+    def duplicate_object_id(self, object_id: ObjectID, new_object_id: ObjectID) -> None:
+        """
+        Link an object's content to a new object ID on the object storage server.
+        """
+
+        object_id_payload = to_capnp_object_id(object_id).to_bytes()
+
+        with self._socket_lock:
+            self.__send_request(
+                new_object_id,
+                len(object_id_payload),
+                ObjectRequestHeader.ObjectRequestType.DuplicateObjectID,
+                object_id_payload,
+            )
+            response_header, response_payload = self.__receive_response()
+
+        self.__ensure_response_type(response_header, [ObjectResponseHeader.ObjectResponseType.DuplicateOK])
+        self.__ensure_empty_payload(response_payload)
+
+    def __ensure_is_connected(self):
+        if self._io_socket is None:
+            raise ObjectStorageException("connector is closed.")
+
+    def __ensure_response_type(
+        self, header: ObjectResponseHeader, valid_response_types: Iterable[ObjectResponseHeader.ObjectResponseType]
+    ):
+        if header.response_type not in valid_response_types:
+            raise RuntimeError(f"unexpected object storage response_type={header.response_type}.")
+
+    def __ensure_empty_payload(self, payload: bytearray):
+        if len(payload) != 0:
+            raise RuntimeError(f"unexpected response payload_length={len(payload)}, expected 0.")
+
+    def __send_request(
+        self,
+        object_id: ObjectID,
+        payload_length: int,
+        request_type: ObjectRequestHeader.ObjectRequestType,
+        payload: Optional[bytes] = None,
+    ):
+        self.__ensure_is_connected()
+        assert self._io_socket is not None
+
+        request_id = self._next_request_id
+        self._next_request_id += 1
+        self._next_request_id %= 2**64 - 1  # UINT64_MAX
+
+        header = ObjectRequestHeader.new_msg(object_id, payload_length, request_id, request_type)
+        header_bytes = header.get_message().to_bytes()
+
+        if payload is not None:
+            self._io_socket.send_sync(Message(address=None, payload=header_bytes))
+            self._io_socket.send_sync(Message(address=None, payload=payload))
+        else:
+            self._io_socket.send_sync(Message(address=None, payload=header_bytes))
+
+    def __receive_response(self):
+        assert self._io_socket is not None
+
+        try:
+            header = self.__read_response_header()
+            payload = self.__read_response_payload(header)
+            return header, payload
+        except YMQException:
+            self.__raise_connection_failure()
+
+    def __read_response_header(self) -> ObjectResponseHeader:
+        assert self._io_socket is not None
+
+        header_bytes = self._io_socket.recv_sync().payload.data
+        if header_bytes is None:
+            self.__raise_connection_failure()
+
+        # pycapnp does not like to read from a bytearray object. This look like an not-yet-resolved issue.
+        # That's is annoying because it leads to an unnecessary copy of the header's buffer.
+        # See https://github.com/capnproto/pycapnp/issues/153
+        # header_bytes = bytes(header_bytearray)
+
+        with _object_storage.ObjectResponseHeader.from_bytes(header_bytes) as header_message:
+            return ObjectResponseHeader(header_message)
+
+    def __read_response_payload(self, header: ObjectResponseHeader) -> bytearray:
+        if header.payload_length > 0:
+            res = self._io_socket.recv_sync().payload.data
+            if res is None:
+                self.__raise_connection_failure()
+            assert len(res) == header.payload_length
+            return bytearray(res)
+        else:
+            return bytearray()
+
+    @staticmethod
+    def __raise_connection_failure():
+        raise ObjectStorageException("connection failure to object storage server.")

scaler/protocol/capnp/_python.py

@@ -0,0 +1,6 @@
+import capnp  # noqa
+
+import scaler.protocol.capnp.common_capnp as _common  # noqa
+import scaler.protocol.capnp.message_capnp as _message  # noqa
+import scaler.protocol.capnp.object_storage_capnp as _object_storage  # noqa
+import scaler.protocol.capnp.status_capnp as _status  # noqa

scaler/protocol/capnp/common.capnp

@@ -0,0 +1,68 @@
+@0xf57f79ac88fab620;
+
+enum TaskResultType {
+    success @0;           # if submit and task is done and get result
+    failed @1;            # if submit and task is failed on worker
+    failedWorkerDied @2;  # if submit and worker died (only happened when scheduler keep_task=False)
+}
+
+enum TaskCancelConfirmType {
+    canceled @0;               # if cancel success
+    cancelFailed @1;           # if cancel failed, this might happened if the task is in process
+    cancelNotFound @2;         # if cancel cannot find such task
+}
+
+enum TaskTransition {
+    hasCapacity @0;
+    taskResultSuccess @1;
+    taskResultFailed @2;
+    taskResultWorkerDied @3;
+    taskCancel @4;
+    taskCancelConfirmCanceled @5;
+    taskCancelConfirmFailed @6;
+    taskCancelConfirmNotFound @7;
+    balanceTaskCancel @8;
+    workerDisconnect @9;
+    schedulerHasTask @10;
+    schedulerHasNoTask @11;
+}
+
+enum TaskState {
+    inactive @0;
+    running @1;
+    canceling @2;
+    balanceCanceling @3;
+    success @4;
+    failed @5;
+    failedWorkerDied @6;
+    canceled @7;
+    canceledNotFound @8;
+    balanceCanceled @9;
+    workerDisconnecting @10;
+}
+
+enum WorkerState {
+    connected @0;
+    disconnected @1;
+}
+
+struct TaskCapability {
+    name @0 :Text;      # the name of the capability provided by the worker/required by the task (e.g. "gpu" or "linux")
+    value @1 :Int64;    # the quantity of the capability provided/required. Use -1 for quantity-less capabilities
+}
+
+struct ObjectMetadata {
+    objectIds @0 :List(Data);
+    objectTypes @1 :List(ObjectContentType);
+    objectNames @2 :List(Data);
+
+    enum ObjectContentType {
+        serializer @0;
+        object @1;
+    }
+}
+
+struct ObjectStorageAddress {
+    host @0 :Text;
+    port @1 :UInt16;
+}

scaler/protocol/capnp/message.capnp

@@ -0,0 +1,218 @@
+@0xaf44f44ea94a4675;
+
+using CommonType = import "common.capnp";
+using Status = import "status.capnp";
+
+struct Task {
+    taskId @0 :Data;
+    source @1 :Data;
+    metadata @2 :Data;
+    funcObjectId @3 :Data;
+    functionArgs @4 :List(Argument);
+    capabilities @5 :List(CommonType.TaskCapability);
+
+    struct Argument {
+        type @0 :ArgumentType;
+        data @1 :Data;
+
+        enum ArgumentType {
+            task @0;
+            objectID @1;
+        }
+    }
+}
+
+struct TaskCancel {
+    struct TaskCancelFlags {
+        force @0 :Bool;
+    }
+
+    taskId @0 :Data;
+    flags @1 :TaskCancelFlags;
+}
+
+struct TaskLog {
+    taskId @0 :Data;
+    logType @1 :LogType;
+    content @2 :Text;
+
+    enum LogType {
+        stdout @0;
+        stderr @1;
+    }
+}
+
+struct TaskResult {
+    taskId @0 :Data;
+    resultType @1 :CommonType.TaskResultType;
+    metadata @2 :Data;
+    results @3 :List(Data);
+}
+
+struct TaskCancelConfirm {
+    taskId @0 :Data;
+    cancelConfirmType @1 :CommonType.TaskCancelConfirmType;
+}
+
+struct GraphTask {
+    taskId @0 :Data;
+    source @1 :Data;
+    targets @2 :List(Data);
+    graph @3 :List(Task);
+}
+
+struct ClientHeartbeat {
+    resource @0 :Status.Resource;
+    latencyUS @1 :UInt32;
+}
+
+struct ClientHeartbeatEcho {
+    objectStorageAddress @0 :CommonType.ObjectStorageAddress;
+}
+
+struct WorkerHeartbeat {
+    agent @0 :Status.Resource;
+    rssFree @1 :UInt64;
+    queueSize @2 :UInt32;
+    queuedTasks @3 :UInt32;
+    latencyUS @4 :UInt32;
+    taskLock @5 :Bool;
+    processors @6 :List(Status.ProcessorStatus);
+    capabilities @7 :List(CommonType.TaskCapability);
+}
+
+struct WorkerHeartbeatEcho {
+    objectStorageAddress @0 :CommonType.ObjectStorageAddress;
+}
+
+struct ObjectInstruction {
+    instructionType @0 :ObjectInstructionType;
+    objectUser @1 :Data;
+    objectMetadata @2 :CommonType.ObjectMetadata;
+
+    enum ObjectInstructionType {
+        create @0;
+        delete @1;
+        clear @2;
+    }
+}
+
+struct DisconnectRequest {
+    worker @0 :Data;
+}
+
+struct DisconnectResponse {
+    worker @0 :Data;
+}
+
+struct ClientDisconnect {
+    disconnectType @0 :DisconnectType;
+
+    enum DisconnectType {
+        disconnect @0;
+        shutdown @1;
+    }
+}
+
+struct ClientShutdownResponse {
+    accepted @0 :Bool;
+}
+
+struct StateClient {
+}
+
+struct StateObject {
+}
+
+struct StateBalanceAdvice {
+    workerId @0 :Data;
+    taskIds @1 :List(Data);
+}
+
+struct StateScheduler {
+    binder @0 :Status.BinderStatus;
+    scheduler @1 :Status.Resource;
+    rssFree @2 :UInt64;
+    clientManager @3 :Status.ClientManagerStatus;
+    objectManager @4 :Status.ObjectManagerStatus;
+    taskManager @5 :Status.TaskManagerStatus;
+    workerManager @6 :Status.WorkerManagerStatus;
+    scalingManager @7 :Status.ScalingManagerStatus;
+}
+
+struct StateWorker {
+    workerId @0 :Data;
+    state@1 :CommonType.WorkerState;
+    capabilities @2 :List(CommonType.TaskCapability);
+}
+
+struct StateTask {
+    taskId @0 :Data;
+    functionName @1 :Data;
+    state @2 :CommonType.TaskState;
+    worker @3 :Data;
+    capabilities @4 :List(CommonType.TaskCapability);
+    metadata @5 :Data;
+}
+
+struct StateGraphTask {
+    enum NodeTaskType {
+        normal @0;
+        target @1;
+    }
+
+    graphTaskId @0 :Data;
+    taskId @1 :Data;
+    nodeTaskType @2 :NodeTaskType;
+    parentTaskIds @3 :List(Data);
+}
+
+struct ProcessorInitialized {
+}
+
+struct InformationRequest {
+    request @0 :Data;
+}
+
+struct InformationResponse {
+    response @0 :Data;
+}
+
+struct Message {
+    union {
+        task @0 :Task;
+        taskCancel @1 :TaskCancel;
+        taskCancelConfirm @2 :TaskCancelConfirm;
+        taskResult @3 :TaskResult;
+        taskLog @4 :TaskLog;
+
+        graphTask @5 :GraphTask;
+
+        objectInstruction @6 :ObjectInstruction;
+
+        clientHeartbeat @7 :ClientHeartbeat;
+        clientHeartbeatEcho @8 :ClientHeartbeatEcho;
+
+        workerHeartbeat @9 :WorkerHeartbeat;
+        workerHeartbeatEcho @10 :WorkerHeartbeatEcho;
+
+        disconnectRequest @11 :DisconnectRequest;
+        disconnectResponse @12 :DisconnectResponse;
+
+        stateClient @13 :StateClient;
+        stateObject @14 :StateObject;
+        stateBalanceAdvice @15 :StateBalanceAdvice;
+        stateScheduler @16 :StateScheduler;
+        stateWorker @17 :StateWorker;
+        stateTask @18 :StateTask;
+        stateGraphTask @19 :StateGraphTask;
+
+        clientDisconnect @20 :ClientDisconnect;
+        clientShutdownResponse @21 :ClientShutdownResponse;
+
+        processorInitialized @22 :ProcessorInitialized;
+
+        informationRequest @23 :InformationRequest;
+        informationResponse @24 :InformationResponse;
+    }
+}

scaler/protocol/capnp/object_storage.capnp

@@ -0,0 +1,57 @@
+@0xc2a14174aa42a12a;
+
+using Cxx = import "/capnp/c++.capnp";
+$Cxx.namespace("scaler::protocol");
+
+struct ObjectRequestHeader {
+    objectID @0: ObjectID; # 32 bytes
+    payloadLength @1: UInt64; # 8 bytes
+    requestID @2: UInt64; # 8 bytes
+    requestType @3: ObjectRequestType; # 2 bytes
+
+    enum ObjectRequestType {
+        # Set or override an object to the message's payload.
+        # Overrides the object's content if it already exists
+        # Always immediately answers with a setOK message.
+        setObject @0;
+
+        # Get an object's content.
+        # If the object does not exist, delays the getOk response until the object is created.
+        getObject @1;
+
+        # Remove the object.
+        deleteObject @2;
+
+        # Creates the provided object ID by linking it to the content of the object ID provided in payload.
+        # Overrides the object content if the new object ID already exists.
+        # If the referenced object does not exist, delays the duplicateOK response until the original object is created.
+        duplicateObjectID @3;
+
+        # Request the server to give back internal information, result is returned as payload.
+        # schema: three uint64_t tuple (number of ids, number of objects (hashes), total actual object size in bytes)
+        infoGetTotal @4;
+    }
+}
+
+struct ObjectID {
+    field0 @0: UInt64;
+    field1 @1: UInt64;
+    field2 @2: UInt64;
+    field3 @3: UInt64;
+}
+
+struct ObjectResponseHeader {
+    objectID @0: ObjectID;
+    payloadLength @1: UInt64;
+    responseID @2: UInt64; # 8 bytes
+    responseType @3: ObjectResponseType;
+
+    enum ObjectResponseType {
+        setOK @0;
+        getOK @1;
+        delOK @2;
+        delNotExists @3;
+        duplicateOK @4;
+        infoGetTotalOK @5;
+    }
+}

scaler/protocol/capnp/status.capnp

@@ -0,0 +1,73 @@
+@0xa4dfa1212ad2d0f0;
+
+struct Resource {
+    cpu @0 :UInt16;   # 99.2% will be represented as 992 as integer
+    rss @1 :UInt64;  # 32bit is capped to 4GB, so use 64bit to represent
+}
+
+struct ObjectManagerStatus {
+    numberOfObjects @0 :UInt32;
+}
+
+struct ClientManagerStatus {
+    clientToNumOfTask @0 :List(Pair);
+
+    struct Pair {
+        client @0 :Data;
+        numTask @1 :UInt32;
+    }
+}
+
+struct TaskManagerStatus {
+    stateToCount @0 :List(Pair);
+
+    struct Pair {
+        state @0 :UInt8;
+        count @1 :UInt32;
+    }
+}
+
+struct ProcessorStatus {
+    pid @0 :UInt32;
+    initialized @1 :Bool;
+    hasTask @2 :Bool;
+    suspended @3 :Bool;
+    resource @4 :Resource;
+}
+
+struct WorkerStatus {
+    workerId @0 :Data;
+    agent @1 :Resource;
+    rssFree @2 :UInt64;
+    free @3 :UInt32;
+    sent @4 :UInt32;
+    queued @5 :UInt32;
+    suspended @6: UInt8;
+    lagUS @7 :UInt64;
+    lastS @8 :UInt8;
+    itl @9 :Text;
+    processorStatuses @10 :List(ProcessorStatus);
+}
+
+struct WorkerManagerStatus {
+    workers @0 :List(WorkerStatus);
+}
+
+struct ScalingManagerStatus {
+    workerGroups @0 :List(Pair);
+
+    struct Pair {
+        workerGroupID @0 :Data;
+        workerIDs @1 :List(Data);
+    }
+}
+
+struct BinderStatus {
+    received @0 :List(Pair);
+    sent @1 :List(Pair);
+
+    struct Pair {
+        client @0 :Text;
+        number @1 :UInt32;
+    }
+}