annotations/Azaion.Inference/loader_client.pyx

import zmq
from remote_command cimport RemoteCommand, CommandType
from file_data cimport FileData, UploadFileData

cdef class LoaderClient:
    def __init__(self, str zmq_host, int zmq_port):
        self._context = zmq.Context().instance()
        self._socket = self._context.socket(zmq.DEALER)
        self._socket.connect(f'tcp://{zmq_host}:{zmq_port}')

    cdef RemoteCommand _send_receive_command(self, RemoteCommand command):
        self._socket.send(command.serialize())
        return RemoteCommand.from_msgpack(self._socket.recv())

    cdef load_big_small_resource(self, str filename, str directory):
        cdef FileData file_data = FileData(folder=directory, filename=filename)
        cdef RemoteCommand response = self._send_receive_command(RemoteCommand(CommandType.LOAD_BIG_SMALL, data=file_data.serialize()))
        if response.command_type == CommandType.DATA_BYTES:
            return response.data
        elif response.command_type == CommandType.ERROR:
            raise Exception(f"Error from server: {response.message}")
        else:
            raise Exception(f"Unexpected response command type: {response.command_type}")

    cdef upload_big_small_resource(self, bytes content, str filename, str directory):
        cdef UploadFileData upload_file_data = UploadFileData(content, filename, directory)
        cdef RemoteCommand upload_resp = self._send_receive_command(RemoteCommand(CommandType.UPLOAD_BIG_SMALL, data=upload_file_data.serialize()))
        if upload_resp.command_type == CommandType.OK:
            return
        elif upload_resp.command_type == CommandType.ERROR:
            raise Exception(f"Error from server: {upload_resp.message}")
        else:
            raise Exception(f"Unexpected response command type: {upload_resp.command_type}")

    cdef close(self):
        if self._socket and not self._socket.closed:
            self._socket.close()
        if self._context and not self._context.closed:
            self._context.term()