atopgpud

#!/usr/bin/python3 -Es

# ==============================================================
# Daemon that gathers statistical information from all
# NVIDIA GPUs in the current system. Since the GPU maintains the
# busy percentages only over the last second, this daemon gathers
# the statistics each second and maintains cumulative counters,
# globally and per process.
#
# Client processes can connect to this daemon via a UNIX domain
# socket with the abstract path name "atopgpud_nvidia" and send
# requests of two bytes, consisting of one byte
# request code followed by one byte API version number.
# The request code can be 'T' to obtain the GPU types or 'S' to
# obtain all statistical counters.
# The response of the daemon starts with a 4-byte integer. The
# first byte is the version of the response format and the
# subsequent three bytes indicate the length (big endian) of the
# response string that follows. See the formatters for the layout
# of the response string, later on in this source code.
#
# Dependencies: pip/pip3 install nvidia-ml-py
# --------------------------------------------------------------
# Author: Gerlof Langeveld
# Date:   July   2018 - initial
#         August 2025 - UNIX domain socket instead of IPv6 socket
#         August 2025 - redesign of GPU busy percentage gathering
#         August 2025 - no need to run with specific privileges
#
# Copyright (C) 2018,2025 Gerlof Langeveld
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 2, or (at your option) any
# later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
# ==============================================================

import os
import sys
import time
import socket
import struct
import logging
import logging.handlers as loghand
import threading

SOCKPATH     = b"\0atopgpud_nvidia"

COMPUTE      = 1		# task support bit value
GRAPHIC      = 2        # task support bit value

HIGHSTAMP    = 1000

# =================================
# GPU related bookkeeping
# =================================
gpulist = []    # list with per-GPU bookkeeping
cliterm = {}    # dict with one entry per client (client
                # socket as key), that contains a dict with
                # the terminated per-process bookkeepings
                # that still have to be received by this client
                # (pid of terminated process as key)

gpulock = None	# mutex for access to gpulist/cliterm


# =================================
# per-GPU class
# =================================
class Stats(object): pass		# generic statistics container

class GpuProp(object):
    ###############################
    # initialization method to setup
    # properties
    ###############################
    def __init__(self, num):
        gpuhandle = pynvml.nvmlDeviceGetHandleByIndex(num)
        pciInfo   = pynvml.nvmlDeviceGetPciInfo(gpuhandle)

        self.gpuhandle         = gpuhandle
        self.stats             = Stats()

        self.stats.busid       = pciInfo.busId
        self.stats.devname     = pynvml.nvmlDeviceGetName(gpuhandle)

        # nvml backward compatibility
        if type(self.stats.busid) == bytes:
            self.stats.busid = self.stats.busid.decode('ascii', errors='replace')
        if type(self.stats.devname) == bytes:
            self.stats.devname = self.stats.devname.decode('ascii', errors='replace')

        self.stats.devname = self.stats.devname.replace(' ', '_')

        self.stats.tasksupport = 0		# process stats support

        try:
            procinfo = pynvml.nvmlDeviceGetComputeRunningProcesses(gpuhandle)
            self.stats.tasksupport |= COMPUTE # compute proces support
        except Exception:
            pass                              # no compute process support

        try:
            procinfo = pynvml.nvmlDeviceGetGraphicsRunningProcesses(gpuhandle)
            self.stats.tasksupport |= GRAPHIC # graphic process support
        except Exception:
            pass                              # no graphic process support

        self.stats.gpupercnow   = 0	# perc of time that GPU was busy
        self.stats.mempercnow   = 0	# perc of time that memory was rd/wr

        self.stats.memtotalnow  = 0	# in Kb
        self.stats.memusednow   = 0	# in Kb

        self.stats.gpusamples   = 0
        self.stats.gpuperccum   = 0	# perc of time that GPU was busy
        self.stats.memperccum   = 0	# perc of time that memory was rd/wr
        self.stats.memusedcum   = 0	# in KiB

        self.stats.procstats    = {}	# stats of active processes (key = pid)

    ###############################
    # method to fetch counters and values
    ###############################
    def readstats(self):
        global HIGHSTAMP

        self.stats.gpusamples += 1

        # -----------------------------
        # get rates (utilization percentages)
        # -----------------------------
        try:
            rates = pynvml.nvmlDeviceGetUtilizationRates(self.gpuhandle)

            self.stats.gpupercnow  = rates.gpu
            self.stats.mempercnow  = rates.memory
            self.stats.gpuperccum += rates.gpu
            self.stats.memperccum += rates.memory
        except pynvml.NVMLError as err:
            self.stats.gpupercnow  = -1
            self.stats.mempercnow  = -1
            self.stats.gpuperccum  = -1
            self.stats.memperccum  = -1

        # -----------------------------
        # get memory occupation GPU-wide
        # -----------------------------
        try:
            meminfo = pynvml.nvmlDeviceGetMemoryInfo(self.gpuhandle)

            self.stats.memtotalnow = meminfo.total // 1024
            self.stats.memusednow  = meminfo.used  // 1024
            self.stats.memusedcum += meminfo.used  // 1024	# in KiB
        except pynvml.NVMLError as err:
            pass

        # -----------------------------
        # get per-process statistics
        # -----------------------------
        try:
            if self.stats.tasksupport & COMPUTE:
                cprocinfo = pynvml.nvmlDeviceGetComputeRunningProcesses(self.gpuhandle)

                for proc in cprocinfo:
                    proc.type = 'C'

            if self.stats.tasksupport & GRAPHIC:
                gprocinfo = pynvml.nvmlDeviceGetGraphicsRunningProcesses(self.gpuhandle)

                for proc in gprocinfo:
                    proc.type = 'G'

            procinfo = cprocinfo + gprocinfo

            # -------------------------
            # build list with pids from
            # the previous interval
            # -------------------------
            actprocs = list(self.stats.procstats.keys())

            # -------------------------
            # handle proc stats of this
            # interval
            # -------------------------
            for proc in procinfo:
                pid = proc.pid

                # ---------------------
                # new process?
                #     create new stats
                # ---------------------
                if pid not in actprocs:
                    self.stats.procstats[pid] = Stats()

                    self.stats.procstats[pid].samples    = 0
                    self.stats.procstats[pid].type       = 'U'  # unknown

                    self.stats.procstats[pid].memnow     = 0	# in KiB
                    self.stats.procstats[pid].memcum     = 0	# in KiB

                    self.stats.procstats[pid].gpubusy    = -1
                    self.stats.procstats[pid].membusy    = -1
                    self.stats.procstats[pid].gpubusycum = -1
                    self.stats.procstats[pid].membusycum = -1
                else:
                    actprocs.remove(pid) # remove active pid from list

                # ---------------------
                # maintain proc stats
                # ---------------------
                self.stats.procstats[pid].samples    += 1
                self.stats.procstats[pid].type        = proc.type

                if proc.usedGpuMemory:
                    self.stats.procstats[pid].memnow  = proc.usedGpuMemory // 1024
                    self.stats.procstats[pid].memcum +=	proc.usedGpuMemory // 1024

            # -------------------------
            # get average GPU utilization
            # and memory utiliation for
            # all active processes on GPU
            # -------------------------
            try:
                procutils = pynvml.nvmlDeviceGetProcessUtilization(self.gpuhandle, HIGHSTAMP)

                if procutils and len(procutils) > 0:
                    for proc in procutils:
                        if HIGHSTAMP < proc.timeStamp:
                            HIGHSTAMP = proc.timeStamp

                        if proc.pid in self.stats.procstats:
                            self.stats.procstats[proc.pid].gpubusy     = proc.smUtil
                            self.stats.procstats[proc.pid].membusy     = proc.memUtil

                            self.stats.procstats[proc.pid].gpubusycum += proc.smUtil
                            self.stats.procstats[proc.pid].membusycum += proc.memUtil
            except Exception as e:
                pass

            # -------------------------
            # determine which processes
            # have terminated since
            # previous sample
            # -------------------------
            for pid in actprocs:
                 for client in cliterm:
                     cliterm[client][pid] = self.stats.procstats[pid]

                 del self.stats.procstats[pid]

        except pynvml.NVMLError as err:
            pass


    ###############################
    # obtain current statistics
    ###############################
    def getstats(self):
        return self.stats


# =================================
# Main function
# =================================
def main():
    global gpulock
    gpulock = threading.Lock()	# mutex for access to gpulist/cliterm

    # -----------------------------
    # initialize GPU access,
    # specifically to detect of it
    # succeeds
    # -----------------------------
    try:
        pynvml.nvmlInit()
    except Exception:
        logging.error("Shared lib 'libnvidia-ml' probably not installed!")
        sys.exit()

    # -----------------------------
    #  open UNIX domain stream socket
    # -----------------------------
    try:
        mysock = socket.socket(socket.AF_UNIX,  socket.SOCK_STREAM, 0)
    except Exception as sockex:
        logging.error("Socket creation fails")
        sys.exit(1)

    # -----------------------------
    #  bind to abstract pathname and
    #  make socket passive
    # -----------------------------
    try:
        mysock.bind(SOCKPATH)
        mysock.listen(32)
    except Exception as sockex:
        logging.error("Socket binding to path %s fails", SOCKPATH)
        sys.exit(1)

    # -----------------------------
    # initialize GPU access for the
    # child process
    # -----------------------------
    try:
        pynvml.nvmlInit()
    except Exception:
        pass

    # -----------------------------
    # determine number of GPUs in
    # this system
    # -----------------------------
    gpunum = pynvml.nvmlDeviceGetCount()
    logging.info("Number of GPUs: %d", gpunum)

    if gpunum == 0:
        logging.info("Terminated (no GPUs available)")
        sys.exit()

    # -----------------------------
    # initialize per-GPU bookkeeping
    # -----------------------------
    for i in range(gpunum):
        gpulist.append( GpuProp(i) )

    # -----------------------------
    # kick off new thread to fetch
    # statistics periodically
    # -----------------------------
    t = threading.Thread(target=gpuscanner, args=(1,))
    t.daemon = True
    t.start()

    logging.info("Initialization succeeded")

    # -----------------------------
    # main thread:
    #   await connect of client
    # -----------------------------
    while True:
        newsock, peeraddr = mysock.accept()

        # -------------------------
        # create new thread to
        # serve this client
        # -------------------------
        t = threading.Thread(target=serveclient, args=(newsock, peeraddr))
        t.daemon = True
        t.start()


# ===========================================
# Thread start function:
# Serve new client that has just connected.
# -------------------------------------------
# Protocol between client and server:
#
# - client transmits request
#   consisting of two bytes
#
#     byte 0: type of request
#             'S' get statistical counters
#             'T' get type of each GPU
#
#     byte 1: integer version number
#             response layout might change
#             so the client asks for a
#             specific response version
#
# - server transmits response
#   consisting of a four bytes integer
#   in big endian byte order
#
#     byte 0:   version number, preferably
#               as requested by the client
#
#     byte 1-3: length of the response string
#               that follows
#
#   followed by the response string that is
#   version specific (see gpuformatters)
# ===========================================
def serveclient(sock, peer):
    # -----------------------------
    # create per client bookkeeping
    # for terminated processes
    # -----------------------------
    with gpulock:
        cliterm[sock] = {}

    # -----------------------------
    # main loop
    # -----------------------------
    while True:
        # -------------------------
        # wait for request
        # -------------------------
        try:
            rcvbuf = sock.recv(20)
        except Exception as sockex:
            logging.error("Receive error: %s", sockex)
            sock.close()
            break

        # -------------------------
        # connection closed by peer?
        # -------------------------
        if not rcvbuf:	
            sock.close()
            break

        logging.debug("Received: %s", rcvbuf)

        # -------------------------
        # request has wrong length?
        # -------------------------
        if len(rcvbuf) != 2:
            logging.error('Wrong request length: %d', len(rcvbuf))
            sock.close()
            break

        # -------------------------
        # valid request:
        #     get statistical counters?
        # -------------------------
        command = chr(rcvbuf[0])
        version = rcvbuf[1]

        if command == 'S':
             if version == 0 or version >= len(gpuformatters):
                 version = len(gpuformatters)-1

             xmitbuf = gpuformatters[version](sock).encode('ascii',
                                                    errors='replace')

        # -------------------------
        # valid request:
        #     get GPU types?
        # -------------------------
        elif command == 'T':		
             if version == 0 or version >= len(gpudevnames):
                 version = len(gpudevnames)-1

             xmitbuf = gpudevnames[version]().encode('ascii', errors='replace')

        # -------------------------
        # invalid request!
        # -------------------------
        else:
            logging.error('Wrong request from client: %s', command)
            sock.close()
            break

        # -------------------------
        # transmit GPU statistics
        # as bytes
        # -------------------------
        prelude = struct.pack(">I", (version << 24) + len(xmitbuf))

        try:
            sock.send(prelude)
            sock.send(xmitbuf)

            logging.debug("Send: %s", prelude + xmitbuf)
        except Exception as sockex:
            logging.error("Send error: %s", sockex)
            sock.close()
            break

    # -----------------------------
    # delete per client bookkeeping
    # of terminated processes
    # -----------------------------
    with gpulock:
        del cliterm[sock]

    # -----------------------------
    # END OF CLIENT THREAD
    # -----------------------------


# =================================
# Generate sequence of device names
# =================================
def gpudevname_v1():
    # -----------------------------
    # main loop:
    # - get device name of every GPU
    # - convert into one string
    #   with format:
    #      numgpus@busid devname tasksupport@busid devname tasksupport@...
    # -----------------------------
    strbuf = str( len(gpulist) )

    with gpulock:
        for i, gpu in enumerate(gpulist):
            s = gpu.getstats()
            strbuf += "@{:s} {:s} {:d}".format(s.busid, s.devname, s.tasksupport)

    return strbuf

gpudevnames = [None, gpudevname_v1, gpudevname_v1]


# =================================
# Convert statistics of all GPUs
# into parseable string
# =================================
def gpuformatter_v1(clisock):
    # -----------------------------
    # main loop:
    # - get statistics for every GPU
    # - convert stats to one string
    #   with format:
    #      numgpus@gpu0 stats#pid stats#pid stats@gpu1 stats#pid stats@....
    # -----------------------------
    strbuf = ""

    with gpulock:
        for i, gpu in enumerate(gpulist):
            s = gpu.getstats()

            # ---------------------
            # generic GPU stats
            # ---------------------
            strbuf += "@{:d} {:d} {:d} {:d} {:d} {:d} {:d} {:d}".format(
                       s.gpupercnow,  s.mempercnow,
                       s.memtotalnow, s.memusednow, s.gpusamples,
                       s.gpuperccum,  s.memperccum, s.memusedcum);

            # ---------------------
            # active processes for
            # this GPU
            # ---------------------
            for pid, stat in s.procstats.items():
                strbuf += "#A {:d} {:d} {:d} {:d} {:d} {:d} {:d}".format(pid,
                           stat.gpubusy, stat.membusy, 0,
                           stat.memnow,  stat.memcum,  stat.samples)

            # ---------------------
            # terminated processes
            # for this GPU
            # ---------------------
            for pid, stat in cliterm[clisock].items():
                strbuf += "#E {:d} {:d} {:d} {:d} {:d} {:d} {:d}".format(pid,
                           stat.gpubusy, stat.membusy, 0,
                           stat.memnow,  stat.memcum,  stat.samples)

            cliterm[clisock].clear()


    return strbuf


def gpuformatter_v2(clisock):
    # -----------------------------
    # main loop:
    # - get statistics for every GPU
    # - convert stats to one string
    #   with format:
    #      numgpus@gpu0 stats#pid stats#pid stats@gpu1 stats#pid stats@....
    # -----------------------------
    strbuf = ""

    with gpulock:
        for i, gpu in enumerate(gpulist):
            s = gpu.getstats()

            # ---------------------
            # generic GPU stats
            # ---------------------
            strbuf += "@{:d} {:d} {:d} {:d} {:d} {:d} {:d} {:d}".format(
                       s.gpupercnow,  s.mempercnow,
                       s.memtotalnow, s.memusednow, s.gpusamples,
                       s.gpuperccum,  s.memperccum, s.memusedcum);

            # ---------------------
            # active processes for
            # this GPU
            # ---------------------
            for pid, stat in s.procstats.items():
                strbuf += "#A {:d} {:s} {:d} {:d} {:d} {:d} {:d}".format(pid,
                           stat.type,
                           stat.gpubusycum, stat.membusycum,
                           stat.memnow,     stat.memcum,  stat.samples)

            # ---------------------
            # terminated processes
            # for this GPU
            # ---------------------
            for pid, stat in cliterm[clisock].items():
                strbuf += "#E {:d} {:s} {:d} {:d} {:d} {:d} {:d}".format(pid,
                           stat.type,
                           stat.gpubusycum, stat.membusycum,
                           stat.memnow,     stat.memcum,  stat.samples)

            cliterm[clisock].clear()

    return strbuf

gpuformatters = [None, gpuformatter_v1, gpuformatter_v2]


# =================================
# Thread start function:
# Scan all GPUs with a particular
# interval to obtain their stats
# =================================
def gpuscanner(scaninterval):
    realtimenext = time.time() + scaninterval

    # -----------------------------
    # main loop
    # -----------------------------
    while True:
        # -----------------------------
        # gather the metrics per GPU
        # -----------------------------
        with gpulock:
            for gpu in gpulist:
                gpu.readstats()

        # -----------------------------
        # handle wait time:
        # wait before taking next sample
        # unless gathering the metric took
        # more than the required interval itself
        # -----------------------------
        sleeptime = realtimenext - time.time()

        if sleeptime > 0:
            time.sleep(sleeptime)

        realtimenext += scaninterval


# ==========================================================================

# -----------------------------
# initialize logging
# -----------------------------
if '-v' in sys.argv:
    loglevel = logging.DEBUG
else:
    loglevel = logging.INFO

fm = logging.Formatter('atopgpud %(levelname)s: %(message)s')

fh = loghand.SysLogHandler('/dev/log',
             facility=loghand.SysLogHandler.LOG_DAEMON)
fh.setFormatter(fm)
fh.setLevel(loglevel)

lg = logging.getLogger()		# root logger
lg.addHandler(fh)
lg.setLevel(loglevel)

if '-f' not in sys.argv:
    # -----------------------------
    # release parent process
    # (daemonize)
    # -----------------------------
    try:
        if os.fork():
            sys.exit(0)	# parent process exits; child continues...
    except Exception:
        logging.error("Failed to fork child")

# -----------------------------
# load module pynvml
# -----------------------------
try:
    import pynvml
except Exception:
    logging.error("Python module 'pynvml' not installed!")
    sys.exit(1)


try:
   # -----------------------------
   # call main function
   # -----------------------------
   main()

finally:
    # -----------------------------
    # shutdown GPU access
    # -----------------------------
    try:
        pynvml.nvmlShutdown()
    except Exception:
        pass