def entry_func(args=None): # Get the script to execute, parse only first input parser = get_argparser() args = parser.parse_args(args) # Here, we wrap the training in a try/except block to ensure that we # stop the GPUMonitor process after training, even if an error occurred from mpunet.utils.system import GPUMonitor gpu_mon = GPUMonitor() try: run(args=args, gpu_mon=gpu_mon) except Exception as e: gpu_mon.stop() raise e
def set_gpu_vis(num_GPUs, force_GPU, logger=None): """ Helper function that sets the GPU visibility as per parsed args """ if force_GPU: from mpunet.utils.system import set_gpu set_gpu(force_GPU) else: # Automatically determine GPUs to use from mpunet.utils.system import GPUMonitor GPUMonitor(logger).await_and_set_free_GPU(num_GPUs, stop_after=True)
def get_free_GPU_sets(num_GPUs, ignore_gpus=None): from mpunet.utils.system import GPUMonitor mon = GPUMonitor() ignore_gpus = _gpu_string_to_list(ignore_gpus or "", as_int=True) free_gpus = sorted(mon.free_GPUs, key=lambda x: int(x)) mon.stop() free_gpus = list(filter(lambda gpu: gpu not in ignore_gpus, free_gpus)) total_GPUs = len(free_gpus) if total_GPUs % num_GPUs or not free_gpus: if total_GPUs < num_GPUs: raise ValueError( "Invalid number of GPUs per process '%i' for total " "GPU count of '%i' - must be evenly divisible." % (num_GPUs, total_GPUs)) else: raise NotImplementedError else: return _get_GPU_sets(free_gpus, num_GPUs)
def get_gpu_monitor(num_GPUs, logger): """ Args: num_GPUs: Number of GPUs to train on logger: A mpunet logger object that will be passed to the GPUMonitor Returns: If num_GPUs >= 0, returns a GPUMonitor object, otherwise returns None """ if num_GPUs >= 0: # Initialize GPUMonitor in separate fork now before memory builds up from mpunet.utils.system import GPUMonitor gpu_mon = GPUMonitor(logger) else: gpu_mon = None return gpu_mon