File size: 1,064 Bytes
01bb284 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
# decorators.py
import functools
import torch
import threading
import time
class spaces:
@staticmethod
def GPU(duration=0):
def decorator(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):
if not torch.cuda.is_available():
raise RuntimeError("GPU no est谩 disponible.")
result = [None]
exception = [None]
def target():
try:
result[0] = func(*args, **kwargs)
except Exception as e:
exception[0] = e
thread = threading.Thread(target=target)
thread.start()
thread.join(duration)
if thread.is_alive():
raise TimeoutError(f"La ejecuci贸n de la funci贸n excedi贸 {duration} segundos.")
if exception[0]:
raise exception[0]
return result[0]
return wrapper
return decorator
|