--- 
apiVersion: v1
kind: Pod
metadata:
  name: bert-csi-ebs
spec:
  volumes:
  - name: persistent-storage
    persistentVolumeClaim:
      claimName: ebs-claim
  containers: 
    - name: bert
      command: 
        - python
        - ./code/train.py
        - --train_steps_per_epoch=1
        - --epochs=1
        - --learning_rate=0.00001
        - --epsilon=0.00000001
        - --train_batch_size=36
        - --validation_batch_size=18
        - --test_batch_size=18
        - --train_steps_per_epoch=1
        - --validation_steps=1
        - --test_steps=1
        - --use_xla=True
        - --use_amp=False
        - --max_seq_length=64
        - --freeze_bert_layer=True
        - --enable_sagemaker_debugger=False
        - --enable_checkpointing=False
        - --enable_tensorboard=False
        - --run_validation=True
        - --run_test=False
        - --run_sample_predictions=False
      image: "231218423789.dkr.ecr.us-west-2.amazonaws.com/dlc-demo:bert"
      imagePullPolicy: Always
      env: 
        - name: SM_TRAINING_ENV
          value: "{\"is_master\":true}"
        - name: SAGEMAKER_JOB_NAME
          value: "tf-bert-training-eks"
        - name: SM_CURRENT_HOST
          value: "localhost"
        - name: SM_NUM_GPUS
          value: "0"
        - name: SM_HOSTS
          value: "{\"hosts\":\"localhost\"}"
        - name: SM_MODEL_DIR
          value: "/opt/ml/model/"     
        - name: SM_OUTPUT_DIR
          value: "/opt/ml/output/"
        - name: SM_OUTPUT_DATA_DIR
          value: "/opt/ml/output/data/"
        - name: SM_CHANNEL_TRAIN
          value: "/input/data/train"
        - name: SM_CHANNEL_VALIDATION
          value: "/input/data/validation"     
        - name: SM_CHANNEL_TEST
          value: "/input/data/test"
      volumeMounts:
      - mountPath: /opt/ml
        name: persistent-storage
  restartPolicy: Never 
# TODO: Need to map all of the directories above using PVs/PVCs and local host provisioner


# https://docs.aws.amazon.com/sagemaker/latest/dg/amazon-sagemaker-toolkits.html

# /opt/ml
# ├── input
# │   ├── config
# │   │   ├── hyperparameters.json
# │   │   └── resourceConfig.json
# │   └── data
# │       └── <channel_name>
# │           └── <input data>
# ├── model
# │
# ├── code
# │
# ├── output
# │
# └── failure


# https://github.com/aws/sagemaker-tensorflow-training-toolkit/blob/d2e8848e657ebc4a3c2a71a88b31275e1c1de403/src/sagemaker_tensorflow_container/training.py#L130# https://github.com/aws/sagemaker-
# https://github.com/aws/sagemaker-tensorflow-training-toolkit/blob/d2e8848e657ebc4a3c2a71a88b31275e1c1de403/src/sagemaker_tensorflow_container/training.py#L228
