fix: allow PipelineVariable keys in HyperparameterTuner hyperparameter_ranges annotation

[mohamedzeidan2021]

Issue

Fixes #5243

The v3 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

sagemaker-train/tests/unit/train/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] (verified via git stash).
• New test_init_with_pipeline_variable_hyperparameter_key constructs a tuner with a ParameterString key as a runtime sanity check (documented in the test as not being the regression guard, since annotations aren't runtime-enforced).

All 45 tests in test_tuner.py pass; 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.

[github-actions]

Claude finished @mohamedzeidan2021's task in 54s —— View job

---

Review: PR #6314 — allow PipelineVariable keys in hyperparameter_ranges annotation

• Read PR diff
• Review tuner.py annotation change in context
• Review tests
• Post review feedback

Verdict: LGTM. This is a correct, minimal, backward-compatible annotation + docstring fix, backed by a real regression guard.

Correctness

• The widening Dict[str, ParameterRange] → Dict[Union[str, PipelineVariable], ParameterRange] is sound. ParameterString subclasses PipelineVariable, so the union is exactly the right type for the documented pipeline pattern.
• Both Union and PipelineVariable are already imported (tuner.py:19, tuner.py:48), so no new imports are needed.
• Runtime behavior is genuinely unchanged: _validate_parameter_ranges (tuner.py:867) is a no-op for ModelTrainer, and the ranges are stored as-is in self._hyperparameter_ranges / self._hyperparameter_ranges_dict. Annotations aren't enforced at runtime, so nothing downstream depends on the narrower key type.
• Backward compatibility: widening an accepted type is safe — every existing Dict[str, ParameterRange] argument still satisfies the broader union. No public-interface break.

Tests

• test_hyperparameter_ranges_annotation_allows_pipeline_variable_keys is the right kind of guard: it introspects __init__.__annotations__ and asserts both str and PipelineVariable are in the key union, so it fails if the type is ever re-narrowed. The comment explaining why get_type_hints is avoided (the TYPE_CHECKING-only "ModelTrainer" forward ref) is accurate and helpful.
• test_init_with_pipeline_variable_hyperparameter_key is honestly labeled as a runtime sanity check rather than the regression guard — good, since annotations aren't enforced at runtime.

Minor observations (non-blocking)

• The annotation test asserts value_type is ParameterRange, which couples the test to that exact value type. That's fine and intentional here, but if the value type were ever widened similarly the assertion would need updating — worth a mental note only.
• Consider whether the create() classmethod's hyperparameter_ranges_dict value annotation (values are Dict[..., ParameterRange]) deserves the same key widening for consistency, since the docstring at tuner.py:1040 now describes those keys as accepting a PipelineVariable too. Not required for this fix — the reported mypy false positive is on __init__ — but it would keep the annotation and the updated docstring fully aligned.

No correctness, security, or performance concerns. Docs and tests are appropriate for the change.

[github-actions]

Claude finished @mohamedzeidan2021's task in 45s —— View job

---

Review: allow PipelineVariable keys in hyperparameter_ranges annotation

• Read PR diff
• Review tuner.py annotation change in context
• Review tests
• Post review feedback

Verdict: LGTM. Correct, minimal, backward-compatible annotation + docstring fix, backed by a real regression guard.

Correctness

• Widening Dict[str, ParameterRange] → Dict[Union[str, PipelineVariable], ParameterRange] (tuner.py:95) is sound. ParameterString subclasses PipelineVariable, so the union precisely covers the documented in-pipeline pattern from Incorrect type annotation of hyperparameter_ranges argument of HyperparameterTuner constructor #5243.
• Union and PipelineVariable are already imported (tuner.py:19, tuner.py:48) — no new imports.
• Runtime behavior is genuinely unchanged: _validate_parameter_ranges (tuner.py:867) is a no-op for ModelTrainer, and the ranges are stored as-is in self._hyperparameter_ranges / self._hyperparameter_ranges_dict. Annotations aren't enforced at runtime, so nothing downstream depends on the narrower key type.
• Backward compatible: widening an accepted type is safe — every existing Dict[str, ParameterRange] argument still satisfies the broader union. No public-interface break.

Tests

• test_hyperparameter_ranges_annotation_allows_pipeline_variable_keys is the right kind of guard — it introspects __init__.__annotations__ and asserts both str and PipelineVariable are in the key union, so it fails if the type is ever re-narrowed. The comment explaining why get_type_hints is avoided (the TYPE_CHECKING-only "ModelTrainer" forward ref) is accurate.
• test_init_with_pipeline_variable_hyperparameter_key is honestly labeled as a runtime sanity check rather than the regression guard — good, since annotations aren't enforced at runtime.

Notes on the create() docstring change (non-blocking)

• The prior review suggested also widening create()'s hyperparameter_ranges_dict value annotation for consistency with its updated docstring (tuner.py:1043). That doesn't actually apply here: create()'s parameters are entirely unannotated (tuner.py:1001), so there's no signature-level type to keep in sync — the docstring-only update is appropriate and complete.
• Minor: the annotation test asserts value_type is ParameterRange, coupling it to that exact value type. Fine and intentional; just a mental note if the value type is ever widened later.

No correctness, security, or performance concerns. Docs and tests are appropriate for the change.