Test/unit suite acceleration

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."""

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

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

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"]