Move after the find max batch size

Author: knowicki-nvidia

model_navigator/pipelines/builders/__init__.py

@@ -26,7 +26,7 @@
 from model_navigator.pipelines.pipeline import Pipeline  # noqa: F401
 
 if is_torch_available():
-    from .torch import torch_conversion_builder, torch_export_builder  # noqa: F401
+    from .torch import torch_conversion_builder, torch_export_builder, torch_export_onnx_builder  # noqa: F401
     from .torch_tensorrt import torch_tensorrt_conversion_builder  # noqa: F401
 
 if is_tf_available():

model_navigator/pipelines/builders/find_device_max_batch_size.py

@@ -95,24 +95,6 @@ def _find_max_batch_size_config_for_torch(config: CommonConfig, models_config: D
             )
             configurations.append(mbs_config)
 
-    for model_cfg in models_config.get(Format.ONNX, []):
-        runner_cls = {
-            DeviceKind.CUDA: OnnxrtCUDARunner,
-            DeviceKind.CPU: OnnxrtCPURunner,
-        }[config.target_device]
-
-        if model_cfg.format != runner_cls.format():
-            raise ModelNavigatorRuntimeError(
-                f"Model config format `{model_cfg.format}` does not match `{runner_cls.format()}`."
-            )
-        mbs_config = FindMaxBatchSizeConfig(
-            format=Format.ONNX,
-            model_path=model_cfg.path,
-            runner_cls=runner_cls,
-            reproduction_scripts_dir=pathlib.Path(model_cfg.key),
-        )
-        configurations.append(mbs_config)
-
     return configurations
 
 

model_navigator/pipelines/builders/torch.py

@@ -51,6 +51,20 @@ def torch_export_builder(config: CommonConfig, models_config: Dict[Format, List[
     for model_cfg in models_config.get(Format.TORCHSCRIPT, []):
         execution_units.append(ExecutionUnit(command=ExportTorch2TorchScript, model_config=model_cfg))
 
+    return Pipeline(name=PIPELINE_TORCH_EXPORT, execution_units=execution_units)
+
+
+def torch_export_onnx_builder(config: CommonConfig, models_config: Dict[Format, List[ModelConfig]]) -> Pipeline:
+    """Prepare export steps for pipeline.
+
+    Args:
+        config: A configuration for pipelines
+        models_config: List of model configs per format
+
+    Returns:
+        Pipeline with steps for export
+    """
+    execution_units: List[ExecutionUnit] = []
     for model_cfg in models_config.get(Format.ONNX, []):
         if model_cfg.parent_path in (None, Format.TORCH):
             #  If model_path provided in onnx config, copy this onnx instead of exporting.

model_navigator/torch/__init__.py

@@ -42,6 +42,7 @@
     tensorrt_conversion_builder,
     torch_conversion_builder,
     torch_export_builder,
+    torch_export_onnx_builder,
     torch_tensorrt_conversion_builder,
     verify_builder,
 )
@@ -144,6 +145,7 @@ def optimize(
         preprocessing_builder,
         torch_export_builder,
         find_device_max_batch_size_builder,
+        torch_export_onnx_builder,
         torch_exportedprogram_builder,
         torch_conversion_builder,
         torch_tensorrt_conversion_builder,

tests/unit/base/test_pipelines_builders_find_device_max_batch_size.py

@@ -80,17 +80,14 @@ def test_find_device_max_batch_size_builder_return_execution_unit_when_torch_fra
     assert len(pipeline.execution_units) == 1
 
     execution_unit = pipeline.execution_units[0]
-    assert len(execution_unit.kwargs["configurations"]) == 3
+    assert len(execution_unit.kwargs["configurations"]) == 2
 
     configuration = execution_unit.kwargs["configurations"][0]
     assert configuration.runner_cls == TorchCUDARunner
 
     configuration = execution_unit.kwargs["configurations"][1]
     assert configuration.runner_cls == TorchCompileCUDARunner
 
-    configuration = execution_unit.kwargs["configurations"][2]
-    assert configuration.runner_cls == OnnxrtCUDARunner
-
 
 def test_find_device_max_batch_size_builder_return_execution_unit_when_tensorflow_framework_is_used():
     config = CommonConfig(

tests/unit/torch/test_pipelines_builders_graph_surgeon.py

@@ -16,10 +16,13 @@
 from model_navigator.configuration.common_config import CommonConfig
 from model_navigator.configuration.model.model_config import ONNXModelConfig, TorchScriptModelConfig
 from model_navigator.frameworks import Framework
-from model_navigator.pipelines.builders.torch import torch_conversion_builder, torch_export_builder
+from model_navigator.pipelines.builders.torch import (
+    torch_conversion_builder,
+    torch_export_onnx_builder,
+)
 
 
-def test_torch_export_builder_return_graph_surgeon_optimization_when_enabled():
+def test_torch_export_onnx_builder_return_graph_surgeon_optimization_when_enabled():
     config = CommonConfig(
         framework=Framework.TORCH,
         dataloader=[{"input_name": [idx]} for idx in range(10)],
@@ -34,12 +37,12 @@ def test_torch_export_builder_return_graph_surgeon_optimization_when_enabled():
     models_config = {
         Format.ONNX: [ONNXModelConfig(opset=DEFAULT_ONNX_OPSET, graph_surgeon_optimization=True)],
     }
-    pipeline = torch_export_builder(config=config, models_config=models_config)
+    pipeline = torch_export_onnx_builder(config=config, models_config=models_config)
     assert len(pipeline.execution_units) == 2
     assert pipeline.execution_units[-1].command == GraphSurgeonOptimize
 
 
-def test_torch_export_builder_does_not_return_graph_surgeon_optimization_when_disabled():
+def test_torch_export_onnx_builder_does_not_return_graph_surgeon_optimization_when_disabled():
     config = CommonConfig(
         framework=Framework.TORCH,
         dataloader=[{"input_name": [idx]} for idx in range(10)],
@@ -54,7 +57,7 @@ def test_torch_export_builder_does_not_return_graph_surgeon_optimization_when_di
     models_config = {
         Format.ONNX: [ONNXModelConfig(opset=DEFAULT_ONNX_OPSET, graph_surgeon_optimization=False)],
     }
-    pipeline = torch_export_builder(config=config, models_config=models_config)
+    pipeline = torch_export_onnx_builder(config=config, models_config=models_config)
     assert len(pipeline.execution_units) == 1