annotations/Azaion.Inference/loader_client.pyx

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

cdef class LoadResult:
    def __init__(self, str err, bytes data=None):
        self.err = err
        self.data = data

cdef class LoaderClient:
    def __init__(self, str zmq_host, int zmq_port):
        self._loader_context = zmq.Context()
        self._socket = self._loader_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 LoadResult(None, response.data)
        elif response.command_type == CommandType.ERROR:
            return LoadResult(f"Error from server: {response.message}")
        else:
            return LoadResult(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, folder=directory, filename=filename)
        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 LoadResult(None, None)
        elif upload_resp.command_type == CommandType.ERROR:
            return LoadResult(f"Error from server: {upload_resp.message}")
        else:
            return LoadResult(f"Unexpected response command type: {upload_resp.command_type}")

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