[Pipeline Refactor] Migration (#1460)

* update pathways to use new v2 pipeline

* fix image classification

* quality

* leftover fixes

* fix batch sizes:

* quality; use old tasks for server

* fix server tests

* fix captioning

* fix init for haystack

* fix import

* fix another import

* fix test

* fix init file

* fix test

* fix remaining base tests

* fix docstring

* quality

* update codegen alias to use the new registry and text generation pipeline

src/deepsparse/__init__.py

@@ -33,7 +33,6 @@
 from .engine import *
 from .tasks import *
 from .pipeline import *
-from .base_pipeline import *
 from .loggers import *
 from .version import __version__, is_release
 from .analytics import deepsparse_analytics as _analytics

src/deepsparse/clip/captioning_pipeline.py

@@ -27,7 +27,7 @@
 import torch
 import torch.nn.functional as F
 from deepsparse.clip import CLIPDecoderInput, CLIPTextInput, CLIPVisualInput
-from deepsparse.pipeline import BasePipeline, Pipeline
+from deepsparse.legacy.pipeline import BasePipeline, Pipeline
 
 
 __all__ = ["CLIPCaptionInput", "CLIPCaptionOutput", "CLIPCaptionPipeline"]

src/deepsparse/clip/decoder_pipeline.py

@@ -17,7 +17,7 @@
 import numpy as np
 from pydantic import BaseModel, Field
 
-from deepsparse import Pipeline
+from deepsparse.legacy import Pipeline
 from deepsparse.utils import model_to_path
 
 

src/deepsparse/clip/text_pipeline.py

@@ -17,7 +17,7 @@
 import numpy as np
 from pydantic import BaseModel, Field
 
-from deepsparse.pipeline import Pipeline
+from deepsparse.legacy.pipeline import Pipeline
 from deepsparse.utils import model_to_path
 from open_clip.tokenizer import tokenize
 

src/deepsparse/clip/visual_pipeline.py

@@ -22,7 +22,7 @@
 from torchvision.transforms import InterpolationMode
 
 from deepsparse.clip.constants import CLIP_RGB_MEANS, CLIP_RGB_STDS
-from deepsparse.pipeline import Pipeline
+from deepsparse.legacy.pipeline import Pipeline
 from deepsparse.pipelines.computer_vision import ComputerVisionSchema
 from deepsparse.utils import model_to_path
 

src/deepsparse/clip/zeroshot_pipeline.py

@@ -19,7 +19,7 @@
 from pydantic import BaseModel, Field
 
 from deepsparse.clip import CLIPTextInput, CLIPVisualInput
-from deepsparse.pipeline import BasePipeline, Pipeline
+from deepsparse.legacy.pipeline import BasePipeline, Pipeline
 from scipy.special import softmax
 
 

src/deepsparse/evaluation/cli.py

@@ -77,7 +77,11 @@
 )
 from src.deepsparse.evaluation.results import Result, save_result
 from src.deepsparse.evaluation.utils import args_to_dict, get_save_path
-from src.deepsparse.pipeline import DEEPSPARSE_ENGINE, ORT_ENGINE, TORCHSCRIPT_ENGINE
+from src.deepsparse.operators.engine_operator import (
+    DEEPSPARSE_ENGINE,
+    ORT_ENGINE,
+    TORCHSCRIPT_ENGINE,
+)
 
 
 _LOGGER = logging.getLogger(__name__)

src/deepsparse/evaluation/evaluator.py

@@ -19,7 +19,11 @@
 )
 from src.deepsparse.evaluation.registry import EvaluationRegistry
 from src.deepsparse.evaluation.results import Result
-from src.deepsparse.pipeline import DEEPSPARSE_ENGINE, ORT_ENGINE, TORCHSCRIPT_ENGINE
+from src.deepsparse.operators.engine_operator import (
+    DEEPSPARSE_ENGINE,
+    ORT_ENGINE,
+    TORCHSCRIPT_ENGINE,
+)
 
 
 __all__ = ["evaluate"]

src/deepsparse/evaluation/utils.py

@@ -17,7 +17,8 @@
 
 from transformers import AutoModelForCausalLM
 
-from deepsparse import DEEPSPARSE_ENGINE, ORT_ENGINE, Pipeline
+from deepsparse import Pipeline
+from deepsparse.operators.engine_operator import DEEPSPARSE_ENGINE, ORT_ENGINE
 
 
 __all__ = ["text_generation_model_from_target", "get_save_path", "args_to_dict"]

src/deepsparse/image_classification/__init__.py

@@ -25,9 +25,12 @@
         "Please install deepsparse[image_classification] to use this pathway"
     )
 
-
 from .constants import *
-from .pipelines import *
+from .pipeline import *
+
+# flake8: noqa
+from .postprocess_operator import *
+from .preprocess_operator import *
 from .schemas import *
 
 

src/deepsparse/v2/image_classification/pipeline.py → src/deepsparse/image_classification/pipeline.py

@@ -13,40 +13,43 @@
 # limitations under the License.
 
 import logging
-import warnings
 from typing import Dict, Optional, Tuple, Union
 
-from deepsparse.v2.image_classification.postprocess_operator import (
+from deepsparse.image_classification.postprocess_operator import (
     ImageClassificationPostProcess,
 )
-from deepsparse.v2.image_classification.preprocess_operator import (
+from deepsparse.image_classification.preprocess_operator import (
     ImageClassificationPreProcess,
 )
-from deepsparse.v2.operators.engine_operator import EngineOperator
-from deepsparse.v2.pipeline import Pipeline
-from deepsparse.v2.routers.router import LinearRouter
-from deepsparse.v2.schedulers.scheduler import OperatorScheduler
+from deepsparse.operators.engine_operator import EngineOperator
+from deepsparse.operators.registry import OperatorRegistry
+from deepsparse.pipeline import Pipeline
+from deepsparse.routers.router import LinearRouter
+from deepsparse.schedulers.scheduler import OperatorScheduler
 
 
 _LOGGER = logging.getLogger(__name__)
 
 __all__ = ["ImageClassificationPipeline"]
 
 
+@OperatorRegistry.register(name="image_classification")
 class ImageClassificationPipeline(Pipeline):
     def __init__(
         self,
         model_path: str,
-        engine_kwargs: Optional[Dict] = None,
         class_names: Union[None, str, Dict[str, str]] = None,
         image_size: Optional[Tuple[int]] = None,
         top_k: int = 1,
+        **engine_kwargs,
     ):
+
         if not engine_kwargs:
             engine_kwargs = {}
             engine_kwargs["model_path"] = model_path
         elif engine_kwargs.get("model_path") != model_path:
-            warnings.warn(f"Updating engine_kwargs to include {model_path}")
+            _LOGGER.warning(f"Updating engine_kwargs to include {model_path}")
+            engine_kwargs["model_path"] = model_path
 
         engine = EngineOperator(**engine_kwargs)
         preproces = ImageClassificationPreProcess(

src/deepsparse/v2/image_classification/postprocess_operator.py → src/deepsparse/image_classification/postprocess_operator.py

@@ -18,7 +18,7 @@
 import numpy
 from pydantic import BaseModel, Field
 
-from deepsparse.v2.operators import Operator
+from deepsparse.operators import Operator
 
 
 class ImageClassificationOutput(BaseModel):

src/deepsparse/v2/image_classification/preprocess_operator.py → src/deepsparse/image_classification/preprocess_operator.py

@@ -23,8 +23,8 @@
     IMAGENET_RGB_MEANS,
     IMAGENET_RGB_STDS,
 )
+from deepsparse.operators import Operator
 from deepsparse.pipelines.computer_vision import ComputerVisionSchema
-from deepsparse.v2.operators import Operator
 
 
 class ImageClassificationInput(ComputerVisionSchema):

src/deepsparse/image_classification/validation_script.py

@@ -27,13 +27,13 @@
                                   on Imagenette  [default: zoo:cv/classificati
                                   on/resnet_v1-50/pytorch/sparseml/imagenette/
                                   base-none]
+  --image-size, --image_size INTEGER
+                                  integer size to evaluate images at (will be
+                                  reshaped to square shape)  [default: 224]
   --batch-size, --batch_size INTEGER
                                   Test batch size, must divide the dataset
                                   evenly, else last batch will be dropped
                                   [default: 1]
-  --image-size, --image_size INTEGER
-                                  integer size to evaluate images at (will be
-                                  reshaped to square shape)  [default: 224]
   --num-cores, --num_cores INTEGER
                                   Number of CPU cores to run deepsparse with,
                                   default is all available
@@ -213,11 +213,10 @@ def main(
     pipeline = Pipeline.create(
         task="image_classification",
         model_path=model_path,
+        engine_type=engine,
         batch_size=batch_size,
         num_cores=num_cores,
-        engine_type=engine,
     )
-    print(f"engine info: {pipeline.engine}")
     correct = total = 0
     progress_bar = tqdm(data_loader)
 

src/deepsparse/v2/utils/__init__.py → src/deepsparse/legacy/__init__.py

@@ -1,5 +1,3 @@
-# flake8: noqa
-
 # Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,9 +11,9 @@
 # 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 .helpers import *
-from .state import *
-from .types import *
 
+# flake8: noqa
 
-from .data import *  # isort:skip
+from .base_pipeline import *
+from .pipeline import *
+from .tasks import *

src/deepsparse/base_pipeline.py → src/deepsparse/legacy/base_pipeline.py

@@ -19,10 +19,10 @@
 from pydantic import BaseModel
 
 from deepsparse import Context
+from deepsparse.legacy.tasks import SupportedTasks, dynamic_import_task
 from deepsparse.loggers.base_logger import BaseLogger
 from deepsparse.loggers.build_logger import logger_from_config
 from deepsparse.loggers.constants import validate_identifier
-from deepsparse.tasks import SupportedTasks, dynamic_import_task
 
 
 __all__ = [
@@ -166,7 +166,7 @@ def create(
             implementation
         :return: pipeline object initialized for the given task
         """
-        from deepsparse.pipeline import Bucketable, BucketingPipeline, Pipeline
+        from deepsparse.legacy.pipeline import Bucketable, BucketingPipeline, Pipeline
 
         pipeline_constructor = BasePipeline._get_task_constructor(task)
         model_path = kwargs.get("model_path", None)
@@ -278,7 +278,7 @@ def from_config(
             logging. Default is None
         :return: loaded Pipeline object from the config
         """
-        from deepsparse.pipeline import PipelineConfig
+        from deepsparse.legacy.pipeline import PipelineConfig
 
         if isinstance(config, Path) or (
             isinstance(config, str) and os.path.exists(config)
@@ -308,7 +308,7 @@ def to_config(self) -> "PipelineConfig":  # noqa: F821
         """
         :return: PipelineConfig that can be used to reload this object
         """
-        from deepsparse.pipeline import PipelineConfig
+        from deepsparse.legacy.pipeline import PipelineConfig
 
         if not hasattr(self, "task"):
             raise RuntimeError(

src/deepsparse/v2/image_classification/__init__.py → src/deepsparse/legacy/image_classification/__init__.py

@@ -13,8 +13,5 @@
 # limitations under the License.
 
 # flake8: noqa
-from .postprocess_operator import *
-from .preprocess_operator import *
 
-
-from .pipeline import *  # isort:skip
+from .pipelines import *

src/deepsparse/image_classification/pipelines.py → src/deepsparse/legacy/image_classification/pipelines.py

@@ -31,7 +31,7 @@
     ImageClassificationInput,
     ImageClassificationOutput,
 )
-from deepsparse.pipeline import Pipeline
+from deepsparse.legacy.pipeline import Pipeline
 from deepsparse.utils import model_to_path