text: # Copyright 2020 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://github.com/aws/aws-neuron-sdk/blob/master/docs/neuron-container-tools/k8s-neuron-device-plugin-rbac.yml
  # Source: https://github.com/aws/aws-neuron-sdk/blob/master/docs/neuron-container-tools/k8s-neuron-device-plugin.yml
  
  kind: ClusterRole
  apiVersion: rbac.authorization.k8s.io/v1
  metadata:
    name: neuron-device-plugin
  rules:
    - apiGroups:
        - ""
      resources:
        - nodes
      verbs:
        - get
        - list
        - watch
    - apiGroups:
        - ""
      resources:
        - events
      verbs:
        - create
        - patch
    - apiGroups:
        - ""
      resources:
        - pods
      verbs:
        - update
        - patch
        - get
        - list
        - watch
    - apiGroups:
        - ""
      resources:
        - nodes/status
      verbs:
        - patch
        - update
  ---
  apiVersion: v1
  kind: ServiceAccount
  metadata:
    name: neuron-device-plugin
    namespace: kube-system
  ---
  kind: ClusterRoleBinding
  apiVersion: rbac.authorization.k8s.io/v1
  metadata:
    name: neuron-device-plugin
    namespace: kube-system
  roleRef:
    apiGroup: rbac.authorization.k8s.io
    kind: ClusterRole
    name: neuron-device-plugin
  subjects:
    - kind: ServiceAccount
      name: neuron-device-plugin
      namespace: kube-system
  
  ---
  apiVersion: apps/v1
  kind: DaemonSet
  metadata:
    name: neuron-device-plugin-daemonset
    namespace: kube-system
  spec:
    selector:
      matchLabels:
        name: neuron-device-plugin-ds
    updateStrategy:
      type: RollingUpdate
    template:
      metadata:
        annotations:
          scheduler.alpha.kubernetes.io/critical-pod: ""
        labels:
          name: neuron-device-plugin-ds
      spec:
        serviceAccount: neuron-device-plugin
        tolerations:
          - key: CriticalAddonsOnly
            operator: Exists
          - key: aws.amazon.com/neuron
            operator: Exists
            effect: NoSchedule
          - key: workload
            operator: Exists
            effect: NoSchedule
        # Mark this pod as a critical add-on; when enabled, the critical add-on
        # scheduler reserves resources for critical add-on pods so that they can
        # be rescheduled after a failure.
        # See https://kubernetes.io/docs/tasks/administer-cluster/guaranteed-scheduling-critical-addon-pods/
        priorityClassName: "system-node-critical"
        affinity:
          nodeAffinity:
            requiredDuringSchedulingIgnoredDuringExecution:
              nodeSelectorTerms:
                - matchExpressions:
                    - key: "beta.kubernetes.io/instance-type"
                      operator: In
                      values:
                        - inf1.xlarge
                        - inf1.2xlarge
                        - inf1.6xlarge
                        - inf1.4xlarge
                - matchExpressions:
                    - key: "node.kubernetes.io/instance-type"
                      operator: In
                      values:
                        - inf1.xlarge
                        - inf1.2xlarge
                        - inf1.6xlarge
                        - inf1.24xlarge
        containers:
          #Device Plugin containers are available both in us-east and us-west ecr
          #repos
          - image: $CORTEX_IMAGE_INFERENTIA
            imagePullPolicy: Always
            name: neuron-device-plugin
            env:
              - name: KUBECONFIG
                value: /etc/kubernetes/kubelet.conf
              - name: NODE_NAME
                valueFrom:
                  fieldRef:
                    fieldPath: spec.nodeName
            securityContext:
              allowPrivilegeEscalation: false
              capabilities:
                drop: ["ALL"]
            volumeMounts:
              - name: device-plugin
                mountPath: /var/lib/kubelet/device-plugins
              - name: infa-map
                mountPath: /run
            resources:
              requests:
                cpu: 100m
                memory: 100Mi
        nodeSelector:
          workload: "true"
        volumes:
          - name: device-plugin
            hostPath:
              path: /var/lib/kubelet/device-plugins
          - name: infa-map
            hostPath:
              path: /run
  