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