fix(configurator): add AutoParamType so --flag=auto works for Auto*Param (#353)

<!-- SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION
& AFFILIATES. All rights reserved. -->
<!-- SPDX-License-Identifier: Apache-2.0 -->

<!-- Thank you for contributing to Safe Synthesizer! -->

# Summary
- Introduces AutoParamType, a custom Click parameter type that accepts
the "auto" sentinel string or delegates to a wrapped base type
(click.INT, click.FLOAT, click.BOOL) for validation. This replaces the
previous behavior where Auto*Param union fields (Literal["auto"] | int,
etc.) fell through to click.STRING, which gave no numeric validation and
displayed an unhelpful TEXT type in --help.
- Updates _click_type() to detect Literal["auto"] | <numeric> unions and
return AutoParamType(base_type) instead of click.STRING, so --help now
shows descriptive types like integer|auto, float|auto, and boolean|auto.
- Pure Literal["auto"] | str unions (no numeric component) still
correctly resolve to click.STRING.

## Pre-Review Checklist

<!-- These checks should be completed before a PR is reviewed, -->
<!-- but you can submit a draft early to indicate that the issue is
being worked on. -->

Ensure that the following pass:

- [x] `make format && make check` or via prek validation.
- [x] `make test` passes locally
- [x] `make test-e2e` passes locally
- [ ] `make test-ci-container` passes locally (recommended)
- [ ] GPU CI status check passes -- comment `/sync` on this PR to
trigger a run (auto-triggers on ready-for-review)

## Pre-Merge Checklist

<!-- These checks need to be completed before a PR is merged, -->
<!-- but as PRs often change significantly during review, -->
<!-- it's OK for them to be incomplete when review is first requested.
-->

- [ ] New or updated tests for any fix or new behavior
- [ ] Updated documentation for new features and behaviors, including
docstrings for API docs.

## Other Notes

<!-- Please add the issue number that should be closed when this PR is
merged. -->
- Closes #159

---------

Signed-off-by: Sean Yang <seayang@nvidia.com>
Signed-off-by: seayang <seayang@nvidia.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: seayang-nv

src/nemo_safe_synthesizer/configurator/pydantic_click_options.py

@@ -27,7 +27,9 @@
 from pydantic import BaseModel
 from pydantic.fields import FieldInfo
 
-__all__ = ["pydantic_options", "parse_overrides"]
+from ..config.types import AUTO_STR
+
+__all__ = ["pydantic_options", "parse_overrides", "AutoParamType"]
 
 
 def parse_overrides(values: dict[str, Any] | None = None, field_sep: str = "__") -> dict[str, Any]:
@@ -128,26 +130,103 @@ def _nullable_model_arg(union_args: tuple) -> type[BaseModel] | None:
 ]
 
 
+class AutoParamType(click.ParamType):
+    """A Click type that accepts the sentinel ``AUTO_STR`` or a base numeric/bool value.
+
+    Used for ``Auto*Param`` fields (``AutoIntParam``, ``AutoFloatParam``,
+    ``AutoBoolParam``) so that ``--flag auto`` and ``--flag 2`` both work.
+    The ``--help`` display shows ``integer|auto``, ``float|auto``, or
+    ``boolean|auto`` instead of the generic ``TEXT`` label.
+
+    Args:
+        base_type: The underlying Click type (``click.INT``, ``click.FLOAT``,
+            or ``click.BOOL``) used to parse and validate non-``AUTO_STR`` values.
+    """
+
+    def __init__(self, base_type: click.ParamType) -> None:
+        self.base_type = base_type
+        self.name = f"{base_type.name}|{AUTO_STR}"
+
+    def convert(
+        self,
+        value: str,
+        param: click.Parameter | None,
+        ctx: click.Context | None,
+    ) -> str | int | float | bool:
+        """Convert the raw CLI value to ``AUTO_STR`` or the base numeric/bool type.
+
+        The ``value`` parameter is typed ``str`` to match the parent
+        ``click.ParamType.convert`` stub signature; in practice Click may also
+        pass through default values of any type, but the equality check and
+        delegated ``base_type.convert`` both handle that correctly.
+
+        Args:
+            value: Raw value from the CLI or the option default.
+            param: The Click parameter object (passed through to the base type).
+            ctx: The Click context (passed through to the base type).
+
+        Returns:
+            ``AUTO_STR`` if ``value`` equals it, otherwise the result of
+            delegating to ``self.base_type.convert()``.
+        """
+        if value == AUTO_STR:
+            return AUTO_STR
+        return self.base_type.convert(value, param, ctx)  # type: ignore[return-value]
+
+
 def _has_string_literal(args: set) -> bool:
     """Check if any member is a ``Literal`` containing a string value."""
     return any(get_origin(a) is Literal and any(isinstance(v, str) for v in get_args(a)) for a in args)
 
 
+def _is_auto_only_literal_union(args: set) -> bool:
+    """Check that the union's string-valued ``Literal`` members are exactly ``{AUTO_STR}``.
+
+    Returns ``True`` only if every string-valued Literal member contributes the
+    single value ``AUTO_STR`` -- e.g. ``Literal["auto"] | int``. Returns
+    ``False`` for unions like ``Literal["disabled"] | int`` or
+    ``Literal["auto", "manual"] | int``, where ``AutoParamType`` would
+    silently reject the non-``"auto"`` sentinels at parse time.
+    """
+    string_values: set[str] = set()
+    for a in args:
+        if get_origin(a) is Literal:
+            for v in get_args(a):
+                if isinstance(v, str):
+                    string_values.add(v)
+    return string_values == {AUTO_STR}
+
+
 def _click_type(annotation: Any) -> click.ParamType:
     """Map a Pydantic field annotation to a Click type.
 
     Unwraps ``Annotated[T, ...]`` and ``T | None`` unions, then returns the
-    widest Click type that covers any member of the union.  String-valued
-    ``Literal`` members (e.g. ``Literal["auto"]``) force ``click.STRING``
-    so Click won't reject the sentinel before Pydantic validates it.
-    Falls back to ``click.STRING`` for unrecognized types.
+    widest Click type that covers any member of the union. ``Auto*Param``
+    fields (``Literal["auto"] | <numeric|bool>``) get an ``AutoParamType``
+    wrapping the numeric/bool base so Click can validate non-sentinel values
+    while still accepting the ``"auto"`` sentinel. Other string-valued
+    ``Literal`` members fall through to ``click.STRING`` so Click won't
+    reject the sentinel before Pydantic validates it. Falls back to
+    ``click.STRING`` for unrecognized types.
     """
     t = annotation
     if get_origin(t) is Annotated:
         t = get_args(t)[0]
     args = set(get_args(t)) if get_origin(t) in (Union, types.UnionType) else {t}
     args.discard(type(None))
     if _has_string_literal(args):
+        # Auto*Param: Literal["auto"] | <numeric|bool> -- wrap in AutoParamType so
+        # --help shows "integer|auto" instead of "TEXT" and Click validates the
+        # numeric side before handing the value to Pydantic. The detection is
+        # tightened to only fire when the literal values are exactly {AUTO_STR};
+        # other sentinels (e.g. Literal["disabled"] | int) fall through to STRING
+        # so AutoParamType doesn't reject them with a confusing error.
+        if _is_auto_only_literal_union(args):
+            for py_type, click_type in _CLICK_TYPE_PRIORITY:
+                if py_type is str:
+                    continue
+                if py_type in args:
+                    return AutoParamType(click_type)
         return click.STRING
     for py_type, click_type in _CLICK_TYPE_PRIORITY:
         if py_type in args:

tests/cli/test_run.py

@@ -12,6 +12,7 @@
 import nemo_safe_synthesizer.sdk.library_builder  # noqa: F401 - ensure submodule is loaded for mock.patch
 from nemo_safe_synthesizer.cli.run import run
 from nemo_safe_synthesizer.cli.settings import CLISettings
+from nemo_safe_synthesizer.cli.utils import merge_overrides
 from nemo_safe_synthesizer.tooling import PreflightRenderContext
 
 # =============================================================================
@@ -661,3 +662,113 @@ def test_generate_with_dataset_registry_calls_common_setup(
         call_kwargs = mock_common_setup.call_args.kwargs
         settings: CLISettings = call_kwargs["settings"]
         assert settings.dataset_registry == "./registry.yaml"
+
+
+class TestAutoParamCliOverrides:
+    """End-to-end tests for ``Auto*Param`` field CLI overrides (issue #159).
+
+    These tests drive the real ``run`` Click command and capture the parsed
+    ``synthesis_overrides`` reaching ``common_setup`` to verify each CLI value
+    is parsed correctly. A separate test checks that the value also lands on
+    the resolved ``SafeSynthesizerParameters`` object (using fields that pass
+    through Pydantic validation unchanged -- some ``Auto*Param`` fields have
+    model validators that resolve ``"auto"`` to a concrete value).
+    """
+
+    # Note: ``AutoBoolParam`` is defined in ``config.types`` but no field of
+    # ``SafeSynthesizerParameters`` currently uses it (the only such field,
+    # ``training.use_unsloth``, was removed when the Unsloth backend was
+    # dropped). Bool conversion is exercised at the unit level by
+    # ``tests/configurator/test_pydantic_click_options.py``.
+    @pytest.mark.parametrize(
+        "flag,raw_value,nested_path,expected",
+        [
+            # AutoIntParam / OptionalAutoInt
+            ("--training__rope_scaling_factor", "auto", ("training", "rope_scaling_factor"), "auto"),
+            ("--training__rope_scaling_factor", "2", ("training", "rope_scaling_factor"), 2),
+            ("--training__num_input_records_to_sample", "auto", ("training", "num_input_records_to_sample"), "auto"),
+            ("--training__num_input_records_to_sample", "100", ("training", "num_input_records_to_sample"), 100),
+            ("--data__max_sequences_per_example", "auto", ("data", "max_sequences_per_example"), "auto"),
+            ("--data__max_sequences_per_example", "5", ("data", "max_sequences_per_example"), 5),
+            # AutoFloatParam
+            ("--privacy__delta", "auto", ("privacy", "delta"), "auto"),
+            ("--privacy__delta", "0.001", ("privacy", "delta"), 0.001),
+        ],
+    )
+    def test_auto_param_override_is_parsed_into_synthesis_overrides(
+        self,
+        cli_runner: CliRunner,
+        dummy_csv: Path,
+        patched_run_dependencies: dict,
+        flag: str,
+        raw_value: str,
+        nested_path: tuple[str, ...],
+        expected: object,
+    ):
+        """Auto*Param CLI flags accept ``"auto"`` and typed values, and reach ``common_setup``."""
+        result = cli_runner.invoke(
+            run,
+            ["--data-source", str(dummy_csv), flag, raw_value],
+            catch_exceptions=False,
+        )
+
+        assert result.exit_code == 0, result.output
+
+        mock_common_setup = patched_run_dependencies["common_setup"]
+        mock_common_setup.assert_called_once()
+        settings: CLISettings = mock_common_setup.call_args.kwargs["settings"]
+
+        # Click parses the raw CLI string (``"auto"`` stays a string, numbers
+        # and bools are coerced) and ``parse_overrides`` reshapes the flat
+        # kwargs into the nested overrides dict before it reaches settings.
+        node: object = settings.synthesis_overrides
+        for key in nested_path:
+            assert isinstance(node, dict) and key in node, (
+                f"missing {'.'.join(nested_path)} in synthesis_overrides: {settings.synthesis_overrides}"
+            )
+            node = node[key]
+        assert node == expected
+        assert type(node) is type(expected)
+
+    @pytest.mark.parametrize(
+        "flag,raw_value,nested_path,expected",
+        [
+            # rope_scaling_factor, num_input_records_to_sample, and
+            # privacy.delta pass through Pydantic validation unchanged for both
+            # 'auto' and explicit values. max_sequences_per_example is excluded
+            # because its model validator rewrites 'auto' to a concrete default
+            # (10 with DP disabled, 1 with DP enabled).
+            ("--training__rope_scaling_factor", "auto", ("training", "rope_scaling_factor"), "auto"),
+            ("--training__rope_scaling_factor", "2", ("training", "rope_scaling_factor"), 2),
+            ("--training__num_input_records_to_sample", "auto", ("training", "num_input_records_to_sample"), "auto"),
+            ("--training__num_input_records_to_sample", "100", ("training", "num_input_records_to_sample"), 100),
+            ("--privacy__delta", "auto", ("privacy", "delta"), "auto"),
+            ("--privacy__delta", "0.001", ("privacy", "delta"), 0.001),
+        ],
+    )
+    def test_auto_param_override_reaches_params_object(
+        self,
+        cli_runner: CliRunner,
+        dummy_csv: Path,
+        patched_run_dependencies: dict,
+        flag: str,
+        raw_value: str,
+        nested_path: tuple[str, ...],
+        expected: object,
+    ):
+        """The parsed CLI value also lands on the validated ``SafeSynthesizerParameters`` object."""
+        result = cli_runner.invoke(
+            run,
+            ["--data-source", str(dummy_csv), flag, raw_value],
+            catch_exceptions=False,
+        )
+
+        assert result.exit_code == 0, result.output
+        settings: CLISettings = patched_run_dependencies["common_setup"].call_args.kwargs["settings"]
+
+        params = merge_overrides(None, settings.synthesis_overrides)
+        resolved: object = params
+        for key in nested_path:
+            resolved = getattr(resolved, key)
+        assert resolved == expected
+        assert type(resolved) is type(expected)

tests/configurator/test_pydantic_click_options.py

@@ -13,6 +13,7 @@
 
 from nemo_safe_synthesizer.config import SafeSynthesizerParameters
 from nemo_safe_synthesizer.configurator.pydantic_click_options import (
+    AutoParamType,
     FlagParam,
     LeafParam,
     _click_type,
@@ -140,17 +141,98 @@ def test_click_type_str_or_int_returns_string():
     assert _click_type(str | int | None) == click.STRING
 
 
-def test_click_type_literal_auto_int_returns_string():
-    """Literal['auto'] | int must use STRING so Click accepts 'auto'."""
+def test_click_type_literal_auto_int_returns_auto_param_type():
+    """Literal['auto'] | int must use AutoParamType so Click accepts both 'auto' and integers."""
     from typing import Literal
 
-    assert _click_type(Literal["auto"] | int) == click.STRING
+    result = _click_type(Literal["auto"] | int)
+    assert isinstance(result, AutoParamType)
+    assert result.base_type is click.INT
 
 
-def test_click_type_literal_auto_float_returns_string():
+def test_click_type_literal_auto_float_returns_auto_param_type():
     from typing import Literal
 
-    assert _click_type(Literal["auto"] | float) == click.STRING
+    result = _click_type(Literal["auto"] | float)
+    assert isinstance(result, AutoParamType)
+    assert result.base_type is click.FLOAT
+
+
+def test_click_type_literal_auto_bool_returns_auto_param_type():
+    from typing import Literal
+
+    result = _click_type(Literal["auto"] | bool)
+    assert isinstance(result, AutoParamType)
+    assert result.base_type is click.BOOL
+
+
+def test_click_type_literal_auto_str_returns_string():
+    """Literal['auto'] | str has no numeric side -- STRING is correct."""
+    from typing import Literal
+
+    assert _click_type(Literal["auto"] | str) == click.STRING
+
+
+def test_click_type_literal_non_auto_int_returns_string():
+    """Literal['disabled'] | int must NOT use AutoParamType.
+
+    AutoParamType only accepts the AUTO_STR sentinel, so a sentinel like
+    'disabled' would be rejected with a confusing 'not a valid integer' error
+    even though Pydantic would accept it. Falling through to STRING keeps the
+    sentinel routable to Pydantic for validation.
+    """
+    from typing import Literal
+
+    assert _click_type(Literal["disabled"] | int) == click.STRING
+
+
+def test_click_type_literal_auto_plus_other_int_returns_string():
+    """Literal['auto', 'manual'] | int must NOT use AutoParamType (multiple sentinels)."""
+    from typing import Literal
+
+    assert _click_type(Literal["auto", "manual"] | int) == click.STRING
+
+
+# ---------------------------------------------------------------------------
+# AutoParamType
+# ---------------------------------------------------------------------------
+
+
+def test_auto_param_type_accepts_auto_sentinel():
+    assert AutoParamType(click.INT).convert("auto", None, None) == "auto"
+
+
+@pytest.mark.parametrize(
+    "base_type,value,expected",
+    [
+        (click.INT, "2", 2),
+        (click.INT, "100", 100),
+        (click.FLOAT, "1.5", 1.5),
+        (click.FLOAT, "0.001", 0.001),
+        (click.BOOL, "true", True),
+        (click.BOOL, "false", False),
+    ],
+)
+def test_auto_param_type_converts_numeric_values(base_type, value, expected):
+    assert AutoParamType(base_type).convert(value, None, None) == expected
+
+
+def test_auto_param_type_rejects_non_auto_string():
+    runner = CliRunner()
+
+    @click.command()
+    @click.option("--val", type=AutoParamType(click.INT))
+    def cmd(val):
+        pass
+
+    result = runner.invoke(cmd, ["--val", "notanumber"])
+    assert result.exit_code != 0
+
+
+def test_auto_param_type_name():
+    assert AutoParamType(click.INT).name == "integer|auto"
+    assert AutoParamType(click.FLOAT).name == "float|auto"
+    assert AutoParamType(click.BOOL).name == "boolean|auto"
 
 
 def test_click_type_plain_int_returns_int():