training_job.json

{
  "Comment": "Iterator State Machine Example",
  "StartAt": "ExtractExecutionContext",
  "States": {

    "ExtractExecutionContext": {
      "Parameters": {
        "execution_name.$": "$$.Execution.Name",
        "execution_date.$": "$$.Execution.StartTime"
      },
      "Type": "Pass",
      "ResultPath": "$.unique_id",
      "Next": "Train"
    },

    "Train": {
      "Type": "Task",
      "Resource": "arn:aws:states:::sagemaker:createTrainingJob.sync",
      "Parameters": {
        "TrainingJobName.$": "States.Format('Train-{}', $.unique_id.execution_name)",
        "ResourceConfig": {
          "InstanceCount": 1,
          "InstanceType": "ml.m4.xlarge",
          "VolumeSizeInGB": 5
        },
        "HyperParameters": {
          "episodes": "10000",
          "model_dir": "/opt/ml/model/models",
          "games_dir": "/opt/ml/model/games",
          "lr": "0.2",
          "gamma": "0.95",
          "epsilon": "1",
          "epsilon_final": "0.5",
          "epsilon_decay": "0.995",
          "batch_length": "8",
          "epochs": "2"
        },
        "AlgorithmSpecification": {
          "TrainingImage.$": "States.Format('721396535024.dkr.ecr.us-east-1.amazonaws.com/lux-clarai:{}', $.experiment_name)",
          "TrainingInputMode": "File",
          "MetricDefinitions" : [
                {
                    "Name": "train:total_rounds",
                    "Regex": "rounds::(.*?);"
                },
                {
                    "Name": "train:units_created",
                    "Regex": "units::(.*?);"
                },
                {
                    "Name": "train:cities_created",
                    "Regex": "cities::(.*?);"
                },
                {
                    "Name": "train:winner",
                    "Regex": "winner::(.*?);"
                },
                {
                    "Name": "train:loss",
                    "Regex": "epoch \\d+::(.*?);"
              }
            ]
        },
        "OutputDataConfig": {
          "S3OutputPath.$": "States.Format('s3://lux-lucia/clarai/experiments/{}/models', $.experiment_name)"
        },
        "StoppingCondition": {
          "MaxRuntimeInSeconds": 100000
        },
        "RoleArn": "arn:aws:iam::721396535024:role/AmazonSageMaker-ExecutionRole-lux-clarai-trainjob"
      },
      "ResultPath": "$.result",
      "Next": "Done"
    },

    "Done": {
      "Type": "Pass",
      "End": true
    }
  }
}