request_handler.py 2.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364
  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. "RETURN": self.execute_return,
  12. "SLEEP": self.execute_sleep}
  13. def push_generator(self, gen):
  14. self.generator_stack.append(gen)
  15. def handle_request(self, reply):
  16. self.reply = reply
  17. while self.generator_stack:
  18. try:
  19. print("Stack: " + str(self.generator_stack))
  20. requests = self.generator_stack[-1].send(self.reply)
  21. print("Requests: " + str(requests))
  22. # Generated new request, so process
  23. if requests and requests[0][0] in self.handlers:
  24. # Try to add a new generator to branch into
  25. self.generator_stack.append(None)
  26. # This next command potentially raises a finished message already, meaning that we should stop already
  27. # We avoid an extra try/except block by putting the None on the stack already
  28. self.generator_stack[-1] = self.handlers[requests[0][0]](requests[0][1])
  29. self.reply = None
  30. else:
  31. # MvS request, so forward that instead
  32. return requests
  33. except StopIteration:
  34. # Exception, so finished execution of this generator, passing on None to the caller
  35. print("STOPITER")
  36. del self.generator_stack[-1]
  37. self.reply = [None]
  38. except primitive_functions.SleepKernel:
  39. # Processing sleep, so pop its generator and reraise
  40. del self.generator_stack[-1]
  41. raise
  42. except primitive_functions.PrimitiveFinished as ex:
  43. # Exception, so finished execution of this generator, passing on ex.result to the caller
  44. del self.generator_stack[-1]
  45. self.reply = [ex.result]
  46. print("SHOULD NOT HAPPEN")
  47. def execute_sleep(self, request_args):
  48. raise primitive_functions.SleepKernel(request_args[0], request_args[1])
  49. def execute_return(self, request_args):
  50. print("RETURNING")
  51. print(self.generator_stack)
  52. del self.generator_stack[-1]
  53. self.reply = request_args[0]
  54. print(self.generator_stack)