modelverse/kernel/modelverse_kernel/request_handler.py

import sys
import modelverse_kernel.primitives as primitive_functions
import modelverse_kernel.jit as jit
from collections import defaultdict

class RequestHandler(object):
    def __init__(self):
        self.generator_stack = []
        self.handlers = {"CALL": self.execute_call,
                         "CALL_ARGS": self.execute_call_args,
                         "CALL_KWARGS": self.execute_call_kwargs,
                         "RETURN": self.execute_return,
                         "FINISH": self.execute_finish,
                         "SLEEP": self.execute_sleep}

    def push_generator(self, gen):
        self.generator_stack.append(gen)

    def handle_request(self, reply):
        self.reply = reply
        while self.generator_stack:
            requests = self.generator_stack[-1].send(self.reply)

            # Generated new request, so process
            if requests and requests[0][0] in self.handlers:
                # This next command potentially raises a finished message already, meaning that we should stop already
                # We avoid an extra try/except block by putting the None on the stack already
                self.handlers[requests[0][0]](requests[0][1])
            else:
                # MvS request, so forward that instead
                return requests
                
    def execute_finish(self, request_args):
        self.generator_stack = []
        
    def execute_sleep(self, request_args):
        raise primitive_functions.SleepKernel(request_args[0], request_args[1])

    def execute_return(self, request_args):
        del self.generator_stack[-1]
        self.reply = [request_args[0]]

    def execute_call(self, request_args):
        self.generator_stack.append(request_args[0])
        self.reply = None

    def execute_call_args(self, request_args):
        print("CALL " + str(request_args[0]))
        self.generator_stack.append(request_args[0](*(request_args[1])))
        self.reply = None

    def execute_call_kwargs(self, request_args):
        print("KWCALL " + str(request_args[0]))
        self.generator_stack.append(request_args[0](**(request_args[1])))
        self.reply = None