Fixing bug in is_alive.

This commit is contained in:
Marc-Andre Ferland 2022-10-17 01:05:51 -04:00
parent 012243a880
commit 554b67a2f0
2 changed files with 9 additions and 4 deletions

View File

@ -130,7 +130,7 @@ def device_init(device_selection=None):
device_selection = device_selection.lower() device_selection = device_selection.lower()
if device_selection.startswith('gpu:'): if device_selection.startswith('gpu:'):
device_selection = int(device_selection[4:]) device_selection = int(device_selection[4:])
if device_selection != 'cuda' and device_selection != 'current': if device_selection != 'cuda' and device_selection != 'current' and device_selection != 'gpu':
if device_select(device_selection): if device_select(device_selection):
if isinstance(device_selection, int): if isinstance(device_selection, int):
print(f'Setting GPU:{device_selection} as active') print(f'Setting GPU:{device_selection} as active')
@ -152,6 +152,8 @@ def is_first_cuda_device(device):
if device is None: return False if device is None: return False
if device == 0 or device == '0': return True if device == 0 or device == '0': return True
if device == 'cuda' or device == 'cuda:0': return True if device == 'cuda' or device == 'cuda:0': return True
if device == 'gpu' or device == 'gpu:0': return True
if device == 'current': return True
if device == torch.device(0): return True if device == torch.device(0): return True
return False return False

View File

@ -305,9 +305,12 @@ def is_alive(name=None):
nbr_alive = 0 nbr_alive = 0
try: try:
for rthread in render_threads: for rthread in render_threads:
thread_name = rthread.name[len(THREAD_NAME_PREFIX):] thread_name = rthread.name[len(THREAD_NAME_PREFIX):].lower()
if name and thread_name != name: if name is not None:
if not runtime.is_first_cuda_device(name) and not runtime.is_first_cuda_device(thread_name): if runtime.is_first_cuda_device(name):
if not runtime.is_first_cuda_device(thread_name):
continue
elif thread_name != name:
continue continue
if rthread.is_alive(): if rthread.is_alive():
nbr_alive += 1 nbr_alive += 1