apiVersion: sagemaker.aws.amazon.com/v1
kind: TrainingJob
metadata:
  name: xgboost-mnist-debugger
spec:
  hyperParameters:
    - name: max_depth
      value: "5"
    - name: eta
      value: "0.2"
    - name: gamma
      value: "4"
    - name: min_child_weight
      value: "6"
    - name: silent
      value: "0"
    - name: objective
      value: reg:squarederror
    - name: subsample
      value: "0.7"
    - name: num_round
      value: "51"
  algorithmSpecification:
    trainingImage: 246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:0.90-2-cpu-py3
    trainingInputMode: File
  roleArn: arn:aws:iam::123456789012:role/service-role/AmazonSageMaker-ExecutionRole
  region: us-west-2
  outputDataConfig:
    s3OutputPath: s3://my-bucket/xgboost-debugger/output
  resourceConfig:
    instanceCount: 1
    instanceType: ml.m4.xlarge
    volumeSizeInGB: 5
  stoppingCondition:
    maxRuntimeInSeconds: 86400
  inputDataConfig:
    - channelName: train
      dataSource:
        s3DataSource:
          s3DataType: S3Prefix
          s3Uri: s3://my-bucket/xgboost-debugger/train
          s3DataDistributionType: FullyReplicated
      contentType: libsvm
      compressionType: None
    - channelName: validation
      dataSource:
        s3DataSource:
          s3DataType: S3Prefix
          s3Uri: s3://my-bucket/xgboost-debugger/validation
          s3DataDistributionType: FullyReplicated
      contentType: libsvm
      compressionType: None
  debugHookConfig:
    s3OutputPath: s3://my-bucket/xgboost-debugger/hookconfig
    collectionConfigurations:
      - collectionName: feature_importance
        collectionParameters:
          - name: save_interval
            value: "5"
      - collectionName: losses
        collectionParameters:
          - name: save_interval"
            value: "500"
      - collectionName: average_shap
        collectionParameters:
          - name: save_interval
            value: "5"
      - collectionName: metrics
        collectionParameters:
          - name: save_interval
            value: "5"
  debugRuleConfigurations:
    - ruleConfigurationName: LossNotDecreasing
      ruleEvaluatorImage: 895741380848.dkr.ecr.us-west-2.amazonaws.com/sagemaker-debugger-rules:latest
      ruleParameters:
        - name: collection_names
          value: metrics
        - name: num_steps
          value: "10"
        - name: rule_to_invoke
          value: LossNotDecreasing