# This yaml file will launch a container with any NVIDIA Ampere GPU that has more than 16 GB of GPU memory
# Specifying the nvidia.com/gpu.family and the nvidia.com/gpu.memory amount 
#    restricts the type of GPU this container will run on
#
# Specifying both an affinity and resource type allows for some flexibility in the GPU this workload will execute on
#       but it does not leave things to chance if there are many GPU types available
apiVersion: v1
kind: Pod
metadata:
  name: gpu-pod
spec:
  affinity:
    nodeAffinity:
      requiredDuringSchedulingIgnoredDuringExecution:
        nodeSelectorTerms:
        - matchExpressions:
          - key: nvidia.com/gpu.family
            operator: In
            values:
            - ampere
          - key: nvidia.com/gpu.memory
            operator: Gt
            values:
            - "16384"
  containers:
    - name: gpu-pod
      image:  nvcr.io/nvidia/k8s/cuda-sample:nbody
      command: ["/bin/sh"]
      args: ["-c", "nvidia-smi"]
      resources:
        limits:
          nvidia.com/gpu: 1
