# Copyright 2018 Google Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

apiVersion: batch/v1
kind: Job
metadata:
  name: resnet-gpu
spec:
  template:
    spec:
      restartPolicy: Never
      volumes:
      - name: nfs-pvc
        persistentVolumeClaim:
          claimName: nfs
      containers:
      - name: resnet-gpu
        image: gcr.io/vishnuk-cloud/tensorflow@sha256:c5655a73205a447e1552fc833d14ab66eee2304dbf64846edd88c41434c6f5f2
        command:
          - python
          - /tensorflow_models/models/official/resnet/resnet_main.py
          - --use_tpu=False
          - --tpu=
          - --precision=float32
          - --data_dir=$(DATA_BUCKET)
          - --model_dir=$(MODEL_BUCKET)
          - --train_batch_size=$(TRAIN_BATCH_SIZE)
          - --train_steps=$(TRAIN_STEPS)
        volumeMounts:
        - name: nfs-pvc
          mountPath: /home/vishnuk
        env:
          # The Google Cloud Storage location where the fake ImageNet dataset is
          # stored.
        - name: DATA_BUCKET
          value: "/home/vishnuk/tf-demo/fake_imagenet/"
          # [REQUIRED] Must specify the Google Cloud Storage location where your
          # output model will be stored.
        - name: MODEL_BUCKET
          value: "/home/vishnuk/tf-demo/train/resnet-train" # "/home/vishnuk/tf-demo/resnet-train/"
        - name: TRAIN_BATCH_SIZE
          value: "32"
        - name: TRAIN_STEPS
          value: "400365"
        resources:
          limits:
            nvidia.com/gpu: 1
