# Copyright 2021 Cortex Labs, Inc.
#
# 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.

# Source: https://cloud.google.com/kubernetes-engine/docs/how-to/gpus#installing_drivers

apiVersion: apps/v1
kind: DaemonSet
metadata:
  name: nvidia-driver-installer
  namespace: kube-system
  labels:
    k8s-app: nvidia-driver-installer
spec:
  selector:
    matchLabels:
      k8s-app: nvidia-driver-installer
  updateStrategy:
    type: RollingUpdate
  template:
    metadata:
      labels:
        name: nvidia-driver-installer
        k8s-app: nvidia-driver-installer
    spec:
      affinity:
        nodeAffinity:
          requiredDuringSchedulingIgnoredDuringExecution:
            nodeSelectorTerms:
              - matchExpressions:
                  - key: cloud.google.com/gke-accelerator
                    operator: Exists
      tolerations:
        - key: nvidia.com/gpu
          operator: Exists
          effect: NoSchedule
        - key: workload
          operator: Exists
          effect: NoSchedule
      hostNetwork: true
      hostPID: true
      nodeSelector:
        workload: "true"
      volumes:
        - name: dev
          hostPath:
            path: /dev
        - name: vulkan-icd-mount
          hostPath:
            path: /home/kubernetes/bin/nvidia/vulkan/icd.d
        - name: nvidia-install-dir-host
          hostPath:
            path: /home/kubernetes/bin/nvidia
        - name: root-mount
          hostPath:
            path: /
        - name: cos-tools
          hostPath:
            path: /var/lib/cos-tools
      initContainers:
        - image: "cos-nvidia-installer:fixed"
          imagePullPolicy: Never
          name: nvidia-driver-installer
          resources:
            requests:
              cpu: "0.15"
          securityContext:
            privileged: true
          env:
            - name: NVIDIA_INSTALL_DIR_HOST
              value: /home/kubernetes/bin/nvidia
            - name: NVIDIA_INSTALL_DIR_CONTAINER
              value: /usr/local/nvidia
            - name: VULKAN_ICD_DIR_HOST
              value: /home/kubernetes/bin/nvidia/vulkan/icd.d
            - name: VULKAN_ICD_DIR_CONTAINER
              value: /etc/vulkan/icd.d
            - name: ROOT_MOUNT_DIR
              value: /root
            - name: COS_TOOLS_DIR_HOST
              value: /var/lib/cos-tools
            - name: COS_TOOLS_DIR_CONTAINER
              value: /build/cos-tools
          volumeMounts:
            - name: nvidia-install-dir-host
              mountPath: /usr/local/nvidia
            - name: vulkan-icd-mount
              mountPath: /etc/vulkan/icd.d
            - name: dev
              mountPath: /dev
            - name: root-mount
              mountPath: /root
            - name: cos-tools
              mountPath: /build/cos-tools
      containers:
        - image: "$CORTEX_IMAGE_GOOGLE_PAUSE"
          name: pause
