Tuning with a categorical parameter containing values that are pipeline variables does not work for a framework estimator. #5240
Description
Describe the bug
Here: https://github.com/aws/sagemaker-python-sdk/blob/master/src/sagemaker/tuner.py#L1655 in case of framework estimator, the CategoricalParameter is encoded using function as_json_range instead of as_tuning_range. Function as_json_range does not work in case when CategoricalParameter values are pipeline parameters.
I don't know what causes this limitation, but the Pytorch estimator used for tuning in version 2.3, CategoricalParameters, whose values are pipeline parameters, works fine when we get rid of this condition and encode the categorical parameter using the function: as_tuning_range. The estimator-dependent limitation also seems strange, as it seems that the tuner that selects hyperparameters should select parameters regardless of the estimator used.
To reproduce
test.py
from sagemaker.parameter import ParameterRange
from sagemaker.pytorch import PyTorch
from sagemaker.tuner import HyperparameterTuner, CategoricalParameter, ContinuousParameter
from sagemaker.workflow.parameters import ParameterString
from sagemaker.workflow.steps import TuningStep
from sagemaker.workflow.pipeline import Pipeline
from sagemaker.workflow.pipeline_context import (
PipelineSession
)
if __name__ == "__main__":
pipeline_session = PipelineSession()
estimator = PyTorch(
sagemaker_session=pipeline_session,
instance_type='ml.m5.large',
instance_count=1,
framework_version="2.3",
py_version="py311",
source_dir='source',
entry_point='main.py',
metric_definitions=[
{'Name': 'valid:loss', 'Regex': 'valid_loss=([0-9]+\\.?[0-9]*)'}
]
)
hparam1 = ParameterString("hparam1", default_value="1")
hparam2 = ParameterString("hparam2", default_value="2")
tuner = HyperparameterTuner(
estimator=estimator,
objective_metric_name='valid:loss',
objective_type='Minimize',
hyperparameter_ranges={
"hparam": CategoricalParameter(
[hparam1, hparam2]
)
},
max_jobs=2,
max_parallel_jobs=1,
base_tuning_job_name='test-tuning',
strategy='Grid',
metric_definitions=estimator.metric_definitions,
)
tuning_step = TuningStep(
name="Tuning",
step_args=tuner.fit(),
)
pipeline = Pipeline(
name="TestTuningPipeline",
parameters=[
hparam1,
hparam2
],
steps=[tuning_step],
sagemaker_session=pipeline_session
)
pipeline.upsert()
execution = pipeline.start(
execution_display_name="TuningTest",
)
print(execution)source/main.py
import argparse
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--hparam", type=int, required=True)
args, _ = parser.parse_known_args()
print(f"Hparam: {args.hparam}")
print("valid_loss=0.1")test.py ends up with exception:
TypeError: Object of type ParameterString is not JSON serializable
Traceback:
╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│HIDDEN_PATH/test.py:77 in <module> │
│ │
│ 74 │ │ sagemaker_session=pipeline_session │
│ 75 │ ) │
│ 76 │ │
│ ❱ 77 │ pipeline.upsert() │
│ 78 │ execution = pipeline.start( │
│ 79 │ │ execution_display_name="TuningTest", │
│ 80 │ ) │
│ │
│HIDDEN_PATH/.venv/lib/python3.12/site-packages/sagemaker/workflow/pipeline.py:292 in upsert │
│ │
│ 289 │ │ │ # after fetching the config. │
│ 290 │ │ │ raise ValueError("An AWS IAM role is required to create or update a Pipeline │
│ 291 │ │ try: │
│ ❱ 292 │ │ │ response = self.create(role_arn, description, tags, parallelism_config) │
│ 293 │ │ except ClientError as ce: │
│ 294 │ │ │ error_code = ce.response["Error"]["Code"] │
│ 295 │ │ │ error_message = ce.response["Error"]["Message"] │
│ │
│HIDDEN_PATH/.venv/lib/python3.12/site-packages/sagemaker/workflow/pipeline.py:164 in create │
│ │
│ 161 │ │ tags = format_tags(tags) │
│ 162 │ │ tags = _append_project_tags(tags) │
│ 163 │ │ tags = self.sagemaker_session._append_sagemaker_config_tags(tags, PIPELINE_TAGS_ │
│ ❱ 164 │ │ kwargs = self._create_args(role_arn, description, parallelism_config) │
│ 165 │ │ update_args( │
│ 166 │ │ │ kwargs, │
│ 167 │ │ │ Tags=tags, │
│ │
│HIDDEN_PATH/.venv/lib/python3.12/site-packages/sagemaker/workflow/pipeline.py:186 in _create_args │
│ │
│ 183 │ │ Returns: │
│ 184 │ │ │ A keyword argument dict for calling create_pipeline. │
│ 185 │ │ """ │
│ ❱ 186 │ │ pipeline_definition = self.definition() │
│ 187 │ │ kwargs = dict( │
│ 188 │ │ │ PipelineName=self.name, │
│ 189 │ │ │ RoleArn=role_arn, │
│ │
│HIDDEN_PATH/.venv/lib/python3.12/site-packages/sagemaker/workflow/pipeline.py:392 in definition │
│ │
│ 389 │ │ │ sagemaker_session=self.sagemaker_session, │
│ 390 │ │ │ steps=self.steps, │
│ 391 │ │ │ pipeline_definition_config=self.pipeline_definition_config, │
│ ❱ 392 │ │ ).build() │
│ 393 │ │ │
│ 394 │ │ request_dict = { │
│ 395 │ │ │ "Version": self._version, │
│ │
│HIDDEN_PATH/.venv/lib/python3.12/site-packages/sagemaker/workflow/_steps_compiler.py:406 in build │
│ │
│ 403 │ │ if self._build_count > 1: │
│ 404 │ │ │ raise RuntimeError("Cannot build a pipeline more than once with the same com │
│ 405 │ │ │
│ ❱ 406 │ │ return self._initialize_queue_and_build(self._input_steps) │
│ 407 │
│ │
│HIDDEN_PATH/.venv/lib/python3.12/site-packages/sagemaker/workflow/_steps_compiler.py:390 │
│ in _initialize_queue_and_build │
│ │
│ 387 │ │ │ if isinstance(step, ConditionStep): │
│ 388 │ │ │ │ compiled_steps.append(self._build_condition_step(step)) │
│ 389 │ │ │ else: │
│ ❱ 390 │ │ │ │ compiled_steps.append(self._build_step(step)) │
│ 391 │ │ │
│ 392 │ │ self._set_serialize_output_to_json_flag(compiled_steps) │
│ 393 │ │ return compiled_steps │
│ │
│HIDDEN_PATH/.venv/lib/python3.12/site-packages/sagemaker/workflow/_steps_compiler.py:331 │
│ in _build_step │
│ │
│ 328 │ │ │ pipeline_build_time=self.pipeline_build_time, │
│ 329 │ │ │ function_step_secret_token=self._function_step_secret_token, │
│ 330 │ │ ) as context: │
│ ❱ 331 │ │ │ request_dict = step.to_request() │
│ 332 │ │ │ │
│ 333 │ │ │ self.upload_runtime_scripts = context.upload_runtime_scripts │
│ 334 │ │ │ self.upload_workspace = context.upload_workspace │
│ │
│HIDDEN_PATH/.venv/lib/python3.12/site-packages/sagemaker/workflow/steps.py:1125 in to_request │
│ │
│ 1122 │ │
│ 1123 │ def to_request(self) -> RequestType: │
│ 1124 │ │ """Updates the dictionary with cache configuration.""" │
│ ❱ 1125 │ │ request_dict = super().to_request() │
│ 1126 │ │ if self.cache_config: │
│ 1127 │ │ │ request_dict.update(self.cache_config.config) │
│ 1128 │
│ │
│HIDDEN_PATH/.venv/lib/python3.12/site-packages/sagemaker/workflow/steps.py:390 in to_request │
│ │
│ 387 │ │
│ 388 │ def to_request(self) -> RequestType: │
│ 389 │ │ """Gets the request structure for `ConfigurableRetryStep`.""" │
│ ❱ 390 │ │ step_dict = super().to_request() │
│ 391 │ │ if self.retry_policies: │
│ 392 │ │ │ step_dict["RetryPolicies"] = self._resolve_retry_policy(self.retry_policies) │
│ 393 │ │ return step_dict │
│ │
│HIDDEN_PATH/.venv/lib/python3.12/site-packages/sagemaker/workflow/steps.py:147 in to_request │
│ │
│ 144 │ │ request_dict = { │
│ 145 │ │ │ "Name": self.name, │
│ 146 │ │ │ "Type": self.step_type.value, │
│ ❱ 147 │ │ │ "Arguments": self.arguments, │
│ 148 │ │ } │
│ 149 │ │ if self.depends_on: │
│ 150 │ │ │ request_dict["DependsOn"] = list(self.depends_on) │
│ │
│HIDDEN_PATH/.venv/lib/python3.12/site-packages/sagemaker/workflow/steps.py:1091 in arguments │
│ │
│ 1088 │ │ if self.step_args: │
│ 1089 │ │ │ # execute fit function with saved parameters, │
│ 1090 │ │ │ # and store args in PipelineSession's _context │
│ ❱ 1091 │ │ │ execute_job_functions(self.step_args) │
│ 1092 │ │ │ │
│ 1093 │ │ │ # populate request dict with args │
│ 1094 │ │ │ tuner = self.step_args.func_args[0] │
│ │
│HIDDEN_PATH/.venv/lib/python3.12/site-packages/sagemaker/workflow/utilities.py:444 │
│ in execute_job_functions │
│ │
│ 441 │ │ │ a pipeline step, contains the necessary function information │
│ 442 │ """ │
│ 443 │ │
│ ❱ 444 │ chained_args = step_args.func(*step_args.func_args, **step_args.func_kwargs) │
│ 445 │ if isinstance(chained_args, _StepArguments): │
│ 446 │ │ execute_job_functions(chained_args) │
│ 447 │
│ │
│HIDDEN_PATH/.venv/lib/python3.12/site-packages/sagemaker/tuner.py:1034 in fit │
│ │
│ 1031 │ │ │ │ arguments are needed. │
│ 1032 │ │ """ │
│ 1033 │ │ if self.estimator is not None: │
│ ❱ 1034 │ │ │ self._fit_with_estimator(inputs, job_name, include_cls_metadata, **kwargs) │
│ 1035 │ │ else: │
│ 1036 │ │ │ self._fit_with_estimator_dict(inputs, job_name, include_cls_metadata, estima │
│ 1037 │
│ │
│HIDDEN_PATH/.venv/lib/python3.12/site-packages/sagemaker/tuner.py:1045 in _fit_with_estimator │
│ │
│ 1042 │ │ """Start tuning for tuner instances that have the ``estimator`` field set.""" │
│ 1043 │ │ self._prepare_estimator_for_tuning(self.estimator, inputs, job_name, **kwargs) │
│ 1044 │ │ self._prepare_for_tuning(job_name=job_name, include_cls_metadata=include_cls_met │
│ ❱ 1045 │ │ self.latest_tuning_job = _TuningJob.start_new(self, inputs) │
│ 1046 │ │
│ 1047 │ def _fit_with_estimator_dict(self, inputs, job_name, include_cls_metadata, estimator │
│ 1048 │ │ """Start tuning for tuner instances that have the ``estimator_dict`` field set." │
│ │
│HIDDEN_PATH/.venv/lib/python3.12/site-packages/sagemaker/tuner.py:2141 in start_new │
│ │
│ 2138 │ │ │ sagemaker.tuner._TuningJob: Constructed object that captures all │
│ 2139 │ │ │ information about the started job. │
│ 2140 │ │ """ │
│ ❱ 2141 │ │ tuner_args = cls._get_tuner_args(tuner, inputs) │
│ 2142 │ │ │
│ 2143 │ │ tuner.sagemaker_session.create_tuning_job(**tuner_args) │
│ 2144 │
│ │
│HIDDEN_PATH/.venv/lib/python3.12/site-packages/sagemaker/tuner.py:2183 in _get_tuner_args │
│ │
│ 2180 │ │ │ tuning_config["objective_type"] = tuner.objective_type │
│ 2181 │ │ │ tuning_config["objective_metric_name"] = tuner.objective_metric_name │
│ 2182 │ │ │
│ ❱ 2183 │ │ parameter_ranges = tuner.hyperparameter_ranges() │
│ 2184 │ │ if parameter_ranges is not None: │
│ 2185 │ │ │ tuning_config["parameter_ranges"] = parameter_ranges │
│ 2186 │
│ │
│HIDDEN_PATH/.venv/lib/python3.12/site-packages/sagemaker/tuner.py:1628 in hyperparameter_ranges │
│ │
│ 1625 │ │ if self._hyperparameter_ranges is None: │
│ 1626 │ │ │ return None │
│ 1627 │ │ │
│ ❱ 1628 │ │ return self._prepare_parameter_ranges_for_tuning( │
│ 1629 │ │ │ self._hyperparameter_ranges, self.estimator │
│ 1630 │ │ ) │
│ 1631 │
│ │
│HIDDEN_PATH/.venv/lib/python3.12/site-packages/sagemaker/tuner.py:1658 │
│ in _prepare_parameter_ranges_for_tuning │
│ │
│ 1655 │ │ │ │ │ if isinstance(parameter, CategoricalParameter) and isinstance( │
│ 1656 │ │ │ │ │ │ estimator, Framework │
│ 1657 │ │ │ │ │ ): │
│ ❱ 1658 │ │ │ │ │ │ tuning_range = parameter.as_json_range(parameter_name) │
│ 1659 │ │ │ │ │ else: │
│ 1660 │ │ │ │ │ │ tuning_range = parameter.as_tuning_range(parameter_name) │
│ 1661 │ │ │ │ │ hp_ranges.append(tuning_range) │
│ │
│HIDDEN_PATH/.venv/lib/python3.12/site-packages/sagemaker/parameter.py:148 in as_json_range │
│ │
│ 145 │ │ │ dict[str, list[str]]: A dictionary that contains the name and values of the │
│ 146 │ │ │ hyperparameter, where the values are serialized as JSON. │
│ 147 │ │ """ │
│ ❱ 148 │ │ return {"Name": name, "Values": [json.dumps(v) for v in self.values]} │
│ 149 │ │
│ 150 │ def is_valid(self, value): │
│ 151 │ │ """Placeholder docstring""" │
│ │
│ /usr/lib/python3.12/json/__init__.py:231 in dumps │
│ │
│ 228 │ │ check_circular and allow_nan and │
│ 229 │ │ cls is None and indent is None and separators is None and │
│ 230 │ │ default is None and not sort_keys and not kw): │
│ ❱ 231 │ │ return _default_encoder.encode(obj) │
│ 232 │ if cls is None: │
│ 233 │ │ cls = JSONEncoder │
│ 234 │ return cls( │
│ │
│ /usr/lib/python3.12/json/encoder.py:200 in encode │
│ │
│ 197 │ │ # This doesn't pass the iterator directly to ''.join() because the │
│ 198 │ │ # exceptions aren't as detailed. The list call should be roughly │
│ 199 │ │ # equivalent to the PySequence_Fast that ''.join() would do. │
│ ❱ 200 │ │ chunks = self.iterencode(o, _one_shot=True) │
│ 201 │ │ if not isinstance(chunks, (list, tuple)): │
│ 202 │ │ │ chunks = list(chunks) │
│ 203 │ │ return ''.join(chunks) │
│ │
│ /usr/lib/python3.12/json/encoder.py:258 in iterencode │
│ │
│ 255 │ │ │ │ markers, self.default, _encoder, self.indent, floatstr, │
│ 256 │ │ │ │ self.key_separator, self.item_separator, self.sort_keys, │
│ 257 │ │ │ │ self.skipkeys, _one_shot) │
│ ❱ 258 │ │ return _iterencode(o, 0) │
│ 259 │
│ 260 def _make_iterencode(markers, _default, _encoder, _indent, _floatstr, │
│ 261 │ │ _key_separator, _item_separator, _sort_keys, _skipkeys, _one_shot, │
│ │
│ /usr/lib/python3.12/json/encoder.py:180 in default │
│ │
│ 177 │ │ │ │ return super().default(o) │
│ 178 │ │ │
│ 179 │ │ """ │
│ ❱ 180 │ │ raise TypeError(f'Object of type {o.__class__.__name__} ' │
│ 181 │ │ │ │ │ │ f'is not JSON serializable') │
│ 182 │ │
│ 183 │ def encode(self, o): │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
TypeError: Object of type ParameterString is not JSON serializable
However replacing HyperparameterTuner with:
class FixedHyperparameterTuner(HyperparameterTuner):
@classmethod
def _prepare_parameter_ranges_for_tuning(cls, parameter_ranges, estimator):
"""Prepare hyperparameter ranges for tuning"""
processed_parameter_ranges = dict()
for range_type in ParameterRange.__all_types__:
hp_ranges = []
for parameter_name, parameter in parameter_ranges.items():
if parameter is not None and parameter.__name__ == range_type:
tuning_range = parameter.as_tuning_range(parameter_name)
hp_ranges.append(tuning_range)
processed_parameter_ranges[range_type + "ParameterRanges"] = hp_ranges
return processed_parameter_rangesworks properly
Expected behavior
I expect that the given code without Fixed HyperparameterTuner will start a pipeline with a tuning step that will start 2 training jobs: one with hparam=1 and the other with hparam=2
Screenshots or logs
System information
A description of your system. Please provide:
- SageMaker Python SDK version: 2.239.3
- Framework name (eg. PyTorch) or algorithm (eg. KMeans): PyTorch
- Framework version: 2.3
- Python version: local python: 3.12.9 (machine that creates pipeline), estimator python: 3.11
- CPU or GPU: on both
- Custom Docker image (Y/N): N
Additional context
Possible solution:
- Remove or limit condition for frameworks where it really can't work: https://github.com/aws/sagemaker-python-sdk/blob/master/src/sagemaker/tuner.py#L1655