alnn_rewrite/util/control.py

"""Shared pipeline control helpers: cooperative stop/pause and memory release.

These were previously copy-pasted across the training tasks and training loop.
Centralizing them keeps the stop/pause semantics identical everywhere.
"""

import gc
import time
from threading import Event

import torch


def check_control_events(
    stop_event: Event | None,
    pause_event: Event | None,
) -> None:
    """Cooperatively honor stop/pause events.

    Raises:
        InterruptedError: if ``stop_event`` is set (including while paused).
    """
    if stop_event is not None and stop_event.is_set():
        raise InterruptedError("Pipeline execution stopped by user.")
    while pause_event is not None and pause_event.is_set():
        time.sleep(0.5)
        if stop_event is not None and stop_event.is_set():
            raise InterruptedError("Pipeline execution stopped by user while paused.")


def release_torch_memory(device: str) -> None:
    """Free Python and accelerator caches between models to limit VRAM growth."""
    gc.collect()
    if device.startswith("cuda"):
        torch.cuda.empty_cache()
    elif device.startswith("mps") and hasattr(torch, "mps"):
        torch.mps.empty_cache()