1from collections import namedtuple
2
3_MemoryInfo = namedtuple("_MemoryInfo", "free,total")
4
5class FakeCUDAContext(object):
6    '''
7    This stub implements functionality only for simulating a single GPU
8    at the moment.
9    '''
10    def __init__(self, device):
11        self._device = device
12
13    def __enter__(self):
14        pass
15
16    def __exit__(self, exc_type, exc_val, exc_tb):
17       pass
18
19    def __str__(self):
20        return "<Managed Device {self.id}>".format(self=self)
21
22    @property
23    def id(self):
24        return self._device
25
26    @property
27    def compute_capability(self):
28        return (5, 2)
29
30    def reset(self):
31        pass
32
33    def get_memory_info(self):
34        """
35        Cross-platform free / total host memory is hard without external dependencies,
36        e.g. `psutil` - so return infinite memory to maintain API type compatibility
37        """
38        return _MemoryInfo(float('inf'), float('inf'))
39
40
41class FakeDeviceList(object):
42    '''
43    This stub implements a device list containing a single GPU. It also
44    keeps track of the GPU status, i.e. whether the context is closed or not,
45    which may have been set by the user calling reset()
46    '''
47    def __init__(self):
48        self.lst = (FakeCUDAContext(0),)
49        self.closed = False
50
51    def __getitem__(self, devnum):
52        self.closed = False
53        return self.lst[devnum]
54
55    def __str__(self):
56        return ', '.join([str(d) for d in self.lst])
57
58    def __iter__(self):
59        return iter(self.lst)
60
61    def __len__(self):
62        return len(self.lst)
63
64    @property
65    def current(self):
66        if self.closed:
67            return None
68        return self.lst[0]
69
70
71gpus = FakeDeviceList()
72
73
74def reset():
75    gpus[0].closed = True
76
77
78def get_context(devnum=0):
79    return FakeCUDAContext(devnum)
80
81
82def require_context(func):
83    '''
84    In the simulator, a context is always "available", so this is a no-op.
85    '''
86    return func
87