if __name__ == "__main__":
  flags.mark_flag_as_required("data_dir")
  flags.mark_flag_as_required("task_name")
  flags.mark_flag_as_required("vocab_file")
  flags.mark_flag_as_required("bert_config_file")
  flags.mark_flag_as_required("output_dir")
  tf.app.run()

apiVersion: sagemaker.aws.amazon.com/v1
kind: TrainingJob
metadata:
  # Maintan a job name as custom format
  name: tensorflow-1594280036
spec:
  # Maintain a Role which have sagemaker full access
  roleArn: arn:aws:iam::XXXXX:role/sagemaker
  # Change it to China Region (BJ: cn-north-1; NingXia: cn-northwest-1)
  region: cn-north-1
  algorithmSpecification:
    # Maintain a offical address, and you can search available image repo in https://aws.amazon.com/cn/releasenotes/available-deep-learning-containers-images/ , pls pay attention if using "gpu" image the instance for training must us ml.p3.* size.
    trainingImage: 763104351884.dkr.ecr.us-east-1.amazonaws.com/tensorflow-training:1.14-gpu-py3
    # Available option is File or PIPE
    trainingInputMode: File
  outputDataConfig:
    # The Module will upload to this address after training
    s3OutputPath: s3://your_bucket/output/models/
  inputDataConfig:
    # This spec is describe for training stage, and need to maintain where SageMaker can load the training Data. Need to change "S3Uri" in your local account. This Spec is required
    - channelName: training
      dataSource:
        s3DataSource:
          s3DataType: S3Prefix
          s3Uri: s3://your_bucket/DataSet/TraingJobName
          s3DataDistributionType: FullyReplicated
      compressionType: None
  # How Many instances is used for training
  resourceConfig:
    instanceCount: 2
    instanceType: ml.p3.8xlarge
    volumeSizeInGB: 50
  # Maintain the hyperParameters according job requirement
  hyperParameters:
    # This items is required by entrypoint paramters,this file is for bert configuration, the code will be load from s3://your_code_bucket/bert-tf/sourcedir.tar.gz and store in this path
    - name: bert_config_file
      value: "/opt/ml/code/bert_config.json"
    # This item is requred by entry point paramters, the SageMaker will load data form s3://your_bucket/DataSet/TraingJobName and will store in this path for training
    - name: data_dir
      value: "/opt/ml/input/data/training/"
    # How many iterations for eval
    - name: do_eval
      value: "1"
    # How many iterations for train
    - name: do_train
      value: "1"
    # The chekpoint address, the code is loading from s3://your_code_bucket/bert-tf/sourcedir.tar.gz, and the bert_model.ckpt should be store in your code.
    - name: init_checkpoint
      value: "/opt/ml/code/bert_model.ckpt"
    - name: learning_rate
      value: "2e-05"
    - name: max_seq_length
      value: "128"
    - name: model_dir
      value: "/opt/ml/model"
    # How many epochs your job is training
    - name: num_train_epochs
      value: "100.0"
    - name: output_dir
      value: "/opt/ml/model"
    - name: sagemaker_container_log_level
      value: "20"
    - name: sagemaker_enable_cloudwatch_metrics
      value: "true"
    - name: sagemaker_mpi_custom_mpi_options
      value: "-verbose --NCCL_DEBUG=INFO -x OMPI_MCA_btl_vader_single_copy_mechanism=none"
    - name: sagemaker_mpi_enabled
      value: "true"
    - name: sagemaker_mpi_num_of_processes_per_host
      value: "4"
    # The EntryPoint file name.
    - name: sagemaker_program
      value: "run_classifier_hvd.py"
    - name: sagemaker_region
      value: "cn-north-1"
    # Where SageMaker can load your code.
    - name: sagemaker_submit_directory
      value: "s3://your_code_bucket/bert-tf/sourcedir.tar.gz"
    - name: task_name
      value: "MRPC"
    - name: train_batch_size
      value: "32"
    # This is required by entrypoints paramters, and the code is load form s3://your_code_bucket/bert-tf/sourcedir.tar.gz.
    - name: vocab_file
      value: "/opt/ml/code/vocab.txt"
  # Once you use manage spot module, your have to set maxRuntimeInSeconds and maxWaitTimeInSeconds. and WaitTime must >= Runtime
  stoppingCondition:
    maxRuntimeInSeconds: 86400
    maxWaitTimeInSeconds: 172800
  # If using managment spots frame, need to maintain this value and checkpoint address
  enableManagedSpotTraining: true
  # the checkpoint address which located is S3
  checkpointConfig:
s3Uri: s3://your_bucket/bert-hvd-tf/checkpoints/