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):
                try:
                    # Verifica si la GPU está disponible, si no lanza un error controlado
                    if torch.cuda.is_available():
                        device = torch.device('cuda')
                        print("GPU disponible. Ejecutando en GPU.")
                    else:
                        raise RuntimeError("No se detecta la GPU. Asegúrate de que los drivers CUDA están instalados y configurados correctamente.")
                
                    # Pasar el dispositivo a la función como argumento
                    kwargs['device'] = device

                    result = [None]
                    exception = [None]

                    def target():
                        try:
                            result[0] = func(*args, **kwargs)
                        except Exception as e:
                            exception[0] = e

                    # Ejecutar la función en un hilo separado
                    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]
                except RuntimeError as e:
                    print(f"Error: {str(e)}")
                    raise e  # Lanzar el error de GPU para que se maneje adecuadamente
            return wrapper
        return decorator