# for single gpu training
apiVersion: batch/v1
kind: Job
metadata:
  name: dpt-train
spec:
  template:
    spec:
      containers:
      - name: chris-dpt
        image: IMAGE
        command: ["stdbuf", "-i0", "-o0", "-e0", "/usr/bin/time", "-v", "bash", "-c"]
        args:
        - >-
          python3 
          /opt/repo/dynamic-inference/train_dpt.py
          -b 16
          -e 200
          -k
        volumeMounts:
        - name: git-repo
          mountPath: /opt/repo
        - name: christh9-pvc
          mountPath: /christh9-pvc
        - name: dshm
          mountPath: /dev/shm
        resources:
          limits:
            memory: 100Gi
            cpu: "6"
            nvidia.com/gpu: "4"
          requests:
            memory: 72Gi
            cpu: "4"
            nvidia.com/gpu: "4"  
      # nvme not available on A40 gpus
      # nodeSelector:
      #   nautilus.io/disktype: nvme  
      initContainers:
      - name: init-clone-repo
        image: INIT_IMAGE
        env:
        - name: GIT_USERNAME
          valueFrom:
            secretKeyRef:
              name: gitlab-secret
              key: user    
        - name: GIT_PASSWORD
          valueFrom:
            secretKeyRef:
              name: gitlab-secret
              key: password    
        command:
          - git
        args:
          - clone
          - --single-branch
          - --recurse-submodules
          - --branch
          - main
          - --remote-submodules
          - https://$(GIT_USERNAME):$(GIT_PASSWORD)@gitlab.nrp-nautilus.io/christh9/dynamic-inference
          - /opt/repo/dynamic-inference
        volumeMounts:
          - name: git-repo
            mountPath: /opt/repo
          - name: ssh-key
            mountPath: /secrets
      volumes:
      - name: ssh-key
        secret:
          secretName: cho-ssh-key
          defaultMode: 0400
      - name: git-repo
        emptyDir: {}
      - name: christh9-pvc
        persistentVolumeClaim:
          claimName: christh9-pvc
      - name: dshm
        emptyDir:
          medium: Memory
      restartPolicy: Never
      affinity:
        nodeAffinity:
          requiredDuringSchedulingIgnoredDuringExecution:
            nodeSelectorTerms:
              - matchExpressions:
                - key: gpu-type
                  operator: In
                  values:
                  - 'A40'
  backoffLimit: 0