GPU Bastard, control the resource belong to you!

bastard.py

import os
import time
import argparse
import GPUtil

from multiprocessing import Process

parser = argparse.ArgumentParser(description='GPU Bastard, control the resource belong to you!')
parser.add_argument('--index',
                    help='GPU Index',
                    default='0')
parser.add_argument('--fraction',
                    help='Fraction of target GPU memory needed.\n Default is 0, means all memory',
                    default=0)

args = parser.parse_args()
args = vars(args)

os.environ['CUDA_VISIBLE_DEVICES'] = args['index']
target_index = int(args['index'])
fraction_needed = float(args['fraction'])

if not (0 <= fraction_needed < 1):
    print('--fraction should be a float number between 0 and 1.')
    exit(1)

print('\nGPU Bastard Activated! Target GPU:%s\n' % str(target_index))


def take_control():
    print('\n==> [%s] Available GPU Found, taking control....\n' % time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime()))

    import tensorflow as tf
    if not fraction_needed == 0:
        config = tf.ConfigProto()
        config.gpu_options.per_process_gpu_memory_fraction = fraction_needed
        tf.Session(config=config)

    tf.enable_eager_execution()
    tf.add(1, 2)

    print('\n==> [%s] Available GPU taken\n' % time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime()))
    # This almostly equals to forever...
    time.sleep(60 * 60 * 24)


while True:
    available_list = GPUtil.getAvailable(maxLoad=1, maxMemory=min(0.9, 1 - fraction_needed))

    if target_index in available_list:
        p = Process(target=take_control)
        p.start()

    time.sleep(3)