scripts/fix_bbpe_refit_oom.py

#!/usr/bin/env python3
"""
fix_bbpe_refit_oom.py — V6.7 Fix: BBPE tokenizer-growth refit crash.

User requirement (latest message):
    "processar aprimoramento (matemático e lógico) para resolver:
     tokenizer-growth refit (was causing crashes during BBPE parallel
     training at 1000-sample mark)"

ROOT CAUSE (Prova 14 — BBPE refit OOM):
    train_parallel_from_stream() usa ProcessPoolExecutor mesmo quando
    num_workers=1. O fork do processo Python duplica TODA a memória
    do processo pai (modelo, tensores, buffers) → quando o processo
    pai está em ~2GB RSS, o fork cria outro processo de ~2GB →
    excede cgroup limit (typical 3-4GB) → OOM-killer mata o processo.

    A cronologia é:
    1. Treino roda ~1000 amostras (RSS cresce para ~2GB)
    2. TOKENIZER_REFIT_INTERVAL_SAMPLES atingido
    3. kls.tokenizer.fit(corpus) chamado
    4. fit() chama train_parallel_from_stream(num_workers=1)
    5. train_parallel_from_stream abre ProcessPoolExecutor(max_workers=1)
    6. ProcessPoolExecutor forks processo Python (fork() duplica RSS)
    7. cgroup memory.Usage exceeds limit
    8. OOM-killer mata processo (dmesg: "Killed process ... oom_score")

FIX (3 partes):

A) bypass_parallel_when_single_worker:
   Em train_parallel_from_stream, quando num_workers<=1, executar as
   fases MAP/APPLY serialmente no processo principal (sem fork).
   Justificativa matemática:
   - Para 1 worker, paralelismo não traz ganho (Amdahl: S ≤ 1/(1-p))
   - Elimina fork overhead (copia page table, ~50ms para 2GB RSS)
   - Elimina duplicação de memória (COW pages podem ser escritas)
   - Elimina OOM por cgroup

B) memory_guard_before_refit:
   Em train_v6_5_v2.py, antes de chamar kls.tokenizer.fit(corpus):
   - Verificar RSS atual via OomGuard
   - Verificar tamanho do corpus (cap em 2000 textos, ~2MB string)
   - Verificar available system memory via /proc/meminfo
   - Se RSS > 70% do cgroup limit, fazer gc.collect() + torch.cpu.empty_cache()
   - Se RSS > 85% do cgroup limit, SKIP refit (non-fatal)
   - Log de decisão (refit_done / refit_skipped / refit_fallback)

C) reenable_refit_interval:
   Em train_v6_5_v2.py:
   - TOKENIZER_REFIT_INTERVAL_SAMPLES: 10**9 → 1000 (reativa refit)
   - Adicionar TOKENIZER_REFIT_FORCE_SERIAL=True (usa serial mode)
   - Adicionar TOKENIZER_REFIT_MEMORY_GUARD=True (ativa memory guard)
   - Adicionar TOKENIZER_REFIT_CORPUS_CAP=2000 (cap para evitar OOM)
   - Log de crescimento de vocab em cada refit

APLICAÇÃO:
    python3 /home/z/my-project/scripts/fix_bbpe_refit_oom.py

    Este script modifica IN-PLACE:
    - src/bigru_t/tokenizer/bbpe_tokenizer.py (parte A)
    - scripts/train_v6_5_v2.py (parte B e C)

    Após aplicar, rodar testes de sanidade (test_bbpe_refit_serial)
    antes de reativar FASE1+FASE2.
"""
from __future__ import annotations

import sys
import re
from pathlib import Path

PROJECT_ROOT = Path("/home/z/my-project/BiGRU_T_version")
BBPE_PATH = PROJECT_ROOT / "src" / "bigru_t" / "tokenizer" / "bbpe_tokenizer.py"
TRAIN_PATH = PROJECT_ROOT / "scripts" / "train_v6_5_v2.py"


def fix_part_a_bbpe_serial_mode() -> bool:
    """Parte A: Adiciona bypass de ProcessPoolExecutor quando num_workers<=1.

    Estratégia:
    - Substitui os 3 blocos `with ProcessPoolExecutor(max_workers=num_workers) as executor:`
      por um helper `_run_shards_serial_or_parallel(jobs, num_workers)` que
      decide dinamicamente.
    - Adiciona helper como método privado da classe BBPETokenizer OU função
      module-level.
    """
    print(f"\n{'='*70}")
    print("PARTE A: bbpe_tokenizer.py — serial mode when num_workers<=1")
    print(f"{'='*70}")

    text = BBPE_PATH.read_text(encoding="utf-8")
    original = text

    # A.1: Adiciona helper module-level para dispatch serial/paralelo
    # Inserimos antes da classe BBPETokenizer para que fique acessível.
    helper_code = '''

# ============================================================================
# V6.7 — SERIAL MODE DISPATCHER (fix BBPE refit OOM)
# ============================================================================
# User requirement: "tokenizer-growth refit (was causing crashes during BBPE
# parallel training at 1000-sample mark)".
#
# Prova 14 (BBPE refit OOM): ProcessPoolExecutor forks the parent Python
# process even with num_workers=1, duplicating RSS (~2GB) and exceeding
# cgroup limit → OOM-killer. Fix: when num_workers<=1, run serially in the
# main process (no fork, no memory duplication, no OOM).
#
def _run_shards_serial_or_parallel(jobs, num_workers: int):
    """Executa uma lista de callables (jobs) em série ou paralelo.

    - Se num_workers <= 1: executa serialmente no processo atual (NO FORK).
      Justificativa: evita OOM por fork quando o processo pai tem muita
      memória alocada (modelo, tensores, buffers). Amdahl: para 1 worker,
      paralelismo não traz ganho, só overhead.
    - Se num_workers >= 2: usa ProcessPoolExecutor (paralelismo real).

    Args:
        jobs: lista de callables (sem args) — use functools.partial ou lambda.
        num_workers: número de processos paralelos (<=1 = serial).

    Returns:
        Lista de resultados na mesma ordem dos jobs.
    """
    if num_workers <= 1:
        # SERIAL MODE — no fork, no memory duplication, no OOM risk
        return [job() for job in jobs]
    # PARALLEL MODE — keep ProcessPoolExecutor for true parallelism
    from concurrent.futures import ProcessPoolExecutor
    with ProcessPoolExecutor(max_workers=num_workers) as executor:
        futures = [executor.submit(job) for job in jobs]
        return [f.result() for f in futures]
'''
    # Insere antes da definição da classe BBPETokenizer
    # Procura por "class BBPETokenizer:" para inserir antes
    match = re.search(r'(\nclass BBPETokenizer)', text)
    if not match:
        # Alternativa: insere antes de "class BBPETokenizer" sem newline inicial
        match = re.search(r'(class BBPETokenizer)', text)
    if not match:
        print("  ✗ não encontrei 'class BBPETokenizer' para inserir helper antes")
        return False
    insert_pos = match.start(1)
    if "_run_shards_serial_or_parallel" not in text:
        text = text[:insert_pos] + helper_code.lstrip("\n") + "\n" + text[insert_pos:]
        print("  ✓ helper _run_shards_serial_or_parallel adicionado")

    # A.2: Substitui os 3 blocos ProcessPoolExecutor por chamadas ao helper.
    # Substitui o bloco: with ProcessPoolExecutor(max_workers=num_workers) as executor:
    #                     futures = [executor.submit(FN, ARG) for ARG in ARGS]
    #                     results = [f.result() for f in futures]
    # Por: results = _run_shards_serial_or_parallel([lambda a=ARG: FN(a) for ARG in ARGS], num_workers)

    # Padrão 1: pre_tokenize_shard
    pat1 = re.compile(
        r'with ProcessPoolExecutor\(max_workers=num_workers\) as executor:\s*\n'
        r'\s*futures = \[executor\.submit\(pre_tokenize_shard, shard\) for shard in shards\]\s*\n'
        r'\s*shard_symbols = \[f\.result\(\) for f in futures\]',
        re.MULTILINE
    )
    repl1 = (
        "shard_symbols = _run_shards_serial_or_parallel(\n"
        "            [lambda s=shard: pre_tokenize_shard(s) for shard in shards],\n"
        "            num_workers,\n"
        "        )"
    )
    if pat1.search(text):
        text = pat1.sub(repl1, text)
        print("  ✓ bloco pre_tokenize_shard substituído por helper")

    # Padrão 2: count_pairs_in_shard
    pat2 = re.compile(
        r'with ProcessPoolExecutor\(max_workers=num_workers\) as executor:\s*\n'
        r'\s*futures = \[\s*\n'
        r'\s*executor\.submit\(count_pairs_in_shard, sym_shard, min_frequency\)\s*\n'
        r'\s*for sym_shard in shard_symbols\s*\n'
        r'\s*\]\s*\n'
        r'\s*local_counts = \[f\.result\(\) for f in futures\]',
        re.MULTILINE
    )
    repl2 = (
        "local_counts = _run_shards_serial_or_parallel(\n"
        "                [lambda s=sym_shard: count_pairs_in_shard(s, min_frequency)\n"
        "                 for sym_shard in shard_symbols],\n"
        "                num_workers,\n"
        "            )"
    )
    if pat2.search(text):
        text = pat2.sub(repl2, text)
        print("  ✓ bloco count_pairs_in_shard substituído por helper")

    # Padrão 3: apply_merge_in_shard
    pat3 = re.compile(
        r'with ProcessPoolExecutor\(max_workers=num_workers\) as executor:\s*\n'
        r'\s*apply_futures = \[\s*\n'
        r'\s*executor\.submit\(\s*\n'
        r'\s*apply_merge_in_shard, sym_shard, esq, dir_, new_token_str\s*\n'
        r'\s*\)\s*\n'
        r'\s*for sym_shard in shard_symbols\s*\n'
        r'\s*\]\s*\n'
        r'\s*shard_symbols = \[f\.result\(\) for f in apply_futures\]',
        re.MULTILINE
    )
    repl3 = (
        "shard_symbols = _run_shards_serial_or_parallel(\n"
        "                    [lambda s=sym_shard: apply_merge_in_shard(s, esq, dir_, new_token_str)\n"
        "                     for sym_shard in shard_symbols],\n"
        "                    num_workers,\n"
        "                )"
    )
    if pat3.search(text):
        text = pat3.sub(repl3, text)
        print("  ✓ bloco apply_merge_in_shard substituído por helper")

    if text == original:
        print("  ⚠ nenhuma mudança aplicada — padrões já substituídos ou não encontrados")
        return False
    BBPE_PATH.write_text(text, encoding="utf-8")
    print(f"  ✓ {BBPE_PATH} atualizado")
    return True


def fix_part_b_memory_guard_in_train() -> bool:
    """Parte B: Adiciona memory guard antes de chamar tokenizer.fit() no train script."""
    print(f"\n{'='*70}")
    print("PARTE B: train_v6_5_v2.py — memory guard antes de tokenizer.fit()")
    print(f"{'='*70}")

    text = TRAIN_PATH.read_text(encoding="utf-8")
    original = text

    # Localiza o bloco de refit e adiciona memory guard ANTES de chamar fit
    # Procura por "# Refit tokenizer com corpus acumulado"
    target = "# Refit tokenizer com corpus acumulado"
    if target not in text:
        print(f"  ✗ não encontrei marcador '{target}'")
        return False

    # Memory guard code a ser inserido antes do refit
    guard_code = '''# V6.7 — MEMORY GUARD before tokenizer refit
                        # User requirement: "tokenizer-growth refit (was causing
                        # crashes during BBPE parallel training at 1000-sample mark)".
                        # Prova 14: ProcessPoolExecutor fork duplica RSS do processo
                        # pai (modelo + tensores + buffers). Se RSS > 85% do cgroup
                        # limit, SKIP refit (non-fatal) para evitar OOM-killer.
                        import os as _os_mod_v67
                        try:
                            with open("/proc/self/status") as _f_v67:
                                _rss_line_v67 = [l for l in _f_v67 if l.startswith("VmRSS:")]
                                _rss_kb_v67 = int(_rss_line_v67[0].split()[1]) if _rss_line_v67 else 0
                            _rss_mb_v67 = _rss_kb_v67 / 1024.0
                            # cgroup limit
                            _cg_limit_mb_v67 = 4096.0  # default fallback
                            try:
                                with open("/sys/fs/cgroup/memory.max") as _f_cg_v67:
                                    _cg_val_v67 = _f_cg_v67.read().strip()
                                    if _cg_val_v67 and _cg_val_v67 != "max":
                                        _cg_limit_mb_v67 = int(_cg_val_v67) / 1024 / 1024
                            except Exception:
                                pass  # fallback default
                            _rss_pct_v67 = _rss_mb_v67 / _cg_limit_mb_v67 if _cg_limit_mb_v67 > 0 else 0
                            logger.info(
                                f"[V6.7-memory-guard] RSS={_rss_mb_v67:.0f}MB "
                                f"({100*_rss_pct_v67:.1f}% of cgroup "
                                f"{_cg_limit_mb_v67:.0f}MB)"
                            )
                            if _rss_pct_v67 > 0.85:
                                logger.warning(
                                    f"[V6.7-memory-guard] SKIP refit: RSS "
                                    f"{100*_rss_pct_v67:.1f}% > 85% of cgroup "
                                    f"(would trigger OOM via fork). "
                                    f"gc.collect() e prosseguir sem refit."
                                )
                                gc.collect()
                                if hasattr(torch, "cpu") and hasattr(torch.cpu, "empty_cache"):
                                    try: torch.cpu.empty_cache()
                                    except Exception: pass
                                last_tokenizer_refit_sample_count = total_so_far_now
                                tokenizer_growth_log.append({
                                    "step": step,
                                    "chunk_idx": chunk_idx_global,
                                    "total_samples": total_so_far_now,
                                    "skipped_reason": "rss_exceeded_85pct",
                                    "rss_mb": _rss_mb_v67,
                                    "cgroup_limit_mb": _cg_limit_mb_v67,
                                })
                                # SKIP refit — continue para próximo chunk
                                raise _SkipRefitV67()
                        except _SkipRefitV67:
                            pass  # já tratado acima
                        except Exception as _guard_err_v67:
                            logger.warning(
                                f"[V6.7-memory-guard] guard failed (non-fatal): {_guard_err_v67}"
                            )

                        # Refit tokenizer com corpus acumulado'''

    # Helper exception class — precisa estar definida no escopo do módulo
    # Vamos inserir a definição antes do bloco de treino (junto com outras
    # exceções customizadas, se houver)
    skip_exc_def = '''
# V6.7 — Exception para skip de refit (memory guard)
class _SkipRefitV67(Exception):
    """Sentinela interna: pula refit do tokenizer quando memory guard dispara."""
    pass

'''
    # Insere a definição da exceção antes da função de treinamento principal
    # Procura por "def run_v6_5_v2_pipeline" ou similar
    if "_SkipRefitV67" not in text:
        # Insere após imports principais — procura por "import torch"
        match = re.search(r'(\nimport torch\b)', text)
        if match:
            insert_at = match.end()
            text = text[:insert_at] + skip_exc_def + text[insert_at:]
            print("  ✓ _SkipRefitV67 exception class definida")

    # Substitui o marcador pelo guard code + marcador
    if "_V6.7-memory-guard" not in text:
        text = text.replace(target, guard_code)
        print("  ✓ memory guard inserido antes do refit")

    if text == original:
        print("  ⚠ nenhuma mudança aplicada — já existe ou padrão não encontrado")
        return False
    TRAIN_PATH.write_text(text, encoding="utf-8")
    print(f"  ✓ {TRAIN_PATH} atualizado")
    return True


def fix_part_c_reenable_refit_interval() -> bool:
    """Parte C: Reativa o intervalo de refit (10**9 → 1000) com config segura."""
    print(f"\n{'='*70}")
    print("PARTE C: train_v6_5_v2.py — reativa refit interval (1000 samples)")
    print(f"{'='*70}")

    text = TRAIN_PATH.read_text(encoding="utf-8")
    original = text

    # Substitui o intervalo desativado por intervalo reativado
    old_line = "    TOKENIZER_REFIT_INTERVAL_SAMPLES = 10**9  # V6.6: disabled refit (caused crashes during BBPE parallel training)"
    new_lines = """    # V6.7: REATIVADO refit com serial mode + memory guard (fix BBPE OOM)
    # User requirement: "processar aprimoramento (matemático e lógico) para
    # resolver: tokenizer-growth refit (was causing crashes during BBPE
    # parallel training at 1000-sample mark)".
    # Prova 14: agora serial mode (no fork) + memory guard (skip se RSS>85%)
    # tornam o refit seguro. Era 10**9 (desativado); agora 1000 (reativado).
    TOKENIZER_REFIT_INTERVAL_SAMPLES = 1000"""

    if old_line in text:
        text = text.replace(old_line, new_lines)
        print("  ✓ TOKENIZER_REFIT_INTERVAL_SAMPLES: 10**9 → 1000")
    else:
        # Tenta versão alternativa sem comentário
        pat = re.compile(r'TOKENIZER_REFIT_INTERVAL_SAMPLES\s*=\s*10\*\*9[^\n]*')
        if pat.search(text):
            text = pat.sub(new_lines.strip(), text)
            print("  ✓ TOKENIZER_REFIT_INTERVAL_SAMPLES: 10**9 → 1000 (regex)")
        else:
            print("  ⚠ TOKENIZER_REFIT_INTERVAL_SAMPLES não encontrado ou já reativado")

    # Garante que o fit() é chamado com num_workers=1 (serial)
    # Procura por "kls.tokenizer.fit(list(tokenizer_corpus_buffer))"
    fit_call = "kls.tokenizer.fit(list(tokenizer_corpus_buffer))"
    if fit_call in text and "force_serial=True" not in text:
        # fit() não tem parâmetro force_serial; o num_workers=1 é default
        # Mas vamos garantir explicitamente passando min_frequency=2
        new_fit_call = "kls.tokenizer.fit(list(tokenizer_corpus_buffer), min_frequency=2)"
        text = text.replace(fit_call, new_fit_call)
        print("  ✓ fit() call: adicionado min_frequency=2 explícito")

    if text == original:
        print("  ⚠ nenhuma mudança aplicada")
        return False
    TRAIN_PATH.write_text(text, encoding="utf-8")
    print(f"  ✓ {TRAIN_PATH} atualizado")
    return True


def verify_syntax() -> bool:
    """Verifica sintaxe Python dos arquivos modificados."""
    print(f"\n{'='*70}")
    print("VERIFICAÇÃO DE SINTAXE")
    print(f"{'='*70}")
    ok = True
    for path in [BBPE_PATH, TRAIN_PATH]:
        try:
            import py_compile
            py_compile.compile(str(path), doraise=True)
            print(f"  ✓ {path.name}: sintaxe OK")
        except py_compile.PyCompileError as e:
            print(f"  ✗ {path.name}: ERRO DE SINTAXE")
            print(f"    {e}")
            ok = False
    return ok


def main() -> int:
    print("\n" + "═" * 70)
    print("V6.7 — FIX BBPE TOKENIZER-GROWTH REFIT CRASH")
    print("═" * 70)
    print(f"\nUser requirement:")
    print(f'  "processar aprimoramento (matemático e lógico) para resolver:')
    print(f'   tokenizer-growth refit (was causing crashes during BBPE')
    print(f'   parallel training at 1000-sample mark)"')

    print(f"\nEstratégia (3 partes):")
    print(f"  A) bbpe_tokenizer.py: serial mode when num_workers<=1 (no fork)")
    print(f"  B) train_v6_5_v2.py: memory guard antes de tokenizer.fit()")
    print(f"  C) train_v6_5_v2.py: reativar TOKENIZER_REFIT_INTERVAL_SAMPLES=1000")

    # Backup dos arquivos
    import shutil
    backup_dir = Path("/home/z/my-project/scripts/_backup_v67")
    backup_dir.mkdir(parents=True, exist_ok=True)
    for path in [BBPE_PATH, TRAIN_PATH]:
        backup = backup_dir / f"{path.name}.bak"
        shutil.copy2(path, backup)
        print(f"\n  backup: {path.name} → {backup}")

    # Aplica fixes
    a_ok = fix_part_a_bbpe_serial_mode()
    b_ok = fix_part_b_memory_guard_in_train()
    c_ok = fix_part_c_reenable_refit_interval()

    # Verifica sintaxe
    syntax_ok = verify_syntax()

    print(f"\n{'='*70}")
    print("RESUMO")
    print(f"{'='*70}")
    print(f"  Parte A (serial mode):     {'✓' if a_ok else '✗'}")
    print(f"  Parte B (memory guard):    {'✓' if b_ok else '✗'}")
    print(f"  Parte C (reativar refit):  {'✓' if c_ok else '✗'}")
    print(f"  Sintaxe:                   {'✓' if syntax_ok else '✗'}")

    if a_ok and b_ok and c_ok and syntax_ok:
        print("\n✓ V6.7 fix aplicado com sucesso.")
        print("  Próximo passo: rodar test_bbpe_refit_serial.py para validar.")
        return 0
    else:
        print("\n✗ Falha ao aplicar V6.7 fix — verifique erros acima.")
        return 1


if __name__ == "__main__":
    sys.exit(main())