diff --git a/sagemaker/launch/launch_hyperparameter_tuning.py b/sagemaker/launch/launch_hyperparameter_tuning.py index 2e0c24a460..d0047faa08 100644 --- a/sagemaker/launch/launch_hyperparameter_tuning.py +++ b/sagemaker/launch/launch_hyperparameter_tuning.py @@ -105,7 +105,7 @@ def get_metric_definitions( def parse_hyperparameter_ranges( hyperparameter_ranges_json: str, ) -> dict[str, ParameterRange]: - """Parse the hyperparameter ranges from JSON string and determine if autotune is needed. + """Parse the hyperparameter ranges from JSON file or string. Expected JSON structure is at https://docs.aws.amazon.com/sagemaker/latest/dg/automatic-model-tuning-define-ranges.html diff --git a/sagemaker/pipeline/create_sm_pipeline.py b/sagemaker/pipeline/create_sm_pipeline.py index 9819211873..e230a4e544 100644 --- a/sagemaker/pipeline/create_sm_pipeline.py +++ b/sagemaker/pipeline/create_sm_pipeline.py @@ -22,9 +22,18 @@ from typing import List, Optional, Sequence, Union import boto3 -from sagemaker.processing import ScriptProcessor +from sagemaker.processing import ( + ProcessingInput, + ProcessingOutput, + ScriptProcessor, +) from sagemaker.spark.processing import PySparkProcessor from sagemaker.pytorch.estimator import PyTorch +from sagemaker.tuner import ( + HyperbandStrategyConfig, + HyperparameterTuner, + StrategyConfig, +) from sagemaker.workflow.functions import Join from sagemaker.workflow.parameters import ParameterInteger, ParameterString from sagemaker.workflow.pipeline import Pipeline @@ -33,8 +42,7 @@ CacheConfig, ProcessingStep, TrainingStep, - ProcessingInput, - ProcessingOutput, + TuningStep, ) from pipeline_parameters import ( @@ -43,6 +51,17 @@ save_pipeline_args, ) +# TODO: We need this to be a module to be able to import it +# from launch.launch_hyperparameter_tuning import ( +# get_metric_definitions, +# parse_hyperparameter_ranges, +# ) +# For now duplicate code from launch/launch_hyperparameter_tuning.py +from tuning_utils import ( + get_metric_definitions, + parse_hyperparameter_ranges, +) + SCRIPT_PATH = os.path.dirname(os.path.realpath(__file__)) @@ -249,6 +268,24 @@ def _create_pipeline_steps( pipeline_steps: List[Union[ProcessingStep, TrainingStep]] = [] for job in self.args.task_config.jobs_to_run: step = self._create_step(job, args) + # TODO: Make cache invalidation more robust + # Processing steps do not invalidate cache when only + # the output changes, we should add an env var to ensure + # that happens. + + # https://docs.aws.amazon.com/sagemaker/latest/dg/pipelines-default-keys.html + # Processing steps will only invalidate cache when one of the following changes + # * AppSpecification + # * Environment + # * ProcessingInputs. + # Here we propose adding an env var that helps invalidate the cache when the output changes: + # if isinstance(step, ProcessingStep): + # assert step.processor + # if step.processor.env is None: + # step.processor.env = {} + # step.processor.env.update( + # {"SM_PIPELINE_STEP_OUTPUTS": [output.destination or "" for output in step.outputs or []]} + # ) if step: pipeline_steps.append(step) @@ -274,10 +311,12 @@ def _create_step( "gsprocessing": self._create_gsprocessing_step, "dist_part": self._create_dist_part_step, "gb_convert": self._create_gb_convert_step, + "hpo": self._create_tuning_step, "train": self._create_train_step, "inference": self._create_inference_step, } step_generator = step_generators.get(job_type) + assert step_generator, f"Unknown job type: {job_type}" return step_generator(args) if step_generator else None def _create_gconstruct_step(self, args: PipelineArgs) -> ProcessingStep: @@ -287,6 +326,14 @@ def _create_gconstruct_step(self, args: PipelineArgs) -> ProcessingStep: "Use --graph-construction-config-filename" ) + gconstruct_s3_output = Join( + on="/", + values=[ + self.output_subpath, + "gconstruct", + ], + ) + gconstruct_processor = ScriptProcessor( image_uri=args.aws_config.graphstorm_pytorch_cpu_image_uri, role=args.aws_config.execution_role, @@ -295,14 +342,7 @@ def _create_gconstruct_step(self, args: PipelineArgs) -> ProcessingStep: command=["python3"], sagemaker_session=self.pipeline_session, volume_size_in_gb=self.volume_size_gb_param, - ) - - gconstruct_s3_output = Join( - on="/", - values=[ - self.output_subpath, - "gconstruct", - ], + env={"GCONSTRUCT_OUTPUT": gconstruct_s3_output}, ) gc_local_input_path = "/opt/ml/processing/input" @@ -527,6 +567,112 @@ def _create_gb_convert_step(self, args: PipelineArgs) -> ProcessingStep: # No need to update next step input for GB convert return gb_convert_step + def _create_tuning_step(self, args: PipelineArgs) -> TuningStep: + # Implementation for HPO tuning step + hpo_output_path = Join( + on="/", + values=[ + self.output_subpath, + "hpo", + ], + ) + + train_params = { + "eval-metric": args.tuning_config.metric_name, + "graph-data-s3": self.next_step_data_input, + "graph-name": self.graph_name_param, + "log-level": args.task_config.log_level, + "num-trainers": self.num_trainers_param, + "task-type": args.training_config.train_inference_task, + "topk-model-to-save": "1", + "train-yaml-s3": self.train_config_file_param, + "use-graphbolt": args.training_config.use_graphbolt_str, + } + + # Add strategy-specific parameters + if args.tuning_config.strategy == "Hyperband": + # Hyperband must control max epochs and early stopping + # Set num-epochs of GraphStorm to hyperband max epoch + # Disable early-stop of GraphStorm + train_params.update( + { + "num-epochs": args.tuning_config.hb_max_epochs, + "use_early_stop": "false", + } + ) + + # Base train estimator for all HPO jobs + train_estimator = PyTorch( + entry_point=os.path.basename(args.script_paths.train_script), + source_dir=os.path.dirname(args.script_paths.train_script), + image_uri=self.train_infer_image, + role=args.aws_config.execution_role, + instance_count=self.instance_count_param, + instance_type=self.train_infer_instance, + output_path=hpo_output_path, # All tuners should create output under "hpo" + py_version="py3", + hyperparameters=train_params, + sagemaker_session=self.pipeline_session, + disable_profiler=True, + debugger_hook_config=False, + volume_size=self.volume_size_gb_param, # Do not compress model output + disable_output_compression=True + ) + + hyperparameter_ranges = parse_hyperparameter_ranges( + args.tuning_config.hyperparameter_ranges + ) + + # Get metric definitions based on strategy + metric_definitions = get_metric_definitions( + args.tuning_config.metric_name, + args.tuning_config.eval_mask, + args.tuning_config.strategy, + ) + # Configure the tuner + tuner_config = { + "estimator": train_estimator, + # Use first metric as objective + "objective_metric_name": metric_definitions[0]["Name"], + "hyperparameter_ranges": hyperparameter_ranges, + "objective_type": args.tuning_config.objective_type, + "max_jobs": args.tuning_config.max_jobs, + "max_parallel_jobs": args.tuning_config.max_parallel_jobs, + "metric_definitions": metric_definitions, + "strategy": args.tuning_config.strategy, + } + + # Add Hyperband-specific configuration if needed + if args.tuning_config.strategy == "Hyperband": + tuner_config["strategy_config"] = StrategyConfig( + HyperbandStrategyConfig( + max_resource=args.tuning_config.hb_max_epochs, + min_resource=args.tuning_config.hb_min_epochs, + ) + ) + + tuning_step = TuningStep( + name="Tuning", + tuner=HyperparameterTuner(**tuner_config), + cache_config=self.cache_config, + inputs={"train": self.train_config_file_param}, + ) + + # For selecting the best model see + # https://sagemaker.readthedocs.io/en/stable/amazon_sagemaker_model_building_pipeline.html#tuningstep + best_model_path = Join( + on="/", + values=[ + hpo_output_path, + tuning_step.properties.BestTrainingJob.TrainingJobName, + "output/model" + ], + ) + + self.model_input_path = best_model_path + + return tuning_step + def _create_train_step(self, args: PipelineArgs) -> TrainingStep: # Implementation for Training step model_output_path = Join( diff --git a/sagemaker/pipeline/pipeline_parameters.py b/sagemaker/pipeline/pipeline_parameters.py index 7f220e7306..1c61ff1e19 100644 --- a/sagemaker/pipeline/pipeline_parameters.py +++ b/sagemaker/pipeline/pipeline_parameters.py @@ -22,7 +22,7 @@ import logging import os from dataclasses import asdict, dataclass, fields, is_dataclass -from typing import List +from typing import List, Literal JOB_ORDER = { @@ -30,8 +30,9 @@ "gsprocessing": 1, "dist_part": 2, "gb_convert": 3, - "train": 4, - "inference": 5, + "hpo": 4, + "train": 5, + "inference": 6, } SCRIPT_PATH = os.path.dirname(os.path.realpath(__file__)) @@ -239,6 +240,28 @@ class TrainingConfig: def __post_init__(self): self.use_graphbolt_str = self.use_graphbolt_str.lower() +@dataclass +class TuningConfig: + """Configuration for the tuning step. + + """ + + max_jobs: int + max_parallel_jobs: int + hyperparameter_ranges: str + metric_name: str + eval_mask: Literal["test", "val"] + objective_type: Literal["Maximize", "Minimize"] + strategy: Literal["Bayesian", "Random", "Hyperband", "Grid"] + hb_min_epochs: int = 1 + hb_max_epochs: int = 20 + + def __post_init__(self): + assert self.eval_mask in ["test", "val"] + assert self.objective_type in ["Maximize", "Minimize"] + assert self.strategy in ["Bayesian", "Random", "Hyperband", "Grid"] + assert self.hb_min_epochs >= 1 + assert self.hb_max_epochs >= self.hb_min_epochs @dataclass class InferenceConfig: @@ -308,6 +331,8 @@ class PipelineArgs: Partition configuration settings. training_config : TrainingConfig Training configuration settings. + tuning_config: TuningCofig + HPO tuning jobs settings. inference_config : InferenceConfig Inference configuration settings. script_paths : ScriptPaths @@ -324,6 +349,7 @@ class PipelineArgs: task_config: TaskConfig partition_config: PartitionConfig training_config: TrainingConfig + tuning_config: TuningConfig inference_config: InferenceConfig script_paths: ScriptPaths step_cache_expiration: str @@ -755,6 +781,67 @@ def parse_pipeline_args() -> PipelineArgs: help="Whether to use GraphBolt. Default: 'false'", ) + # HPO tuning configuration + hpo_group = parser.add_argument_group("Hyperparameter tuning arguments") + + hpo_group.add_argument( + "--hpo-max-jobs", + type=int, + default=10, + help="Maximum number of training jobs to run", + ) + hpo_group.add_argument( + "--hpo-max-parallel-jobs", + type=int, + default=2, + help="Maximum number of parallel training jobs", + ) + hpo_group.add_argument( + "--hyperparameter-ranges", + type=str, + help="Path to a JSON file, or a JSON string defining hyperparameter ranges. " + "For syntax see 'Dynamic hyperparameters' in " + "https://docs.aws.amazon.com/sagemaker/latest/dg/automatic-model-tuning-define-ranges.html", + ) + hpo_group.add_argument( + "--hpo-metric-name", + type=str, + help="Name of the metric to optimize for (e.g., 'accuracy', 'amri')", + ) + hpo_group.add_argument( + "--hpo-eval-mask", + type=str, + choices=["test", "val"], + help="Whether to use test or validation metrics for HPO.", + ) + hpo_group.add_argument( + "--hpo-objective-type", + type=str, + default="Maximize", + choices=["Maximize", "Minimize"], + help="Type of objective, can be 'Maximize' or 'Minimize'", + ) + hpo_group.add_argument( + "--hpo-strategy", + type=str, + default="Bayesian", + choices=["Bayesian", "Random", "Hyperband", "Grid"], + help="Optimization strategy. Default: 'Bayesian'.", + ) + # Hyperband-specific arguments + hpo_group.add_argument( + "--hpo-hb-min-epochs", + type=int, + default=1, + help="Minimum number of epochs for Hyperband strategy. Default: 1", + ) + hpo_group.add_argument( + "--hpo-hb-max-epochs", + type=int, + default=20, + help="Maximum number of epochs for Hyperband strategy. Default: 20", + ) + # Inference Configuration optional_args.add_argument( "--inference-yaml-s3", @@ -867,6 +954,17 @@ def parse_pipeline_args() -> PipelineArgs: train_yaml_file=args.train_yaml_s3, use_graphbolt_str=args.use_graphbolt, ), + tuning_config=TuningConfig( + args.hpo_max_jobs, + args.hpo_max_parallel_jobs, + args.hyperparameter_ranges, + args.hpo_metric_name, + args.hpo_eval_mask, + args.hpo_objective_type, + args.hpo_strategy, + args.hpo_hb_min_epochs, + args.hpo_hb_max_epochs, + ), inference_config=InferenceConfig( save_predictions=args.save_predictions, save_embeddings=args.save_embeddings, diff --git a/sagemaker/pipeline/tuning_utils.py b/sagemaker/pipeline/tuning_utils.py new file mode 100644 index 0000000000..bade0d9d45 --- /dev/null +++ b/sagemaker/pipeline/tuning_utils.py @@ -0,0 +1,139 @@ +""" +Copyright Contributors + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Duplicates the HPO utilities from launch/launch_hyperparameter_tuning.py +""" + +import json +import logging +import os +from typing import Literal + +from sagemaker.parameter import ( + ParameterRange, + ContinuousParameter, + IntegerParameter, + CategoricalParameter, +) + +def get_metric_definitions( + metric_name: str, + eval_mask: Literal["val", "test"], + strategy: Literal["Grid", "Random", "Bayesian", "Hyperband"], +) -> list[dict]: + """Generate metric definitions based on the HPO strategy. + + Parameters + ---------- + metric_name : str + Name of the metric to optimize, as it appears in GraphStorm logs (e.g., 'roc_auc') + eval_mask : str + Dataset mask to use when collecting metric values ('test' or 'val') + strategy : str + HPO strategy ('Bayesian', 'Random', 'Hyperband', or 'Grid') + + Returns + ------- + list[dict] + List of metric definitions for SageMaker HPO + """ + base_metric_name = f"{eval_mask.lower()}_{metric_name.lower()}" + + if strategy == "Hyperband": + # For Hyperband, capture intermediate metrics during training + test_set_string = "Validation" if eval_mask.lower() == "val" else "Test" + return [ + { + "Name": f"{base_metric_name}", + "Regex": f"INFO:root:Step [0-9]+ \\| {test_set_string} {metric_name}: ([0-9\\.]+)", + } + ] + else: + # For other strategies, only capture the final best metric + return [ + { + "Name": f"best_{base_metric_name}", + "Regex": f"INFO:root:best_{eval_mask.lower()}_score: {{'{metric_name}': ([0-9\\.]+)}}", + } + ] + + +def parse_hyperparameter_ranges( + hyperparameter_ranges_json: str, +) -> dict[str, ParameterRange]: + """Parse the hyperparameter ranges from JSON file or string. + + Expected JSON structure is at + https://docs.aws.amazon.com/sagemaker/latest/dg/automatic-model-tuning-define-ranges.html + + Parameters + ---------- + hyperparameter_ranges_json : str + Path to a JSON file or JSON as a string. + + Returns + ------- + dict + Hyperparameter dict that can be used to create a HyperparameterTuner + object. + + Raises + ------ + ValueError + If the JSON dict contains an invalid parameter type. + + .. versionadded:: 0.4.1 + """ + if os.path.exists(hyperparameter_ranges_json): + with open(hyperparameter_ranges_json, "r", encoding="utf-8") as f: + str_ranges: dict[str, list[dict]] = json.load(f)["ParameterRanges"] + else: + try: + str_ranges = json.loads(hyperparameter_ranges_json)["ParameterRanges"] + except json.decoder.JSONDecodeError as e: + logging.fatal("Invalid JSON string: %s", e) + raise e + + hyperparameter_ranges: dict[str, ParameterRange] = {} + for params_type, config_list in str_ranges.items(): + if params_type == "ContinuousParameterRanges": + for config in config_list: + param_name = config["Name"] + hyperparameter_ranges[param_name] = ContinuousParameter( + config["MinValue"], + config["MaxValue"], + config.get("ScalingType", "Auto"), + ) + elif params_type == "IntegerParameterRanges": + for config in config_list: + param_name = config["Name"] + hyperparameter_ranges[param_name] = IntegerParameter( + config["MinValue"], + config["MaxValue"], + config.get("ScalingType", "Auto"), + ) + elif params_type == "CategoricalParameterRanges": + for config in config_list: + param_name = config["Name"] + hyperparameter_ranges[param_name] = CategoricalParameter( + config["Values"] + ) + else: + raise ValueError( + f"Unknown parameter type {params_type}. " + "Expect one of 'CategoricalParameterRanges', 'ContinuousParameterRanges', " + "'IntegerParameterRanges'" + ) + return hyperparameter_ranges \ No newline at end of file