---
apiVersion: apps/v1
kind: DaemonSet
metadata:
  name: {{ include "tpu-dra-driver.fullname" . }}-kubeletplugin
  namespace: {{ include "tpu-dra-driver.namespace" . }}
  labels:
    {{- include "tpu-dra-driver.labels" . | nindent 4 }}
spec:
  selector:
    matchLabels:
      {{- include "tpu-dra-driver.selectorLabels" . | nindent 6 }}
  {{- with .Values.kubeletPlugin.updateStrategy }}
  updateStrategy:
    {{- toYaml . | nindent 4 }}
  {{- end }}
  template:
    metadata:
      {{- with .Values.kubeletPlugin.podAnnotations }}
      annotations:
        {{- toYaml . | nindent 8 }}
      {{- end }}
      labels:
        {{- include "tpu-dra-driver.templateLabels" . | nindent 8 }}
    spec:
      hostNetwork: true
      {{- if .Values.kubeletPlugin.priorityClassName }}
      priorityClassName: {{ .Values.kubeletPlugin.priorityClassName }}
      {{- end }}
      {{- with .Values.imagePullSecrets }}
      imagePullSecrets:
        {{- toYaml . | nindent 8 }}
      {{- end }}
      serviceAccountName: {{ include "tpu-dra-driver.serviceAccountName" . }}
      securityContext:
        {{- toYaml .Values.kubeletPlugin.podSecurityContext | nindent 8 }}
      initContainers:
      - name: "tpu-network-optimization"
        image: "gcr.io/gke-release/tpu_network_optimizer@sha256:ad1446302c2a7714af62959dd84fc9cd47552085cac4f03580cb4c5eb523e81c"
        securityContext:
          privileged: true
        command:
        - /tpu_network_optimizer.sh
        volumeMounts:
        - name: sys
          mountPath: /sys
        - name: proc
          mountPath: /proc
      containers:
      - name: plugin
        securityContext:
          {{- toYaml .Values.kubeletPlugin.containers.plugin.securityContext | nindent 10 }}
        image: {{ include "tpu-dra-driver.fullimage" . }}
        imagePullPolicy: {{ .Values.image.pullPolicy }}
        command: ["tpu-dra-kubeletplugin"]
        resources:
          {{- toYaml .Values.kubeletPlugin.containers.plugin.resources | nindent 10 }}
        env:
        - name: CDI_ROOT
          value: /var/run/cdi
        - name: NODE_NAME
          valueFrom:
            fieldRef:
              fieldPath: spec.nodeName
        - name: POD_NAMESPACE
          valueFrom:
            fieldRef:
              fieldPath: metadata.namespace
        - name: POD_NAME
          valueFrom:
            fieldRef:
              fieldPath: metadata.name
        - name: NODE_IP
          valueFrom:
            fieldRef:
              fieldPath: status.hostIP
        - name: DEVICE_CLASSES
          value: {{ .Values.deviceClasses | join "," }}
        volumeMounts:
        - name: plugins-registry
          mountPath: /var/lib/kubelet/plugins_registry
        - name: plugins
          mountPath: /var/lib/kubelet/plugins
        - name: cdi
          mountPath: /var/run/cdi
        - name: sys
          mountPath: /sys
        - name: proc
          mountPath: /proc
      - image: "gcr.io/gke-release/vbar_control_agent@sha256:4bd33c9d448a446b8a1893f1caee93f0206674faeda82047491d144416c42c88"
        name: vbar-control-agent
        env:
        - name: TPU_SKIP_MDS_QUERY
          value: 'true'
        securityContext:
          privileged: true
          readOnlyRootFilesystem: true
          seccompProfile:
            type: RuntimeDefault
        resources:
          requests:
            memory: 150Mi
            cpu: 10m
          limits:
            memory: 150Mi
        command: ["vbar_control_agent_files/bin/vbar_control_agent"]
        args: ['--logtostderr', '--census_enabled=false', '--gid=', '--uid=', '--chroot=']
      volumes:
      - name: plugins-registry
        hostPath:
          path: /var/lib/kubelet/plugins_registry
      - name: plugins
        hostPath:
          path: /var/lib/kubelet/plugins
      - name: cdi
        hostPath:
          path: /var/run/cdi
      - name: dev
        hostPath:
          path: /dev
          type: DirectoryOrCreate
      - name: pod-resources
        hostPath:
          path: /var/lib/kubelet/pod-resources
          type: DirectoryOrCreate
      - name: tpu-env
        hostPath:
          path: /etc/tpu
          type: DirectoryOrCreate
      - name: tpu-logs
        hostPath:
          path: /tmp/tpu_logs
          type: DirectoryOrCreate
      - name: sys
        hostPath:
          path: /sys
          type: Directory
      - name: proc
        hostPath:
          path: /proc
          type: Directory
      {{- with .Values.kubeletPlugin.nodeSelector }}
      nodeSelector:
        {{- toYaml . | nindent 8 }}
      {{- end }}
      {{- with .Values.kubeletPlugin.affinity }}
      affinity:
        {{- toYaml . | nindent 8 }}
      {{- end }}
      {{- with .Values.kubeletPlugin.tolerations }}
      tolerations:
        {{- toYaml . | nindent 8 }}
      {{- end }}
