From 1be38738c1811ffcb64bb22c56161cba44c1389a Mon Sep 17 00:00:00 2001 From: Sylvester Kaczmarek <16242628+sylvesterkaczmarek@users.noreply.github.com> Date: Thu, 3 Sep 2026 18:51:22 +0100 Subject: [PATCH 1/3] Fix boolean parsing in capture-onnx CLI --- olive/cli/capture_onnx.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/olive/cli/capture_onnx.py b/olive/cli/capture_onnx.py index 3966b798c..55449296d 100644 --- a/olive/cli/capture_onnx.py +++ b/olive/cli/capture_onnx.py @@ -37,6 +37,19 @@ def parse_dim_dict(s): raise argparse.ArgumentTypeError("Format must be key=value,... with positive integers as values") from exc +def parse_bool(value): + if isinstance(value, bool): + return value + + normalized = value.lower() + if normalized in {"true", "1", "yes", "on"}: + return True + if normalized in {"false", "0", "no", "off"}: + return False + + raise argparse.ArgumentTypeError(f"invalid boolean value: {value!r}") + + class CaptureOnnxGraphCommand(BaseOliveCLICommand): @staticmethod def register_subcommand(parser: ArgumentParser): @@ -147,21 +160,21 @@ def register_subcommand(parser: ArgumentParser): ) mb_group.add_argument( "--exclude_embeds", - type=bool, + type=parse_bool, default=False, required=False, help="Remove embedding layer from your ONNX model.", ) mb_group.add_argument( "--exclude_lm_head", - type=bool, + type=parse_bool, default=False, required=False, help="Remove language modeling head from your ONNX model.", ) mb_group.add_argument( "--enable_cuda_graph", - type=bool, + type=parse_bool, default=None, # Explicitly setting to None to differentiate between user intent and default. required=False, help=( From 89bb22626d1ad6c1325da3962079a31103ac62b6 Mon Sep 17 00:00:00 2001 From: Sylvester Kaczmarek <16242628+sylvesterkaczmarek@users.noreply.github.com> Date: Thu, 3 Sep 2026 18:52:21 +0100 Subject: [PATCH 2/3] Add capture-onnx boolean argument regression tests --- test/cli/test_capture_onnx_args.py | 49 ++++++++++++++++++++++++++++++ 1 file changed, 49 insertions(+) create mode 100644 test/cli/test_capture_onnx_args.py diff --git a/test/cli/test_capture_onnx_args.py b/test/cli/test_capture_onnx_args.py new file mode 100644 index 000000000..2aca9b872 --- /dev/null +++ b/test/cli/test_capture_onnx_args.py @@ -0,0 +1,49 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +import argparse + +import pytest + +from olive.cli.capture_onnx import CaptureOnnxGraphCommand + + +def _parse_capture_args(*args): + parser = argparse.ArgumentParser() + commands = parser.add_subparsers() + CaptureOnnxGraphCommand.register_subcommand(commands) + return parser.parse_args(["capture-onnx-graph", *args]) + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("true", True), + ("1", True), + ("yes", True), + ("on", True), + ("false", False), + ("0", False), + ("no", False), + ("off", False), + ], +) +def test_capture_onnx_boolean_arguments(value, expected): + args = _parse_capture_args( + "--exclude_embeds", + value, + "--exclude_lm_head", + value, + "--enable_cuda_graph", + value, + ) + + assert args.exclude_embeds is expected + assert args.exclude_lm_head is expected + assert args.enable_cuda_graph is expected + + +def test_capture_onnx_boolean_arguments_reject_invalid_value(): + with pytest.raises(SystemExit): + _parse_capture_args("--exclude_embeds", "not-a-bool") From 0cf06cc887b6d2c62601402ca2baea1847d5941b Mon Sep 17 00:00:00 2001 From: Sylvester Kaczmarek <16242628+sylvesterkaczmarek@users.noreply.github.com> Date: Tue, 8 Sep 2026 23:52:47 +0100 Subject: [PATCH 3/3] Use argparse boolean flags for capture ONNX options Signed-off-by: Sylvester Kaczmarek <16242628+sylvesterkaczmarek@users.noreply.github.com> --- olive/cli/capture_onnx.py | 24 +++--------------- test/cli/test_capture_onnx_args.py | 39 +++++++++--------------------- 2 files changed, 15 insertions(+), 48 deletions(-) diff --git a/olive/cli/capture_onnx.py b/olive/cli/capture_onnx.py index 55449296d..b7f2150f8 100644 --- a/olive/cli/capture_onnx.py +++ b/olive/cli/capture_onnx.py @@ -37,19 +37,6 @@ def parse_dim_dict(s): raise argparse.ArgumentTypeError("Format must be key=value,... with positive integers as values") from exc -def parse_bool(value): - if isinstance(value, bool): - return value - - normalized = value.lower() - if normalized in {"true", "1", "yes", "on"}: - return True - if normalized in {"false", "0", "no", "off"}: - return False - - raise argparse.ArgumentTypeError(f"invalid boolean value: {value!r}") - - class CaptureOnnxGraphCommand(BaseOliveCLICommand): @staticmethod def register_subcommand(parser: ArgumentParser): @@ -160,23 +147,18 @@ def register_subcommand(parser: ArgumentParser): ) mb_group.add_argument( "--exclude_embeds", - type=parse_bool, - default=False, - required=False, + action="store_true", help="Remove embedding layer from your ONNX model.", ) mb_group.add_argument( "--exclude_lm_head", - type=parse_bool, - default=False, - required=False, + action="store_true", help="Remove language modeling head from your ONNX model.", ) mb_group.add_argument( "--enable_cuda_graph", - type=parse_bool, + action="store_true", default=None, # Explicitly setting to None to differentiate between user intent and default. - required=False, help=( "The model can use CUDA graph capture for CUDA execution provider. " "If enabled, all nodes being placed on the CUDA EP is the prerequisite " diff --git a/test/cli/test_capture_onnx_args.py b/test/cli/test_capture_onnx_args.py index 2aca9b872..cba94d030 100644 --- a/test/cli/test_capture_onnx_args.py +++ b/test/cli/test_capture_onnx_args.py @@ -4,8 +4,6 @@ # -------------------------------------------------------------------------- import argparse -import pytest - from olive.cli.capture_onnx import CaptureOnnxGraphCommand @@ -16,34 +14,21 @@ def _parse_capture_args(*args): return parser.parse_args(["capture-onnx-graph", *args]) -@pytest.mark.parametrize( - ("value", "expected"), - [ - ("true", True), - ("1", True), - ("yes", True), - ("on", True), - ("false", False), - ("0", False), - ("no", False), - ("off", False), - ], -) -def test_capture_onnx_boolean_arguments(value, expected): +def test_capture_onnx_boolean_flags_default_to_disabled(): + args = _parse_capture_args() + + assert args.exclude_embeds is False + assert args.exclude_lm_head is False + assert args.enable_cuda_graph is None + + +def test_capture_onnx_boolean_flags_enable_on_presence(): args = _parse_capture_args( "--exclude_embeds", - value, "--exclude_lm_head", - value, "--enable_cuda_graph", - value, ) - assert args.exclude_embeds is expected - assert args.exclude_lm_head is expected - assert args.enable_cuda_graph is expected - - -def test_capture_onnx_boolean_arguments_reject_invalid_value(): - with pytest.raises(SystemExit): - _parse_capture_args("--exclude_embeds", "not-a-bool") + assert args.exclude_embeds is True + assert args.exclude_lm_head is True + assert args.enable_cuda_graph is True