#!/usr/bin/python

"""
Report users making use of each NVIDIA board on the system.
Also provides the first free board for use with cudaSetDevice().

Iain Murray, March 2008.
"""

import sys
from glob import glob
from subprocess import Popen, PIPE
import os.path

lsof_wrapper = '/pkgs/bin/gpu_usage'
if not os.path.exists(lsof_wrapper):
    print 'gpu_usage is not (yet) supported on this machine'
    sys.exit(1)

# Get ID's of NVIDIA boards. Should do this through a CUDA call, but this is
# a quick and dirty way that works for now:
dev_prefix = '/dev/nvidia'
board_devs = glob(dev_prefix + '[0-9]*')
def dev_to_id(str):
    return int(str[len(dev_prefix):])
board_ids = map(dev_to_id, board_devs)
board_ids.sort()

# Find out which boards are in use
#cmd = ["lsof", "-w", "--"]
#cmd.extend(board_devs)
#lsof_output = Popen(cmd, stdout=PIPE).communicate()[0]
lsof_output = Popen(lsof_wrapper, stdout=PIPE).communicate()[0]
lines = map(lambda x: x.split(), lsof_output.split('\n'))
lines = [x for x in lines if x] # filter out empty lines
lines = [x for x in lines if x[3] == 'mem'] # seems to imply serious use
def board_users(id):
    return ", ".join([x[2] for x in lines if x[7] == dev_prefix + str(id)])

# Report
if '--id' in sys.argv:
    ###
    print 'gpu_usage is no longer a valid way to pick an id for your job.'
    print 'Use /u/murray/bin/gpu_lock instead (run with no args for help).\n'
    sys.exit(1)
    ###
    for id in board_ids:
        if not board_users(id):
            print id
            break
    else:
        print '-1'
else:
    ###
    print "\nNOTE: This script no longer works correctly, users are spuriously"
    print "reported on all boards for each single board usage.\n"
    ###
    print "\nNVIDIA board users:"
    for id in board_ids:
        print "    Board %d: %s" % (id, board_users(id))
    #print '\nTo get the first free id for use in code, run: gpu_usage --id\n'
    ###
    print
    print "Please use the discretionary locking system: /u/murray/bin/gpu_lock\n"