annotations/Azaion.Inference/main_inference.pyx

import queue
import traceback
from queue import Queue
cimport constants_inf
from threading import Thread

from inference cimport Inference
from loader_client cimport LoaderClient
from remote_command_inf cimport RemoteCommand, CommandType
from remote_command_handler_inf cimport RemoteCommandHandler

cdef class CommandProcessor:

    cdef RemoteCommandHandler remote_handler
    cdef object inference_queue
    cdef bint running
    cdef Inference inference
    cdef LoaderClient loader_client

    def __init__(self, int zmq_port, str loader_zmq_host, int loader_zmq_port, str api_url):
        self.remote_handler = RemoteCommandHandler(zmq_port, self.on_command)
        self.inference_queue = Queue(maxsize=constants_inf.QUEUE_MAXSIZE)
        self.loader_client = LoaderClient(loader_zmq_host, loader_zmq_port)
        self.inference = Inference(self.loader_client, self.remote_handler)
        self.running = True
        self.remote_handler.start()


    def start(self):
        while self.running:
            try:
                command = self.inference_queue.get(timeout=0.5)
                self.inference.run_inference(command)
                self.remote_handler.send(command.client_id, <RemoteCommand>RemoteCommand(CommandType.INFERENCE_DONE))
            except queue.Empty:
                continue
            except Exception as e:
                traceback.print_exc()
        constants_inf.log(<str>'EXIT!')

    cdef on_command(self, RemoteCommand command):
        try:
            if command.command_type == CommandType.INFERENCE:
                self.inference_queue.put(command)
            elif command.command_type == CommandType.AI_AVAILABILITY_CHECK:
                status = self.inference.ai_availability_status.serialize()
                self.remote_handler.send(command.client_id, <RemoteCommand>RemoteCommand(CommandType.AI_AVAILABILITY_RESULT, status))
            elif command.command_type == CommandType.STOP_INFERENCE:
                self.inference.stop()
            elif command.command_type == CommandType.EXIT:
                t = Thread(target=self.stop) # non-block worker:
                t.start()
            else:
                pass
        except Exception as e:
            constants_inf.logerror(<str>f"Error handling client: {str(e)}")

    def stop(self):
        self.inference.stop()
        self.remote_handler.stop()
        self.loader_client.stop()
        self.running = False