decorators.py

# 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