apiVersion: v1
kind: ServiceAccount
metadata:
  name: gpu-taint-controller
  namespace: kube-system
---
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRole
metadata:
  name: gpu-taint-controller
rules:
  - apiGroups: [""]
    resources: ["nodes"]
    verbs: ["get", "list", "watch", "patch"]
---
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRoleBinding
metadata:
  name: gpu-taint-controller
subjects:
  - kind: ServiceAccount
    name: gpu-taint-controller
    namespace: kube-system
roleRef:
  kind: ClusterRole
  name: gpu-taint-controller
  apiGroup: rbac.authorization.k8s.io
---
# DaemonSet that runs on GPU nodes and taints them on startup.
# This ensures new GPU nodes auto-scaled by the cluster autoscaler
# get the nvidia.com/gpu taint, preventing non-GPU workloads from
# scheduling on expensive GPU instances.
apiVersion: apps/v1
kind: DaemonSet
metadata:
  name: gpu-taint-controller
  namespace: kube-system
  labels:
    app.kubernetes.io/name: gpu-taint-controller
spec:
  selector:
    matchLabels:
      app.kubernetes.io/name: gpu-taint-controller
  template:
    metadata:
      labels:
        app.kubernetes.io/name: gpu-taint-controller
    spec:
      serviceAccountName: gpu-taint-controller
      # Only schedule on nodes that have GPU pool labels
      affinity:
        nodeAffinity:
          requiredDuringSchedulingIgnoredDuringExecution:
            nodeSelectorTerms:
              - matchExpressions:
                  - key: k8s.scaleway.com/pool-name
                    operator: In
                    values:
                      - gpu
                      - gpu-training
                      - gpu-inference
      # Tolerate the taint we're about to set (otherwise the DaemonSet
      # gets evicted by its own taint)
      tolerations:
        - key: nvidia.com/gpu
          operator: Exists
          effect: NoSchedule
      containers:
        - name: taint-setter
          image: bitnami/kubectl:latest
          env:
            - name: NODE_NAME
              valueFrom:
                fieldRef:
                  fieldPath: spec.nodeName
          command:
            - /bin/sh
            - -c
            - |
              echo "Applying GPU taint to node ${NODE_NAME}"
              kubectl taint nodes "${NODE_NAME}" nvidia.com/gpu=true:NoSchedule --overwrite || true
              echo "Taint applied. Sleeping."
              while true; do sleep 86400; done
          resources:
            requests:
              cpu: 10m
              memory: 16Mi
            limits:
              cpu: 50m
              memory: 32Mi