diff --git a/CHANGELOG.md b/CHANGELOG.md index 023e0b76..5f1a8050 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,10 @@ e este projeto adere ao [Versionamento Semântico](https://semver.org/lang/pt-BR ## [Unreleased] +### Corrigido + +- Com `checkpoint_path`, um Ctrl-C (`KeyboardInterrupt`) grava no checkpoint as linhas concluídas desde a última gravação e relança a interrupção, nos modos sequencial e paralelo. Antes, o arquivo ficava na última gravação de `batch_size`, e a retomada pagava de novo as linhas concluídas depois dela. No paralelo, entram as chamadas já em andamento, que terminam enquanto a interrupção espera as threads do executor, também quando o Ctrl-C sai de dentro do despacho; a linha que ainda está na fila, sem thread que a tenha começado, é cancelada e fica para a retomada, e as que um segundo Ctrl-C deixa sem terminar ficam pendentes. A espera pelas linhas corre em fatias curtas, porque o handler de SIGINT que o polars instala ao ser importado faz o kernel reiniciar a espera sem prazo: com o polars instalado, o Ctrl-C do paralelo só chegava quando uma linha terminava, e o segundo não cortava a espera. A falha ao gravar vira aviso e não troca a interrupção (#213). + ## [0.11.0] - 2026-10-29 ### Adicionado diff --git a/docs/en/guides/faq.md b/docs/en/guides/faq.md index 29f505f8..8c4418cd 100644 --- a/docs/en/guides/faq.md +++ b/docs/en/guides/faq.md @@ -64,3 +64,5 @@ Run a sample first (`df.sample(30)`) with `track_tokens=True`. The summary at th ### A long run was interrupted Use `batch_size` and `checkpoint_path` to save progress during the run, and resume by running the same call again, which continues from the file; see [Checkpoints for Long Runs](performance.md#checkpoints-for-long-runs). + +A Ctrl-C saves to the checkpoint the rows completed since the last save before stopping, and resuming does not call them again. The row whose call the Ctrl-C cut off stays pending. With `parallel_requests`, the interruption waits for the calls in progress to finish, and they go into the save; a second Ctrl-C cuts that wait short, and the calls that did not finish stay pending. diff --git a/docs/guides/faq.md b/docs/guides/faq.md index e71ec335..9fa98d7b 100644 --- a/docs/guides/faq.md +++ b/docs/guides/faq.md @@ -64,3 +64,5 @@ Rode primeiro uma amostra (`df.sample(30)`) com `track_tokens=True`. O resumo ao ### Uma execução longa foi interrompida Use `batch_size` e `checkpoint_path` para gravar o progresso durante a execução, e retome rodando a mesma chamada de novo, que continua do arquivo; ver [Checkpoints em Execuções Longas](performance.md#checkpoints-em-execucoes-longas). + +Um Ctrl-C grava no checkpoint as linhas concluídas desde a última gravação antes de interromper, e a retomada não as chama de novo. A linha cuja chamada o Ctrl-C cortou fica pendente. Com `parallel_requests`, a interrupção espera as chamadas em andamento terminarem, e elas entram na gravação; um segundo Ctrl-C corta essa espera, e as chamadas que não terminaram ficam pendentes. diff --git a/src/dataframeit/core.py b/src/dataframeit/core.py index ed009f50..1fd79ba4 100644 --- a/src/dataframeit/core.py +++ b/src/dataframeit/core.py @@ -6,6 +6,7 @@ import datetime import hashlib import importlib.util +import itertools import json import logging import numbers @@ -157,6 +158,14 @@ def _install_fallback_handler(logger: logging.Logger) -> None: # Chamadas de busca por campo acima das quais vale avisar sobre rate limit. _PER_FIELD_SEARCH_WARNING_CALLS = 100 +# Prazo de cada espera do laço paralelo pelas linhas em andamento; ver o wait +# de _process_rows_parallel. +_WAIT_SLICE_SECONDS = 0.1 + +# Numera os executores do laço paralelo, um por execução, para que o nome das +# threads de cada um não se confunda com o das de outro; ver _closing_in_slices. +_executor_ids = itertools.count() + # Segundos que o aviso de rate limit supõe para processar uma linha com busca, # além do rate_limit_delay. A estimativa do aviso é um piso calculado só da # configuração: supõe uma busca por execução do agente. Um agente que faz mais @@ -2535,7 +2544,7 @@ def _progress_suffix( return "" -def _process_rows( # noqa: PLR0913 (estado da execução repassado por dataframeit) +def _process_rows( # noqa: C901, PLR0913 (estado da execução repassado por dataframeit, e um ramo por desfecho da linha) df: pd.DataFrame, *, text_column: str, @@ -2588,16 +2597,11 @@ def _process_rows( # noqa: PLR0913 (estado da execução repassado por datafram rows_processed_this_run = 0 rows_saved = 0 - # Processar cada linha - for i, (idx, row) in enumerate(tqdm(df.iterrows(), total=len(df), desc=desc)): + def process_row(i: int, idx: Hashable, row: pd.Series) -> tuple[bool, dict[str, Any]]: + """Registra o desfecho da linha e devolve se a chamada deu certo e os campos gravados.""" # Verificar se linha já foi processada row_already_processed = pd.notna(row[status_col]) and row[status_col] == "processed" - # Com reprocess_columns, todas as linhas são processadas. - if not reprocess_columns and not is_pending[i]: - continue - - started = time.monotonic() success = False values: dict[str, Any] = {} if _is_missing_text(row[text_column]): @@ -2640,7 +2644,6 @@ def _process_rows( # noqa: PLR0913 (estado da execução repassado por datafram status_col=status_col, reprocess_columns=reprocess_columns, ) - break except Exception as e: # noqa: BLE001 (qualquer falha da linha vira status 'error') _record_error_safely( df, @@ -2652,28 +2655,53 @@ def _process_rows( # noqa: PLR0913 (estado da execução repassado por datafram reprocess_columns=reprocess_columns, token_stats=token_stats, ) + return success, values - _notify_row_done( - on_row_done, - RowResult( - index=idx, - position=rows_processed_this_run, - success=success, - error=_recorded_error(df, idx), - values=values, - elapsed_seconds=time.monotonic() - started, - ), - ) - rows_processed_this_run += 1 - if ( - checkpoint - and rows_processed_this_run % checkpoint.batch_size == 0 - and _try_save_checkpoint(df, checkpoint) - ): - rows_saved = rows_processed_this_run + # Processar cada linha. A gravação do checkpoint fica neste laço, e não em + # process_row, para que o aviso de falha aponte a chamada do usuário (ver o + # stacklevel de _try_save_checkpoint). + try: + for i, (idx, row) in enumerate(tqdm(df.iterrows(), total=len(df), desc=desc)): + # Com reprocess_columns, todas as linhas são processadas. + if not reprocess_columns and not is_pending[i]: + continue + started = time.monotonic() + success, values = process_row(i, idx, row) + if "abort" in token_stats: + break - if success and config.rate_limit_delay > 0: - time.sleep(config.rate_limit_delay) + # A linha conta antes do callback: um Ctrl-C dentro dele chega com a + # linha já gravada no DataFrame, e a gravação do Ctrl-C a leva. + position = rows_processed_this_run + rows_processed_this_run += 1 + _notify_row_done( + on_row_done, + RowResult( + index=idx, + position=position, + success=success, + error=_recorded_error(df, idx), + values=values, + elapsed_seconds=time.monotonic() - started, + ), + ) + if ( + checkpoint + and rows_processed_this_run % checkpoint.batch_size == 0 + and _try_save_checkpoint(df, checkpoint) + ): + rows_saved = rows_processed_this_run + + if success and config.rate_limit_delay > 0: + time.sleep(config.rate_limit_delay) + except KeyboardInterrupt: + # Ctrl-C: grava as linhas registradas desde a última gravação, que a + # retomada pagaria de novo, e relança a interrupção. A linha cuja chamada + # foi cortada não chegou a ser registrada e fica pendente. A falha de + # gravação vira aviso em _try_save_checkpoint e não troca a interrupção. + if checkpoint and rows_processed_this_run > rows_saved: + _try_save_checkpoint(df, checkpoint) + raise # Save final: a cauda (< batch_size), o que uma gravação que falhou deixou de # fora e as marcas de uma interrupção, que não contam como linha processada. @@ -2683,6 +2711,62 @@ def _process_rows( # noqa: PLR0913 (estado da execução repassado por datafram return token_stats +@contextmanager +def _on_ctrl_c(callback: Callable[[], None]) -> Generator[None, None, None]: + """Chama `callback` quando um KeyboardInterrupt atravessa o bloco, e o relança.""" + try: + yield + except KeyboardInterrupt: + callback() + raise + + +@contextmanager +def _closing_in_slices(max_workers: int) -> Generator[ThreadPoolExecutor, None, None]: + """Abre o executor da execução e, na saída, o encerra e espera as threads dele em fatias. + + Faz as vezes do `with ThreadPoolExecutor`, cuja saída espera as threads sem + prazo: com prazo em cada espera, um segundo Ctrl-C a corta; ver o `wait` de + _process_rows_parallel. O encerramento cancela as linhas que ainda estão na + fila do executor, e nenhuma thread começa outra depois dele: a linha + cancelada não é paga e fica sem status, para a retomada. Isso cobre o + Ctrl-C que sai de dentro do `Thread.start()`, com a thread já lançada e + ainda sem marcar o início; sem o cancelamento, ela começaria depois da + gravação do checkpoint e pagaria fora dele a linha da fila. Fora do Ctrl-C, + a saída chega com a fila vazia, porque o laço espera todos os futuros em + voo. + + A espera é pelas threads, e não pelos futuros em voo, porque o Ctrl-C pode + sair de dentro do `submit` depois de uma thread tirar a linha da fila. A + linha roda e é paga sem que o futuro dela chegue ao laço, e a gravação do + checkpoint precisa encontrá-la registrada. A espera cobre também a linha + cujo `on_row_done` ainda roda: o callback vem da thread do worker, que não + recebe o KeyboardInterrupt, e a linha já está contada quando ele começa. + + As threads são achadas pelo nome, com o prefixo deste executor, em + `threading.enumerate()`. O conjunto `_threads` do executor é privado e só + recebe a thread depois de `Thread.start()` voltar; `enumerate` a lista + desde antes. `is_alive` deixa de fora a thread que ainda não marcou o + início, que `join` recusaria: ela nunca começa ou, ao começar, encontra a + fila vazia e sai sem rodar linha. A lista se refaz a cada fatia e pega + também a thread que começa durante a espera. O número no prefixo separa as + threads de duas chamadas concorrentes, e o `_` depois dele impede que o + executor 1 tome por suas as threads do 10. + """ + prefix = f"dataframeit-paralelo{next(_executor_ids)}" + executor = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix=prefix) + try: + yield executor + finally: + executor.shutdown(wait=False, cancel_futures=True) + while alive := [ + thread + for thread in threading.enumerate() + if thread.name.startswith(f"{prefix}_") and thread.is_alive() + ]: + alive[0].join(_WAIT_SLICE_SECONDS) + + def _process_rows_parallel( # noqa: C901, PLR0912, PLR0913, PLR0915 (estado compartilhado entre threads) df: pd.DataFrame, *, @@ -2918,14 +3002,31 @@ def process_single_row(row_data: tuple) -> dict: time.sleep(config.rate_limit_delay) return {"success": True, "idx": idx} + def save_registered_rows() -> None: + """Grava o checkpoint no Ctrl-C. + + O executor espera as linhas em andamento antes de chegar aqui, e elas + chegam registradas. Um segundo Ctrl-C durante essa espera deixa threads + rodando: a cópia sob `lock` leva só as linhas que terminaram, e as + demais ficam pendentes para a retomada. Pelo rótulo, o _SnapshotWriter + não regrava o arquivo sem linha contada desde o último snapshot gravado. + """ + if writer: + with lock: + frame, label = df.copy(), checkpoint_counter + writer.save(frame, label) + # Um executor para a execução inteira. O teto de linhas em voo é # current_workers, lido a cada reposição: a redução por rate limit só deixa - # de despachar, sem cancelar o que já está em voo, e vale até o fim. + # de despachar, sem cancelar o que já está em voo, e vale até o fim. Na + # saída por Ctrl-C, o executor fecha e espera as threads antes de + # save_registered_rows gravar: os `with` saem na ordem inversa. pending_rows = iter(rows_to_process) in_flight: dict = {} with ( + _on_ctrl_c(save_registered_rows), tqdm(total=len(rows_to_process), desc=desc) as pbar, - ThreadPoolExecutor(max_workers=parallel_requests) as executor, + _closing_in_slices(parallel_requests) as executor, ): while True: with lock: @@ -2947,7 +3048,18 @@ def process_single_row(row_data: tuple) -> dict: if not in_flight: break - done, _ = wait(in_flight, return_when=FIRST_COMPLETED) + # A espera tem prazo por causa do polars, que, ao ser importado, + # instala um handler de SIGINT com SA_RESTART. Com ele, o kernel + # reinicia a espera sem prazo de um lock em vez de interrompê-la, e + # o Ctrl-C só viraria KeyboardInterrupt quando uma linha terminasse. + # O que corta a espera é haver prazo, e não a duração da fatia: no + # Linux, a espera com prazo de um semáforo não é reiniciada nem com + # SA_RESTART (ver signal(7)) e volta com EINTR assim que o sinal + # chega, e o interpretador levanta o KeyboardInterrupt na hora. A + # fatia curta só limita a demora numa plataforma que reinicie também + # a espera com prazo. Sem linha terminada no prazo, `done` volta + # vazio e o laço só dá a volta: relê o teto e espera de novo. + done, _ = wait(in_flight, timeout=_WAIT_SLICE_SECONDS, return_when=FIRST_COMPLETED) for future in done: row, started = in_flight.pop(future) try: diff --git a/tests/test_interrupcao.py b/tests/test_interrupcao.py index d70c9f0b..6316021b 100644 --- a/tests/test_interrupcao.py +++ b/tests/test_interrupcao.py @@ -1,7 +1,12 @@ """Falha que impede as linhas seguintes interrompe a execução sem gastá-las.""" +import signal +import sys import threading +import time +from collections.abc import Generator from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager, nullcontext from unittest.mock import patch import pandas as pd @@ -229,3 +234,734 @@ def test_interrupcao_com_texto_ilegivel_ainda_interrompe(paralelo): saida = _roda(df, llm, parallel_requests=paralelo) assert saida["_dataframeit_status"].isna().all() + + +# ============================================================================= +# Ctrl-C +# ============================================================================= + + +def _ctrl_c(): + """Manda SIGINT ao thread principal, como o Ctrl-C do terminal. + + No modo sequencial a chamada roda no thread principal, e o + `KeyboardInterrupt` sai do próprio `pthread_kill`, antes de a chamada + terminar. No paralelo, ele chega ao thread principal enquanto este espera as + linhas, e a chamada do worker termina e é paga. + """ + assert signal.getsignal(signal.SIGINT) is signal.default_int_handler + principal = threading.main_thread().ident + assert principal is not None + signal.pthread_kill(principal, signal.SIGINT) + + +def _llm_com_ctrl_c(na_chamada: int | None = None): + """Devolve a lista dos textos cujas chamadas terminaram (pagas) e o falso LLM.""" + trava = threading.Lock() + iniciadas: list[str] = [] + pagas: list[str] = [] + + def llm(texto, *args, **kwargs): + with trava: + iniciadas.append(texto) + vez = len(iniciadas) + if vez == na_chamada: + _ctrl_c() + with trava: + pagas.append(texto) + return {"data": {"campo1": f"r-{texto}"}, "usage": {}} + + return pagas, llm + + +@contextmanager +def _ctrl_c_depois(gatilho: threading.Event, *folgas: float) -> Generator[None, None, None]: + """Manda um Ctrl-C a cada folga, contada desde `gatilho`, enquanto o bloco roda. + + Na saída do bloco, o Ctrl-C que falta é cancelado e o envio em curso é + esperado: o `KeyboardInterrupt` cai dentro do teste, nunca num teste adiante. + """ + saiu = threading.Event() + trava = threading.Lock() + + def envia() -> None: + while not gatilho.wait(0.05): + if saiu.is_set(): + return + inicio = time.monotonic() + for folga in folgas: + if saiu.wait(max(0.0, inicio + folga - time.monotonic())): + return + with trava: + if saiu.is_set(): + return + _ctrl_c() + + remetente = threading.Thread(target=envia, daemon=True) + remetente.start() + try: + yield + finally: + # `saiu` vem antes da trava: um KeyboardInterrupt que chegue enquanto o + # thread principal espera a trava já não deixa outro envio para depois. + saiu.set() + with trava: + pass + remetente.join() + + +def _entrada(biblioteca: str, textos: list[str]): + if biblioteca == "polars": + pl = pytest.importorskip("polars") + return pl.DataFrame({"texto": textos}) + return pd.DataFrame({"texto": textos}) + + +requer_pthread_kill = pytest.mark.skipif( + not hasattr(signal, "pthread_kill"), reason="Ctrl-C simulado com signal.pthread_kill" +) + + +@requer_pthread_kill +@pytest.mark.parametrize("paralelo", [1, 3]) +@pytest.mark.parametrize("biblioteca", ["pandas", "polars"]) +def test_ctrl_c_grava_as_linhas_pagas_e_a_retomada_nao_as_chama_de_novo( + tmp_path, biblioteca, paralelo +): + """12 linhas, batch_size=5 e Ctrl-C durante a 9ª chamada. + + No sequencial, a 9ª chamada não termina, e 8 linhas estão pagas. Com 3 + workers, a 9ª chamada começa com a 7ª e a 8ª em voo, presas até o executor + ser encerrado: as três terminam enquanto a interrupção espera as threads, e + as linhas 10 a 12, que nenhuma thread começou, ficam pendentes. A gravação + a cada 5 linhas deixaria as últimas 3 ou 4 de fora. + """ + textos = [f"t{i:02d}" for i in range(1, 13)] + entrada = _entrada(biblioteca, textos) + ckpt = tmp_path / "ckpt.csv" + pagas, llm_base = _llm_com_ctrl_c() + trava = threading.Lock() + comecadas: list[str] = [] + encerrando = threading.Event() + encerra = ThreadPoolExecutor.shutdown + + def encerra_observado(self, *args, **kwargs): + encerrando.set() + return encerra(self, *args, **kwargs) + + def llm(texto, *args, **kwargs): + # A 7ª e a 8ª chamadas seguram as threads delas até o encerramento, e + # com as três em voo o laço não despacha a 10ª linha antes do Ctrl-C. + with trava: + comecadas.append(texto) + vez = len(comecadas) + if vez == 9: + _ctrl_c() + elif paralelo > 1 and vez in {7, 8}: + assert encerrando.wait(timeout=10) + return llm_base(texto, *args, **kwargs) + + with ( + patch.object(ThreadPoolExecutor, "shutdown", encerra_observado), + pytest.raises(KeyboardInterrupt), + ): + _roda( + entrada, + llm, + parallel_requests=paralelo, + batch_size=5, + checkpoint_path=ckpt, + ) + + assert sorted(pagas) == textos[: 8 if paralelo == 1 else 9] + gravado = pd.read_csv(ckpt) + processadas = gravado.loc[gravado["_dataframeit_status"].eq("processed"), "texto"] + assert sorted(processadas) == sorted(pagas) + + retomada, llm_ok = _llm_com_ctrl_c() + saida = _roda( + _entrada(biblioteca, textos), + llm_ok, + parallel_requests=paralelo, + batch_size=5, + checkpoint_path=ckpt, + ) + + assert sorted(retomada) == sorted(set(textos) - set(pagas)) + assert list(saida["campo1"]) == [f"r-{texto}" for texto in textos] + + +@requer_pthread_kill +@pytest.mark.parametrize("paralelo", [1, 3]) +def test_falha_ao_gravar_no_ctrl_c_avisa_e_mantem_o_keyboard_interrupt(tmp_path, paralelo): + """A falha vira aviso, como nas gravações a cada batch_size.""" + _, llm = _llm_com_ctrl_c(na_chamada=5) + + with ( + patch("dataframeit.core._save_checkpoint", side_effect=OSError("disco cheio")) as grava, + pytest.warns(UserWarning, match="Falha ao gravar o checkpoint.*OSError: disco cheio"), + pytest.raises(KeyboardInterrupt), + ): + _roda( + pd.DataFrame({"texto": [f"t{i}" for i in range(9)]}), + llm, + parallel_requests=paralelo, + batch_size=100, + checkpoint_path=tmp_path / "ckpt.csv", + ) + + assert grava.call_count == 1 + + +def _sigint_com_restart(reinicia: bool) -> Generator[None, None, None]: + """Liga ou desliga o SA_RESTART do SIGINT durante o teste. + + Na volta, o SA_RESTART fica como o processo o tinha: ligado com o polars + importado, que instala um handler com ele, e desligado sem o polars. + """ + com_polars = "polars" in sys.modules + signal.siginterrupt(signal.SIGINT, not reinicia) + yield + signal.siginterrupt(signal.SIGINT, not com_polars) + + +@pytest.fixture +def sigint_reinicia_a_espera(): + """Põe SA_RESTART no SIGINT, como o handler que o polars instala ao ser importado. + + Com SA_RESTART, o kernel reinicia a espera sem prazo de um lock quando o + SIGINT chega, e o thread principal só levanta o `KeyboardInterrupt` quando a + espera termina. Sem isso, o SIGINT interromperia também a espera sem prazo, + e o teste não a distinguiria da espera em fatias. + """ + yield from _sigint_com_restart(reinicia=True) + + +@pytest.fixture +def sigint_interrompe_a_espera(): + """Tira o SA_RESTART do SIGINT, como fica o processo sem o polars importado. + + Sem SA_RESTART, o SIGINT interrompe também a espera sem prazo do + `Thread.start()` pelo início da thread nova, e o `KeyboardInterrupt` sai de + dentro dele. + """ + yield from _sigint_com_restart(reinicia=False) + + +@requer_pthread_kill +@pytest.mark.usefixtures("sigint_reinicia_a_espera") +def test_segundo_ctrl_c_corta_a_espera_e_a_chamada_lenta_fica_pendente(tmp_path): + """O segundo Ctrl-C chega enquanto a interrupção espera a chamada lenta. + + A espera acaba sem a chamada terminar, o checkpoint leva só a linha + concluída, e a retomada chama a lenta de novo. + """ + textos = ["lenta", "ctrl-c"] + ckpt = tmp_path / "ckpt.csv" + encerrando = threading.Event() + liberada = threading.Event() + terminou = threading.Event() + encerra = ThreadPoolExecutor.shutdown + + def encerra_observado(self, *args, **kwargs): + encerrando.set() + return encerra(self, *args, **kwargs) + + def llm(texto, *args, **kwargs): + if texto == "lenta": + liberada.wait(timeout=10) + terminou.set() + else: + _ctrl_c() + return {"data": {"campo1": f"r-{texto}"}, "usage": {}} + + try: + # O primeiro Ctrl-C encerra o executor antes de esperar a chamada + # lenta; a folga deixa a espera começar antes do segundo. + with ( + patch.object(ThreadPoolExecutor, "shutdown", encerra_observado), + _ctrl_c_depois(encerrando, 0.2), + pytest.raises(KeyboardInterrupt), + ): + _roda( + pd.DataFrame({"texto": textos}), + llm, + parallel_requests=2, + batch_size=100, + checkpoint_path=ckpt, + ) + cortada = not terminou.is_set() + gravado = pd.read_csv(ckpt) + finally: + liberada.set() + + assert cortada + status = dict(zip(gravado["texto"], gravado["_dataframeit_status"], strict=True)) + assert status["ctrl-c"] == "processed" + assert pd.isna(status["lenta"]) + + retomada, llm_ok = _llm_com_ctrl_c() + _roda( + pd.DataFrame({"texto": textos}), + llm_ok, + parallel_requests=2, + batch_size=100, + checkpoint_path=ckpt, + ) + assert retomada == ["lenta"] + + +@requer_pthread_kill +@pytest.mark.usefixtures("sigint_reinicia_a_espera") +def test_ctrl_c_com_todas_as_linhas_lentas_sai_sem_esperar_uma_terminar(tmp_path): + """Os dois Ctrl-C chegam com as três linhas em andamento e nenhuma termina. + + O primeiro vira `KeyboardInterrupt` sem esperar uma linha terminar, e o + segundo corta a espera da saída: a execução sai antes da duração das linhas. + """ + duracao = 5.0 + textos = ["a", "b", "c"] + trava = threading.Lock() + comecadas: list[str] = [] + terminadas: list[str] = [] + todas_comecaram = threading.Event() + liberada = threading.Event() + + def llm(texto, *args, **kwargs): + with trava: + comecadas.append(texto) + if len(comecadas) == len(textos): + todas_comecaram.set() + liberada.wait(timeout=duracao) + with trava: + terminadas.append(texto) + return {"data": {"campo1": f"r-{texto}"}, "usage": {}} + + inicio = time.monotonic() + try: + # A folga do primeiro Ctrl-C deixa o thread principal entrar na espera + # das linhas: o que chega antes dela vira KeyboardInterrupt em qualquer + # espera, e o teste não distinguiria a espera sem prazo. + with _ctrl_c_depois(todas_comecaram, 0.1, 0.4), pytest.raises(KeyboardInterrupt): + _roda( + pd.DataFrame({"texto": textos}), + llm, + parallel_requests=3, + batch_size=100, + checkpoint_path=tmp_path / "ckpt.csv", + ) + saida = time.monotonic() - inicio + with trava: + terminadas_na_saida = list(terminadas) + finally: + liberada.set() + + assert saida < duracao + assert terminadas_na_saida == [] + + +@requer_pthread_kill +def test_ctrl_c_espera_a_linha_em_andamento_e_nao_deixa_thread_do_executor(tmp_path): + """As linhas lentas em andamento no Ctrl-C terminam na espera e entram no checkpoint. + + Elas só terminam depois de o executor ser encerrado, e nenhuma thread do + executor sobrevive à interrupção. + """ + textos = ["lenta1", "ctrl-c", "lenta2"] + ckpt = tmp_path / "ckpt.csv" + antes = set(threading.enumerate()) + trava = threading.Lock() + comecadas: list[str] = [] + todas_comecaram = threading.Event() + encerrando = threading.Event() + encerra = ThreadPoolExecutor.shutdown + + def encerra_observado(self, *args, **kwargs): + encerrando.set() + return encerra(self, *args, **kwargs) + + def llm(texto, *args, **kwargs): + with trava: + comecadas.append(texto) + if len(comecadas) == len(textos): + todas_comecaram.set() + if texto == "ctrl-c": + todas_comecaram.wait(timeout=10) + _ctrl_c() + else: + encerrando.wait(timeout=10) + time.sleep(0.2) + return {"data": {"campo1": f"r-{texto}"}, "usage": {}} + + with ( + patch.object(ThreadPoolExecutor, "shutdown", encerra_observado), + pytest.raises(KeyboardInterrupt), + ): + _roda( + pd.DataFrame({"texto": textos}), + llm, + parallel_requests=3, + batch_size=100, + checkpoint_path=ckpt, + ) + sobreviventes = [t for t in set(threading.enumerate()) - antes if not t.daemon] + + gravado = pd.read_csv(ckpt) + status = dict(zip(gravado["texto"], gravado["_dataframeit_status"], strict=True)) + assert status == dict.fromkeys(textos, "processed") + assert sobreviventes == [] + + retomada, llm_ok = _llm_com_ctrl_c() + _roda( + pd.DataFrame({"texto": textos}), + llm_ok, + parallel_requests=3, + batch_size=100, + checkpoint_path=ckpt, + ) + assert retomada == [] + + +@requer_pthread_kill +def test_ctrl_c_dentro_do_submit_espera_a_linha_que_a_thread_ja_comecou(tmp_path): + """O Ctrl-C chega dentro do `submit` da última linha, que a thread já começou. + + A linha roda e é paga, embora o laço não tenha guardado o futuro dela. A + interrupção espera as threads do executor, e a linha entra no checkpoint. + """ + textos = ["a", "b", "lenta"] + ckpt = tmp_path / "ckpt.csv" + pagas, llm_rapido = _llm_com_ctrl_c() + lenta_comecou = threading.Event() + + def llm(texto, *args, **kwargs): + if texto == "lenta": + lenta_comecou.set() + # Termina depois da gravação que esperasse só os futuros guardados. + time.sleep(0.5) + return llm_rapido(texto, *args, **kwargs) + + ajusta = ThreadPoolExecutor._adjust_thread_count + despachos: list[int] = [] + + def ajusta_e_interrompe(self): + # O submit põe o item na fila e chama este método, que inicia a thread + # quando falta uma; o KeyboardInterrupt do pthread_kill sai do submit + # antes de o laço guardar o futuro. O Ctrl-C espera a thread tirar a + # linha da fila: a que ainda estivesse nela seria cancelada. + ajusta(self) + despachos.append(1) + if len(despachos) == len(textos): + assert lenta_comecou.wait(timeout=10) + _ctrl_c() + + with ( + patch.object(ThreadPoolExecutor, "_adjust_thread_count", ajusta_e_interrompe), + pytest.raises(KeyboardInterrupt), + ): + _roda( + pd.DataFrame({"texto": textos}), + llm, + parallel_requests=3, + batch_size=100, + checkpoint_path=ckpt, + ) + + gravado = pd.read_csv(ckpt) + status = dict(zip(gravado["texto"], gravado["_dataframeit_status"], strict=True)) + assert status == dict.fromkeys(textos, "processed") + assert sorted(pagas) == sorted(textos) + + retomada, llm_ok = _llm_com_ctrl_c() + _roda( + pd.DataFrame({"texto": textos}), + llm_ok, + parallel_requests=3, + batch_size=100, + checkpoint_path=ckpt, + ) + assert retomada == [] + + +@requer_pthread_kill +@pytest.mark.usefixtures("sigint_interrompe_a_espera") +def test_ctrl_c_dentro_do_start_da_thread_deixa_a_linha_da_fila_pendente(tmp_path): + """O Ctrl-C sai do `Thread.start()` da primeira thread do executor, antes de ela começar. + + A thread já foi lançada, mas ainda não marcou o início, e a linha que ela + pegaria está na fila. O encerramento cancela a linha, e a thread, ao + começar, não a roda: a linha não é paga e a retomada a chama. Nenhuma + linha é paga fora do checkpoint. + """ + textos = ["a", "b", "c"] + ckpt = tmp_path / "ckpt.csv" + pagas, llm = _llm_com_ctrl_c() + inicia = threading.Thread._bootstrap_inner # ty: ignore[unresolved-attribute] + lancadas: list[threading.Thread] = [] + saiu = threading.Event() + + def inicia_depois_do_ctrl_c(self): + # Roda na thread nova, antes de ela marcar o início: o thread principal + # está no `_started.wait()` do `start()`, e o Ctrl-C o interrompe ali. + if not self.name.startswith("dataframeit-paralelo") or lancadas: + inicia(self) + return + lancadas.append(self) + _ctrl_c() + time.sleep(0.2) + try: + inicia(self) + finally: + saiu.set() + + with ( + patch.object(threading.Thread, "_bootstrap_inner", inicia_depois_do_ctrl_c), + pytest.raises(KeyboardInterrupt), + ): + _roda( + pd.DataFrame({"texto": textos}), + llm, + parallel_requests=3, + batch_size=100, + checkpoint_path=ckpt, + ) + assert len(lancadas) == 1 + # A thread ainda não começou, e `join` a recusaria. + assert saiu.wait(timeout=10) + + # Sem linha registrada, a interrupção não grava o arquivo. + processadas = set() + if ckpt.exists(): + gravado = pd.read_csv(ckpt) + processadas = set(gravado.loc[gravado["_dataframeit_status"].eq("processed"), "texto"]) + assert set(pagas) <= processadas + + retomada, llm_ok = _llm_com_ctrl_c() + _roda( + pd.DataFrame({"texto": textos}), + llm_ok, + parallel_requests=3, + batch_size=100, + checkpoint_path=ckpt, + ) + assert sorted(retomada) == sorted(set(textos) - set(pagas)) + + +@pytest.mark.parametrize( + "numeros", + [None, (10, 1)], + ids=["numeracao_do_modulo", "executor_1_ao_lado_do_executor_10"], +) +def test_execucao_nao_espera_as_threads_de_outra_chamada_concorrente(numeros): + """Duas chamadas ao mesmo tempo: a curta termina com a longa em andamento. + + A saída de cada execução espera só as threads do próprio executor. A longa + roda num thread do usuário, com as linhas presas até a curta voltar. Com os + números fixados, a longa abre o executor 10 e a curta o 1, e o nome + `dataframeit-paralelo10_0` começa com `dataframeit-paralelo1`. + """ + prazo = 10.0 + trava = threading.Lock() + longas = [f"longa{i}" for i in range(3)] + comecadas: list[str] = [] + terminadas: list[str] = [] + longas_comecaram = threading.Event() + liberada = threading.Event() + + def llm(texto, *args, **kwargs): + if texto in longas: + with trava: + comecadas.append(texto) + if len(comecadas) == len(longas): + longas_comecaram.set() + liberada.wait(timeout=prazo) + with trava: + terminadas.append(texto) + return {"data": {"campo1": f"r-{texto}"}, "usage": {}} + + saidas: list[pd.DataFrame] = [] + + def roda_longa() -> None: + saidas.append( + dataframeit( + pd.DataFrame({"texto": longas}), + questions=Modelo, + prompt="{texto}", + parallel_requests=3, + ) + ) + + numeracao = patch("dataframeit.core._executor_ids", iter(numeros)) if numeros else nullcontext() + longa = threading.Thread(target=roda_longa) + with ( + patch("dataframeit.core.call_langchain", side_effect=llm), + patch("dataframeit.core.validate_provider_dependencies"), + numeracao, + ): + try: + longa.start() + assert longas_comecaram.wait(timeout=prazo) + inicio = time.monotonic() + curta = dataframeit( + pd.DataFrame({"texto": ["a", "b"]}), + questions=Modelo, + prompt="{texto}", + parallel_requests=2, + ) + duracao = time.monotonic() - inicio + with trava: + terminadas_na_volta = list(terminadas) + finally: + liberada.set() + longa.join(timeout=prazo) + + assert duracao < 2.0 + assert terminadas_na_volta == [] + assert list(curta["campo1"]) == ["r-a", "r-b"] + assert [list(saida["campo1"]) for saida in saidas] == [[f"r-{t}" for t in longas]] + + +@pytest.mark.parametrize("paralelo", [1, 3]) +def test_ctrl_c_antes_de_terminar_uma_linha_nao_reescreve_o_checkpoint(tmp_path, paralelo): + """Sem linha registrada desde a última gravação, o arquivo fica como estava.""" + ckpt = tmp_path / "ckpt.csv" + df = pd.DataFrame({"texto": ["a", "b", "c"]}) + _roda( + df.copy(), + _llm_com_ctrl_c()[1], + model_kwargs={"temperature": 0.5}, + batch_size=10, + checkpoint_path=ckpt, + ) + anterior = ckpt.read_bytes() + + def interrompe(*args, **kwargs): + raise KeyboardInterrupt + + with pytest.warns(UserWarning, match="outra configuração"), pytest.raises(KeyboardInterrupt): + _roda( + df.copy(), + interrompe, + parallel_requests=paralelo, + batch_size=10, + checkpoint_path=ckpt, + ) + + assert ckpt.read_bytes() == anterior + + +@requer_pthread_kill +@pytest.mark.parametrize("paralelo", [1, 3]) +def test_ctrl_c_sem_checkpoint_deixa_as_linhas_pagas_na_entrada_pandas(paralelo): + df = pd.DataFrame({"texto": [f"t{i}" for i in range(9)]}) + pagas, llm = _llm_com_ctrl_c(na_chamada=5) + + with pytest.raises(KeyboardInterrupt): + _roda(df, llm, parallel_requests=paralelo) + + assert sorted(df.loc[df["_dataframeit_status"].eq("processed"), "texto"]) == sorted(pagas) + + +@requer_pthread_kill +def test_ctrl_c_espera_o_on_row_done_em_andamento_e_grava_a_linha(tmp_path): + """O Ctrl-C chega com o `on_row_done` de uma linha ainda rodando na thread do worker. + + O callback não recebe o KeyboardInterrupt, e a linha já está contada quando + ele começa: a interrupção espera a thread dele, o callback termina antes de + a interrupção sair da chamada, e a linha entra no checkpoint. A linha que + não foi despachada fica pendente. + """ + textos = ["callback", "ctrl-c", "fila"] + ckpt = tmp_path / "ckpt.csv" + callback_comecou = threading.Event() + encerrando = threading.Event() + trava = threading.Lock() + entregues: list[str] = [] + encerra = ThreadPoolExecutor.shutdown + + def encerra_observado(self, *args, **kwargs): + encerrando.set() + return encerra(self, *args, **kwargs) + + def llm(texto, *args, **kwargs): + if texto == "ctrl-c": + assert callback_comecou.wait(timeout=10) + _ctrl_c() + # Termina só depois do encerramento, como a linha em andamento. + assert encerrando.wait(timeout=10) + time.sleep(0.2) + return {"data": {"campo1": f"r-{texto}"}, "usage": {}} + + def ao_terminar(resultado): + if resultado.index == 0: + callback_comecou.set() + # Termina depois de a interrupção começar a esperar as threads. + encerrando.wait(timeout=10) + time.sleep(0.2) + with trava: + entregues.append(textos[resultado.index]) + + with ( + patch.object(ThreadPoolExecutor, "shutdown", encerra_observado), + pytest.raises(KeyboardInterrupt), + ): + _roda( + pd.DataFrame({"texto": textos}), + llm, + parallel_requests=2, + batch_size=100, + checkpoint_path=ckpt, + on_row_done=ao_terminar, + ) + with trava: + entregues_na_saida = sorted(entregues) + + assert entregues_na_saida == ["callback", "ctrl-c"] + gravado = pd.read_csv(ckpt) + status = dict(zip(gravado["texto"], gravado["_dataframeit_status"], strict=True)) + assert status["callback"] == "processed" + assert status["ctrl-c"] == "processed" + assert pd.isna(status["fila"]) + + retomada, llm_ok = _llm_com_ctrl_c() + _roda( + pd.DataFrame({"texto": textos}), + llm_ok, + parallel_requests=2, + batch_size=100, + checkpoint_path=ckpt, + ) + assert retomada == ["fila"] + + +def test_ctrl_c_dentro_do_on_row_done_no_sequencial_grava_a_linha_entregue(tmp_path): + """No sequencial, o callback roda no thread principal e pode receber o Ctrl-C. + + A linha entregue ao callback já está gravada no DataFrame e contada, e a + gravação do Ctrl-C a leva, também quando é a primeira linha da execução. A + retomada não a chama de novo. + """ + textos = ["a", "b", "c"] + ckpt = tmp_path / "ckpt.csv" + pagas, llm = _llm_com_ctrl_c() + + def interrompe(resultado): + raise KeyboardInterrupt + + with pytest.raises(KeyboardInterrupt): + _roda( + pd.DataFrame({"texto": textos}), + llm, + batch_size=100, + checkpoint_path=ckpt, + on_row_done=interrompe, + ) + + assert pagas == ["a"] + gravado = pd.read_csv(ckpt) + assert gravado["_dataframeit_status"].tolist()[0] == "processed" + assert gravado["_dataframeit_status"].isna().tolist()[1:] == [True, True] + + retomada, llm_ok = _llm_com_ctrl_c() + _roda(pd.DataFrame({"texto": textos}), llm_ok, batch_size=100, checkpoint_path=ckpt) + assert retomada == ["b", "c"] diff --git a/tests/test_retomada_assinada.py b/tests/test_retomada_assinada.py index 01bbe546..495cf82d 100644 --- a/tests/test_retomada_assinada.py +++ b/tests/test_retomada_assinada.py @@ -73,21 +73,30 @@ def test_execucao_recusada_e_interrompida_nao_adota_o_checkpoint_antigo(tmp_path df = pd.DataFrame({"texto": ["a", "b", "c"]}) _roda(df, _llm(100, SystemExit())[1], prompt="P1 {texto}", batch_size=10, checkpoint_path=ckpt) - with ( - pytest.warns(UserWarning, match="outra configuração"), - contextlib.suppress(KeyboardInterrupt), - ): - _roda(df, _llm(0, interrupcao)[1], prompt="P2 {texto}", batch_size=10, checkpoint_path=ckpt) + # Os avisos se conferem pelo registro: um pytest.warns dentro do outro engole + # os avisos do de fora no pytest 7, e o pytest 8 reemite o que não casa. + with warnings.catch_warnings(record=True) as avisos, contextlib.suppress(KeyboardInterrupt): + warnings.simplefilter("always") + _roda( + df, + _llm(1, interrupcao, prefixo="B")[1], + prompt="P2 {texto}", + batch_size=10, + checkpoint_path=ckpt, + ) + mensagens = [str(aviso.message) for aviso in avisos] + assert any("outra configuração" in mensagem for mensagem in mensagens) + assert any("Execução interrompida" in mensagem for mensagem in mensagens) == isinstance( + interrupcao, ProviderUsageLimitError + ) - # Com o limite de uso, B grava o próprio checkpoint, com as linhas pendentes, - # e C retoma dele; com Ctrl-C, o arquivo ainda é o de A, e C recomeça. - chamadas, llm = _llm(100, SystemExit(), prefixo="P2-") - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - saida = _roda(df, llm, prompt="P2 {texto}", batch_size=10, checkpoint_path=ckpt) + # Nas duas interrupções, B grava o próprio checkpoint, com a linha que pagou e + # as pendentes, e C retoma dele, sem nenhuma resposta de A. + chamadas, llm = _llm(100, SystemExit(), prefixo="C") + saida = _roda(df, llm, prompt="P2 {texto}", batch_size=10, checkpoint_path=ckpt) - assert len(chamadas) == 3 - assert saida["campo1"].tolist() == ["P2-1", "P2-2", "P2-3"] + assert len(chamadas) == 2 + assert saida["campo1"].tolist() == ["B1", "C1", "C2"] @pytest.mark.parametrize("outra", [["x", "y", "z"], ["x", "y"]])