[Pipeline Refactor] Split/Join Functionality for multiple prompts

src/deepsparse/v2/operators/engine_operator.py

@@ -20,7 +20,7 @@
 from deepsparse import Context as EngineContext
 from deepsparse import Engine, MultiModelEngine, Scheduler
 from deepsparse.benchmark import ORTEngine
-from deepsparse.utils import join_engine_outputs, model_to_path, split_engine_inputs
+from deepsparse.utils import model_to_path
 from deepsparse.v2.operators import Operator
 
 
@@ -145,18 +145,6 @@ def run(self, inp: EngineOperatorInputs, **kwargs) -> Dict:
             # planned refactor
             engine_outputs = inp.engine(inp.engine_inputs)
             return {"engine_outputs": engine_outputs}
-        inp = inp.engine_inputs
-        batches, orig_batch_size = self.expand_inputs(engine_inputs=inp)
-        batches_outputs = list(map(self.engine, batches))
-        engine_outputs = self.condense_inputs(
-            batch_outputs=batches_outputs, orig_batch_size=orig_batch_size
-        )
-        return {"engine_outputs": engine_outputs}
 
-    def expand_inputs(self, **kwargs):
-        return split_engine_inputs(kwargs["engine_inputs"], self._batch_size)
-
-    def condense_inputs(self, **kwargs):
-        batch_outputs = kwargs["batch_outputs"]
-        orig_batch_size = kwargs["orig_batch_size"]
-        return join_engine_outputs(batch_outputs, orig_batch_size)
+        engine_outputs = self.engine(inp.engine_inputs)
+        return {"engine_outputs": engine_outputs}

src/deepsparse/v2/operators/operator.py

@@ -99,7 +99,6 @@ def __call__(
                 pipeline_state=pipeline_state,
                 **kwargs,
             )
-
         if self.has_output_schema():
             return self.output_schema(**run_output)
         return run_output
@@ -117,18 +116,6 @@ def can_operate(self, inp: Any) -> bool:
         """
         return True
 
-    def expand_inputs(self, **kwargs):
-        """
-        Generic function to handle expanding values.
-        """
-        raise NotImplementedError
-
-    def condense_inputs(self, **kwargs):
-        """
-        Generic function to handle condensing values.
-        """
-        raise NotImplementedError
-
     def yaml(self):
         pass
 

src/deepsparse/v2/pipeline.py

@@ -13,7 +13,9 @@
 # limitations under the License.
 
 
-from typing import Dict, List, Union
+import copy
+from functools import partial
+from typing import Any, Dict, List, Union
 
 from deepsparse.v2.operators import Operator
 from deepsparse.v2.routers import Router
@@ -59,6 +61,55 @@ def __init__(
         # SchedulerGroup handles running all schedulers in order of priority
         self._scheduler_group = SchedulerGroup(self.schedulers)
 
+    def _run_sequential(
+        self,
+        inp: Any,
+        inference_state: InferenceState,
+        pipeline_state: PipelineState,
+        start: str,
+        end: str,
+    ):
+        # TODO: somehow refactor to prevent repeat code.
+        next_step = start
+        while next_step != end:
+            operator = self.ops[next_step]
+            if isinstance(inp, dict):
+                operator_output = operator(
+                    pipeline_state=pipeline_state,
+                    inference_state=inference_state,
+                    **inp,
+                )
+            else:
+                operator_output = operator(
+                    inp, pipeline_state=pipeline_state, inference_state=inference_state
+                )
+            if isinstance(operator_output, tuple):
+                state_update = operator_output[-1]
+                operator_output = operator_output[0]
+                inference_state.update_state(state_update)
+
+            next_step = self.router.next(next_step, self.ops, operator_output)
+            inp = operator_output
+        return inp
+
+    def _apply_split(self, inp: Any, inference_state: InferenceState):
+
+        batches, orig_batch_size = self.expand_inputs(inp, 1)
+        run_with_state = partial(
+            self._run_sequential,
+            pipeline_state=self.pipeline_state,
+            start=self.router.route[self.router.SPLIT_ROUTE],
+            end=self.router.END_SPLIT,
+        )
+        inference_state_list = [
+            copy.deepcopy(inference_state) for x in range(len(batches))
+        ]
+        outputs = self._scheduler_group.map(
+            batches, inference_state_list, func=run_with_state
+        )
+        outputs = self.condense_inputs(outputs)
+        return outputs
+
     def run(
         self,
         *args,
@@ -78,7 +129,11 @@ def run(
         operator_output = None
 
         while next_step != self.router.END_ROUTE:
-            # Either a dictionary key or valid index
+            # Split_Route should be after Start_Route
+            if next_step == self.router.SPLIT_ROUTE:
+                operator_output = self._apply_split(operator_output, inference_state)
+                next_step = self.router.route[self.router.END_SPLIT]
+
             operator = self.ops[next_step]
             if next_step == self.router.START_ROUTE:
                 output_future = self._scheduler_group.submit(
@@ -136,6 +191,18 @@ def __call__(self, *args, **kwargs):
 
         return self.run(*args, **kwargs)
 
+    def expand_inputs(self, *args, **kwargs):
+        """
+        Generic function to handle expanding values.
+        """
+        raise NotImplementedError
+
+    def condense_inputs(self, *args, **kwargs):
+        """
+        Generic function to handle condensing values.
+        """
+        raise NotImplementedError
+
     def validate(self):
         """
         Validate that compatability of the router and operators provided.

src/deepsparse/v2/routers/router.py

@@ -128,8 +128,10 @@ class GraphRouter(Router):
     where `can_operate` returns True will run. Paths should be deterministic.
     """
 
-    def __init__(self, end_route: str, start_route: str, route: Dict):
+    def __init__(self, end_route: str, start_route: str, route: Dict, **kwargs):
         super().__init__(end_route=end_route, start_route=start_route, route=route)
+        self.SPLIT_ROUTE = kwargs.get("split_route")
+        self.END_SPLIT = kwargs.get("end_split")
 
     def next(
         self,

src/deepsparse/v2/schedulers/scheduler.py

@@ -14,6 +14,7 @@
 
 
 from concurrent.futures import Future, ThreadPoolExecutor
+from typing import Callable
 
 from deepsparse.v2.operators import Operator
 
@@ -52,6 +53,22 @@ def submit(
             **kwargs,
         )
 
+    def can_map(self, *args):
+        """
+        args containing list of inputs to be used for each worker. This function if we
+        have sufficient workes available
+        """
+        if len(args[0]) <= self._threadpool._max_workers:
+            return True
+        return False
+
+    def map(self, *args, func: Callable):
+        """
+        :param func: Callable to run as part of the map function
+        args containing a list of function variables to map
+        """
+        return list(self._threadpool.map(func, *args))
+
     def can_process(
         self,
         *args,

src/deepsparse/v2/schedulers/scheduler_group.py

@@ -14,7 +14,7 @@
 
 
 from concurrent.futures import Future
-from typing import List
+from typing import Callable, List
 
 from deepsparse.v2.operators import Operator
 from deepsparse.v2.schedulers.scheduler import OperatorScheduler
@@ -56,22 +56,13 @@ def submit(
                     **kwargs,
                 )
 
-    def can_process(
-        self,
-        *args,
-        operator: Operator,
-        **kwargs,
-    ) -> bool:
+    def map(self, *args, func: Callable):
         """
-        :param operator: operator to check
-        :return: True if this Operator can process the given operator and input.
-            SchedulerGroup always returns True
+        :param operator: operator to run
+        :return: list of outputs from multiple workers
         """
-        return any(
-            scheduler.can_process(
-                *args,
-                operator=operator,
-                **kwargs,
-            )
-            for scheduler in self.schedulers
-        )
+        for scheduler in self.schedulers:
+            if scheduler.can_map(
+                args[0],
+            ):
+                return scheduler.map(*args, func=func)

src/deepsparse/v2/text_generation/__init__.py

@@ -17,6 +17,7 @@
 from .compile_generations import *
 from .compile_logits import *
 from .generate_new_token import *
+from .join_output import *
 from .kv_cache_operator import *
 from .multi_engine_prefill_operator import *
 from .nl_engine_operator import *

src/deepsparse/v2/text_generation/join_output.py

@@ -0,0 +1,70 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List
+
+import numpy
+
+from deepsparse.transformers.utils.helpers import pad_to_fixed_length
+from deepsparse.v2.operators import Operator
+from deepsparse.v2.text_generation.compile_generations import CompileGenerationsOutput
+
+
+__all__ = ["JoinOutput"]
+
+
+class JoinOutput(Operator):
+    """
+    Run this operator to combine the results from multiple prompts.
+    """
+
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+
+    def run(self, inp: List[CompileGenerationsOutput], **kwargs):
+        batch_outputs = [x for x in inp[0]]
+        generated_tokens = [x.generated_tokens for x in batch_outputs]
+        generated_logits = [x.generated_logits for x in batch_outputs]
+        finished_reason = [x.finished_reason for x in batch_outputs]
+
+        max_len = max(token.shape[1] for token in generated_tokens)
+
+        # pad all tokens to the same length
+        tokens = [
+            pad_to_fixed_length(
+                array=prediction,
+                max_len=max_len,
+                value=self.tokenizer.pad_token_id,
+                axis=1,
+            )
+            for prediction in generated_tokens
+        ]
+
+        # find the longest sequence in the batch of logits
+        max_len = max(logits.shape[1] for logits in generated_logits)
+
+        # pad all logits to the same length
+        logits = [
+            pad_to_fixed_length(array=single_logits, max_len=max_len, axis=1)
+            for single_logits in generated_logits
+        ]
+
+        tokens = numpy.concatenate(tokens)
+        logits = numpy.concatenate(logits)
+
+        return {
+            "generated_tokens": tokens,
+            "generated_logits": logits,
+            "finished_reason": finished_reason,
+        }

src/deepsparse/v2/text_generation/pipeline.py

@@ -15,6 +15,7 @@
 from typing import Dict
 
 from deepsparse.transformers.utils.helpers import process_generation_config
+from deepsparse.utils import split_engine_inputs
 from deepsparse.v2.pipeline import Pipeline
 from deepsparse.v2.routers import GraphRouter
 from deepsparse.v2.schedulers import OperatorScheduler
@@ -24,6 +25,7 @@
     CompileGenerations,
     CompilePromptLogits,
     GenerateNewTokenOperator,
+    JoinOutput,
     KVCacheCreator,
     MultiEnginePrefill,
     NLEngineOperator,
@@ -131,6 +133,7 @@ def __init__(
         process_output = ProcessOutputs(tokenizer=self.tokenizer)
         compile_generations = CompileGenerations()
         compile_generated_tokens = CompileGeneratedTokens()
+        join_output = JoinOutput(tokenizer=self.tokenizer)
 
         ops = {
             "process_input": process_inputs,
@@ -146,10 +149,12 @@ def __init__(
             "process_outputs": process_output,
             "compile_generations": compile_generations,
             "compile_generated_tokens": compile_generated_tokens,
+            "join_output": join_output,
         }
 
         routes = {
-            "process_input": "prepare_prefill",
+            "process_input": "SPLIT",
+            "SPLIT": "prepare_prefill",
             "prepare_prefill": ["multi_engine_prefill", "autoregressive_preprocess"],
             "multi_engine_prefill": "multi_engine",
             "multi_engine": "compile_logits",
@@ -169,18 +174,33 @@ def __init__(
                 "autoregressive_preprocess",
                 "compile_generations",
             ],
-            "compile_generations": "process_outputs",
+            "compile_generations": "JOIN",
+            "JOIN": "join_output",
+            "join_output": "process_outputs",
             "process_outputs": "STOP",
         }
 
         router = GraphRouter(
-            end_route="STOP", start_route="process_input", route=routes
+            end_route="STOP",
+            start_route="process_input",
+            route=routes,
+            split_route="SPLIT",
+            end_split="JOIN",
         )
-        scheduler = [OperatorScheduler()]
+        scheduler = [OperatorScheduler(), OperatorScheduler(max_workers=4)]
         super().__init__(
             ops=ops, router=router, schedulers=scheduler, pipeline_state=pipeline_state
         )
 
+    def expand_inputs(self, items, batch_size):
+        items = [items.get(key) for key in items.keys()]
+        out, orig_batch_size = split_engine_inputs(items, batch_size)
+        combined_batches = [{"input_ids": b[0], "attention_mask": b[1]} for b in out]
+        return combined_batches, orig_batch_size
+
+    def condense_inputs(self, *args, **kwargs):
+        return args[0], kwargs
+
     # TODO: Move to be part of a generic transformers set-up Operator.
     def setup_onnx_file_path(self, model_path, sequence_length) -> str:
         import logging