GPU holder

gistfile1.txt

import os
import time
import torchvision as tv
import torch
import threading

class GPU(object):
    def __init__(self):
        self.header = "nvidia-smi"
        self.tail = "--format=csv,nounits,noheader"

    def query(self, q):
        cmd = " ".join([self.header, "--query-gpu="+q, self.tail])
        return os.popen(cmd)

    def utilization_memory(self):
        return [float(line.strip()) for line in self.query("utilization.memory")]

    def utilization_gpu(self):
        return [float(line.strip()) for line in self.query("utilization.gpu")]

    def memory_free(self):
        return [int(line.strip())for line in self.query("memory.free")]


class Boarder(threading.Thread):
    def __init__(self, dev_id):
        super().__init__()
        self.m = tv.models.resnet34(True) # .cuda()
        self.m.train()
        self.gpu = torch.device("cuda:"+str(dev_id))
        self.m = self.m.to(self.gpu)
        self.inputs = torch.randn((256, 3, 224, 224)) # .cuda()
        self.inputs = self.inputs.to(self.gpu)

    def on(self):
        print("on", self.gpu)
        self.m(self.inputs)

    def run(self):
        while True:
            self.on()
            torch.cuda.synchronize()
            time.sleep(.5)
        print("fini", self.gpu)


class Scout(object):
    def __init__(self, mission, seconds=1):
        self.mission = mission
        self.seconds = seconds
        self.boarders = {}

    def search(self):
        for info in self.tick():
            for idx, usage in enumerate(info):
                print(idx, usage)
                if usage > 10000 and idx not in self.boarders:
                    print(idx, usage)
                    b = Boarder(idx)
                    self.boarders[idx] = b
                    b.start()

    def tick(self):
        while True:
            time.sleep(self.seconds)
            yield self.mission()


def main():
    p = GPU()
    s = Scout(p.memory_free)
    s.search()


if __name__ == "__main__":
    main()