request_handler.py 2.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152
  1. import sys
  2. import modelverse_kernel.primitives as primitive_functions
  3. import modelverse_kernel.jit as jit
  4. from collections import defaultdict
  5. class RequestHandler(object):
  6. def __init__(self):
  7. self.generator_stack = []
  8. self.handlers = {"CALL": lambda i: i[0],
  9. "CALL_ARGS": lambda i: i[0](*(i[1])),
  10. "CALL_KWARGS": lambda i: i[0](**(i[1])),
  11. "SLEEP": self.execute_sleep}
  12. def push_generator(self, gen):
  13. self.generator_stack.append(gen)
  14. def handle_request(self, reply):
  15. while self.generator_stack:
  16. try:
  17. gen = self.generator_stack[-1]
  18. requests = gen.send(reply)
  19. # Generated new request, so process
  20. if requests and requests[0][0] in self.handlers:
  21. # Try to add a new generator to branch into
  22. reply = None
  23. self.generator_stack.append(None)
  24. # This next command potentially raises a finished message already, meaning that we should stop already
  25. # We avoid an extra try/except block by putting the None on the stack already
  26. self.generator_stack[-1] = self.handlers[requests[0][0]](requests[0][1])
  27. else:
  28. # MvS request, so forward that instead
  29. return requests
  30. except StopIteration:
  31. # Exception, so finished execution of this generator, passing on None to the caller
  32. del self.generator_stack[-1]
  33. reply = [None]
  34. except primitive_functions.PrimitiveFinished as ex:
  35. # Exception, so finished execution of this generator, passing on ex.result to the caller
  36. del self.generator_stack[-1]
  37. reply = [ex.result]
  38. except primitive_functions.SleepKernel:
  39. # Processing sleep, so pop its generator and reraise
  40. del self.generator_stack[-1]
  41. raise
  42. def execute_sleep(self, request_args):
  43. raise primitive_functions.SleepKernel(request_args[0], request_args[1])