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:
            try:
                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.reply = None
                    self.handlers[requests[0][0]](requests[0][1])
                else:
                    # MvS request, so forward that instead
                    return requests

            except StopIteration:
                # Exception, so finished execution of this generator, passing on None to the caller
                #print("STOP for " + str(self.generator_stack))
                del self.generator_stack[-1]
                self.reply = [None]
                
    def execute_finish(self, request_args):
        del self.generator_stack[-1]
        self.reply = [None]
        
    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])

    def execute_call_args(self, request_args):
        self.generator_stack.append(request_args[0](*(request_args[1])))

    def execute_call_kwargs(self, request_args):
        self.generator_stack.append(request_args[0](**(request_args[1])))