Improve pipeline stop logic to ensure join is called exactly once for all stages

morpheus/pipeline/pipeline.py

@@ -75,10 +75,10 @@ def __init__(self, config: Config):
         self._num_threads = config.num_threads
 
         # Complete set of nodes across segments in this pipeline
-        self._stages: typing.Set[Stage] = set()
+        self._stages: typing.List[Stage] = []
 
         # Complete set of sources across segments in this pipeline
-        self._sources: typing.Set[SourceStage] = set()
+        self._sources: typing.List[SourceStage] = []
 
         # Dictionary containing segment information for this pipeline
         self._segments: typing.Dict = defaultdict(lambda: {"nodes": set(), "ingress_ports": [], "egress_ports": []})
@@ -94,6 +94,7 @@ def __init__(self, config: Config):
 
         self._loop: asyncio.AbstractEventLoop = None
 
+        # Future that allows post_start to propagate exceptions back to pipeline
         self._post_start_future: asyncio.Future = None
 
     @property
@@ -124,10 +125,10 @@ def add_stage(self, stage: StageT, segment_id: str = "main") -> StageT:
         # Add to list of stages if it's a stage, not a source
         if (isinstance(stage, Stage)):
             segment_nodes.add(stage)
-            self._stages.add(stage)
+            self._stages.append(stage)
         elif (isinstance(stage, SourceStage)):
             segment_nodes.add(stage)
-            self._sources.add(stage)
+            self._sources.append(stage)
         else:
             raise NotImplementedError(f"add_stage() failed. Unknown node type: {type(stage)}")
 
@@ -425,7 +426,6 @@ async def post_start(executor):
                 with self._mutex:
                     self._state = PipelineState.COMPLETED
 
-        # asyncio.create_task(post_start(self._mrc_executor))
         self._post_start_future = asyncio.create_task(post_start(self._mrc_executor))
 
     def stop(self):
@@ -448,9 +448,11 @@ def stop(self):
 
     async def join(self):
         """
-        Suspend execution all currently running stages and the MRC pipeline.
-        Typically called after `stop`.
+        Wait until pipeline completes upon which join methods of sources and stages will be called.
         """
+        assert self._post_start_future is not None, "Pipeline must be started before joining"
+        assert self._state != PipelineState.COMPLETED, "Cannot join a pipeline that has already completed"
+
         await self._post_start_future
 
     def _on_stop(self):

tests/pipeline/test_pipeline_state.py

@@ -130,6 +130,29 @@ async def test_stop_after_stop(config: Config):
     await pipeline.join()
 
 
+async def test_join_without_start(config: Config):
+
+    pipeline = LinearPipeline(config)
+    assert pipeline.state == PipelineState.INITIALIZED
+    pipeline.set_source(source_test_stage(config))
+    with pytest.raises(Exception) as excinfo:
+        await pipeline.join()
+    assert "must be started" in str(excinfo.value)
+
+
+async def test_join_after_join(config: Config):
+    pipeline = LinearPipeline(config)
+    assert pipeline.state == PipelineState.INITIALIZED
+    pipeline.set_source(source_test_stage(config))
+    await pipeline.build_and_start()
+    assert pipeline.state == PipelineState.STARTED
+    await pipeline.join()
+    assert pipeline.state == PipelineState.COMPLETED
+    with pytest.raises(Exception) as excinfo:
+        await pipeline.join()
+    assert "has already completed" in str(excinfo.value)
+
+
 @mock.patch('morpheus.stages.preprocess.deserialize_stage.DeserializeStage.join')
 @mock.patch('morpheus.stages.preprocess.deserialize_stage.DeserializeStage.stop')
 @mock.patch('morpheus.stages.input.in_memory_source_stage.InMemorySourceStage.join')
@@ -242,3 +265,50 @@ async def test_stage_methods_called_stop_after_stop(mock_source_stage_stop,
     mock_source_stage_join.assert_called_once()
     mock_deserialize_stage_stop.assert_called_once()
     mock_deserialize_stage_join.assert_called_once()
+
+
+@mock.patch('morpheus.stages.preprocess.deserialize_stage.DeserializeStage.join')
+@mock.patch('morpheus.stages.preprocess.deserialize_stage.DeserializeStage.stop')
+@mock.patch('morpheus.stages.input.in_memory_source_stage.InMemorySourceStage.join')
+@mock.patch('morpheus.stages.input.in_memory_source_stage.InMemorySourceStage.stop')
+async def test_stage_methods_called_join_without_start(mock_source_stage_stop,
+                                                       mock_source_stage_join,
+                                                       mock_deserialize_stage_stop,
+                                                       mock_deserialize_stage_join,
+                                                       config: Config,
+                                                       filter_probs_df: DataFrameType):
+    pipeline = LinearPipeline(config)
+    pipeline.set_source(InMemorySourceStage(config, [filter_probs_df]))
+    pipeline.add_stage(DeserializeStage(config))
+    with pytest.raises(Exception) as excinfo:
+        await pipeline.join()
+    assert "must be started" in str(excinfo.value)
+    mock_source_stage_stop.assert_not_called()
+    mock_source_stage_join.assert_not_called()
+    mock_deserialize_stage_stop.assert_not_called()
+    mock_deserialize_stage_join.assert_not_called()
+
+
+@mock.patch('morpheus.stages.preprocess.deserialize_stage.DeserializeStage.join')
+@mock.patch('morpheus.stages.preprocess.deserialize_stage.DeserializeStage.stop')
+@mock.patch('morpheus.stages.input.in_memory_source_stage.InMemorySourceStage.join')
+@mock.patch('morpheus.stages.input.in_memory_source_stage.InMemorySourceStage.stop')
+async def test_stage_methods_called_join_after_join(mock_source_stage_stop,
+                                                    mock_source_stage_join,
+                                                    mock_deserialize_stage_stop,
+                                                    mock_deserialize_stage_join,
+                                                    config: Config,
+                                                    filter_probs_df: DataFrameType):
+    pipeline = LinearPipeline(config)
+    pipeline.set_source(InMemorySourceStage(config, [filter_probs_df]))
+    pipeline.add_stage(DeserializeStage(config))
+    await pipeline.build_and_start()
+    await pipeline.join()
+    assert pipeline.state == PipelineState.COMPLETED
+    with pytest.raises(Exception) as excinfo:
+        await pipeline.join()
+    assert "has already completed" in str(excinfo.value)
+    mock_source_stage_stop.assert_not_called()
+    mock_source_stage_join.assert_called_once()
+    mock_deserialize_stage_stop.assert_not_called()
+    mock_deserialize_stage_join.assert_called_once()

tests/test_monitor_stage.py

@@ -83,7 +83,7 @@ async def test_join(mock_morph_tqdm: mock.MagicMock, config: Config):
     await stage.join()
     mock_morph_tqdm.assert_not_called()
 
-    stage.on_start()
+    await stage.start_async()
     await stage.join()
     mock_morph_tqdm.close.assert_called_once()