apiVersion: v1
kind: ServiceAccount
metadata:
  name: gpu-labeler-sa
  namespace: kube-system
  labels:
    parent: skypilot
    job: sky-gpu-labeler

---
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRole
metadata:
  name: node-patcher-role
  namespace: kube-system
  labels:
    parent: skypilot
    job: sky-gpu-labeler
rules:
- apiGroups: [""]
  resources: ["nodes"]
  verbs: ["patch"]

---
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRoleBinding
metadata:
  name: node-patcher-rolebinding
  labels:
    parent: skypilot
    job: sky-gpu-labeler
roleRef:
  apiGroup: rbac.authorization.k8s.io
  kind: ClusterRole
  name: node-patcher-role
subjects:
- kind: ServiceAccount
  name: gpu-labeler-sa
  namespace: kube-system

---
apiVersion: v1
kind: ConfigMap
metadata:
  name: gpu-labeler-script
  namespace: kube-system
  labels:
    parent: skypilot
    job: sky-gpu-labeler
data:
  label_gpus.py: |
    #!/usr/bin/env python3
    import os
    import subprocess
    import sys
    from typing import Optional

    from kubernetes import client
    from kubernetes import config

    # Canonical GPU names - injected from sky.provision.kubernetes.constants
    canonical_gpu_names = {{ canonical_gpu_names }}


    def get_gpu_name() -> Optional[str]:
        try:
            result = subprocess.run(
                ['nvidia-smi', '--query-gpu=name', '--format=csv,noheader,nounits'],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                check=True,
                timeout=30)
            gpu_name = result.stdout.decode('utf-8').strip()
            if not gpu_name:
                print('nvidia-smi returned empty output. GPU driver may not be loaded.')
                return None
            # In the case of multi-gpu nodes, we assume the node is homogenous and
            # just use the first GPU name.
            gpu_name = gpu_name.split('\n')[0]
            return gpu_name.lower()
        except FileNotFoundError:
            print('nvidia-smi not found. Ensure NVIDIA drivers are installed.')
            return None
        except subprocess.CalledProcessError as e:
            stderr = e.stderr.decode('utf-8') if e.stderr else ''
            print(f'nvidia-smi failed with exit code {e.returncode}: {stderr}')
            return None
        except subprocess.TimeoutExpired:
            print('nvidia-smi timed out. GPU may be unresponsive.')
            return None
        except Exception as e:
            print(f'Unexpected error getting GPU name: {type(e).__name__}: {e}')
            return None


    def label_node(gpu_name: str) -> None:
        try:
            config.load_incluster_config()  # Load in-cluster configuration
            v1 = client.CoreV1Api()

            # Fetch the current node's name from the environment variable
            node_name = os.environ.get('MY_NODE_NAME')
            if not node_name:
                raise ValueError('Failed to get node name from environment')

            # Label the node with the GPU name
            body = {'metadata': {'labels': {'skypilot.co/accelerator': gpu_name}}}
            v1.patch_node(node_name, body)

            print(f'Labeled node {node_name} with GPU {gpu_name}')

        except Exception as e:
            print(f'Error labeling node: {e}')
            # Re-raise so the script exits non-zero and the job is marked as failed
            raise


    def main():
        gpu_name = get_gpu_name()
        if gpu_name is not None:
            labelled = False
            for canonical_name in canonical_gpu_names:
                if canonical_name.lower() in gpu_name.lower():
                    label_node(canonical_name.lower())
                    labelled = True
                    break
            if not labelled:
                # If we didn't find a canonical name:
                # 1. remove 'NVIDIA ' if present (e.g., 'NVIDIA RTX A6000' -> 'RTX A6000')
                # 2. remove 'GeForce ' if present (e.g., 'NVIDIA GeForce RTX 3070' -> 'RTX 3070')
                # 3. replace 'RTX ' with 'RTX' (without spaces) (e.g., 'RTX 6000' -> 'RTX6000')
                # 4. replace any other spaces with dashes (e.g. 'RTX 2080 Ti' -> 'RTX2080-Ti')
                gpu_name = gpu_name.lower().replace('nvidia ', '').replace('geforce ', '').replace('rtx ', 'rtx').replace(' ', '-')
                gpu_label = gpu_name
                label_node(gpu_label)
                labelled = True
        else:
            print('No GPU detected. Try running nvidia-smi in the container.')
            sys.exit(1)


    if __name__ == '__main__':
        main()
