nzfc_gram_runtime/cache_profiles.py

from __future__ import annotations

import gc
import time
import types
import traceback
from typing import Any, Dict, Optional

import torch

try:
    from .runtime import sanitize_model_answer
except Exception:
    def sanitize_model_answer(x):
        return x


def clear_cuda():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def vram_snapshot() -> Dict[str, float]:
    if not torch.cuda.is_available():
        return {}
    snap = {}
    allocs = []
    reserved = []
    peaks = []
    for i in range(torch.cuda.device_count()):
        alloc = torch.cuda.memory_allocated(i) / 1e9
        resv = torch.cuda.memory_reserved(i) / 1e9
        peak = torch.cuda.max_memory_allocated(i) / 1e9
        snap[f'gpu{i}_alloc_gb'] = float(alloc)
        snap[f'gpu{i}_reserved_gb'] = float(resv)
        snap[f'gpu{i}_peak_gb'] = float(peak)
        allocs.append(alloc)
        reserved.append(resv)
        peaks.append(peak)
    snap['sum_alloc_gb'] = float(sum(allocs))
    snap['sum_reserved_gb'] = float(sum(reserved))
    snap['max_alloc_gb'] = float(max(allocs)) if allocs else 0.0
    snap['max_reserved_gb'] = float(max(reserved)) if reserved else 0.0
    snap['max_peak_gb'] = float(max(peaks)) if peaks else 0.0
    return snap


def is_oom_like_error(e: BaseException) -> bool:
    s = (type(e).__name__ + ' ' + str(e)).lower()
    markers = ['out of memory', 'cuda out of memory', 'cuda error', 'illegal memory access', 'cublas', 'cudnn', 'memory', 'oom']
    return any(m in s for m in markers)


def infer_input_device(bot) -> torch.device:
    dev = getattr(bot, 'input_device', None)
    if dev is not None:
        return torch.device(dev)
    model = getattr(bot, 'model', None)
    if model is not None:
        try:
            emb = model.get_input_embeddings()
            if emb is not None:
                return next(emb.parameters()).device
        except Exception:
            pass
        try:
            for p in model.parameters():
                if not getattr(p, 'is_meta', False):
                    return p.device
        except Exception:
            pass
    return torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')


def build_messages(bot, system_prompt: str, user_prompt: str):
    if hasattr(bot, '_build_messages'):
        try:
            return bot._build_messages(system_prompt, user_prompt)
        except Exception:
            pass
    return [
        {'role': 'system', 'content': str(system_prompt)},
        {'role': 'user', 'content': str(user_prompt)},
    ]


def encode_messages(bot, messages, device: torch.device):
    if hasattr(bot, '_encode_messages'):
        try:
            encoded = bot._encode_messages(messages)
            return {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in encoded.items()}
        except Exception:
            pass
    processor = getattr(bot, 'processor', None)
    tokenizer = getattr(bot, 'tokenizer', None)
    prompt_text = None
    if processor is not None and hasattr(processor, 'apply_chat_template'):
        try:
            prompt_text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        except Exception:
            prompt_text = None
    if prompt_text is None and tokenizer is not None and hasattr(tokenizer, 'apply_chat_template'):
        try:
            prompt_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        except Exception:
            prompt_text = None
    if prompt_text is None:
        system_text = str(messages[0].get('content', '')) if messages else ''
        user_text = str(messages[1].get('content', '')) if len(messages) > 1 else ''
        prompt_text = '<system>\n' + system_text + '\n</system>\n\n<user>\n' + user_text + '\n</user>\n\n<assistant>\n'
    if processor is not None:
        try:
            encoded = processor(text=prompt_text, return_tensors='pt')
            return {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in encoded.items()}
        except Exception:
            pass
    if tokenizer is None:
        raise RuntimeError('No tokenizer or processor available for encoding.')
    encoded = tokenizer(prompt_text, return_tensors='pt')
    return {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in encoded.items()}


def decode_generated(bot, out_ids, input_len: int) -> str:
    if hasattr(bot, '_decode_ids'):
        try:
            return str(bot._decode_ids(out_ids[input_len:])).strip()
        except Exception:
            pass
    tokenizer = getattr(bot, 'tokenizer', None)
    processor = getattr(bot, 'processor', None)
    if tokenizer is None and processor is not None:
        tokenizer = getattr(processor, 'tokenizer', None)
    if tokenizer is None:
        raise RuntimeError('No tokenizer available for decoding.')
    return str(tokenizer.decode(out_ids[input_len:], skip_special_tokens=True, clean_up_tokenization_spaces=True)).strip()


@torch.inference_mode()
def generate_once(bot, *, system_prompt: str, user_prompt: str, max_new_tokens: int, use_cache: bool, do_sample: bool = False, temperature: float = 0.0, top_p: Optional[float] = None, repetition_penalty: Optional[float] = None, context_hard_cap: int = 16000) -> Dict[str, Any]:
    model = getattr(bot, 'model', None)
    if model is None:
        return {'ran': False, 'answer': None, 'reason': 'model_not_loaded', 'use_cache': use_cache}
    device = infer_input_device(bot)
    messages = build_messages(bot, system_prompt, user_prompt)
    encoded = encode_messages(bot, messages, device)
    input_ids = encoded.get('input_ids')
    input_len = int(input_ids.shape[-1]) if input_ids is not None else 0
    if input_len > context_hard_cap:
        return {'ran': False, 'answer': None, 'reason': f'context_hard_cap_exceeded:{input_len}>{context_hard_cap}', 'input_tokens': input_len, 'use_cache': use_cache}
    tokenizer = getattr(bot, 'tokenizer', None)
    processor = getattr(bot, 'processor', None)
    if tokenizer is None and processor is not None:
        tokenizer = getattr(processor, 'tokenizer', None)
    gen_kwargs = dict(encoded)
    gen_kwargs['max_new_tokens'] = int(max_new_tokens)
    gen_kwargs['do_sample'] = bool(do_sample)
    gen_kwargs['use_cache'] = bool(use_cache)
    if do_sample and temperature and temperature > 0:
        gen_kwargs['temperature'] = float(temperature)
    if do_sample and top_p is not None:
        gen_kwargs['top_p'] = float(top_p)
    if repetition_penalty is not None:
        gen_kwargs['repetition_penalty'] = float(repetition_penalty)
    if tokenizer is not None:
        if getattr(tokenizer, 'pad_token_id', None) is not None:
            gen_kwargs['pad_token_id'] = tokenizer.pad_token_id
        elif getattr(tokenizer, 'eos_token_id', None) is not None:
            gen_kwargs['pad_token_id'] = tokenizer.eos_token_id
        if getattr(tokenizer, 'eos_token_id', None) is not None:
            gen_kwargs['eos_token_id'] = tokenizer.eos_token_id
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    t0 = time.perf_counter()
    out = model.generate(**gen_kwargs)
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    t1 = time.perf_counter()
    out_ids = out[0] if isinstance(out, torch.Tensor) else out.sequences[0]
    raw = decode_generated(bot, out_ids, input_len=input_len)
    answer = sanitize_model_answer(raw)
    return {'ran': True, 'answer': answer, 'answer_raw': raw, 'input_tokens': input_len, 'new_tokens': int(out_ids.shape[-1] - input_len), 'latency_s': float(t1 - t0), 'use_cache': bool(use_cache), 'vram': vram_snapshot()}


def adaptive_generate_answer(self, system_prompt: str, user_prompt: str, max_new_tokens: int = 128, do_sample: bool = False, temperature: float = 0.0, top_p: Optional[float] = None, repetition_penalty: Optional[float] = None, cache_policy: str = 'adaptive', prefer_cache: bool = True, oom_retry_tokens: int = 32, context_hard_cap: int = 16000, verbose: bool = False, **kwargs) -> Dict[str, Any]:
    cache_policy = str(cache_policy or 'adaptive').lower().strip()
    if cache_policy in ['off', 'safe', 'false', '0', 'no_cache']:
        attempts = [False]
    elif cache_policy in ['on', 'true', '1', 'cache']:
        attempts = [True]
    else:
        attempts = [True, False] if prefer_cache else [False, True]
    errors = []
    lock = getattr(self, 'model_lock', None)
    def run(use_cache_value):
        return generate_once(self, system_prompt=system_prompt, user_prompt=user_prompt, max_new_tokens=max_new_tokens, use_cache=use_cache_value, do_sample=do_sample, temperature=temperature, top_p=top_p, repetition_penalty=repetition_penalty, context_hard_cap=context_hard_cap)
    for idx, use_cache_value in enumerate(attempts):
        clear_cuda()
        try:
            if verbose:
                print(f'[NZFC adaptive cache] attempt={idx + 1} use_cache={use_cache_value} max_new_tokens={max_new_tokens}')
            if lock is None:
                out = run(use_cache_value)
            else:
                with lock:
                    out = run(use_cache_value)
            out['cache_policy'] = cache_policy
            out['attempt_index'] = idx + 1
            out['fallback_used'] = idx > 0
            out['errors_before_success'] = errors
            return out
        except Exception as e:
            err = {'attempt_index': idx + 1, 'use_cache': bool(use_cache_value), 'type': type(e).__name__, 'message': str(e)[:1000], 'is_oom_like': is_oom_like_error(e)}
            errors.append(err)
            if verbose:
                print('[NZFC adaptive cache][WARN]', err)
                print(traceback.format_exc()[:2000])
            clear_cuda()
            if cache_policy in ['on', 'off', 'safe', 'true', 'false', '1', '0', 'cache', 'no_cache']:
                break
            continue
    return {'ran': False, 'answer': None, 'reason': 'adaptive_generation_failed', 'cache_policy': cache_policy, 'errors': errors, 'vram': vram_snapshot()}


def attach_adaptive_kv_cache_generation(bot, *, default_cache_policy: str = 'adaptive', default_prefer_cache: bool = True, default_oom_retry_tokens: int = 32, default_context_hard_cap: int = 16000, verbose: bool = True):
    def bound_generate_answer(self, system_prompt: str, user_prompt: str, max_new_tokens: int = 128, do_sample: bool = False, temperature: float = 0.0, top_p: Optional[float] = None, repetition_penalty: Optional[float] = None, cache_policy: Optional[str] = None, prefer_cache: Optional[bool] = None, oom_retry_tokens: Optional[int] = None, context_hard_cap: Optional[int] = None, verbose: Optional[bool] = None, **kwargs):
        return adaptive_generate_answer(self, system_prompt=system_prompt, user_prompt=user_prompt, max_new_tokens=max_new_tokens, do_sample=do_sample, temperature=temperature, top_p=top_p, repetition_penalty=repetition_penalty, cache_policy=cache_policy or default_cache_policy, prefer_cache=default_prefer_cache if prefer_cache is None else bool(prefer_cache), oom_retry_tokens=default_oom_retry_tokens if oom_retry_tokens is None else int(oom_retry_tokens), context_hard_cap=default_context_hard_cap if context_hard_cap is None else int(context_hard_cap), verbose=bool(verbose) if verbose is not None else False, **kwargs)
    bot.generate_answer = types.MethodType(bound_generate_answer, bot)
    bot.nzfc_cache_profile = {'default_cache_policy': default_cache_policy, 'default_prefer_cache': default_prefer_cache, 'default_oom_retry_tokens': default_oom_retry_tokens, 'default_context_hard_cap': default_context_hard_cap, 'description': 'Adaptive KV-cache generation: try use_cache=True first, fallback to use_cache=False on failure.'}
    if verbose:
        print('[NZFC adaptive cache][OK] attached')
        print(bot.nzfc_cache_profile)
    return bot