Add todo comments

morpheus/stages/input/rest_source_stage.py

@@ -45,6 +45,7 @@ class RestSourceStage(PreallocatorMixin, SingleOutputSource):
         Pipeline configuration instance.
     """
 
+    # TODO add configs, and pass-thru for both gunicorn and flask configs
     def __init__(self, c: Config, sleep_time: float = 0.1):
         super().__init__(c)
         self._sleep_time = sleep_time
@@ -56,14 +57,14 @@ def name(self) -> str:
         return "from-rest"
 
     def supports_cpp_node(self) -> bool:
-        return False
+        return False  # TODO Add C++ impl
 
     def _generate_frames(self) -> typing.Iterator[MessageMeta]:
         self._is_running.value = 1
         self._server_proc.start()
 
         processing = True
-        while (processing and self._server_proc.is_alive()):
+        while (processing):
             # Read as many messages as we can from the queue if it's empty check to see if we should be shutting down
             # It is important that any messages we received that are in the queue are processed before we shutdown since
             # we already returned an OK response to the client.
@@ -72,7 +73,7 @@ def _generate_frames(self) -> typing.Iterator[MessageMeta]:
             try:
                 data = self._queue.get_nowait()
             except queue.Empty:
-                if self._is_running.value == 0:
+                if (self._is_running.value == 0 or not self._server_proc.is_alive()):
                     processing = False
                 else:
                     time.sleep(self._sleep_time)
@@ -90,13 +91,13 @@ def _generate_frames(self) -> typing.Iterator[MessageMeta]:
             if df is not None:
                 yield MessageMeta(df)
 
+        self._queue.close()
+
+    def _stop(self):
         logger.debug("Stopping REST server ...")
         self._server_proc.terminate()
         self._server_proc.join()
         logger.debug("REST server stopped")
-        self._queue.close()
-
-    def _stop(self):
         self._is_running.value = 0
 
     def _build_source(self, builder: mrc.Builder) -> StreamPair:
@@ -105,6 +106,7 @@ def _build_source(self, builder: mrc.Builder) -> StreamPair:
 
     def _post_build_single(self, builder: mrc.Builder, out_pair: StreamPair) -> StreamPair:
         src_node = out_pair[0]
+        # TODO: This doesn't work its called when the source exits, not when the pipeline is shutting down
         post_node = builder.make_node(self.unique_name + "-post", ops.on_completed(self._stop))
         builder.make_edge(src_node, post_node)
 

morpheus/utils/rest_server.py

@@ -16,7 +16,6 @@
 import multiprocessing as mp
 import queue
 import typing
-from functools import partial
 
 import gunicorn.app.base
 from flask import Flask
@@ -29,7 +28,7 @@
 DEFAULT_OPTIONS = {
     'bind': '127.0.0.1:8080',
     'keepalive': 0,
-    'loglevel': 'DEBUG',  # TODO adjust this
+    'loglevel': 'WARNING',
     'preload_app': True,
     'proc_name': 'morpheus_rest_server',
     'reuse_port': True,
@@ -54,6 +53,8 @@ def dispatch_request(self) -> ResponseReturnValue:
         Note this does not perform any validation on the request pyaload, it is possible that incoming data
         which would later fail to be processed by Morpheus is accepted by this endpoint.
         """
+
+        # This should work the same for both POST, PUT & PATCH
         if request.is_json and request.content_length is not None and request.content_length > 0:
             try:
                 self._logger.debug("Received request with content length %s", request.content_length)
@@ -87,6 +88,8 @@ def load(self):
 
 def _start_rest_server(options: dict, queue: mp.Queue):
     app = Flask('morpheus_rest_server')
+
+    # TODO make this configurable
     app.logger.setLevel(logging.DEBUG)
     app.add_url_rule('/submit',
                      methods=["POST", "PUT"],
@@ -95,10 +98,16 @@ def _start_rest_server(options: dict, queue: mp.Queue):
     server.run()
 
 
+# TODO: Currently `options` configures the gunicorn server, we need a second config to configure Flask specifically the
+# endpoint(s) and supported methods.
 def start_rest_server(options: dict = None) -> typing.Tuple[mp.Process, mp.Queue]:
     """
     Starts a REST server.
     """
+
+    # gunicorn is opinionated and doesn't really support being embedded in another app, so we need to start it in a
+    # separate process.
+    # TODO: Although gunicorn is a good choice we should look into cherrypy and other options.
     server_options = DEFAULT_OPTIONS.copy()
     server_options.update(options or {})
     queue = mp.Queue()