remove context, clean-up args, remove prefill_preprocess_operaator

Author: dsikka

src/deepsparse/v2/operators/engine_operator.py

@@ -51,7 +51,7 @@ def __init__(
         num_streams: int = None,
         scheduler: Scheduler = None,
         input_shapes: List[List[int]] = None,
-        engine_context: Optional[Context] = None,
+        engine_context: Optional[EngineContext] = None,
         engine_kwargs: Dict = None,
     ):
         self.model_path = model_to_path(model_path)
@@ -76,7 +76,7 @@ def __init__(
             engine_args["scheduler"] = scheduler
             engine_args["num_streams"] = num_streams
 
-        engine_args.updated(engine_kwargs)
+        engine_args.update(engine_kwargs)
         self.engine = self._create_engine(self.model_path, engine_type, engine_args)
 
     def _create_engine(
@@ -95,12 +95,12 @@ def _create_engine(
 
         if engine_type == DEEPSPARSE_ENGINE:
             if self.engine_context is not None and isinstance(
-                self.engine_context, Context
+                self.engine_context, EngineContext
             ):
                 engine_args.pop("num_cores", None)
                 engine_args.pop("scheduler", None)
                 engine_args.pop("num_streams", None)
-                engine_args["context"] = self.engien_context
+                engine_args["context"] = self.engine_context
                 return MultiModelEngine(
                     model=onnx_file_path,
                     **engine_args,
@@ -116,7 +116,7 @@ def _create_engine(
             f"{SUPPORTED_PIPELINE_ENGINES}"
         )
 
-    def run(self, inp: EngineOperatorInputs) -> Dict:
+    def run(self, inp: EngineOperatorInputs, **kwargs) -> Dict:
         inp = inp.engine_inputs
         batches, orig_batch_size = self.expand_inputs(engine_inputs=inp)
         batches_outputs = list(map(self.engine, batches))

src/deepsparse/v2/operators/operator.py

@@ -17,6 +17,8 @@
 
 from pydantic import BaseModel
 
+from deepsparse.v2.utils import InferenceState, PipelineState
+
 
 __all__ = ["Operator"]
 
@@ -54,6 +56,8 @@ def has_output_schema(cls) -> bool:
     def __call__(
         self,
         *args,
+        inference_state: InferenceState,
+        pipeline_state: PipelineState,
         **kwargs,
     ) -> Any:
         """
@@ -81,10 +85,18 @@ def __call__(
                     "in the form of a dictionary or an instance of the input_schema"
                     "object"
                 )
-
-            run_output = self.run(inference_input)
+            run_output = self.run(
+                inference_input,
+                inference_state=inference_state,
+                pipeline_state=pipeline_state,
+            )
         else:
-            run_output = self.run(*args, **kwargs)
+            run_output = self.run(
+                *args,
+                inference_state=inference_state,
+                pipeline_state=pipeline_state,
+                **kwargs,
+            )
 
         if self.has_output_schema():
             return self.output_schema(**run_output)
@@ -99,13 +111,11 @@ def run(self, *args, **kwargs) -> Any:
         """
         raise NotImplementedError
 
-    def can_operate(
-        self, inp: Any, context: Context, inference_state: InferenceState
-    ) -> bool:
+    def can_operate(self, inp: Any) -> bool:
         """
-        Whether or not the given operator can run, based on input, context, or state
+        Whether or not the given operator can run, based on input
         """
-        raise NotImplementedError
+        return True
 
     def expand_inputs(self, **kwargs):
         """

src/deepsparse/v2/pipeline.py

@@ -18,7 +18,7 @@
 from deepsparse.v2.operators import Operator
 from deepsparse.v2.routers import Router
 from deepsparse.v2.schedulers import OperatorScheduler, SchedulerGroup
-from deepsparse.v2.utils import Context
+from deepsparse.v2.utils import InferenceState, PipelineState
 
 
 __all__ = ["Pipeline"]
@@ -40,71 +40,76 @@ class Pipeline(Operator):
     :param schedulers: A list of schedulers to run operators.
 
     """
-    
+
     def __init__(
         self,
         ops: Union[Dict[str, Operator], List[Operator]],
         router: Router,
         schedulers: List[OperatorScheduler],
+        pipeline_state: PipelineState = None,
     ):
 
         self.ops = ops
         self.router = router
         self.schedulers = schedulers
+        self.pipeline_state = pipeline_state
         self.validate()
 
         # SchedulerGroup handles running all schedulers in order of priority
         self._scheduler_group = SchedulerGroup(self.schedulers)
 
-    def run(self, *args, **kwargs):
+    def run(
+        self,
+        *args,
+        inference_state: InferenceState,
+        pipeline_state: PipelineState,
+        **kwargs,
+    ):
         """
-        Run through the operators using the provided router and scheduler. Update the
-        context to reflect each step of the router. The input to a given operator is the
-        output of the previous operator.
+        Run through the operators using the provided router and scheduler.
+        The input to a given operator is the output of the previous operator.
 
         :param inp: input to the operator. expected to be of any type that is
         expected by the operator.
 
         """
         next_step = self.router.START_ROUTE
         operator_output = None
+
         while next_step != self.router.END_ROUTE:
             # Either a dictionary key or valid index
             operator = self.ops[next_step]
             if next_step == self.router.START_ROUTE:
                 output_future = self._scheduler_group.submit(
-                    *args, operator=operator, **kwargs
+                    *args,
+                    inference_state=inference_state,
+                    operator=operator,
+                    pipeline_state=pipeline_state,
+                    **kwargs,
                 )
             else:
                 if isinstance(operator_output, dict):
                     output_future = self._scheduler_group.submit(
-                        operator=operator, **operator_output
+                        inference_state=inference_state,
+                        operator=operator,
+                        pipeline_state=pipeline_state,
+                        **operator_output,
                     )
                 else:
                     output_future = self._scheduler_group.submit(
-                        operator_output, operator=operator
+                        operator_output,
+                        inference_state=inference_state,
+                        pipeline_state=pipeline_state,
+                        operator=operator,
                     )
-                    
-            # print("Current State", inference_state.current_state)
-
-            """
-            output_future = self._scheduler_group.submit(
-                operator=operator,
-                operator_input=inp,
-                context=context,
-                pipeline_state=self.pipeline_state,
-                inference_state=inference_state,
-            )
-            """
-
-            # wait for future to resolve
-            operator_output, state_update = output_future.result()
-            inference_state.update_state(state_update)
-
-            next_step = self.router.next(
-                next_step, self.ops, context, operator_output, inference_state
-            )
-            inp = operator_output
+
+            operator_output = output_future.result()
+            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)
 
         return operator_output
 
@@ -113,6 +118,18 @@ def __call__(self, *args, **kwargs):
         :return: output of the pipeline operators ran with the router for the given
         input
         """
+        if kwargs.get("inference_state"):
+            inference_state = kwargs.pop("inference_state")
+        else:
+            inference_state = InferenceState()
+            inference_state.create_state({})
+
+        if "pipeline_state" in kwargs:
+            self.pipeline_state = kwargs.get("pipeline_state")
+
+        kwargs["inference_state"] = inference_state
+        kwargs["pipeline_state"] = self.pipeline_state
+
         return self.run(*args, **kwargs)
 
     def validate(self):

src/deepsparse/v2/routers/router.py

@@ -15,15 +15,14 @@
 
 import logging
 from abc import abstractmethod
-from typing import Dict, List, Union
+from typing import Any, Dict, List, Optional, Union
 
 from deepsparse.v2.operators import Operator
-from deepsparse.v2.utils import Context, InferenceState
 
 
 _LOGGER = logging.getLogger(__name__)
 
-__all__ = ["Router", "LinearRouter"]
+__all__ = ["Router", "LinearRouter", "GraphRouter"]
 
 
 class Router:
@@ -36,14 +35,22 @@ class Router:
 
     """
 
-    def __init__(self, end_route: Union[str, int], start_route: Union[str, int]):
+    def __init__(
+        self,
+        end_route: Union[str, int],
+        start_route: Union[str, int],
+        route: Optional[Dict] = None,
+    ):
         self.START_ROUTE = start_route
         self.END_ROUTE = end_route
         self.route = route
 
     @abstractmethod
     def next(
-        self, past: Union[str, int], ops: Union[List[Operator], Dict[str, Operator]]
+        self,
+        past: Union[str, int],
+        ops: Optional[Union[List[Operator], Dict[str, Operator]]],
+        inp: Optional[Any],
     ) -> Union[str, int]:
         """
         Determines the index or dictionary key for the next operator which should run.
@@ -73,7 +80,9 @@ class LinearRouter(Router):
     def __init__(self, end_route: int, start_route: int = 0):
         super().__init__(end_route=end_route, start_route=start_route)
 
-    def next(self, past: int, ops: List[Operator]) -> int:
+    def next(
+        self, past: int, ops: Optional[List[Operator]] = None, inp: Optional[Any] = None
+    ) -> int:
         new_index = past + 1
         if new_index < self.END_ROUTE:
             return new_index
@@ -111,7 +120,7 @@ def validate(operators: List[Operator]) -> bool:
         return True
 
 
-class TextGenerationRouter(Router):
+class GraphRouter(Router):
     """
     Router for a DAG. Expects graphs be presented in the form of a dictionary, where
     keys are the nodes of the graph and the values are the connected nodes. For
@@ -126,17 +135,15 @@ def next(
         self,
         past: str,
         ops: Dict[str, Operator],
-        context: Context,
         inp: Any,
-        inference_state: InferenceState,
     ) -> int:
         node = past
         if isinstance(self.route[node], str):
             return self.route[node]
         else:
             for neighbour_node in self.route[node]:
                 neighbour_node_op = ops[neighbour_node]
-                if neighbour_node_op.can_operate(inp, context, inference_state):
+                if neighbour_node_op.can_operate(inp):
                     return neighbour_node
             raise ValueError("Cannot operate on any of the nodes")
 

src/deepsparse/v2/schedulers/scheduler.py

@@ -14,7 +14,6 @@
 
 
 from concurrent.futures import Future, ThreadPoolExecutor
-from typing import Any
 
 from deepsparse.v2.operators import Operator
 
@@ -37,19 +36,30 @@ class OperatorScheduler:
     def __init__(self, max_workers: int = 1):
         self._threadpool = ThreadPoolExecutor(max_workers=max_workers)
 
-    def submit(self, *args, operator: Operator, **kwargs) -> Future:
+    def submit(
+        self,
+        *args,
+        operator: Operator,
+        **kwargs,
+    ) -> Future:
         """
         :param operator: operator to run
-        :param operator_input: input schema to the operator
-        :param context: context of already run operators
         :return: future referencing the asynchronously run output of the operator
         """
-        return self._threadpool.submit(operator, *args, **kwargs)
+        return self._threadpool.submit(
+            operator,
+            *args,
+            **kwargs,
+        )
 
-    def can_process(self, *args, operator: Operator, **kwargs) -> bool:
+    def can_process(
+        self,
+        *args,
+        operator: Operator,
+        **kwargs,
+    ) -> bool:
         """
         :param operator: operator to check
-        :param operator_input: operator_input to check
         :return: True if this Operator can process the given operator and input.
             Base OperatorScheduler always returns True
         """