How to configure dynamically a AWS Sage Maker task with AWS Step Function - amazon-web-services

Question

I'm trying to build an ML pipeline using AWS Step Function.
I would like to configure the 'CreateHyperParameterTuningJob' dynamically depending on the input of the task.
Here is a screenshot of the State Machine that I'm trying to build:
ML State Machine
When I try to create this State Machine, I got the following error:
The value for the field 'MaxParallelTrainingJobs' must be an INTEGER
I'm struggling to figure out what is the issue here.
Do you have any suggestion to make the SM configuration dynamic with Step Function? Is it even possible?
Here is the input data passed to the 'Run training job' task:
{
"client_id": "test",
"training_job_definition": {
"AlgorithmSpecification": {
"TrainingImage": "433757028032.dkr.ecr.us-west-2.amazonaws.com/xgboost:latest",
"TrainingInputMode": "File"
},
"ResourceConfig": {
"InstanceCount": 1,
"InstanceType": "ml.m5.large",
"VolumeSizeInGB": 5
},
"StaticHyperParameters": {
"num_round": 750
},
"StoppingCondition": {
"MaxRuntimeInSeconds": 900
},
"InputDataConfig": [
{
"ChannelName": "train",
"CompressionType": "None",
"ContentType": "csv",
"DataSource": {
"S3DataSource": {
"S3DataDistributionType": "FullyReplicated",
"S3DataType": "S3Prefix",
"S3Uri": "..."
}
}
},
{
"ChannelName": "validation",
"CompressionType": "None",
"ContentType": "csv",
"DataSource": {
"S3DataSource": {
"S3DataDistributionType": "FullyReplicated",
"S3DataType": "S3Prefix",
"S3Uri": "..."
}
}
}
],
"OutputDataConfig": {
"S3OutputPath": "..."
},
"RoleArn": "arn:aws:iam::679298748479:role/landingzone_sagemaker_role"
},
"hyper_parameter_tuning_job_config": {
"HyperParameterTuningJobObjective": {
"MetricName": "validation:rmse",
"Type": "Minimize"
},
"Strategy": "Bayesian",
"ResourceLimits": {
"MaxParallelTrainingJobs": 2,
"MaxNumberOfTrainingJobs": 10
},
"ParameterRanges": {
"ContinuousParameterRanges": [
{
"Name": "eta",
"MinValue": 0.01,
"MaxValue": 0.04
},
{
"Name": "gamma",
"MinValue": 0,
"MaxValue": 100
},
{
"Name": "subsample",
"MinValue": 0.6,
"MaxValue": 1
},
{
"Name": "lambda",
"MinValue": 0,
"MaxValue": 5
},
{
"Name": "alpha",
"MinValue": 0,
"MaxValue": 2
}
],
"IntegerParameterRanges": [
{
"Name": "max_depth",
"MinValue": 5,
"MaxValue": 10
}
]
}
}
}
Here is JSON file that describes the State Machine:
{
"StartAt": "Generate Training Dataset",
"States": {
"Generate Training Dataset": {
"Resource": "arn:aws:lambda:uswest-2:012345678912:function:StepFunctionsSample-SageMaLambdaForDataGeneration-1TF67BUE5A12U",
"Type": "Task",
"Next": "Run training job"
},
"Run training job": {
"Resource": "arn:aws:states:::sagemaker:createHyperParameterTuningJob.sync",
"Parameters": {
"HyperParameterTuningJobName.$": "$.execution_date",
"HyperParameterTuningJobConfig": {
"HyperParameterTuningJobObjective": {
"MetricName": "$.hyper_parameter_tuning_job_config.HyperParameterTuningJobObjective.MetricName",
"Type": "Minimize"
},
"Strategy": "$.hyper_parameter_tuning_job_config.Strategy",
"ResourceLimits": {
"MaxParallelTrainingJobs": "$.hyper_parameter_tuning_job_config.ResourceLimits.MaxParallelTrainingJobs",
"MaxNumberOfTrainingJobs": "$.hyper_parameter_tuning_job_config.ResourceLimits.MaxNumberOfTrainingJobs"
},
"ParameterRanges": "$.hyper_parameter_tuning_job_config.ParameterRanges"
},
"TrainingJobDefinition": {
"AlgorithmSpecification": "$.training_job_definition.AlgorithmSpecification",
"StoppingCondition": "$.training_job_definition.StoppingCondition",
"ResourceConfig": "$.training_job_definition.ResourceConfig",
"RoleArn": "$.training_job_definition.RoleArn",
"InputDataConfig": "$.training_job_definition.InputDataConfig",
"OutputDataConfig": "$.training_job_definition.OutputDataConfig",
"StaticHyperParameters": "$.training_job_definition.StaticHyperParameters"
},
"HyperParameterTuningJobConfig.ResourceLimits": ""
},
"Type": "Task",
"End": true
}
}
}