Changes related to #1234

nv-morpheus · Oct 4, 2023 · bd09ea6 · bd09ea6
1 parent 7655b93
commit bd09ea6
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/morpheus/stages/input/http_client_source_stage.py b/morpheus/stages/input/http_client_source_stage.py
@@ -24,6 +24,7 @@
 import cudf
 
 from morpheus.cli.register_stage import register_stage
+from morpheus.config import PipelineModes
 from morpheus.config import Config
 from morpheus.messages import MessageMeta
 from morpheus.pipeline.preallocator_mixin import PreallocatorMixin
@@ -34,7 +35,7 @@
 logger = logging.getLogger(__name__)
 
 
-@register_stage("from-http-client", ignore_args=["query_params", "headers", "**request_kwargs"])
+@register_stage("from-http-client", modes=[PipelineModes.FIL, PipelineModes.NLP, PipelineModes.OTHER],ignore_args=["query_params", "headers", "**request_kwargs"])
 class HttpClientSourceStage(PreallocatorMixin, SingleOutputSource):
     """
     Source stage that polls a remote HTTP server for incoming data.

diff --git a/morpheus/stages/input/http_server_source_stage.py b/morpheus/stages/input/http_server_source_stage.py
@@ -25,6 +25,7 @@
 import cudf
 
 from morpheus.cli.register_stage import register_stage
+from morpheus.config import PipelineModes
 from morpheus.config import Config
 from morpheus.messages import MessageMeta
 from morpheus.pipeline.preallocator_mixin import PreallocatorMixin
@@ -40,7 +41,7 @@
 SUPPORTED_METHODS = (HTTPMethod.POST, HTTPMethod.PUT)
 
 
-@register_stage("from-http")
+@register_stage("from-http", modes=[PipelineModes.FIL, PipelineModes.NLP, PipelineModes.OTHER])
 class HttpServerSourceStage(PreallocatorMixin, SingleOutputSource):
     """
     Source stage that starts an HTTP server and listens for incoming requests on a specified endpoint.