Skip to content

[SageMaker] Add support for HPO step in SM Pipelines #1336

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
thvasilo wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

[SageMaker] Add support for HPO step in SM Pipelines #1336

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion sagemaker/launch/launch_hyperparameter_tuning.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
168 changes: 157 additions & 11 deletions sagemaker/pipeline/create_sm_pipeline.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -33,8 +42,7 @@
CacheConfig,
ProcessingStep,
TrainingStep,
ProcessingInput,
ProcessingOutput,
TuningStep,
)

from pipeline_parameters import (
Expand All @@ -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__))


Expand Down Expand Up @@ -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)

Expand All @@ -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:
Expand All @@ -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,
Expand All @@ -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"
Expand Down Expand Up @@ -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(
Expand Down
104 changes: 101 additions & 3 deletions sagemaker/pipeline/pipeline_parameters.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,16 +22,17 @@
import logging
import os
from dataclasses import asdict, dataclass, fields, is_dataclass
from typing import List
from typing import List, Literal


JOB_ORDER = {
"gconstruct": 0,
"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__))
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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,
Expand Down
Loading