diff --git a/olive/telemetry/library/telemetry_logger.py b/olive/telemetry/library/telemetry_logger.py index 7eb236e759..32562e1d3c 100644 --- a/olive/telemetry/library/telemetry_logger.py +++ b/olive/telemetry/library/telemetry_logger.py @@ -113,17 +113,23 @@ def register_payload_transmitted_callback( return self._logger_exporter.register_payload_transmitted_callback(callback, include_failures) return None - def log(self, event_name: str, attributes: Optional[dict[str, Any]] = None) -> None: + def log(self, event_name: str, attributes: Optional[dict[str, Any]] = None) -> bool: """Log a telemetry event. Args: event_name: Name of the event attributes: Optional event attributes + Returns: + Whether the event was submitted to the logger. + """ - if self._logger: - extra = attributes if attributes else {} - self._logger.info(event_name, extra=extra) + if not self._logger or self._logger.disabled: + return False + + extra = attributes if attributes else {} + self._logger.info(event_name, extra=extra) + return True def disable_telemetry(self) -> None: """Disable telemetry logging.""" diff --git a/olive/telemetry/telemetry.py b/olive/telemetry/telemetry.py index 0ddb690e2a..0a326a6c98 100644 --- a/olive/telemetry/telemetry.py +++ b/olive/telemetry/telemetry.py @@ -548,8 +548,8 @@ def log( attrs = _merge_metadata(attributes, metadata) if self._logger is None: return - self._logger.log(event_name, attrs) - if self._cache_handler: + event_logged = self._logger.log(event_name, attrs) + if event_logged and self._cache_handler: self._cache_handler.record_event_logged() except Exception: # Fail silently — telemetry must never crash the host application diff --git a/test/cli/test_cli.py b/test/cli/test_cli.py index 3db547f63d..fb986fb306 100644 --- a/test/cli/test_cli.py +++ b/test/cli/test_cli.py @@ -14,7 +14,6 @@ from olive.cli.launcher import main as cli_main -@pytest.mark.parametrize("console_script", [True, False]) @pytest.mark.parametrize( "command", [ @@ -28,40 +27,44 @@ "diffusion-lora", ], ) -def test_valid_command(console_script, command): - # setup - command_args = [] - if console_script: - command_args.append("olive") - else: - command_args.extend([sys.executable, "-m", "olive"]) - if command: - command_args.append(command) - command_args.append("--help") +def test_command_help(command, capsys): + with pytest.raises(SystemExit) as exc_info: + cli_main([command, "--help"]) - # execute + assert exc_info.value.code == 0 + assert f"usage: olive {command}" in capsys.readouterr().out + + +@pytest.mark.parametrize( + ("command_args", "expected_usage"), + [ + ( + [str(Path(sys.executable).with_name("olive.exe" if sys.platform == "win32" else "olive")), "run", "--help"], + "usage: olive run", + ), + ([sys.executable, "-m", "olive", "run", "--help"], "usage: python -m olive run"), + ], +) +def test_cli_entrypoint(command_args, expected_usage): out = subprocess.run(command_args, check=True, capture_output=True) - # assert - if not console_script: - # the help message only says `python` when running as a module - command_args[0] = "python" - assert f"usage: {' '.join(command_args[:-1])}" in out.stdout.decode("utf-8") + assert expected_usage in out.stdout.decode("utf-8") -@pytest.mark.parametrize("console_script", [True, False]) -def test_invalid_command(console_script): - # setup - command_args = [] - if console_script: - command_args.append("olive") - else: - command_args.extend([sys.executable, "-m", "olive"]) - command_args.append("invalid-command") +def test_cli_top_level_help(capsys): + with pytest.raises(SystemExit) as exc_info: + cli_main(["--help"]) + + assert exc_info.value.code == 0 + assert "usage: olive" in capsys.readouterr().out + + +def test_invalid_command(capsys): + with pytest.raises(SystemExit) as exc_info: + cli_main(["invalid-command"]) - # execute and assert - with pytest.raises(subprocess.CalledProcessError): - subprocess.run(command_args, check=True, capture_output=True) + assert exc_info.value.code == 2 + assert "invalid choice: 'invalid-command'" in capsys.readouterr().err @pytest.mark.parametrize("deprecated_module", ["olive.workflows.run", "olive.platform_sdk.qualcomm.configure"]) @@ -76,15 +79,14 @@ def test_legacy_call(deprecated_module): ) -def test_unknown_args(): - # setup +def test_unknown_args(capsys): command_args = ["olive", "run", "--config", "config.json", "--unknown-arg", "-u"] - # execute and assert - with pytest.raises(subprocess.CalledProcessError) as exc_info: - subprocess.run(command_args, check=True, capture_output=True) + with pytest.raises(SystemExit) as exc_info: + cli_main(command_args[1:]) - error_message = exc_info.value.stderr.decode("utf-8") + assert exc_info.value.code == 2 + error_message = capsys.readouterr().err assert "Unknown arguments:" in error_message assert "--unknown-arg" in error_message assert "-u" in error_message diff --git a/test/cli/test_run_pass.py b/test/cli/test_run_pass.py index 12776cc443..d4c4c24fe1 100644 --- a/test/cli/test_run_pass.py +++ b/test/cli/test_run_pass.py @@ -2,21 +2,23 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- -import subprocess -import sys from argparse import ArgumentParser +import pytest -def test_run_pass_command_help(): - """Test that the run-pass command shows help properly.""" - # setup - command_args = [sys.executable, "-m", "olive", "run-pass", "--help"] - # execute - out = subprocess.run(command_args, check=True, capture_output=True) +def test_run_pass_command_help(capsys): + parser = ArgumentParser() + sub_parsers = parser.add_subparsers() + + from olive.cli.run_pass import RunPassCommand + + RunPassCommand.register_subcommand(sub_parsers) + with pytest.raises(SystemExit) as exc_info: + parser.parse_args(["run-pass", "--help"]) - # assert - help_text = out.stdout.decode("utf-8") + assert exc_info.value.code == 0 + help_text = capsys.readouterr().out assert "usage:" in help_text assert "run-pass" in help_text assert "--pass-name" in help_text @@ -114,79 +116,6 @@ def test_run_pass_command_missing_model(): assert getattr(args, "model_name_or_path", None) is None -def test_run_pass_command_config_generation(): - """Test the configuration generation logic.""" - from copy import deepcopy - - # Test the configuration template and generation logic - template = { - "input_model": {"type": "HfModel", "load_kwargs": {"attn_implementation": "eager"}}, - "systems": { - "local_system": { - "type": "LocalSystem", - "accelerators": [{"device": "cpu", "execution_providers": ["CPUExecutionProvider"]}], - } - }, - "output_dir": "models", - "host": "local_system", - "target": "local_system", - "no_artifacts": True, - } - - # Simulate the _get_run_config logic - config = deepcopy(template) - pass_name = "OnnxConversion" - - # Add the pass configuration - pass_config = {"type": pass_name} - config["passes"] = {pass_name.lower(): pass_config} - - # Update output directory - config["output_dir"] = "/tmp/test_output" - - # Verify the structure - assert "passes" in config - assert "onnxconversion" in config["passes"] - assert config["passes"]["onnxconversion"]["type"] == "OnnxConversion" - assert config["output_dir"] == "/tmp/test_output" - assert config["host"] == "local_system" - assert config["target"] == "local_system" - - -def test_run_pass_command_config_generation_with_pass_config(): - """Test the configuration generation with additional pass config.""" - from copy import deepcopy - - template = { - "input_model": {"type": "HfModel", "load_kwargs": {"attn_implementation": "eager"}}, - "systems": { - "local_system": { - "type": "LocalSystem", - "accelerators": [{"device": "cpu", "execution_providers": ["CPUExecutionProvider"]}], - } - }, - "output_dir": "models", - "host": "local_system", - "target": "local_system", - "no_artifacts": True, - } - - # Simulate enhanced configuration - config = deepcopy(template) - pass_name = "OnnxConversion" - pass_config = {"type": pass_name} - - # Add additional configuration - additional_config = {"convert_attribute": True} - pass_config.update(additional_config) - - config["passes"] = {pass_name.lower(): pass_config} - - # Verify the enhanced structure - assert config["passes"]["onnxconversion"]["type"] == "OnnxConversion" - assert config["passes"]["onnxconversion"]["convert_attribute"] is True - - def test_run_pass_command_device_options(): """Test the --device and accelerator argument parsing.""" parser = ArgumentParser() @@ -223,138 +152,3 @@ def test_run_pass_command_device_options(): assert args_default.device == "cpu" assert args_default.provider == "CPUExecutionProvider" - - -def test_run_pass_command_accelerator_config_integration(): - """Test that accelerator options are integrated into the configuration properly.""" - from copy import deepcopy - - from olive.common.utils import set_nested_dict_value - - # Test template - template = { - "input_model": {"type": "HfModel", "load_kwargs": {"attn_implementation": "eager"}}, - "systems": { - "local_system": { - "type": "LocalSystem", - "accelerators": [{"device": "cpu", "execution_providers": ["CPUExecutionProvider"]}], - } - }, - "output_dir": "models", - "host": "local_system", - "target": "local_system", - "no_artifacts": True, - } - - # Simulate update_accelerator_options logic - config = deepcopy(template) - - # Simulate args with GPU device - class MockArgs: - device = "gpu" - provider = "CUDAExecutionProvider" - memory = None - - args = MockArgs() - - # Apply the accelerator updates - execution_providers = [args.provider] - to_replace = [ - (("systems", "local_system", "accelerators", 0, "device"), args.device), - (("systems", "local_system", "accelerators", 0, "execution_providers"), execution_providers), - (("systems", "local_system", "accelerators", 0, "memory"), args.memory), - ] - for k, v in to_replace: - if v is not None: - set_nested_dict_value(config, k, v) - - # Verify the configuration was updated correctly - accelerator = config["systems"]["local_system"]["accelerators"][0] - assert accelerator["device"] == "gpu" - assert accelerator["execution_providers"] == ["CUDAExecutionProvider"] - - -def test_run_pass_command_device_provider_consistency(): - """Test that provider/device consistency is enforced.""" - from copy import deepcopy - - # Test template - template = { - "input_model": {"type": "HfModel", "load_kwargs": {"attn_implementation": "eager"}}, - "systems": { - "local_system": { - "type": "LocalSystem", - "accelerators": [{"device": "cpu", "execution_providers": ["CPUExecutionProvider"]}], - } - }, - "output_dir": "models", - "host": "local_system", - "target": "local_system", - "no_artifacts": True, - } - - # Test the consistency enforcement logic - config = deepcopy(template) - - # Simulate a case where user specifies device=cpu but provider=CUDAExecutionProvider - accelerator = config["systems"]["local_system"]["accelerators"][0] - accelerator["device"] = "cpu" - accelerator["execution_providers"] = ["CUDAExecutionProvider"] - - # Define test constants (from olive.hardware.constants) - device_to_ep = { - "cpu": {"CPUExecutionProvider", "OpenVINOExecutionProvider"}, - "gpu": { - "DmlExecutionProvider", - "CUDAExecutionProvider", - "ROCMExecutionProvider", - "MIGraphXExecutionProvider", - "TensorrtExecutionProvider", - "NvTensorRTRTXExecutionProvider", - "OpenVINOExecutionProvider", - "JsExecutionProvider", - }, - "npu": { - "DmlExecutionProvider", - "QNNExecutionProvider", - "VitisAIExecutionProvider", - "OpenVINOExecutionProvider", - }, - } - - # Apply consistency logic (simulate _ensure_device_provider_consistency) - providers = accelerator.get("execution_providers", []) - current_device = accelerator.get("device", "cpu") - - if providers: - provider = providers[0] - - # Define provider-specific device preferences - provider_device_preference = { - "CPUExecutionProvider": "cpu", - "CUDAExecutionProvider": "gpu", - "ROCMExecutionProvider": "gpu", - "TensorrtExecutionProvider": "gpu", - "NvTensorRTRTXExecutionProvider": "gpu", - "MIGraphXExecutionProvider": "gpu", - "JsExecutionProvider": "gpu", - "DmlExecutionProvider": "gpu", - "QNNExecutionProvider": "npu", - "VitisAIExecutionProvider": "npu", - "OpenVINOExecutionProvider": "cpu", - } - - # Check if current device is valid for the provider - valid_devices = [] - for device, device_providers in device_to_ep.items(): - if provider in device_providers: - valid_devices.append(device) - - if current_device not in valid_devices and valid_devices: - # Current device is not valid for the provider, use the preferred device - preferred_device = provider_device_preference.get(provider, valid_devices[0]) - accelerator["device"] = preferred_device - - # Verify that device was corrected to match the provider - assert accelerator["device"] == "gpu" # CUDAExecutionProvider requires gpu - assert accelerator["execution_providers"] == ["CUDAExecutionProvider"] diff --git a/test/common/quant/test_hf_utils.py b/test/common/quant/test_hf_utils.py index b5799fb4c5..83e845df77 100644 --- a/test/common/quant/test_hf_utils.py +++ b/test/common/quant/test_hf_utils.py @@ -37,64 +37,48 @@ def test_enum_value(self): class TestOliveHfQuantizationOverrideConfig: - def test_default_initialization(self): - """Test default initialization of override config.""" - config = OliveHfQuantizationOverrideConfig() - assert config.bits is None - assert config.symmetric is None - assert config.group_size is None - - def test_custom_initialization(self): - """Test custom initialization of override config.""" - config = OliveHfQuantizationOverrideConfig(bits=8, symmetric=False, group_size=32) - assert config.bits == 8 - assert config.symmetric is False - assert config.group_size == 32 - - def test_partial_initialization(self): - """Test partial initialization of override config.""" - config = OliveHfQuantizationOverrideConfig(bits=4) - assert config.bits == 4 - assert config.symmetric is None - assert config.group_size is None + @pytest.mark.parametrize( + ("kwargs", "expected"), + [ + pytest.param({}, (None, None, None), id="default"), + pytest.param({"bits": 8, "symmetric": False, "group_size": 32}, (8, False, 32), id="custom"), + pytest.param({"bits": 4}, (4, None, None), id="partial"), + ], + ) + def test_initialization(self, kwargs, expected): + config = OliveHfQuantizationOverrideConfig(**kwargs) + + assert (config.bits, config.symmetric, config.group_size) == expected class TestOliveHfQuantizationConfig: - def test_basic_initialization(self): - """Test basic initialization with required parameters.""" - config = OliveHfQuantizationConfig(bits=4, symmetric=True, group_size=128) - assert config.bits == 4 - assert config.symmetric is True - assert config.group_size == 128 - assert config.lm_head is False - assert config.embeds is False - assert config.modules_to_not_convert is None - assert config.overrides == {} - assert config.quant_method == OliveHfQuantizationMethod.OLIVE + @pytest.mark.parametrize("full_config", [False, True], ids=["basic", "full"]) + def test_initialization(self, full_config): + kwargs = {"bits": 4, "symmetric": True, "group_size": 128} + if full_config: + kwargs.update( + { + "lm_head": True, + "embeds": True, + "modules_to_not_convert": ["layer3"], + "overrides": {"layer1": {"bits": 8}, "layer2": {"symmetric": False, "group_size": 64}}, + "tie_word_embeddings": True, + } + ) + config = OliveHfQuantizationConfig(**kwargs) - def test_full_initialization(self): - """Test initialization with all parameters.""" - overrides = {"layer1": {"bits": 8}, "layer2": {"symmetric": False, "group_size": 64}} - config = OliveHfQuantizationConfig( - bits=4, - symmetric=True, - group_size=128, - lm_head=True, - embeds=True, - modules_to_not_convert=["layer3"], - overrides=overrides, - tie_word_embeddings=True, - ) assert config.bits == 4 assert config.symmetric is True assert config.group_size == 128 - assert config.lm_head is True - assert config.embeds is True - assert config.modules_to_not_convert == ["layer3"] - assert len(config.overrides) == 2 - assert isinstance(config.overrides["layer1"], OliveHfQuantizationOverrideConfig) - assert config.overrides["layer1"].bits == 8 - assert config.tie_word_embeddings is True + assert config.lm_head is full_config + assert config.embeds is full_config + assert config.modules_to_not_convert == (["layer3"] if full_config else None) + assert len(config.overrides) == (2 if full_config else 0) + assert config.tie_word_embeddings is full_config + assert config.quant_method == OliveHfQuantizationMethod.OLIVE + if full_config: + assert isinstance(config.overrides["layer1"], OliveHfQuantizationOverrideConfig) + assert config.overrides["layer1"].bits == 8 def test_invalid_bits(self): """Test that invalid bits raise ValueError.""" diff --git a/test/common/quant/test_utils.py b/test/common/quant/test_utils.py index 8e9ea928b1..73cf5bab9e 100644 --- a/test/common/quant/test_utils.py +++ b/test/common/quant/test_utils.py @@ -11,33 +11,21 @@ class TestGetMaxqMinq: - def test_unsigned_4bit(self): - """Test 4-bit unsigned quantization range.""" - maxq, minq = get_maxq_minq(bits=4, signed=False) - assert minq == 0 - assert maxq == 15 - assert maxq - minq + 1 == 16 - - def test_unsigned_8bit(self): - """Test 8-bit unsigned quantization range.""" - maxq, minq = get_maxq_minq(bits=8, signed=False) - assert minq == 0 - assert maxq == 255 - assert maxq - minq + 1 == 256 - - def test_signed_4bit(self): - """Test 4-bit signed quantization range.""" - maxq, minq = get_maxq_minq(bits=4, signed=True) - assert minq == -8 - assert maxq == 7 - assert maxq - minq + 1 == 16 - - def test_signed_8bit(self): - """Test 8-bit signed quantization range.""" - maxq, minq = get_maxq_minq(bits=8, signed=True) - assert minq == -128 - assert maxq == 127 - assert maxq - minq + 1 == 256 + @pytest.mark.parametrize( + ("bits", "signed", "expected_maxq", "expected_minq"), + [ + pytest.param(4, False, 15, 0, id="unsigned-4bit"), + pytest.param(8, False, 255, 0, id="unsigned-8bit"), + pytest.param(4, True, 7, -8, id="signed-4bit"), + pytest.param(8, True, 127, -128, id="signed-8bit"), + ], + ) + def test_quantization_range(self, bits, signed, expected_maxq, expected_minq): + maxq, minq = get_maxq_minq(bits=bits, signed=signed) + + assert minq == expected_minq + assert maxq == expected_maxq + assert maxq - minq + 1 == 2**bits class TestWeightQuantizer: @@ -322,18 +310,6 @@ def test_pack_with_padding(self, bits): # Should still work correctly assert torch.all(tensor == unpacked) - def test_unpack_shape_trimming(self): - """Test that unpacking trims to the correct shape.""" - bits = 4 - original_shape = (8, 17) # Not divisible by packing factor - tensor = torch.randint(0, 2**bits, original_shape, dtype=torch.uint8) - - packed = pack_to_uint8(tensor, bits) - unpacked = unpack_from_uint8(packed, bits, original_shape) - - assert unpacked.shape == original_shape - assert torch.all(tensor == unpacked) - @pytest.mark.parametrize("bits", [2, 4, 8]) def test_pack_all_zeros(self, bits): """Test packing tensor with all zeros.""" diff --git a/test/engine/packaging/test_packaging_generator.py b/test/engine/packaging/test_packaging_generator.py index 7ccc850560..09dcde1cda 100644 --- a/test/engine/packaging/test_packaging_generator.py +++ b/test/engine/packaging/test_packaging_generator.py @@ -6,194 +6,102 @@ import shutil import zipfile from pathlib import Path -from unittest.mock import patch -import mlflow import onnx import pytest -from olive.engine import Engine -from olive.engine.footprint import Footprint, FootprintNode +from olive.engine.footprint import Footprint, FootprintNode, FootprintNodeMetric from olive.engine.output import WorkflowOutput from olive.engine.packaging.packaging_config import ( PackagingConfig, PackagingType, + ZipfilePackagingConfig, ) from olive.engine.packaging.packaging_generator import generate_output_artifacts -from olive.evaluator.metric import AccuracySubType -from olive.evaluator.olive_evaluator import OliveEvaluatorConfig +from olive.evaluator.metric_result import MetricResult, SubMetricResult from olive.hardware import DEFAULT_CPU_ACCELERATOR -from olive.passes.onnx.conversion import OnnxConversion -from test.utils import get_accuracy_metric, get_pytorch_model_config +from test.utils import ONNX_MODEL_PATH -# TODO(team): no engine API envolved, use generate_output_artifacts API directly -@patch("onnx.external_data_helper.sys.getsizeof") @pytest.mark.parametrize( - ("save_as_external_data", "mocked_size_value"), - [(True, 2048), (False, 100)], + "save_as_external_data", + [ + pytest.param(False, id="inline-onnx-with-metrics"), + pytest.param(True, id="external-data-onnx-with-metrics"), + ], ) -def test_generate_zipfile_artifacts(mock_sys_getsizeof, save_as_external_data, mocked_size_value, tmp_path): - # setup - # onnx will save with external data when tensor size is greater than 1024(default threshold) - mock_sys_getsizeof.return_value = mocked_size_value - metric = get_accuracy_metric(AccuracySubType.ACCURACY_SCORE) - evaluator_config = OliveEvaluatorConfig(metrics=[metric]) - options = { - "cache_config": { - "cache_dir": tmp_path, - "clean_cache": True, - "clean_evaluation_cache": True, - }, - "search_strategy": { - "execution_order": "joint", - "sampler": "random", - }, - "evaluator": evaluator_config, - } - engine = Engine(**options) - # Use TorchScript because dynamo export creates models with strict input shape requirements - # that don't match the dummy data used for evaluation - engine.register(OnnxConversion, {"save_as_external_data": save_as_external_data, "use_dynamo_exporter": False}) - - input_model_config = get_pytorch_model_config() - - packaging_config = PackagingConfig() - packaging_config.type = PackagingType.Zipfile - packaging_config.name = "OutputModels" - +def test_generate_zipfile_artifacts_with_metrics(tmp_path, save_as_external_data): + packaging_config = PackagingConfig(type=PackagingType.Zipfile, name="OutputModels") + workflow_output = create_workflow_output( + create_test_model_config(tmp_path, save_as_external_data=save_as_external_data), + include_metrics=True, + ) output_dir = tmp_path / "outputs" + output_dir.mkdir() - # execute - engine.run( - input_model_config=input_model_config, - accelerator_spec=DEFAULT_CPU_ACCELERATOR, - packaging_config=packaging_config, - output_dir=output_dir, - ) + generate_output_artifacts(packaging_config, workflow_output, output_dir) - # assert artifacts_path = output_dir / "OutputModels.zip" assert artifacts_path.exists() with zipfile.ZipFile(artifacts_path) as zip_ref: zip_ref.extractall(output_dir) + verify_output_artifacts(output_dir) models_rank_path = output_dir / "models_rank.json" verify_models_rank_json_file(output_dir, models_rank_path, save_as_external_data=save_as_external_data) - # contain the evaluation result candidate_model_path = output_dir / "CandidateModels" / "cpu-cpu" / "BestCandidateModel_1" if save_as_external_data: assert (candidate_model_path / "model.onnx.data").exists() - try: - model_path = candidate_model_path / "model.onnx" - onnx.load(str(model_path)) - except Exception as e: - pytest.fail(f"Failed to load the model: {e}") + assert_onnx_loads(candidate_model_path / "model.onnx") metrics_file = candidate_model_path / "metrics.json" with metrics_file.open() as f: metrics = json.load(f) - assert "input_model_metrics" in metrics - assert "candidate_model_metrics" in metrics - - # clean up - shutil.rmtree(output_dir) - + assert "input_model_metrics" in metrics + assert "candidate_model_metrics" in metrics -# TODO(team): no engine API envolved, use generate_output_artifacts API directly -def test_generate_zipfile_artifacts_no_search(tmp_path): - # setup - options = { - "cache_config": { - "cache_dir": tmp_path, - "clean_cache": True, - "clean_evaluation_cache": True, - }, - } - engine = Engine(**options) - engine.register(OnnxConversion, {"use_dynamo_exporter": True}) - - input_model_config = get_pytorch_model_config() - - packaging_config = PackagingConfig() - packaging_config.type = PackagingType.Zipfile - packaging_config.name = "OutputModels" - output_dir = tmp_path / "outputs" - - # execute - engine.run( - input_model_config=input_model_config, - accelerator_spec=DEFAULT_CPU_ACCELERATOR, - packaging_config=packaging_config, - output_dir=output_dir, - evaluate_input_model=False, +@pytest.mark.parametrize( + "export_in_mlflow_format", + [ + pytest.param(False, id="zip-without-metrics-or-search"), + pytest.param(True, id="mlflow-without-metrics"), + ], +) +def test_generate_zipfile_artifacts_without_metrics(tmp_path, export_in_mlflow_format): + packaging_config = PackagingConfig( + type=PackagingType.Zipfile, + name="OutputModels", + config=ZipfilePackagingConfig(export_in_mlflow_format=export_in_mlflow_format), ) - - # assert - artifacts_path = output_dir / "OutputModels.zip" - assert artifacts_path.exists() - with zipfile.ZipFile(artifacts_path) as zip_ref: - zip_ref.extractall(output_dir) - verify_output_artifacts(output_dir) - models_rank_path = output_dir / "models_rank.json" - verify_models_rank_json_file(output_dir, models_rank_path) - - # clean up - shutil.rmtree(output_dir) - - -# TODO(team): no engine API envolved, use generate_output_artifacts API directly -def test_generate_zipfile_artifacts_mlflow(tmp_path): - # setup - options = { - "cache_config": { - "cache_dir": tmp_path, - "clean_cache": True, - "clean_evaluation_cache": True, - }, - } - engine = Engine(**options) - engine.register(OnnxConversion, {"use_dynamo_exporter": True}) - - input_model_config = get_pytorch_model_config() - - packaging_config = PackagingConfig() - packaging_config.type = PackagingType.Zipfile - packaging_config.name = "OutputModels" - # Initialize config with ZipfilePackagingConfig if it's None - if packaging_config.config is None: - from olive.engine.packaging.packaging_config import ZipfilePackagingConfig - - packaging_config.config = ZipfilePackagingConfig() - packaging_config.config.export_in_mlflow_format = True - + workflow_output = create_workflow_output(create_test_model_config(tmp_path), include_metrics=False) output_dir = tmp_path / "outputs" + output_dir.mkdir() - # execute - engine.run( - input_model_config=input_model_config, - accelerator_spec=DEFAULT_CPU_ACCELERATOR, - packaging_config=packaging_config, - output_dir=output_dir, - evaluate_input_model=False, - ) + generate_output_artifacts(packaging_config, workflow_output, output_dir) - # assert artifacts_path = output_dir / "OutputModels.zip" assert artifacts_path.exists() with zipfile.ZipFile(artifacts_path) as zip_ref: zip_ref.extractall(output_dir) + verify_output_artifacts(output_dir) models_rank_path = output_dir / "models_rank.json" - verify_models_rank_json_file(output_dir, models_rank_path, export_in_mlflow_format=True) - assert (output_dir / "CandidateModels" / "cpu-cpu" / "BestCandidateModel_1" / "mlflow_model").exists() + verify_models_rank_json_file( + output_dir, + models_rank_path, + export_in_mlflow_format=export_in_mlflow_format, + ) - # clean up - shutil.rmtree(output_dir) - if Path("mlruns").exists(): - shutil.rmtree("mlruns") + candidate_model_path = output_dir / "CandidateModels" / "cpu-cpu" / "BestCandidateModel_1" + assert not (candidate_model_path / "metrics.json").exists() + if export_in_mlflow_format: + assert (candidate_model_path / "mlflow_model").exists() + if Path("mlruns").exists(): + shutil.rmtree("mlruns") + else: + assert_onnx_loads(candidate_model_path / "model.onnx") def test_generate_zipfile_artifacts_no_output_models(tmp_path): @@ -235,6 +143,63 @@ def test__package_dockerfile(tmp_path): assert dockerfile_path.exists() +def create_test_model_config(tmp_path, save_as_external_data=False): + if not save_as_external_data: + return {"type": "ONNXModel", "config": {"model_path": str(ONNX_MODEL_PATH)}} + + model_proto = onnx.load(ONNX_MODEL_PATH) + model_dir = tmp_path / "external_data_model" + model_dir.mkdir() + onnx.save_model( + model_proto, + model_dir / "model.onnx", + save_as_external_data=True, + all_tensors_to_one_file=True, + location="model.onnx.data", + size_threshold=0, + ) + return { + "type": "ONNXModel", + "config": {"model_path": str(model_dir), "onnx_file_name": "model.onnx"}, + } + + +def create_workflow_output(model_config, include_metrics): + input_model_id = "input_model" + output_model_id = "candidate_model" + input_node = FootprintNode( + model_id=input_model_id, + model_config={"type": "ONNXModel", "config": {"model_path": str(ONNX_MODEL_PATH)}}, + metrics=create_accuracy_metrics(0.80) if include_metrics else None, + pass_run_config={"type": "input_model"}, + ) + output_node = FootprintNode( + model_id=output_model_id, + parent_model_id=input_model_id, + model_config=model_config, + metrics=create_accuracy_metrics(0.90) if include_metrics else None, + from_pass="test_pass", + pass_run_config={"type": "test_pass"}, + is_pareto_frontier=True, + ) + objective_dict = {"accuracy": {"higher_is_better": True, "goal": None, "priority": 1}} if include_metrics else {} + footprint = Footprint( + nodes={input_model_id: input_node, output_model_id: output_node}, + objective_dict=objective_dict, + is_marked_pareto_frontier=True, + ) + footprint.input_model_id = input_model_id + footprint.output_model_ids = [output_model_id] + return WorkflowOutput(DEFAULT_CPU_ACCELERATOR, footprint) + + +def create_accuracy_metrics(value): + return FootprintNodeMetric( + value=MetricResult(root={"accuracy": SubMetricResult(value=value, priority=1, higher_is_better=True)}), + cmp_direction={"accuracy": 1}, + ) + + def get_footprint(model_id, model_path): model_config = {"config": {"model_path": model_path}, "type": "ONNXModel"} footprint_node = FootprintNode(model_id=model_id, is_pareto_frontier=True, model_config=model_config) @@ -249,8 +214,15 @@ def verify_output_artifacts(output_dir): assert (output_dir / "models_rank.json").exists() +def assert_onnx_loads(model_path): + try: + onnx.load(str(model_path)) + except Exception as e: + pytest.fail(f"Failed to load the model: {e}") + + def verify_models_rank_json_file(output_dir, file_path, save_as_external_data=False, export_in_mlflow_format=False): - with Path.open(file_path) as file: + with file_path.open() as file: data = json.load(file) assert data is not None @@ -259,12 +231,12 @@ def verify_models_rank_json_file(output_dir, file_path, save_as_external_data=Fa model_path = output_dir / Path(model_data["model_config"]["config"]["model_path"]) assert model_path.exists(), "Model path in model rank file does not exist." if export_in_mlflow_format: + import mlflow + assert mlflow.onnx.load_model(str(model_path)), ( "Model path in model rank file is not a valid MLflow model path." ) elif save_as_external_data: - assert onnx.load(str(model_path / "model.onnx")), ( - "With external data, model path in model rank file is not a valid ONNX model path." - ) + assert_onnx_loads(model_path / "model.onnx") else: - assert onnx.load(str(model_path)), "Model path in model rank file is not a valid ONNX model path." + assert_onnx_loads(model_path) diff --git a/test/model/test_onnx_model.py b/test/model/test_onnx_model.py index c6dc59b034..efd7850f89 100644 --- a/test/model/test_onnx_model.py +++ b/test/model/test_onnx_model.py @@ -8,9 +8,21 @@ from olive.exception import OliveEvaluationError from olive.hardware.accelerator import Device from olive.model import ONNXModelHandler +from olive.model.config.model_config import ModelConfig from test.utils import get_onnx_model +def test_model_identifier_is_consistent(): + input_model = get_onnx_model() + + model_config = ModelConfig.model_validate(input_model.to_json()) + model_identifier = model_config.get_model_identifier() + repeated_identifier = ModelConfig.model_validate(input_model.to_json()).get_model_identifier() + + assert model_identifier == repeated_identifier + assert len(model_identifier) == 64 + + @patch("onnxruntime.InferenceSession") @patch("onnxruntime.get_available_providers") def test_model_prepare_session(get_available_providers_mock, inference_session_mock): diff --git a/test/passes/onnx/test_add_metadata.py b/test/passes/onnx/test_add_metadata.py index f4db0e290b..e6f75d210e 100644 --- a/test/passes/onnx/test_add_metadata.py +++ b/test/passes/onnx/test_add_metadata.py @@ -272,50 +272,6 @@ def test_add_metadata_with_model_hashes(self, tmp_path): assert metadata_dict["test_key"] == "test_value" assert "olive_version" in metadata_dict - def test_add_metadata_no_custom_config(self, tmp_path): - """Test that hashes are always included even with minimal config.""" - # Setup - input_model = get_onnx_model() - config = { - "graph_name": "test_graph_minimal", - } - p = create_pass_from_dict(AddOliveMetadata, config, disable_search=True) - output_folder = str(tmp_path / "onnx") - - # Execute - output_model = p.run(input_model, output_folder) - - # Assert - onnx_model = onnx.load_model(output_model.model_path) - metadata_dict = {entry.key: entry.value for entry in onnx_model.metadata_props} - - # Verify hash fields are always present even with minimal config - assert "model_hash" in metadata_dict - - # Verify olive_version is always present - assert "olive_version" in metadata_dict - - def test_add_metadata_always_includes_hashes(self, tmp_path): - """Test that hashes are always included by default.""" - # Setup - input_model = get_onnx_model() - config = { - "graph_name": "test_graph_default_hashes" - # No hash-related config specified - should be included by default - } - p = create_pass_from_dict(AddOliveMetadata, config, disable_search=True) - output_folder = str(tmp_path / "onnx") - - # Execute - output_model = p.run(input_model, output_folder) - - # Assert - onnx_model = onnx.load_model(output_model.model_path) - metadata_dict = {entry.key: entry.value for entry in onnx_model.metadata_props} - - # Verify hash fields are always present (default behavior) - assert "model_hash" in metadata_dict - @patch("olive.passes.onnx.add_metadata.logger") def test_add_metadata_hash_calculation_error(self, mock_logger, tmp_path): """Test handling of hash calculation errors.""" @@ -344,25 +300,6 @@ def test_add_metadata_hash_calculation_error(self, mock_logger, tmp_path): # Verify warning was logged mock_logger.warning.assert_called() - def test_calculate_model_hash_consistency(self, tmp_path): - """Test that hash calculation is consistent for the same model.""" - # Setup - input_model = get_onnx_model() - - # Calculate hash twice for same model using the actual implementation approach - from olive.model.config.model_config import ModelConfig - - model_config1 = ModelConfig.model_validate(input_model.to_json()) - hash1 = model_config1.get_model_identifier() - - model_config2 = ModelConfig.model_validate(input_model.to_json()) - hash2 = model_config2.get_model_identifier() - - # Hashes should be identical - assert hash1 == hash2 - assert len(hash1) == 64 # SHA256 length - assert isinstance(hash1, str) - def test_add_metadata_with_hf_model_name(self, tmp_path): """Test that HF model name is automatically included when model type is HfModel.""" # Setup - Mock model.to_json() to return HfModel config @@ -398,150 +335,50 @@ def test_add_metadata_with_hf_model_name(self, tmp_path): assert "hf_model_name" in metadata_dict assert metadata_dict["hf_model_name"] == "microsoft/Phi-3.5-mini-instruct" - def test_add_metadata_without_hf_model_name(self, tmp_path): + @pytest.mark.parametrize( + "model_config", + [ + pytest.param({"type": "ONNXModel", "model_path": "/path/to/model.onnx"}, id="onnx-model"), + pytest.param({"type": "PyTorchModel", "model_path": "microsoft/Phi-3.5-mini-instruct"}, id="pytorch-model"), + ], + ) + def test_add_metadata_without_hf_model_name(self, tmp_path, model_config): """Test that non-HF models don't include HF model name.""" - # Setup - Mock model.to_json() to return non-HfModel config input_model = get_onnx_model() - # Patch the to_json method to return ONNXModel configuration - with patch.object( - input_model, "to_json", return_value={"config": {"type": "ONNXModel", "model_path": "/path/to/model.onnx"}} - ): + with patch.object(input_model, "to_json", return_value={"config": model_config}): config = {"graph_name": "non_hf_model_graph"} p = create_pass_from_dict(AddOliveMetadata, config, disable_search=True) output_folder = str(tmp_path / "onnx") - # Execute output_model = p.run(input_model, output_folder) - # Assert onnx_model = onnx.load_model(output_model.model_path) metadata_dict = {entry.key: entry.value for entry in onnx_model.metadata_props} - # Verify HF model name is not included for non-HF models assert "hf_model_name" not in metadata_dict - def test_add_metadata_hf_model_with_pytorch_type(self, tmp_path): - """Test that PyTorchModel with HF model path doesn't include HF model name.""" - # Setup - Mock model.to_json() to return PyTorchModel config with HF model path + @pytest.mark.parametrize( + "model_config", + [ + pytest.param({"type": "HfModel", "task": "text-generation"}, id="missing"), + pytest.param({"type": "HfModel", "model_path": "", "task": "text-generation"}, id="empty"), + pytest.param({"type": "HfModel", "model_path": 12345, "task": "text-generation"}, id="non-string"), + ], + ) + def test_add_metadata_hf_model_invalid_model_path(self, tmp_path, model_config): input_model = get_onnx_model() - # Patch the to_json method to return PyTorchModel configuration - with patch.object( - input_model, - "to_json", - return_value={ - "config": { - "type": "PyTorchModel", - "model_path": "microsoft/Phi-3.5-mini-instruct", # HF path but wrong type - } - }, - ): - config = {"graph_name": "pytorch_with_hf_path_graph"} + with patch.object(input_model, "to_json", return_value={"config": model_config}): + config = {"graph_name": "hf_model_invalid_path_graph"} p = create_pass_from_dict(AddOliveMetadata, config, disable_search=True) output_folder = str(tmp_path / "onnx") - # Execute output_model = p.run(input_model, output_folder) - # Assert - onnx_model = onnx.load_model(output_model.model_path) - metadata_dict = {entry.key: entry.value for entry in onnx_model.metadata_props} - - # Verify HF model name is not included since type is not HfModel - assert "hf_model_name" not in metadata_dict - - def test_add_metadata_hf_model_missing_model_path(self, tmp_path): - """Test that HfModel without model_path doesn't include HF model name.""" - # Setup - Mock model.to_json() to return HfModel config without model_path - input_model = get_onnx_model() - - # Patch the to_json method to return HfModel configuration without model_path - with patch.object( - input_model, - "to_json", - return_value={ - "config": { - "type": "HfModel", - "task": "text-generation", - # No model_path field - } - }, - ): - config = {"graph_name": "hf_model_no_path_graph"} - p = create_pass_from_dict(AddOliveMetadata, config, disable_search=True) - output_folder = str(tmp_path / "onnx") - - # Execute - output_model = p.run(input_model, output_folder) - - # Assert - onnx_model = onnx.load_model(output_model.model_path) - metadata_dict = {entry.key: entry.value for entry in onnx_model.metadata_props} - - # Verify HF model name is not included since model_path is missing - assert "hf_model_name" not in metadata_dict - - def test_add_metadata_hf_model_empty_model_path(self, tmp_path): - """Test that HfModel with empty model_path doesn't include HF model name.""" - # Setup - Mock model.to_json() to return HfModel config with empty model_path - input_model = get_onnx_model() - - # Patch the to_json method to return HfModel configuration with empty model_path - with patch.object( - input_model, - "to_json", - return_value={ - "config": { - "type": "HfModel", - "model_path": "", # Empty string - "task": "text-generation", - } - }, - ): - config = {"graph_name": "hf_model_empty_path_graph"} - p = create_pass_from_dict(AddOliveMetadata, config, disable_search=True) - output_folder = str(tmp_path / "onnx") - - # Execute - output_model = p.run(input_model, output_folder) - - # Assert - onnx_model = onnx.load_model(output_model.model_path) - metadata_dict = {entry.key: entry.value for entry in onnx_model.metadata_props} - - # Verify HF model name is not included since model_path is empty - assert "hf_model_name" not in metadata_dict - - def test_add_metadata_hf_model_non_string_model_path(self, tmp_path): - """Test that HfModel with non-string model_path doesn't include HF model name.""" - # Setup - Mock model.to_json() to return HfModel config with non-string model_path - input_model = get_onnx_model() - - # Patch the to_json method to return HfModel configuration with non-string model_path - with patch.object( - input_model, - "to_json", - return_value={ - "config": { - "type": "HfModel", - "model_path": 12345, # Non-string type - "task": "text-generation", - } - }, - ): - config = {"graph_name": "hf_model_non_string_path_graph"} - p = create_pass_from_dict(AddOliveMetadata, config, disable_search=True) - output_folder = str(tmp_path / "onnx") - - # Execute - output_model = p.run(input_model, output_folder) - - # Assert onnx_model = onnx.load_model(output_model.model_path) metadata_dict = {entry.key: entry.value for entry in onnx_model.metadata_props} - # Verify HF model name is not included since model_path is not a string assert "hf_model_name" not in metadata_dict def test_add_metadata_hf_model_no_config(self, tmp_path): diff --git a/test/passes/onnx/test_common.py b/test/passes/onnx/test_common.py index 9498cd9c39..a100884a05 100644 --- a/test/passes/onnx/test_common.py +++ b/test/passes/onnx/test_common.py @@ -7,14 +7,13 @@ import pytest from olive.common.utils import is_hardlink -from olive.passes.olive_pass import create_pass_from_dict +from olive.model import ONNXModelHandler from olive.passes.onnx.common import ( add_version_metadata_to_model_proto, model_proto_to_olive_model, resave_model, ) -from olive.passes.onnx.conversion import OnnxConversion -from test.utils import ONNX_MODEL_PATH, get_hf_model +from test.utils import ONNX_MODEL_PATH @pytest.mark.parametrize( @@ -22,13 +21,6 @@ [ {}, {"save_as_external_data": True}, - { - "save_as_external_data": False, - "all_tensors_to_one_file": True, - "external_data_name": None, - "size_threshold": 1024, - "convert_attribute": False, - }, ], ) def test_model_proto_to_olive_model(external_data_config, tmp_path): @@ -39,14 +31,23 @@ def test_model_proto_to_olive_model(external_data_config, tmp_path): @pytest.mark.parametrize("has_external_data", [True, False]) def test_resave_model(has_external_data, tmp_path): - # setup - from transformers.cache_utils import DynamicLayer - - original_lazy_initialization = DynamicLayer.lazy_initialization - input_model = create_pass_from_dict( - OnnxConversion, {"save_as_external_data": has_external_data, "use_dynamo_exporter": True}, disable_search=True - ).run(get_hf_model(), str(tmp_path / "input")) - assert DynamicLayer.lazy_initialization is original_lazy_initialization + input_dir = tmp_path / "input" + input_dir.mkdir() + input_path = input_dir / "input.onnx" + model_proto = onnx.load(ONNX_MODEL_PATH) + if has_external_data: + onnx.save_model( + model_proto, + input_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location="input.onnx.data", + size_threshold=0, + convert_attribute=True, + ) + else: + onnx.save(model_proto, input_path) + input_model = ONNXModelHandler(input_path) # execute resave_path = tmp_path / "resave" / "resave.onnx" diff --git a/test/passes/onnx/test_extract_adapters.py b/test/passes/onnx/test_extract_adapters.py index 20b98fbff3..3cd3888751 100644 --- a/test/passes/onnx/test_extract_adapters.py +++ b/test/passes/onnx/test_extract_adapters.py @@ -72,16 +72,18 @@ def input_model_info_fixture(tmp_path_factory, request): } +@pytest.mark.parametrize("adapter_type", [AdapterType.LORA, AdapterType.DORA, AdapterType.LOHA]) +def test_model_without_adapters(adapter_type): + assert not model_has_adapters(get_onnx_model().model_path, adapter_type) + + @pytest.mark.parametrize("input_model_info", [AdapterType.LORA, AdapterType.DORA, AdapterType.LOHA], indirect=True) -@pytest.mark.parametrize("model_type", [None, "float", "int4"]) +@pytest.mark.parametrize("model_type", ["float", "int4"]) def test_model_has_adapters(input_model_info, model_type): model_info = input_model_info adapter_type = model_info["adapter_type"] - if model_type is None: - assert not model_has_adapters(get_onnx_model().model_path, adapter_type) - else: - assert model_has_adapters(model_info[model_type]["onnx_model"].model_path, adapter_type) + assert model_has_adapters(model_info[model_type]["onnx_model"].model_path, adapter_type) @pytest.mark.parametrize("input_model_info", [AdapterType.LORA], indirect=True) diff --git a/test/passes/onnx/test_mnb_to_qdq.py b/test/passes/onnx/test_mnb_to_qdq.py index 299e5cf1a5..66d43562de 100644 --- a/test/passes/onnx/test_mnb_to_qdq.py +++ b/test/passes/onnx/test_mnb_to_qdq.py @@ -19,6 +19,7 @@ @pytest.fixture( + scope="module", params=[ pytest.param( (True, 2), marks=pytest.mark.skipif(SKIP_2BIT, reason="2-bit not supported in this version of ONNX Runtime") @@ -35,8 +36,9 @@ ids=["symmetric-2bit", "asymmetric-2bit", "symmetric-4bit", "asymmetric-4bit", "symmetric-8bit", "asymmetric-8bit"], name="create_mnb_model", ) -def create_mnb_model_fixture(request, tmp_path): +def create_mnb_model_fixture(request, tmp_path_factory): symmetric, bits = request.param + tmp_path = tmp_path_factory.mktemp(f"mnb-{bits}bit-{'symmetric' if symmetric else 'asymmetric'}") if version.parse("1.22.0") > ORT_VERSION: if bits == 8: pytest.skip("MatMulNBitsQuantizer doesn't support 8 bits in this version of ONNX Runtime") @@ -100,6 +102,7 @@ def forward(self, x): quant.process() onnx.save(quant.model.model, mnb_path) + # The generated model is read-only and shared by all 16 pass configurations. return mnb_path, in_dim, symmetric, bits diff --git a/test/passes/onnx/test_mobius_model_builder.py b/test/passes/onnx/test_mobius_model_builder.py index 92a925a47a..21cc8ee615 100644 --- a/test/passes/onnx/test_mobius_model_builder.py +++ b/test/passes/onnx/test_mobius_model_builder.py @@ -241,19 +241,6 @@ def test_single_component_returns_onnx_handler(tmp_path): assert call_kwargs["dtype"] == "f32" -def test_model_onnx_exists_after_run(tmp_path): - """The saved model.onnx file must exist on disk.""" - out = tmp_path / "out" - pkg = _fake_pkg(["model"], out) - - with _patch_build(pkg): - p = _make_pass() - result = p.run(_make_hf_model("org/model"), out) - - # ONNXModelHandler.model_path already points to the .onnx file - assert Path(result.model_path).exists() - - def test_genai_artifacts_in_single_component(tmp_path): """ORT GenAI artifacts must be included in single-component model's additional_files.""" out = tmp_path / "out" diff --git a/test/passes/onnx/test_model_builder.py b/test/passes/onnx/test_model_builder.py index 3de060b106..57889d1e4c 100644 --- a/test/passes/onnx/test_model_builder.py +++ b/test/passes/onnx/test_model_builder.py @@ -95,19 +95,24 @@ def test_model_builder_olive_quant(tmp_path, embeds, group_size): assert Path(output_folder / "genai_config.json").exists() -@pytest.mark.parametrize("layer_annotations", [True, False]) -def test_model_builder_layer_annotations(tmp_path, layer_annotations): +def test_model_builder_layer_annotations(tmp_path, monkeypatch): """Test that layer annotations are correctly applied to the output ONNX model.""" - input_model = make_local_tiny_llama(tmp_path / "input_model", "hf") - - if layer_annotations: - # Create layer annotations to be applied - # Keys are layer names, values are lists of node-name substrings to match - annotations = { - "embedding_layer": ["embed_tokens"], - "norm_layer": ["norm"], - } - input_model.model_attributes = {"layer_annotations": annotations} + + def fake_create_model( + model_name, input_path, output_dir, precision, execution_provider, cache_dir, filename, **kwargs + ): + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + _create_test_onnx_model(output_dir / "model.onnx", "model.layers.0.embed_tokens") + (output_dir / "genai_config.json").write_text("{}") + + _mock_genai_builder(monkeypatch, fake_create_model) + input_model = Mock(spec=HfModelHandler) + input_model.model_name_or_path = "dummy-model" + input_model.adapter_path = None + input_model.test_model_config = None + input_model.test_model_path = None + input_model.model_attributes = {"layer_annotations": {"embedding_layer": ["embed_tokens"]}} p = create_pass_from_dict( ModelBuilder, @@ -123,13 +128,9 @@ def test_model_builder_layer_annotations(tmp_path, layer_annotations): assert isinstance(output_model, ONNXModelHandler) assert Path(output_model.model_path).exists() - if layer_annotations: - # Verify that metadata properties were applied to nodes - model_proto = onnx.load(output_model.model_path, load_external_data=False) - node_names_with_metadata = {node.name for node in model_proto.graph.node if node.metadata_props} - assert len(node_names_with_metadata) > 0, ( - "Expected nodes with metadata_props when layer_annotations are provided" - ) + model_proto = onnx.load(output_model.model_path, load_external_data=False) + annotated_nodes = {node.name for node in model_proto.graph.node if node.metadata_props} + assert annotated_nodes == {"model.layers.0.embed_tokens"} def test_model_builder_uses_saved_test_model_path(tmp_path): diff --git a/test/passes/onnx/test_peephole_optimizer.py b/test/passes/onnx/test_peephole_optimizer.py index 7d54b69094..c726853e95 100644 --- a/test/passes/onnx/test_peephole_optimizer.py +++ b/test/passes/onnx/test_peephole_optimizer.py @@ -14,7 +14,6 @@ from olive.hardware import DEFAULT_CPU_ACCELERATOR from olive.model import ONNXModelHandler from olive.passes.olive_pass import create_pass_from_dict -from olive.passes.onnx.common import model_proto_to_olive_model from olive.passes.onnx.peephole_optimizer import ModelOptimizer, OnnxPeepholeOptimizer from test.utils import get_onnx_model @@ -83,7 +82,9 @@ def test_onnx_peephole_optimizer_pass_fuse_reshape_operations(tmp_path, external opset_imports=opset_imports, ) - m = model_proto_to_olive_model(model, str(tmp_path / "input.onnx"), external_data_config) + input_model_path = tmp_path / "input.onnx" + onnx.save(model, input_model_path) + m = ONNXModelHandler(input_model_path) p = create_pass_from_dict( OnnxPeepholeOptimizer, external_data_config, disable_search=True, accelerator_spec=DEFAULT_CPU_ACCELERATOR ) @@ -109,7 +110,7 @@ def test_onnx_peephole_optimizer_pass_fuse_reshape_operations(tmp_path, external @patch("olive.passes.onnx.peephole_optimizer.model_proto_to_olive_model") @patch("onnxoptimizer.optimize") @patch("onnxscript.optimizer.optimize") -def test_onnxscript(mock_onnxscript, mock_onnxoptimizer, mock_model_proto_to_olive_model, tmp_path): +def test_optimizers(mock_onnxscript, mock_onnxoptimizer, mock_model_proto_to_olive_model, tmp_path): # setup input_model = get_onnx_model() p = create_pass_from_dict(OnnxPeepholeOptimizer, {}, disable_search=True) @@ -122,23 +123,6 @@ def test_onnxscript(mock_onnxscript, mock_onnxoptimizer, mock_model_proto_to_oli # assert mock_onnxscript.assert_called_once_with(input_model.load_model()) - - -@patch("olive.passes.onnx.peephole_optimizer.model_proto_to_olive_model") -@patch("onnxoptimizer.optimize") -@patch("onnxscript.optimizer.optimize") -def test_onnxoptimizer(mock_onnxscript, mock_onnxoptimizer, mock_model_proto_to_olive_model, tmp_path): - # setup - input_model = get_onnx_model() - p = create_pass_from_dict(OnnxPeepholeOptimizer, {}, disable_search=True) - mock_onnxscript.return_value = input_model.load_model() - mock_onnxoptimizer.return_value = input_model.load_model() - output_folder = str(tmp_path / "onnx") - - # execute - p.run(input_model, output_folder) - - # assert mock_onnxoptimizer.assert_called_once() diff --git a/test/passes/onnx/test_qairt_mha2sha.py b/test/passes/onnx/test_qairt_mha2sha.py index b3e113aba4..2f7f73130c 100644 --- a/test/passes/onnx/test_qairt_mha2sha.py +++ b/test/passes/onnx/test_qairt_mha2sha.py @@ -137,27 +137,6 @@ def test_mha2sha_for_onnx_model_handler(qairt_pass_instance, tmp_output_dir, moc loaded_qairt_instance.mock_export.assert_called_once_with(tmp_output_dir, prefix=input_model.onnx_file_name) -def test_mha2sha_v1_fallback(qairt_pass_instance, tmp_output_dir, mock_qairt_sdk_classes): - """Test that the pass falls back to mha2sha (v1) if v2 is not available.""" - dummy_model = get_onnx_model() - - original_side_effect_func = mock_qairt_sdk_classes.load.side_effect - - def custom_side_effect_func(model_path): - instance = original_side_effect_func(model_path) - instance.has_mha2sha_v2 = False - return instance - - mock_qairt_sdk_classes.load.side_effect = custom_side_effect_func - - _ = qairt_pass_instance.run(dummy_model, output_model_path=tmp_output_dir) - - # Verify V1 was called and V2 was not - loaded_qairt_instance = mock_qairt_sdk_classes.load.call_args - loaded_qairt_instance.mock_mha2sha.assert_called_once_with() - loaded_qairt_instance.mock_mha2sha_v2.assert_not_called() - - def test_mha2sha_kwargs_passed(tmp_output_dir, mock_qairt_sdk_classes): """Test that additional kwargs are passed to mha2sha_v2/mha2sha.""" dummy_model = get_onnx_model() @@ -177,8 +156,7 @@ def test_mha2sha_kwargs_passed(tmp_output_dir, mock_qairt_sdk_classes): loaded_qairt_instance_v2.mock_mha2sha_v2.reset_mock() loaded_qairt_instance_v2.mock_mha2sha.reset_mock() - # Configure the mock OnnxModel instance to NOT have mha2sha_v2 for this part - # We reuse the logic from test_mha2sha_v1_fallback + # Configure the mock OnnxModel instance to use the v1 fallback. original_side_effect_func = mock_qairt_sdk_classes.load.side_effect def custom_side_effect_func_v1(model_path): diff --git a/test/passes/onnx/test_transformer_optimization.py b/test/passes/onnx/test_transformer_optimization.py index e6ae228dd8..269bbbd045 100644 --- a/test/passes/onnx/test_transformer_optimization.py +++ b/test/passes/onnx/test_transformer_optimization.py @@ -54,85 +54,60 @@ def test_ort_transformer_optimization_pass(tmp_path): p.run(input_model, output_folder) -@pytest.mark.parametrize("use_gpu", [True, False]) -@pytest.mark.parametrize("fp16", [True, False]) @pytest.mark.parametrize( - "accelerator_spec", [DEFAULT_CPU_ACCELERATOR, DEFAULT_GPU_CUDA_ACCELERATOR, DEFAULT_GPU_TRT_ACCELERATOR] + ("use_gpu", "fp16", "accelerator_spec", "expected_message"), + [ + pytest.param( + False, True, DEFAULT_CPU_ACCELERATOR, "CPUExecutionProvider does not support float16", id="cpu-fp16" + ), + pytest.param(True, False, DEFAULT_CPU_ACCELERATOR, "CPUExecutionProvider does not support GPU", id="cpu-gpu"), + pytest.param( + False, True, DEFAULT_GPU_TRT_ACCELERATOR, "TensorRT has its own float16 implementation", id="trt-fp16" + ), + pytest.param(True, True, DEFAULT_GPU_CUDA_ACCELERATOR, None, id="cuda-valid"), + ], ) -@pytest.mark.parametrize("mock_inferece_session", [True, False]) -def test_invalid_ep_config(use_gpu, fp16, accelerator_spec, mock_inferece_session, tmp_path, caplog): - import onnxruntime as ort - from onnxruntime.transformers.onnx_model import OnnxModel - from packaging import version - - if accelerator_spec == DEFAULT_GPU_TRT_ACCELERATOR and not mock_inferece_session: - pytest.skip("Skipping test: TRT EP does not support compiled nodes when mock_inferece_session=False") - - input_model = get_onnx_model() +def test_invalid_ep_config(use_gpu, fp16, accelerator_spec, expected_message, caplog): config = {"model_type": "bert", "use_gpu": use_gpu, "float16": fp16} - # Use caplog.at_level to capture logs from the "olive" logger specifically, - # bypassing the propagate=False issue in olive/__init__.py with caplog.at_level(logging.INFO, logger="olive"): pass_config = OrtTransformersOptimization.generate_config(accelerator_spec, config, disable_search=True) p = OrtTransformersOptimization(accelerator_spec, pass_config, True) - is_pruned = not p.validate_config(pass_config, accelerator_spec) - - if accelerator_spec.execution_provider == "CPUExecutionProvider": - if fp16 and use_gpu: - assert is_pruned - assert "CPUExecutionProvider does not support float16" in caplog.text - elif use_gpu: - assert is_pruned - assert "CPUExecutionProvider does not support GPU inference" in caplog.text - - if accelerator_spec.execution_provider == "TensorrtExecutionProvider" and fp16: - assert is_pruned - assert "TensorRT has its own float16 implementation" in caplog.text - - if not is_pruned: - inference_session_mock_call_count = 0 - - def inference_session_init( - self, - path_or_bytes, - sess_options=None, - providers=None, - provider_options=None, - **kwargs, - ): - nonlocal inference_session_mock_call_count - inference_session_mock_call_count += 1 - shutil.copyfile(ONNX_MODEL_PATH, sess_options.optimized_model_filepath) - - with patch("onnxruntime.transformers.optimizer.optimize_by_fusion") as optimize_by_fusion_mock: - optimize_by_fusion_mock.return_value = OnnxModel(input_model.load_model()) - output_folder = str(tmp_path / "onnx") - if mock_inferece_session: - with patch.object(ort.InferenceSession, "__init__", new=inference_session_init): - p.run(input_model, output_folder) - else: - p.run(input_model, output_folder) - optimize_by_fusion_mock.assert_called() - - if accelerator_spec.execution_provider == "TensorrtExecutionProvider": - if accelerator_spec.execution_provider not in ort.get_available_providers(): - if use_gpu: - # the use_gpu will be ignored by optimize_model, please refef to the following links for more info. - # https://github.com/microsoft/onnxruntime/blob/v1.15.1/onnxruntime/python/tools/transformers/optimizer.py#L280 - if version.parse(ort.__version__) >= version.parse("1.16.0"): - # for TensorRT EP, the graph optimization will be skipped but the fusion will be applied. - with caplog.at_level(logging.INFO, logger="olive"): - assert "There is no gpu for onnxruntime to do optimization." in caplog.text - if mock_inferece_session: - assert inference_session_mock_call_count == 0 - else: - # for cpu graph optimization, the graph optimization will always be run. So there is not need check - if mock_inferece_session: - assert inference_session_mock_call_count > 0 - else: - if mock_inferece_session: - assert inference_session_mock_call_count > 0 + is_valid = p.validate_config(pass_config, accelerator_spec) + + assert is_valid is (expected_message is None) + if expected_message: + assert expected_message in caplog.text + + +def test_transformer_optimization_valid_cuda_config_runs(tmp_path): + import onnxruntime as ort + from onnxruntime.transformers.onnx_model import OnnxModel + + input_model = get_onnx_model() + config = {"model_type": "bert", "use_gpu": True, "float16": True} + pass_config = OrtTransformersOptimization.generate_config(DEFAULT_GPU_CUDA_ACCELERATOR, config, disable_search=True) + p = OrtTransformersOptimization(DEFAULT_GPU_CUDA_ACCELERATOR, pass_config, True) + + def inference_session_init( + self, + path_or_bytes, + sess_options=None, + providers=None, + provider_options=None, + **kwargs, + ): + shutil.copyfile(ONNX_MODEL_PATH, sess_options.optimized_model_filepath) + + with ( + patch("onnxruntime.transformers.optimizer.optimize_by_fusion") as optimize_by_fusion_mock, + patch.object(ort.InferenceSession, "__init__", new=inference_session_init), + ): + optimize_by_fusion_mock.return_value = OnnxModel(input_model.load_model()) + output_model = p.run(input_model, str(tmp_path / "onnx")) + + optimize_by_fusion_mock.assert_called_once() + assert output_model.model_path def test_transformer_optimization_invalid_model_type(tmp_path): diff --git a/test/passes/pytorch/test_lora.py b/test/passes/pytorch/test_lora.py index 11cf1effb3..fa43eff57c 100644 --- a/test/passes/pytorch/test_lora.py +++ b/test/passes/pytorch/test_lora.py @@ -67,7 +67,7 @@ def test_lora(tmp_path): # TODO(team): Failed in pipeline (linux gpu). Need to investigate. @pytest.mark.skipif( - platform.system() == OS.WINDOWS or not torch.cuda.is_available() or True, + platform.system() == OS.WINDOWS or not torch.cuda.is_available(), reason="bitsandbytes requires Linux GPU.", ) def test_qlora(tmp_path): @@ -81,7 +81,7 @@ def test_qlora(tmp_path): # TODO(team): Failed in pipeline (linux gpu). Need to investigate. @pytest.mark.skipif( - platform.system() == OS.WINDOWS or not torch.cuda.is_available() or True, + platform.system() == OS.WINDOWS or not torch.cuda.is_available(), reason="bitsandbytes requires Linux GPU.", ) def test_loftq(tmp_path): diff --git a/test/passes/pytorch/test_selective_mixed_precision.py b/test/passes/pytorch/test_selective_mixed_precision.py index 53f4493dfc..db7055aeef 100644 --- a/test/passes/pytorch/test_selective_mixed_precision.py +++ b/test/passes/pytorch/test_selective_mixed_precision.py @@ -288,23 +288,16 @@ def test_selective_mixed_precision_scored_run_co_promotes_qkv(tmp_path): ) -def test_selective_mixed_precision_scored_rejects_per_tensor_group_size_when_qkv_grouped(input_model, tmp_path): - """Per-tensor (group_size=0) is rejected when QKV groups exist because per-member aggregation isn't exact.""" - p = create_pass_from_dict( - SelectiveMixedPrecision, - {"algorithm": "iqe", "ratio": 0.5, "group_size": 0, "high_group_size": 16}, - disable_search=True, - ) - - with pytest.raises(ValueError, match="per-tensor"): - p.run(input_model, str(tmp_path)) - - -def test_selective_mixed_precision_scored_rejects_per_tensor_high_group_size_when_qkv_grouped(input_model, tmp_path): - """Per-tensor high precision (high_group_size=0) is also rejected for the same reason.""" +@pytest.mark.parametrize( + ("group_size", "high_group_size"), + [pytest.param(0, 16, id="base-precision"), pytest.param(16, 0, id="high-precision")], +) +def test_selective_mixed_precision_scored_rejects_per_tensor_when_qkv_grouped( + input_model, tmp_path, group_size, high_group_size +): p = create_pass_from_dict( SelectiveMixedPrecision, - {"algorithm": "iqe", "ratio": 0.5, "group_size": 16, "high_group_size": 0}, + {"algorithm": "iqe", "ratio": 0.5, "group_size": group_size, "high_group_size": high_group_size}, disable_search=True, ) @@ -742,26 +735,8 @@ def _kld_unit_scores(model, memory_mode, quantizer, high_quantizer, device="cpu" ) -def test_kld_strategy_low_memory_matches_legacy_grad_accum(monkeypatch): - """LOW_MEMORY KLD scoring is numerically equivalent to the legacy gradient-accumulation path.""" - data = get_kld_gradient_test_data() - patch_kld_calibration_data(monkeypatch, data) - quantizer, high_quantizer = get_kld_gradient_quantizers() - model = KldGradientTestModel() - - expected_numels, expected_scores = get_legacy_kld_scores( - deepcopy(model), data, quantizer, high_quantizer, device="cpu" - ) - actual_numels, actual_scores = _kld_unit_scores( - deepcopy(model), KldMemoryMode.LOW_MEMORY, quantizer, high_quantizer - ) - - assert actual_numels == expected_numels - assert_scores_close(actual_scores, expected_scores) - - -def test_kld_strategy_full_matches_legacy_grad_accum(monkeypatch): - """FULL KLD scoring is numerically equivalent to the legacy gradient-accumulation path.""" +@pytest.mark.parametrize("memory_mode", [KldMemoryMode.LOW_MEMORY, KldMemoryMode.FULL]) +def test_kld_strategy_matches_legacy_grad_accum(monkeypatch, memory_mode): data = get_kld_gradient_test_data() patch_kld_calibration_data(monkeypatch, data) quantizer, high_quantizer = get_kld_gradient_quantizers() @@ -770,7 +745,7 @@ def test_kld_strategy_full_matches_legacy_grad_accum(monkeypatch): expected_numels, expected_scores = get_legacy_kld_scores( deepcopy(model), data, quantizer, high_quantizer, device="cpu" ) - actual_numels, actual_scores = _kld_unit_scores(deepcopy(model), KldMemoryMode.FULL, quantizer, high_quantizer) + actual_numels, actual_scores = _kld_unit_scores(deepcopy(model), memory_mode, quantizer, high_quantizer) assert actual_numels == expected_numels assert_scores_close(actual_scores, expected_scores) diff --git a/test/passes/pytorch/test_slicegpt.py b/test/passes/pytorch/test_slicegpt.py index eb460b2f75..8f999408ee 100644 --- a/test/passes/pytorch/test_slicegpt.py +++ b/test/passes/pytorch/test_slicegpt.py @@ -13,9 +13,7 @@ # TODO(team): Failed in pipeline (linux gpu). Need to investigate. -@pytest.mark.skipif( - (sys.version_info < (3, 10) and not torch.cuda.is_available()) or True, reason="requires python3.10 or higher" -) +@pytest.mark.skipif(sys.version_info < (3, 10) or not torch.cuda.is_available(), reason="requires Python 3.10+ and GPU") def test_slicegpt(tmp_path): from olive.passes.pytorch.slicegpt import SliceGPT diff --git a/test/passes/qairt/test_preparation.py b/test/passes/qairt/test_preparation.py index 90a9db1509..b8aefd7a1a 100644 --- a/test/passes/qairt/test_preparation.py +++ b/test/passes/qairt/test_preparation.py @@ -4,6 +4,7 @@ # -------------------------------------------------------------------------- import json +from pathlib import Path from unittest.mock import MagicMock, Mock, patch import pytest @@ -308,6 +309,7 @@ def stderr_generator(): with ( patch("subprocess.Popen", return_value=mock_process), patch("tempfile.NamedTemporaryFile") as mock_temp, + patch("olive.passes.qairt.preparation.Path.unlink", autospec=True) as mock_unlink, ): temp_file_path = tmp_path / "olive_qairt_prep_test.json" @@ -325,8 +327,7 @@ def stderr_generator(): prep_pass.run(mock_hf_model, str(output_path)) - # Verify temp file would be cleaned up (unlink called) - # Note: In actual implementation, cleanup happens in finally block + mock_unlink.assert_called_once_with(Path(temp_file_path)) def test_preparation_uses_sys_executable_and_env(tmp_path, mock_hf_model, mock_qairt_modules): diff --git a/test/requirements-test-gpu.txt b/test/requirements-test-gpu.txt index 8946682889..9a4a9e8dc9 100644 --- a/test/requirements-test-gpu.txt +++ b/test/requirements-test-gpu.txt @@ -1,6 +1,7 @@ -r requirements-test.txt bitsandbytes onnxruntime-genai-cuda +slicegpt torch<2.11.0 # torch 2.11.0 is not compatible with CI v100 GPU triton # auto-gptq: installed from source in run_test.sh (no Python 3.12 wheels on PyPI) diff --git a/test/systems/python_environment/test_python_environment_system.py b/test/systems/python_environment/test_python_environment_system.py index 2826b3a9e1..587423a090 100644 --- a/test/systems/python_environment/test_python_environment_system.py +++ b/test/systems/python_environment/test_python_environment_system.py @@ -7,7 +7,6 @@ import shutil import tempfile import venv -from pathlib import Path from unittest.mock import MagicMock, patch import pytest @@ -36,26 +35,21 @@ class TestPythonEnvironmentSystem: @pytest.fixture(autouse=True) def setup(self, tmp_path): - # create a virtual environment with no packages installed - venv_path = tmp_path / "venv" - venv.create(venv_path, with_pip=True) - # python path - if platform.system() == OS.WINDOWS: - self.python_environment_path = Path(venv_path) / "Scripts" - else: - self.python_environment_path = Path(venv_path) / "bin" - # use the current python environment as the test environment + self.python_environment_path = tmp_path / "python-env" + self.python_environment_path.mkdir() self.system = PythonEnvironmentSystem(self.python_environment_path) - yield - shutil.rmtree(venv_path) - def test_get_supported_execution_providers(self): - python_path = shutil.which("python", path=self.python_environment_path) + def test_get_supported_execution_providers(self, tmp_path): + venv_path = tmp_path / "venv" + venv.create(venv_path, with_pip=True) + python_environment_path = venv_path / ("Scripts" if platform.system() == OS.WINDOWS else "bin") + system = PythonEnvironmentSystem(python_environment_path) + python_path = shutil.which("python", path=python_environment_path) # install only onnxruntime - run_subprocess([python_path, "-m", "pip", "install", "onnxruntime"], env=self.system.environ) + run_subprocess([python_path, "-m", "pip", "install", "onnxruntime"], env=system.environ) # for GPU ort, the get_available_providers will return ["CUDAExecutionProvider", "DmlExecutionProvider"] - assert set(self.system.get_supported_execution_providers()) == { + assert set(system.get_supported_execution_providers()) == { "AzureExecutionProvider", "CPUExecutionProvider", } @@ -82,9 +76,8 @@ def test__run_command(self, mock_temp_dir, mock_run_subprocess, tmp_path): # assert assert res == dummy_output - python_path = shutil.which("python", path=self.python_environment_path) expected_command = [ - python_path, + self.system.executable, str(script_path), "--dummy_config", str(tmp_path / "dummy_config.json"), diff --git a/test/workflows/test_setup.py b/test/workflows/test_setup.py index e9f21767b7..4996f3e7af 100644 --- a/test/workflows/test_setup.py +++ b/test/workflows/test_setup.py @@ -4,13 +4,13 @@ # -------------------------------------------------------------------------- import json import platform -import shutil from pathlib import Path +from unittest.mock import patch import pytest from olive.common.constants import OS -from olive.common.utils import run_subprocess +from olive.workflows.run.run import run as olive_run # pylint: disable=redefined-outer-name @@ -33,29 +33,14 @@ def config_json(tmp_path): return str(config_json_file) -def test_dependency_setup(tmp_path, config_json): - cmd = [ - "olive", - "run", - "--config", - str(config_json), - "--list_required_packages", - ] +def test_dependency_setup(config_json): + with patch("olive.workflows.run.run.generate_files_from_packages") as mock_generate: + olive_run(config_json, list_required_packages=True) - return_code, _, stderr = run_subprocess(cmd, check=False) - if return_code != 0: - pytest.fail(stderr) - - output_filepath = Path("olive_requirements.txt") - assert output_filepath.exists() - - required_packages = [] - with output_filepath.open() as strm: - required_packages = [_.strip() for _ in strm.readlines()] + required_packages, output_path = mock_generate.call_args.args ort_extra = "onnxruntime-directml" if platform.system() == OS.WINDOWS else "onnxruntime-gpu" + assert output_path == "olive_requirements.txt" assert ort_extra in required_packages assert "psutil" in required_packages - output_filepath.unlink() - shutil.rmtree(tmp_path, ignore_errors=True)