#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
"""
@File    :   nvidia_smi_query.py
@Time    :   2021/6/19 14:30
@Author  :   boomyoung
@Version :   1.0
@Contact :   boomyoung@***.com
@Desc    :   
"""
import json
import logging
import os
import subprocess

import docker


class NvidiaSmiQuery(object):

    def __init__(self):
        self.log_to = logging.getLogger()

        self._QUERY_GPU = '/usr/bin/nvidia-smi --query-gpu="uuid,index" --format=csv,noheader'
        self._QUERY_APP = '/usr/bin/nvidia-smi --query-compute-apps="gpu_uuid,pid,used_memory" --format=csv,noheader'

    def queryGpu(self) -> dict:
        """
        获取gpu_uuid和gpu索引的对应关系
        :return:
        """
        status, output = subprocess.getstatusoutput(self._QUERY_GPU)
        # status, output = (0, 'GPU-9ac93f4d-726e-c371-f189-7ed073d96fad, 0\nGPU-0822505e-f54d-898a-66eb-16ec1a79ff45, 1')
        gpu_dict = {}
        if status == 0:
            gpu_list = output.split('\n')

            for i in gpu_list:
                try:
                    gpu_uuid = i.split(',')[0].strip()
                    gpu_index = i.split(',')[1].strip()
                except:
                    continue
                gpu_dict[gpu_uuid] = gpu_index
        return gpu_dict

    def getContainerId(self, pid: str) -> str:
        """
        通过pid获取容器id
        :param pid:
        :return:
        """
        cgroup_file = "/home/proc/{}/cgroup".format(pid)
        # 1:name=systemd:/kubepods/burstable/pod079552fe-c445-11eb-90a6-b8ca3af096f1/173939f128334165c78443e4045d3e7cee45795f3c7b3029c72b3388215bfe2e
        container_id = ''
        if os.path.exists(cgroup_file):
            status, output = subprocess.getstatusoutput(
                "/usr/bin/tail -n 1 %s | awk -F'/' '{print $NF}' | cut -d'-' -f 2 | cut -c 1-12" % cgroup_file)
            if status == 0:
                container_id = output.strip()
        return container_id

    def getContainerName(self, container_id: str) -> str:
        """
        容器内连接宿主机docker api,
        通过容器ID获取容器名
        :param container_id:
        :return:
        """
        try:
            client = docker.from_env(version='auto', timeout=5)
            container_name = client.containers.get(container_id).name
        except docker.errors.DockerException:
            try:
                client = docker.DockerClient(base_url="tcp://127.0.0.1:2375", version="auto", timeout=5)
                container_name = client.containers.get(container_id).name
            except:
                try:
                    client = docker.DockerClient(base_url="unix://var/run/docker.sock", version="auto", timeout=5)
                    container_name = client.containers.get(container_id).name
                except:
                    container_name = ''
        except Exception as e:
            # self.log_to.error('Port 2375 Connection refused! detail: {}'.format(e))
            container_name = ''
        return container_name

    def queryComputeApp(self) -> list:
        status, output = subprocess.getstatusoutput(self._QUERY_APP)

        data_list = []
        if status == 0:
            app_list = output.split('\n')

            for i in app_list:
                try:
                    gpu_uuid = i.split(',')[0].strip()
                    pid = i.split(',')[1].strip()
                    used_memory = i.split(',')[2].replace(' ', '')
                except:
                    continue

                data_dict = {
                    'pid': pid,
                    'used_memory': used_memory,
                    'gpu_uuid': gpu_uuid
                }
                data_list.append(data_dict)
        return data_list

    def getQueryResult(self) -> tuple:
        """
        返回查询结果
        pid: 进程id
        used_memory：已使用的显存
        gpu_index：gpu索引
        container_name：容器名
        container_id：容器ID
        :return:
        """

        status, _ = subprocess.getstatusoutput('/usr/bin/nvidia-smi')
        if status != 0:
            return 100, _

        gpu_list = []

        gpu_info = self.queryGpu()

        gpu_compute = self.queryComputeApp()
        if gpu_compute:
            for i in gpu_compute:
                gpu_uuid = i.get('gpu_uuid')
                pid = i.get('pid')
                used_memory = i.get('used_memory')

                gpu_index = gpu_info.get(gpu_uuid, '')

                container_id = self.getContainerId(pid)
                container_name = self.getContainerName(container_id)

                data_dict = {
                    'pid': pid,
                    'used_memory': used_memory,
                    'gpu_index': gpu_index,
                    'container_name': container_name,
                    'container_id': container_id
                }
                gpu_list.append(data_dict)
            return 0, gpu_list
        else:
            return 101, "nvidia-smi query fails."
            

# !/usr/bin/env python3
# -*- encoding: utf-8 -*-
"""
@File    :   gpu-exporter.py
@Time    :   2021.01.15 15:32:40
@Author  :   boomyoung
@Version :   1.0
@Contact :   boomyoung@***.com
@Desc    :   Prometheus Exporter, Container GPU memory occupancy monitoring.
             Python, Shell, Flask, Prometheus, Docker API.
"""
import sys

import prometheus_client

from prometheus_client import Gauge
from prometheus_client.core import CollectorRegistry

from flask import Response, Flask

import logging
from logging.handlers import RotatingFileHandler

from nvidia_smi_query import NvidiaSmiQuery

app = Flask(__name__)
app.debug = True

used_gpu_memory = Gauge("used_gpu_memory", "GPU memory usage of container",
                        ["container_name", "gpu", "pid", "container_id", "unit"])

BYTE_UNITS = {
    'b': 1,
    'k': 1024,
    'm': 1024 * 1024,
    'g': 1024 * 1024 * 1024
}


def setupLogger(log_file, level=logging.INFO):
    log = logging.getLogger()
    formatter = logging.Formatter('%(asctime)s -- %(levelname)s -- %(message)s')
    log_file_handler = RotatingFileHandler(filename=log_file, maxBytes=1024 * 1024, backupCount=3)
    log_file_handler.setFormatter(formatter)
    log.setLevel(level)
    log.addHandler(log_file_handler)


@app.route("/metrics")
def index():
    # setupLogger('/var/log/gpu-exporter.log')
    log_to = logging.getLogger()

    unit = 'bits'

    status, result = NvidiaSmiQuery().getQueryResult()
    if status == 0:
        # log_to.info('nvidia-smi query SUCCESS!output: \n{}'.format(result))
        for i in result:
            pid = i.get('pid', '')
            gpu_mem_used = i.get('used_memory', '')
            gpu_index = i.get('gpu_index', '')
            container_name = i.get('container_name', '')
            container_id = i.get('container_id', '')
            if "KiB" in gpu_mem_used:
                gpu_mem_used = gpu_mem_used.strip('KiB')
                conversion = BYTE_UNITS.get('k')
            elif "MiB" in gpu_mem_used:
                gpu_mem_used = gpu_mem_used.strip('MiB')
                conversion = BYTE_UNITS.get('m')
            elif "GiB" in gpu_mem_used:
                gpu_mem_used = gpu_mem_used.strip('GiB')
                conversion = BYTE_UNITS.get('g')
            else:
                gpu_mem_used = 0
                conversion = BYTE_UNITS.get('b')

            used_gpu_memory.labels(container_name, gpu_index, pid, container_id, unit).set(
                float(gpu_mem_used) * conversion)

        return Response(prometheus_client.generate_latest(used_gpu_memory),
                        mimetype="text/plain")
    elif status == 101:
        # log_to.warning(result)
        used_gpu_memory.labels('None', 'None', 'None', 'None', unit).set(int(status))
        return Response(prometheus_client.generate_latest(used_gpu_memory),
                        mimetype="text/plain")
    else:
        # log_to.error(result)
        used_gpu_memory.labels('ERROR', 'ERROR', 'ERROR', 'ERROR', unit).set(int(status))
        return Response(prometheus_client.generate_latest(used_gpu_memory),
                        mimetype="text/plain")


if __name__ == "__main__":
    app.run(host="0.0.0.0", port="44110")