Use pass.run and automatic tmpdir in pass unit tests

test/unit_test/passes/inc/test_inc_quantization.py

@@ -3,7 +3,6 @@
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
 import platform
-import tempfile
 from pathlib import Path
 
 import pytest
@@ -15,62 +14,57 @@
 from olive.passes.olive_pass import create_pass_from_dict
 from olive.passes.onnx.conversion import OnnxConversion
 from olive.passes.onnx.inc_quantization import IncDynamicQuantization, IncQuantization, IncStaticQuantization
-from olive.systems.local import LocalSystem
 
 
 @pytest.mark.skipif(
     platform.system() == "Windows", reason="Skip test on Windows. neural-compressor import is hanging on Windows."
 )
-def test_inc_quantization():
-    with tempfile.TemporaryDirectory() as tempdir:
-        # setup
-        ov_model = get_onnx_model(tempdir)
-        local_system = LocalSystem()
-        data_dir = Path(tempdir) / "data"
-        data_dir.mkdir(exist_ok=True)
-        config = {"data_dir": data_dir, "dataloader_func": create_dataloader}
-        output_folder = str(Path(tempdir) / "quantized")
-
-        # create IncQuantization pass
-        p = create_pass_from_dict(IncQuantization, config, disable_search=True)
-        # execute
-        quantized_model = local_system.run_pass(p, ov_model, None, output_folder)
-        # assert
-        assert quantized_model.model_path.endswith(".onnx")
-        assert Path(quantized_model.model_path).exists()
-        assert Path(quantized_model.model_path).is_file()
-        assert "QLinearConv" in [node.op_type for node in quantized_model.load_model().graph.node]
-
-        # clean
-        del p
-        # create IncDynamicQuantization pass
-        p = create_pass_from_dict(IncDynamicQuantization, config, disable_search=True)
-        # execute
-        quantized_model = local_system.run_pass(p, ov_model, None, output_folder)
-        # assert
-        assert quantized_model.model_path.endswith(".onnx")
-        assert Path(quantized_model.model_path).exists()
-        assert Path(quantized_model.model_path).is_file()
-        assert "DynamicQuantizeLinear" in [node.op_type for node in quantized_model.load_model().graph.node]
-
-        # clean
-        del p
-        # create IncStaticQuantization pass
-        p = create_pass_from_dict(IncStaticQuantization, config, disable_search=True)
-        # execute
-        quantized_model = local_system.run_pass(p, ov_model, None, output_folder)
-        # assert
-        assert quantized_model.model_path.endswith(".onnx")
-        assert Path(quantized_model.model_path).exists()
-        assert Path(quantized_model.model_path).is_file()
-        assert "QLinearConv" in [node.op_type for node in quantized_model.load_model().graph.node]
-
-
-def get_onnx_model(tempdir):
-    local_system = LocalSystem()
+def test_inc_quantization(tmpdir):
+    ov_model = get_onnx_model(tmpdir)
+    data_dir = Path(tmpdir) / "data"
+    data_dir.mkdir(exist_ok=True)
+    config = {"data_dir": data_dir, "dataloader_func": create_dataloader}
+    output_folder = str(Path(tmpdir) / "quantized")
+
+    # create IncQuantization pass
+    p = create_pass_from_dict(IncQuantization, config, disable_search=True)
+    # execute
+    quantized_model = p.run(ov_model, None, output_folder)
+    # assert
+    assert quantized_model.model_path.endswith(".onnx")
+    assert Path(quantized_model.model_path).exists()
+    assert Path(quantized_model.model_path).is_file()
+    assert "QLinearConv" in [node.op_type for node in quantized_model.load_model().graph.node]
+
+    # clean
+    del p
+    # create IncDynamicQuantization pass
+    p = create_pass_from_dict(IncDynamicQuantization, config, disable_search=True)
+    # execute
+    quantized_model = p.run(ov_model, None, output_folder)
+    # assert
+    assert quantized_model.model_path.endswith(".onnx")
+    assert Path(quantized_model.model_path).exists()
+    assert Path(quantized_model.model_path).is_file()
+    assert "DynamicQuantizeLinear" in [node.op_type for node in quantized_model.load_model().graph.node]
+
+    # clean
+    del p
+    # create IncStaticQuantization pass
+    p = create_pass_from_dict(IncStaticQuantization, config, disable_search=True)
+    # execute
+    quantized_model = p.run(ov_model, None, output_folder)
+    # assert
+    assert quantized_model.model_path.endswith(".onnx")
+    assert Path(quantized_model.model_path).exists()
+    assert Path(quantized_model.model_path).is_file()
+    assert "QLinearConv" in [node.op_type for node in quantized_model.load_model().graph.node]
+
+
+def get_onnx_model(tmpdir):
     torch_hub_model_path = "chenyaofo/pytorch-cifar-models"
     pytorch_hub_model_name = "cifar10_mobilenetv2_x1_0"
-    torch.hub.set_dir(tempdir)
+    torch.hub.set_dir(tmpdir)
     pytorch_model = PyTorchModel(
         model_loader=lambda torch_hub_model_path: torch.hub.load(torch_hub_model_path, pytorch_hub_model_name),
         model_path=torch_hub_model_path,
@@ -79,10 +73,10 @@ def get_onnx_model(tempdir):
     onnx_conversion_config = {}
 
     p = create_pass_from_dict(OnnxConversion, onnx_conversion_config, disable_search=True)
-    output_folder = str(Path(tempdir) / "onnx")
+    output_folder = str(Path(tmpdir) / "onnx")
 
     # execute
-    onnx_model = local_system.run_pass(p, pytorch_model, None, output_folder)
+    onnx_model = p.run(pytorch_model, None, output_folder)
     return onnx_model
 
 

test/unit_test/passes/onnx/pipeline/test_step_utils.py

@@ -3,29 +3,26 @@
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
 import json
-import tempfile
 from pathlib import Path
 from test.unit_test.passes.onnx.test_pre_post_processing_op import (
     convert_superresolution_model,
     get_superresolution_model,
 )
 
 from olive.passes.onnx.pipeline.step_utils import parse_steps
-from olive.systems.local import LocalSystem
 
 
 class CustomizedParam:
     def __init__(self, params: dict):
         self.params = params
 
 
-def test_step_parser():
+def test_step_parser(tmpdir):
     from onnxruntime_extensions.tools.pre_post_processing import TokenizerParam
 
     pytorch_model = get_superresolution_model()
-    with tempfile.TemporaryDirectory() as tempdir:
-        input_model = convert_superresolution_model(pytorch_model, tempdir, LocalSystem())
-        model = input_model.load_model()
+    input_model = convert_superresolution_model(pytorch_model, tmpdir)
+    model = input_model.load_model()
 
     step_config = Path(__file__).parent / "step_config.json"
     with step_config.open() as f:

test/unit_test/passes/onnx/test_conversion.py

@@ -1,26 +1,21 @@
-import tempfile
 from pathlib import Path
 from test.unit_test.utils import get_hf_model_with_past, get_pytorch_model
 
 import pytest
 
 from olive.passes.olive_pass import create_pass_from_dict
 from olive.passes.onnx.conversion import OnnxConversion
-from olive.systems.local import LocalSystem
 
 
 @pytest.mark.parametrize("input_model", [get_pytorch_model(), get_hf_model_with_past()])
-def test_onnx_conversion_pass(input_model):
+def test_onnx_conversion_pass(input_model, tmpdir):
     # setup
-    local_system = LocalSystem()
-
     p = create_pass_from_dict(OnnxConversion, {}, disable_search=True)
-    with tempfile.TemporaryDirectory() as tempdir:
-        output_folder = str(Path(tempdir) / "onnx")
+    output_folder = str(Path(tmpdir) / "onnx")
 
-        # The conversion need torch version > 1.13.1, otherwise, it will complain
-        # Unsupported ONNX opset version: 18
-        onnx_model = local_system.run_pass(p, input_model, None, output_folder)
+    # The conversion need torch version > 1.13.1, otherwise, it will complain
+    # Unsupported ONNX opset version: 18
+    onnx_model = p.run(input_model, None, output_folder)
 
-        # assert
-        assert Path(onnx_model.model_path).exists()
+    # assert
+    assert Path(onnx_model.model_path).exists()

test/unit_test/passes/onnx/test_insert_beam_search.py

@@ -1,16 +1,13 @@
-import tempfile
 from pathlib import Path
 from test.unit_test.utils import get_onnx_model
 
 from olive.model import CompositeOnnxModel
 from olive.passes.olive_pass import create_pass_from_dict
 from olive.passes.onnx.insert_beam_search import InsertBeamSearch
-from olive.systems.local import LocalSystem
 
 
-def test_insert_beam_search_pass():
+def test_insert_beam_search_pass(tmpdir):
     # setup
-    local_system = LocalSystem()
     input_models = []
     input_models.append(get_onnx_model())
     input_models.append(get_onnx_model())
@@ -21,8 +18,7 @@ def test_insert_beam_search_pass():
     )
 
     p = create_pass_from_dict(InsertBeamSearch, {}, disable_search=True)
-    with tempfile.TemporaryDirectory() as tempdir:
-        output_folder = str(Path(tempdir) / "onnx")
+    output_folder = str(Path(tmpdir) / "onnx")
 
-        # execute
-        local_system.run_pass(p, composite_model, None, output_folder)
+    # execute
+    p.run(composite_model, None, output_folder)

test/unit_test/passes/onnx/test_mixed_precision.py

@@ -1,19 +1,15 @@
-import tempfile
 from pathlib import Path
 from test.unit_test.utils import get_onnx_model
 
 from olive.passes.olive_pass import create_pass_from_dict
 from olive.passes.onnx.mixed_precision import OrtMixedPrecision
-from olive.systems.local import LocalSystem
 
 
-def test_ort_mixed_precision_pass():
+def test_ort_mixed_precision_pass(tmpdir):
     # setup
-    local_system = LocalSystem()
     input_model = get_onnx_model()
     p = create_pass_from_dict(OrtMixedPrecision, {}, disable_search=True)
-    with tempfile.TemporaryDirectory() as tempdir:
-        output_folder = str(Path(tempdir) / "onnx")
+    output_folder = str(Path(tmpdir) / "onnx")
 
-        # execute
-        local_system.run_pass(p, input_model, None, output_folder)
+    # execute
+    p.run(input_model, None, output_folder)

test/unit_test/passes/onnx/test_model_optimizer.py

@@ -2,22 +2,18 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
-import tempfile
 from pathlib import Path
 from test.unit_test.utils import get_onnx_model
 
 from olive.passes.olive_pass import create_pass_from_dict
 from olive.passes.onnx import OnnxModelOptimizer
-from olive.systems.local import LocalSystem
 
 
-def test_onnx_model_optimizer_pass():
+def test_onnx_model_optimizer_pass(tmpdir):
     # setup
-    local_system = LocalSystem()
     input_model = get_onnx_model()
     p = create_pass_from_dict(OnnxModelOptimizer, {}, disable_search=True)
-    with tempfile.TemporaryDirectory() as tempdir:
-        output_folder = str(Path(tempdir) / "onnx")
+    output_folder = str(Path(tmpdir) / "onnx")
 
-        # execute
-        local_system.run_pass(p, input_model, None, output_folder)
+    # execute
+    p.run(input_model, None, output_folder)

test/unit_test/passes/onnx/test_optimum_conversion.py

@@ -1,22 +1,20 @@
-import tempfile
 from pathlib import Path
 from test.unit_test.utils import get_optimum_model_by_hf_config, get_optimum_model_by_model_path
 
 import pytest
 
 from olive.passes.olive_pass import create_pass_from_dict
 from olive.passes.onnx.optimum_conversion import OptimumConversion
-from olive.systems.local import LocalSystem
 
 
 @pytest.mark.parametrize("input_model", [get_optimum_model_by_hf_config(), get_optimum_model_by_model_path()])
-def test_optimum_conversion_pass(input_model):
+def test_optimum_conversion_pass(input_model, tmpdir):
     # setup
-    local_system = LocalSystem()
-
     p = create_pass_from_dict(OptimumConversion, {}, disable_search=True)
-    with tempfile.TemporaryDirectory() as tempdir:
-        output_folder = Path(tempdir)
-        onnx_model = local_system.run_pass(p, input_model, None, output_folder)
-        # assert
-        assert Path(onnx_model.model_path).exists()
+    output_folder = Path(tmpdir)
+
+    # execute
+    onnx_model = p.run(input_model, None, output_folder)
+
+    # assert
+    assert Path(onnx_model.model_path).exists()

test/unit_test/passes/onnx/test_perf_tuning.py

@@ -2,7 +2,6 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
-import tempfile
 from pathlib import Path
 from test.unit_test.utils import get_onnx_model
 from unittest.mock import patch
@@ -11,24 +10,21 @@
 
 from olive.passes.olive_pass import create_pass_from_dict
 from olive.passes.onnx import OrtPerfTuning
-from olive.systems.local import LocalSystem
 
 
 @pytest.mark.parametrize("config", [{"input_names": ["input"], "input_shapes": [[1, 1]]}, {}])
-def test_ort_perf_tuning_pass(config):
+def test_ort_perf_tuning_pass(config, tmpdir):
     # setup
-    local_system = LocalSystem()
     input_model = get_onnx_model()
     p = create_pass_from_dict(OrtPerfTuning, config, disable_search=True)
-    with tempfile.TemporaryDirectory() as tempdir:
-        output_folder = str(Path(tempdir) / "onnx")
+    output_folder = str(Path(tmpdir) / "onnx")
 
-        # execute
-        local_system.run_pass(p, input_model, None, output_folder)
+    # execute
+    p.run(input_model, None, output_folder)
 
 
 @patch("olive.model.ONNXModel.get_io_config")
-def test_ort_perf_tuning_pass_with_dynamic_shapes(mock_get_io_config):
+def test_ort_perf_tuning_pass_with_dynamic_shapes(mock_get_io_config, tmpdir):
     mock_get_io_config.return_value = {
         "input_names": ["input"],
         "input_shapes": [["input_0", "input_1"]],
@@ -38,13 +34,11 @@ def test_ort_perf_tuning_pass_with_dynamic_shapes(mock_get_io_config):
         "output_types": ["float32", "float32"],
     }
 
-    local_system = LocalSystem()
     input_model = get_onnx_model()
     p = create_pass_from_dict(OrtPerfTuning, {}, disable_search=True)
-    with tempfile.TemporaryDirectory() as tempdir:
-        output_folder = str(Path(tempdir) / "onnx")
+    output_folder = str(Path(tmpdir) / "onnx")
 
-        with pytest.raises(TypeError) as e:
-            # execute
-            local_system.run_pass(p, input_model, None, output_folder)
-            assert "ones() received an invalid combination of arguments" in str(e.value)
+    with pytest.raises(TypeError) as e:
+        # execute
+        p.run(input_model, None, output_folder)
+        assert "ones() received an invalid combination of arguments" in str(e.value)