apiVersion: batch/v1
kind: Job
metadata:
  name: simple-training
spec:
  backoffLimit: 1
  template:
    metadata:
      annotations:
        k8s.v1.cni.cncf.io/networks: rdma-net
    spec:
      securityContext:
        runAsUser: 0
      containers:
      - name: mnist-training
        image: graphcore/pytorch-tutorials:latest
        workingDir: "/opt/tutorials/simple_applications/pytorch/mnist"
        command:
        - "bash"
        args:
        - "-c"
        - "pip3 install -r requirements.txt && python3 mnist_poptorch.py --epochs=1"
        env:
        - name: IPUOF_VIPU_API_HOST
          value: "localhost"
        - name: IPUOF_VIPU_API_PORT
          value: "8090"
        - name: IPUOF_VIPU_API_PARTITION_ID
          value: "training_partition"
        securityContext:
          capabilities:
            add: [ "IPC_LOCK" ]
        resources:
          limits:
            rdma/rdma_shared_device_a: 1
          requests:
            rdma/rdma_shared_device_a: 1
      restartPolicy: OnFailure
