apiVersion: sagemaker.aws.amazon.com/v1
kind: TrainingJob
metadata:
  name: fsx-kmeans-mnist
spec:
  hyperParameters:
    - name: k
      value: "10"
    - name: feature_dim
      value: "784"
  algorithmSpecification:
    trainingImage: 174872318107.dkr.ecr.us-west-2.amazonaws.com/kmeans:1
    trainingInputMode: File
  roleArn: {ROLE_ARN}
  region: us-west-2
  outputDataConfig:
    s3OutputPath: s3://{DATA_BUCKET}/kmeans_mnist_example/output
  resourceConfig:
    instanceCount: 1
    instanceType: ml.m4.xlarge
    volumeSizeInGB: 5
  stoppingCondition:
    maxRuntimeInSeconds: 86400
  inputDataConfig:
    - channelName: train
      dataSource:
        fileSystemDataSource:
          fileSystemType: FSxLustre
          fileSystemAccessMode: ro
          fileSystemId: "{FSX_ID}"
          directoryPath: /fsx/kmeans_mnist_example/input
      contentType: "text/csv;label_size=0"
      compressionType: None
  vpcConfig:
    securityGroupIds: 
      - sg-e6b84ca2
    subnets: 
      - subnet-187e9960
      - subnet-21d8d86a
      - subnet-91b1d8ba