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