fix: allow PipelineVariable keys in HyperparameterTuner hyperparameter_ranges annotation (v2)

[mohamedzeidan2021]

Issue

Fixes #5243 (v2 backport — companion to #6314 which fixes the same annotation in v3/sagemaker-train).

The HyperparameterTuner constructor annotates hyperparameter_ranges as Dict[str, ParameterRange]. However, the tuner accepts a PipelineVariable (e.g. a pipeline ParameterString) as a dict key — the hyperparameter name — which is a documented, working pattern when building a tuning step inside a pipeline. Because the annotation only allows str keys, mypy reports a false positive on valid code:

error: Dict entry 0 has incompatible type "ParameterString": "CategoricalParameter";
       expected "str": "ParameterRange"  [dict-item]

Fix

Broaden the key type to Union[str, PipelineVariable]:

hyperparameter_ranges: Dict[Union[str, PipelineVariable], ParameterRange],

Union and PipelineVariable are already imported in the module. The __init__ and create() docstrings are updated to note that keys may be a str or a PipelineVariable. This is an annotation/docstring-only change — runtime behavior is unchanged (annotations are not enforced, and _validate_parameter_ranges already accepts any key), and ParameterString subclasses PipelineVariable, so the union is exactly right.

Testing

tests/unit/test_tuner.py:

• New test_hyperparameter_ranges_annotation_allows_pipeline_variable_keys inspects the resolved __init__ annotation and asserts the key union contains both str and PipelineVariable. This is the regression guard — it fails if the type is narrowed back to Dict[str, ParameterRange].
• New test_init_with_pipeline_variable_hyperparameter_key constructs a tuner with a ParameterString key as a runtime sanity check (documented as not being the regression guard, since annotations aren't runtime-enforced).

All 107 tests in test_tuner.py pass. Reproduced the mypy dict-item error against the unpatched annotation and confirmed it is cleared by the fix. black (line-length 100) and flake8 (line-length 120) clean.

Backwards compatibility

Widening a parameter's accepted type is backwards compatible — every existing Dict[str, ParameterRange] argument still satisfies the broader union. No runtime effect. This is the only signature-level annotation of this parameter in the repo.