src/pipeline/gpu_runtime.py

"""GPU lease runtime — lease-decorated SDK stage calls, duration estimators,
per-request state reset, and the startup AOTI compilation probe."""
import time

import torch

from config import get_vad_duration, get_asr_duration, ZEROGPU_MAX_DURATION
from qua_sdk.components.recognition.spec import RecognitionParams
from qua_sdk.components.segmentation.runtimes.recitation_params import RecitationSegmenterParams
from qua_sdk.observe import collect_stage_metrics
from qua_sdk.registry import resolve
from qua_sdk.schemas import Audio, Region, Regions

from src.core.deploy_select import gpu_deploy, select_deploy
from src.core.request_stats import get_request_stats, reset_request_stats
from src.core.zero_gpu import gpu_with_fallback


def _reset_request_state():
    """Drop stale per-request stats and start a fresh DebugCollector.

    Called at every pipeline entry so v3 log rows always carry populated
    `events` / `anchor` / per-segment `dp_debug` — not only `/debug_process`
    runs. A previous collector on the thread is replaced (stale state dropped).
    """
    reset_request_stats()
    try:
        from src.core.debug_collector import start_debug_collection
        start_debug_collection()
    except Exception:
        pass


SEGMENTER_KEY = "segmentation.recitation_v2@v1"
RECOGNIZER_KEY = "recognition.w2v2_ctc@v1"

# Stale-lease cleanup: drop the SDK runtimes' cached models inside the next
# lease. Registered at import so every process that can hold a lease has it.
from src.core.zero_gpu import register_stale_invalidator as _register_stale_invalidator
_register_stale_invalidator(lambda: resolve(SEGMENTER_KEY).invalidate())
_register_stale_invalidator(lambda: resolve(RECOGNIZER_KEY).invalidate())


def _segmentation_params(min_silence_ms, min_speech_ms, pad_ms) -> RecitationSegmenterParams:
    """UI slider values → the Space profile's segmentation knobs (single pad)."""
    return RecitationSegmenterParams(
        cleaning="segmenter_native",
        min_silence_ms=int(min_silence_ms), min_speech_ms=int(min_speech_ms),
        pad_left_ms=int(pad_ms), pad_right_ms=int(pad_ms),
    )

_gpu_info_logged = False
_gpu_info_cache = {}


def _log_gpu_info():
    """Print GPU device info once per lease and cache for logging."""
    global _gpu_info_logged
    if _gpu_info_logged or not torch.cuda.is_available():
        return
    _gpu_info_logged = True
    props = torch.cuda.get_device_properties(0)
    _gpu_info_cache["name"] = props.name
    _gpu_info_cache["total_vram_gb"] = round(props.total_memory / (1024**3), 1)
    _gpu_info_cache["sms"] = props.multi_processor_count
    _gpu_info_cache["compute"] = f"{props.major}.{props.minor}"
    print(f"[GPU LEASE] {props.name} | "
          f"VRAM: {_gpu_info_cache['total_vram_gb']:.1f} GB | "
          f"SMs: {props.multi_processor_count} | "
          f"Compute: {props.major}.{props.minor}")


def _capture_vram_safely():
    """Read CUDA peak VRAM stats — returns (0.0, 0.0) when not on GPU.

    Defensive against the CPU-subprocess path where `torch.cuda.is_available()`
    can deceptively report True (because spaces' patches or stray
    CUDA_VISIBLE_DEVICES handling let the subprocess see the parent's GPU).
    Calling `max_memory_allocated()` in that situation can hang because the
    subprocess has no actual GPU lease — the C-level CUDA query waits forever
    on a context that will never be granted.
    """
    from src.core.zero_gpu import is_user_forced_cpu
    if is_user_forced_cpu() or not torch.cuda.is_available():
        return 0.0, 0.0
    try:
        peak_vram = torch.cuda.max_memory_allocated() / (1024 * 1024)
        reserved_vram = torch.cuda.max_memory_reserved() / (1024 * 1024)
        torch.cuda.reset_peak_memory_stats()
        return peak_vram, reserved_vram
    except RuntimeError:
        return 0.0, 0.0


def _combined_duration(audio, sample_rate, *_args, **_kwargs):
    """Lease duration for VAD+ASR: sum of independent estimates, capped at ZeroGPU max."""
    minutes = len(audio) / sample_rate / 60
    model_name = _args[3] if len(_args) > 3 else _kwargs.get("model_name", "Base")
    uncapped = get_vad_duration(minutes) + get_asr_duration(minutes, model_name)
    capped = min(uncapped, ZEROGPU_MAX_DURATION)
    get_request_stats().lease = {
        "lease_type": "combined",
        "requested_s": round(capped, 3),
        "uncapped_s": round(uncapped, 3),
        "cap_hit": uncapped > ZEROGPU_MAX_DURATION,
        "cap_s": ZEROGPU_MAX_DURATION,
    }
    return capped


def _asr_only_duration(audio, sample_rate, intervals, *_args, **_kwargs):
    """Lease duration for standalone ASR, capped at ZeroGPU max."""
    minutes = sum(e - s for s, e in intervals) / 60
    model_name = _args[0] if _args else _kwargs.get("model_name", "Base")
    uncapped = get_asr_duration(minutes, model_name)
    capped = min(uncapped, ZEROGPU_MAX_DURATION)
    get_request_stats().lease = {
        "lease_type": "asr_only",
        "requested_s": round(capped, 3),
        "uncapped_s": round(uncapped, 3),
        "cap_hit": uncapped > ZEROGPU_MAX_DURATION,
        "cap_s": ZEROGPU_MAX_DURATION,
    }
    return capped


@gpu_with_fallback(duration=_combined_duration)
def run_vad_and_asr_gpu(audio, sample_rate, min_silence_ms, min_speech_ms, pad_ms, model_name="Base"):
    """Single GPU lease: SDK segmentation + recognition.

    Returns (regions, emissions, stage_metrics, vad_gpu_time, asr_gpu_time,
    peak_vram, reserved_vram) — all picklable for the CPU dispatch paths.
    """
    _log_gpu_info()
    deploy = select_deploy()
    audio_obj = Audio.from_array(audio, sample_rate)
    t_lease_start = time.time()

    segmenter = resolve(SEGMENTER_KEY)
    with collect_stage_metrics() as seg_metrics:
        regions = segmenter.segment(
            audio_obj, _segmentation_params(min_silence_ms, min_speech_ms, pad_ms),
            deploy=deploy,
        )
    vad_gpu_time = time.time() - t_lease_start

    if len(regions) == 0:
        return regions, None, {"segmentation": seg_metrics}, vad_gpu_time, 0.0, 0.0, 0.0

    recognizer = resolve(RECOGNIZER_KEY)
    t_asr_start = time.time()
    with collect_stage_metrics() as rec_metrics:
        emissions = recognizer.transcribe(
            audio_obj, regions, RecognitionParams(model=model_name), deploy=deploy,
        )
    asr_gpu_time = time.time() - t_asr_start

    peak_vram, reserved_vram = _capture_vram_safely()

    return (regions, emissions, {"segmentation": seg_metrics, "recognition": rec_metrics},
            vad_gpu_time, asr_gpu_time, peak_vram, reserved_vram)


@gpu_with_fallback(duration=_asr_only_duration)
def run_phoneme_asr_gpu(audio, sample_rate, intervals, model_name="Base"):
    """Standalone recognition lease (resegment/retranscribe/realign paths).

    Returns (emissions, rec_metrics, asr_gpu_time, peak_vram, reserved_vram).
    """
    _log_gpu_info()
    deploy = select_deploy()
    audio_obj = Audio.from_array(audio, sample_rate)
    regions = Regions(
        regions=[Region(start_s=float(s), end_s=float(e)) for s, e in intervals],
        audio_duration_s=len(audio) / sample_rate,
    )

    t_asr_start = time.time()
    recognizer = resolve(RECOGNIZER_KEY)
    with collect_stage_metrics() as rec_metrics:
        emissions = recognizer.transcribe(
            audio_obj, regions, RecognitionParams(model=model_name), deploy=deploy,
        )
    asr_gpu_time = time.time() - t_asr_start

    peak_vram, reserved_vram = _capture_vram_safely()

    return emissions, rec_metrics, asr_gpu_time, peak_vram, reserved_vram


@gpu_with_fallback(duration=lambda: 300)  # 5 min lease for compilation test
def test_aoti_compilation_gpu():
    """AOTI export/compile (or Hub load) for the segmenter — startup, inside a lease."""
    from qua_sdk.components.segmentation.runtimes import aoti as sdk_aoti

    deploy = gpu_deploy()
    runtime = resolve(SEGMENTER_KEY)
    runtime.preload()
    runtime._ensure_on_device(deploy)
    return sdk_aoti.export_and_compile(runtime, deploy)