compressed_tensors.offload.cache.utils

def catch_cpu_mem_error(func):
    """
    Decorator to catch CPU memory errors and log a remediation warning.
    Prevents duplicate logs if nested functions also use this decorator.
    """

    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except (RuntimeError, OSError) as exc:
            # Prevent duplicate logs when DistributedCPUCache calls super().offload()
            if getattr(exc, "_cpu_mem_logged", False):
                raise

            if _is_cpu_memory_error(exc):
                logger.warning(_CPU_MEMORY_REMEDIATION)
                exc._cpu_mem_logged = True
            raise

    return wrapper

def catch_pinned_mem_error(func):
    """
    Decorator to catch failures from page-locked staging and log a remediation warning.
    """

    @wraps(func)
    def wrapper(*args, **kwargs):
        pin_memory = kwargs.get("pin_memory", False)
        if len(args) >= 3:
            pin_memory = args[2]

        try:
            return func(*args, **kwargs)
        except (RuntimeError, OSError) as exc:
            if getattr(exc, "_cpu_mem_logged", False):
                raise

            if pin_memory and _is_pinned_memory_error(exc):
                logger.warning(_PINNED_MEMORY_REMEDIATION)
                exc._cpu_mem_logged = True
            raise

    return wrapper