diff --git a/apps/desktop/sidecar/device_policy.py b/apps/desktop/sidecar/device_policy.py new file mode 100644 index 0000000..988d98c --- /dev/null +++ b/apps/desktop/sidecar/device_policy.py @@ -0,0 +1,110 @@ +""" +STT 추론 장치 선택 정책 (순수 로직, 외부 IO 없음). + +main.py(FastAPI 어댑터)가 ctranslate2/faster-whisper 호출을 주입하고, 이 모듈은 +"어떤 장치·정밀도로 올릴지"와 "언제 CPU로 내려갈지"만 결정한다. + +배경: ctranslate2 는 CUDA 런타임을 정적 링크하므로 NVIDIA 드라이버만 있어도 +get_cuda_device_count() 가 1 이상을 돌려준다. 그런데 cuBLAS(cublas64_12.dll)가 +없으면 모델 생성은 성공하고 첫 전사에서야 RuntimeError 가 난다. Pascal(CC 6.x) +GPU 는 float16 을 지원하지 않아 모델 생성부터 ValueError 가 난다. 그래서 + 1) 지원 compute type 중 쓸 수 있는 것을 고르고, + 2) GPU 로 만든 직후 짧은 무음 전사로 실제 동작을 검증하고, + 3) 실패하면 CPU(int8)로 다시 만든다. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable, Iterable, Optional, TypeVar + +M = TypeVar("M") + + +@dataclass(frozen=True) +class DeviceChoice: + """모델을 올릴 장치와 연산 정밀도.""" + + device: str + compute_type: str + + @property + def is_gpu(self) -> bool: + return self.device == "cuda" + + +CPU_CHOICE = DeviceChoice(device="cpu", compute_type="int8") + +# 선호 순서. float16 을 못 쓰는 GPU(Pascal 등)는 int8_float32 → float32 로 내려간다. +_CUDA_COMPUTE_PREFERENCE: tuple[str, ...] = ("float16", "int8_float32", "float32") + +# GPU 런타임/라이브러리 문제로 판단할 오류 메시지 조각 (소문자 비교). +_GPU_ERROR_MARKERS: tuple[str, ...] = ( + "cublas", + "cudnn", + "cuda", + "cufft", + "nvrtc", + "float16 compute", + "efficient float16", +) + + +def choose_cuda_compute_type(supported: Iterable[str]) -> Optional[str]: + """GPU 가 지원하는 compute type 중 선호 순서대로 첫 번째를 고른다. 없으면 None.""" + supported_set = {s.lower() for s in supported} + for candidate in _CUDA_COMPUTE_PREFERENCE: + if candidate in supported_set: + return candidate + return None + + +def choose_gpu_device(cuda_count: int, supported: Iterable[str]) -> Optional[DeviceChoice]: + """CUDA 장치 수와 지원 compute type 으로 GPU 후보를 정한다. 쓸 수 없으면 None.""" + if cuda_count <= 0: + return None + compute_type = choose_cuda_compute_type(supported) + if compute_type is None: + return None + return DeviceChoice(device="cuda", compute_type=compute_type) + + +def is_gpu_runtime_error(exc: BaseException) -> bool: + """CUDA 라이브러리 로드 실패·미지원 정밀도 등 CPU 로 내려가면 풀리는 오류인가.""" + if not isinstance(exc, (RuntimeError, ValueError, OSError)): + return False + message = str(exc).lower() + return any(marker in message for marker in _GPU_ERROR_MARKERS) + + +def load_with_fallback( + build: Callable[[DeviceChoice], M], + verify: Callable[[M], None], + preferred: Optional[DeviceChoice], + on_fallback: Optional[Callable[[DeviceChoice, BaseException], None]] = None, +) -> tuple[M, DeviceChoice]: + """preferred(GPU)로 만들어 검증하고, 실패하면 CPU 로 다시 만든다. + + - GPU 생성 실패: GPU 관련 오류일 때만 CPU 로 내려간다. 그 밖의 오류(모델 파일 + 다운로드 실패 등)는 CPU 로 다시 해도 똑같이 실패하므로 그대로 올린다. + - GPU 검증 실패: 모델은 만들어졌으므로 장치 문제다. 오류 종류와 무관하게 CPU 로 간다. + - CPU 후보이거나 후보가 없으면 검증 없이 CPU 로 만든다. + """ + if preferred is not None and preferred.is_gpu: + try: + model = build(preferred) + except Exception as exc: + if not is_gpu_runtime_error(exc): + raise + if on_fallback is not None: + on_fallback(preferred, exc) + else: + try: + verify(model) + return model, preferred + except Exception as exc: + # 검증에 실패한 GPU 모델의 참조를 먼저 끊어 VRAM 을 돌려받는다. + del model + if on_fallback is not None: + on_fallback(preferred, exc) + return build(CPU_CHOICE), CPU_CHOICE diff --git a/apps/desktop/sidecar/main.py b/apps/desktop/sidecar/main.py index 5d3d5fc..4241fa8 100644 --- a/apps/desktop/sidecar/main.py +++ b/apps/desktop/sidecar/main.py @@ -39,6 +39,14 @@ import uvicorn from fastapi import FastAPI, File, Form, UploadFile from fastapi.responses import JSONResponse +from device_policy import ( + CPU_CHOICE, + DeviceChoice, + choose_gpu_device, + is_gpu_runtime_error, + load_with_fallback, +) + # ── 로깅 설정 ────────────────────────────────────────────── # Windows에서 파이프로 연결되면 Python이 로케일(cp949) 인코딩으로 출력해 @@ -64,6 +72,10 @@ _model_id: str | None = None # 보조 모델 (실시간 자막 등 받아쓰기와 다른 모델). 최대 1개. _aux_models: "dict[str, WhisperModel]" = {} _gpu_available: bool = False +# GPU 로 올릴 때 쓸 장치/정밀도. GPU 가 없거나 검증에 실패하면 None. +_gpu_choice: DeviceChoice | None = None +# 로딩된 모델별 실제 장치 (primary/aux 는 같은 model_id 를 동시에 들지 않는다). +_model_devices: dict[str, DeviceChoice] = {} _server: uvicorn.Server | None = None _models_dir: Path | None = None @@ -118,21 +130,44 @@ def _detect_gpu() -> None: torch 의존 제거를 위해 ctranslate2의 네이티브 CUDA 감지를 사용한다. ctranslate2는 faster-whisper의 백엔드이므로 항상 함께 설치된다. """ - global _gpu_available + global _gpu_available, _gpu_choice try: import ctranslate2 cuda_count = ctranslate2.get_cuda_device_count() - _gpu_available = cuda_count > 0 - if _gpu_available: - logger.info("GPU 감지: CUDA 디바이스 %d개", cuda_count) + supported: set[str] = set() + if cuda_count > 0: + try: + supported = set(ctranslate2.get_supported_compute_types("cuda")) + except Exception as exc: + logger.info("CUDA compute type 조회 실패: %s", exc) + _gpu_choice = choose_gpu_device(cuda_count, supported) + _gpu_available = _gpu_choice is not None + if _gpu_choice is not None: + logger.info( + "GPU 감지: CUDA 디바이스 %d개 (compute=%s)", + cuda_count, + _gpu_choice.compute_type, + ) + elif cuda_count > 0: + logger.info("CUDA 디바이스는 있으나 지원 compute type 없음, CPU 모드로 동작") else: logger.info("GPU 미감지, CPU 모드로 동작") except Exception as exc: _gpu_available = False + _gpu_choice = None logger.info("GPU 감지 실패, CPU 모드로 동작: %s", exc) +def _disable_gpu(reason: BaseException) -> None: + """GPU 경로가 실제로 동작하지 않음이 확인되면 이후 모델은 모두 CPU 로 올린다.""" + global _gpu_available, _gpu_choice + if _gpu_available: + logger.warning("GPU 사용 불가 → CPU(int8) 모드로 전환: %s", reason) + _gpu_available = False + _gpu_choice = None + + def _cpu_threads() -> int: """CPU 추론에 사용할 스레드 수 (과도한 점유 방지 위해 8로 상한).""" return max(1, min(8, os.cpu_count() or 4)) @@ -313,6 +348,46 @@ def _build_transcribe_kwargs( return kwargs +def _run_transcription( + model: "WhisperModel", + audio_array: np.ndarray, + transcribe_kwargs: dict, +) -> tuple[list[dict], str, str]: + """모델로 전사하고 세그먼트를 끝까지 소비한다. + + faster-whisper 는 세그먼트를 지연 생성하므로 디코더 오류(GPU 런타임 포함)는 + 반복 중에 난다. 호출부가 전체를 한 단위로 재시도할 수 있도록 여기서 모두 소비한다. + Returns: (segments, full_text, detected_language) + """ + # VAD가 전체 오디오를 제거하면 max() 에러 발생 → VAD 없이 재시도 + try: + segments_iter, info = model.transcribe(audio_array, **transcribe_kwargs) + except ValueError as ve: + if "empty sequence" in str(ve) and transcribe_kwargs.get("vad_filter"): + logger.warning("VAD가 전체 오디오를 제거함 → VAD 없이 재시도") + transcribe_kwargs["vad_filter"] = False + segments_iter, info = model.transcribe(audio_array, **transcribe_kwargs) + else: + raise + + segments_list: list[dict] = [] + full_text_parts: list[str] = [] + + for segment in segments_iter: + seg_dict = { + "text": segment.text.strip(), + "start": round(segment.start, 3), + "end": round(segment.end, 3), + "avg_logprob": round(segment.avg_logprob, 4), + } + segments_list.append(seg_dict) + full_text_parts.append(segment.text.strip()) + + full_text = " ".join(full_text_parts).strip() + detected_language = info.language if info.language else "unknown" + return segments_list, full_text, detected_language + + # ── 엔드포인트 ───────────────────────────────────────────── @@ -331,26 +406,86 @@ async def health() -> JSONResponse: ) -def _create_model(model_id: str) -> "WhisperModel": +def _create_model(model_id: str, choice: DeviceChoice) -> "WhisperModel": """모델을 올린다. /download 로 받아 둔 로컬 디렉토리가 있으면 그것을 쓴다.""" from faster_whisper import WhisperModel - device = "cuda" if _gpu_available else "cpu" - compute_type = "float16" if _gpu_available else "int8" local_dir = _local_model_dir(model_id) model_source = str(local_dir) if local_dir else model_id if local_dir: logger.info("로컬 모델 디렉토리 사용: %s", local_dir) - logger.info("모델 생성: %s (device=%s, compute=%s)", model_id, device, compute_type) + logger.info( + "모델 생성: %s (device=%s, compute=%s)", model_id, choice.device, choice.compute_type + ) return WhisperModel( model_source, - device=device, - compute_type=compute_type, + device=choice.device, + compute_type=choice.compute_type, cpu_threads=_cpu_threads(), num_workers=1, ) +# 0.5초 @16kHz 무음. GPU 모델이 cuBLAS 등을 실제로 로드할 수 있는지 확인하는 데 쓴다. +_PROBE_SAMPLES = 8000 + + +def _probe_model(model: "WhisperModel") -> None: + """짧은 무음 전사로 인코더·디코더를 한 번 돌려 본다 (실패 시 예외).""" + silence = np.zeros(_PROBE_SAMPLES, dtype=np.float32) + segments, _info = model.transcribe( + silence, + language="en", + beam_size=1, + temperature=0.0, + vad_filter=False, + condition_on_previous_text=False, + without_timestamps=True, + ) + # faster-whisper 는 세그먼트를 지연 생성하므로 끝까지 소비해야 디코더가 돈다. + for _ in segments: + pass + + +def _load_on_best_device(model_id: str) -> "WhisperModel": + """GPU 후보가 있으면 GPU 로 올려 검증하고, 안 되면 CPU(int8)로 올린다.""" + + def on_fallback(failed: DeviceChoice, exc: BaseException) -> None: + logger.warning( + "GPU 모델 검증 실패 (%s, compute=%s) → CPU 로 재로딩: %s", + model_id, + failed.compute_type, + exc, + ) + _disable_gpu(exc) + + preferred = _gpu_choice if _gpu_available else None + model, choice = load_with_fallback( + build=lambda c: _create_model(model_id, c), + verify=_probe_model, + preferred=preferred, + on_fallback=on_fallback, + ) + _model_devices[model_id] = choice + return model + + +def _reload_on_cpu(model_id: str, is_primary: bool) -> "WhisperModel": + """전사 중 GPU 런타임 오류가 난 모델을 CPU 로 다시 올려 같은 자리에 둔다.""" + global _model + if is_primary: + _model = None + else: + _aux_models.pop(model_id, None) + model = _create_model(model_id, CPU_CHOICE) + _model_devices[model_id] = CPU_CHOICE + if is_primary: + _model = model + else: + _aux_models[model_id] = model + return model + + @app.post("/load") async def load_model(body: dict) -> JSONResponse: # noqa: ANN001 """Whisper 모델을 로딩한다. @@ -386,11 +521,11 @@ async def load_model(body: dict) -> JSONResponse: # noqa: ANN001 if slot == "aux": # 보조 자리는 하나만 둔다 — 다른 보조 모델은 내려 VRAM 을 돌려받는다. _aux_models.clear() - _aux_models[model_id] = _create_model(model_id) + _aux_models[model_id] = _load_on_best_device(model_id) else: # 모델 교체 시 이전 모델을 먼저 해제해 VRAM/RAM을 회수한다. _model = None - _model = _create_model(model_id) + _model = _load_on_best_device(model_id) _model_id = model_id # 기본 모델이 된 모델은 보조 자리에 중복으로 들고 있지 않는다. _aux_models.pop(model_id, None) @@ -433,7 +568,8 @@ async def transcribe( partial - 녹음 중 미리보기 모드 ('true'면 greedy 디코딩 + 컨텍스트 미사용) model_id - 쓸 모델 (비우면 기본 모델). 올라가 있지 않으면 409 """ - if model_id and model_id != _model_id: + is_primary = not (model_id and model_id != _model_id) + if not is_primary: model = _aux_models.get(model_id) if model is None: return JSONResponse( @@ -442,6 +578,7 @@ async def transcribe( ) else: model = _model + resolved_model_id = model_id if not is_primary else _model_id if model is None: return JSONResponse( status_code=503, @@ -482,35 +619,34 @@ async def transcribe( is_partial=is_partial, ) - # VAD가 전체 오디오를 제거하면 max() 에러 발생 → VAD 없이 재시도 try: - segments_iter, info = model.transcribe(audio_array, **transcribe_kwargs) - except ValueError as ve: - if "empty sequence" in str(ve) and transcribe_kwargs.get("vad_filter"): - logger.warning("VAD가 전체 오디오를 제거함 → VAD 없이 재시도") - transcribe_kwargs["vad_filter"] = False - segments_iter, info = model.transcribe(audio_array, **transcribe_kwargs) - else: + segments_list, full_text, detected_language = _run_transcription( + model, audio_array, dict(transcribe_kwargs) + ) + except Exception as exc: + # cuBLAS 미설치 PC 등: 검증을 통과했더라도 실제 전사에서 GPU 런타임 오류가 + # 나면 이 모델을 CPU 로 다시 올려 한 번만 재시도한다. + device = _model_devices.get(resolved_model_id or "") + if not ( + resolved_model_id + and device is not None + and device.is_gpu + and is_gpu_runtime_error(exc) + ): raise + logger.warning( + "전사 중 GPU 런타임 오류 → CPU 로 재로딩 후 재시도 (%s): %s", + resolved_model_id, + exc, + ) + _disable_gpu(exc) + model = _reload_on_cpu(resolved_model_id, is_primary) + segments_list, full_text, detected_language = _run_transcription( + model, audio_array, dict(transcribe_kwargs) + ) - segments_list: list[dict] = [] - full_text_parts: list[str] = [] - - for segment in segments_iter: - seg_dict = { - "text": segment.text.strip(), - "start": round(segment.start, 3), - "end": round(segment.end, 3), - "avg_logprob": round(segment.avg_logprob, 4), - } - segments_list.append(seg_dict) - full_text_parts.append(segment.text.strip()) - - full_text = " ".join(full_text_parts).strip() processing_time = int((time.monotonic() - start_time) * 1000) - detected_language = info.language if info.language else "unknown" - logger.info( "전사 완료: '%s' (lang=%s, %.1f초, %dms)", full_text[:80], diff --git a/apps/desktop/sidecar/tests/test_device_fallback.py b/apps/desktop/sidecar/tests/test_device_fallback.py new file mode 100644 index 0000000..342c3d6 --- /dev/null +++ b/apps/desktop/sidecar/tests/test_device_fallback.py @@ -0,0 +1,322 @@ +""" +GPU → CPU 폴백 회귀 테스트. + +cuBLAS 가 없는 NVIDIA PC(드라이버만 설치)와 float16 을 지원하지 않는 Pascal GPU 에서 +로컬 전사가 모두 HTTP 500 으로 끝나던 문제를 막는다. + +실행 (apps/desktop/sidecar 에서): + .venv/Scripts/python.exe -m unittest discover -s tests -v +""" + +from __future__ import annotations + +import sys +import types +import unittest +from pathlib import Path +from typing import Iterator +from unittest import mock + +SIDECAR_DIR = Path(__file__).resolve().parent.parent +if str(SIDECAR_DIR) not in sys.path: + sys.path.insert(0, str(SIDECAR_DIR)) + +import numpy as np # noqa: E402 +from fastapi.testclient import TestClient # noqa: E402 + +import device_policy # noqa: E402 +import main # noqa: E402 +from device_policy import CPU_CHOICE, DeviceChoice # noqa: E402 + +CUBLAS_ERROR = "Library cublas64_12.dll is not found or cannot be loaded" +FLOAT16_ERROR = ( + "Requested float16 compute type, but the target device or backend do not " + "support efficient float16 computation." +) +GPU_FLOAT16 = DeviceChoice(device="cuda", compute_type="float16") + + +# ── 순수 정책 ─────────────────────────────────────────────── + + +class ChooseComputeTypeTest(unittest.TestCase): + def test_prefers_float16_when_supported(self) -> None: + choice = device_policy.choose_gpu_device(1, {"float32", "float16", "int8_float16"}) + self.assertEqual(choice, GPU_FLOAT16) + + def test_pascal_without_float16_uses_int8_float32(self) -> None: + choice = device_policy.choose_gpu_device(1, {"float32", "int8", "int8_float32"}) + self.assertEqual(choice, DeviceChoice("cuda", "int8_float32")) + + def test_no_device_or_no_usable_type_means_cpu(self) -> None: + self.assertIsNone(device_policy.choose_gpu_device(0, {"float16"})) + self.assertIsNone(device_policy.choose_gpu_device(1, set())) + + +class GpuErrorClassificationTest(unittest.TestCase): + def test_gpu_library_errors(self) -> None: + self.assertTrue(device_policy.is_gpu_runtime_error(RuntimeError(CUBLAS_ERROR))) + self.assertTrue(device_policy.is_gpu_runtime_error(ValueError(FLOAT16_ERROR))) + self.assertTrue( + device_policy.is_gpu_runtime_error(RuntimeError("CUDA failed with error out of memory")) + ) + + def test_unrelated_errors(self) -> None: + self.assertFalse(device_policy.is_gpu_runtime_error(ValueError("max() arg is an empty sequence"))) + self.assertFalse(device_policy.is_gpu_runtime_error(OSError("model.bin not found"))) + self.assertFalse(device_policy.is_gpu_runtime_error(KeyError("cuda"))) + + +class LoadWithFallbackTest(unittest.TestCase): + def test_gpu_verified_keeps_gpu(self) -> None: + model, choice = device_policy.load_with_fallback( + build=lambda c: c, verify=lambda _m: None, preferred=GPU_FLOAT16 + ) + self.assertEqual(choice, GPU_FLOAT16) + self.assertEqual(model, GPU_FLOAT16) + + def test_verify_failure_falls_back_to_cpu(self) -> None: + built: list[DeviceChoice] = [] + fallbacks: list[BaseException] = [] + + def verify(m: DeviceChoice) -> None: + if m.is_gpu: + raise RuntimeError(CUBLAS_ERROR) + + def build(c: DeviceChoice) -> DeviceChoice: + built.append(c) + return c + + _model, choice = device_policy.load_with_fallback( + build=build, + verify=verify, + preferred=GPU_FLOAT16, + on_fallback=lambda _c, e: fallbacks.append(e), + ) + self.assertEqual(choice, CPU_CHOICE) + self.assertEqual(built, [GPU_FLOAT16, CPU_CHOICE]) + self.assertEqual(len(fallbacks), 1) + + def test_gpu_build_error_falls_back_to_cpu(self) -> None: + def build(c: DeviceChoice) -> DeviceChoice: + if c.is_gpu: + raise ValueError(FLOAT16_ERROR) + return c + + _model, choice = device_policy.load_with_fallback( + build=build, verify=lambda _m: None, preferred=GPU_FLOAT16 + ) + self.assertEqual(choice, CPU_CHOICE) + + def test_non_gpu_build_error_is_not_retried(self) -> None: + calls: list[DeviceChoice] = [] + + def build(c: DeviceChoice) -> DeviceChoice: + calls.append(c) + raise OSError("Unable to download model files") + + with self.assertRaises(OSError): + device_policy.load_with_fallback(build=build, verify=lambda _m: None, preferred=GPU_FLOAT16) + self.assertEqual(calls, [GPU_FLOAT16]) + + def test_cpu_preferred_skips_verify(self) -> None: + verified: list[DeviceChoice] = [] + _model, choice = device_policy.load_with_fallback( + build=lambda c: c, verify=verified.append, preferred=None + ) + self.assertEqual(choice, CPU_CHOICE) + self.assertEqual(verified, []) + + +# ── 사이드카 HTTP 경로 (가짜 faster_whisper) ───────────────── + + +class _Segment: + def __init__(self, text: str) -> None: + self.text = text + self.start = 0.0 + self.end = 1.0 + self.avg_logprob = -0.1 + + +class _Info: + language = "ko" + + +class _FakeGpuEnv: + """가짜 GPU 환경. cublas_missing: GPU 디코딩이 cuBLAS 로드에서 실패한다.""" + + cublas_missing = False + # True 면 짧은 검증 전사는 통과하고 실제(긴) 전사에서만 실패한다. + fail_only_long_audio = False + float16_unsupported = False + + +class FakeWhisperModel: + created: list["FakeWhisperModel"] = [] + + def __init__( + self, + source: str, + device: str, + compute_type: str, + cpu_threads: int, + num_workers: int, + ) -> None: + if device == "cuda" and compute_type == "float16" and _FakeGpuEnv.float16_unsupported: + raise ValueError(FLOAT16_ERROR) + self.source = source + self.device = device + self.compute_type = compute_type + FakeWhisperModel.created.append(self) + + def transcribe(self, audio: np.ndarray, **_kwargs: object) -> tuple[Iterator[_Segment], _Info]: + fails = ( + self.device == "cuda" + and _FakeGpuEnv.cublas_missing + and (not _FakeGpuEnv.fail_only_long_audio or len(audio) > main._PROBE_SAMPLES) + ) + + def gen() -> Iterator[_Segment]: + if fails: + # faster-whisper 처럼 세그먼트 반복 중(디코딩 시점)에 실패한다. + raise RuntimeError(CUBLAS_ERROR) + yield _Segment(" 안녕하세요 ") + + return gen(), _Info() + + +def _pcm(seconds: float) -> bytes: + return np.zeros(int(16000 * seconds), dtype=np.int16).tobytes() + + +class SidecarGpuFallbackTest(unittest.TestCase): + def setUp(self) -> None: + _FakeGpuEnv.cublas_missing = False + _FakeGpuEnv.fail_only_long_audio = False + _FakeGpuEnv.float16_unsupported = False + FakeWhisperModel.created = [] + main._model = None + main._model_id = None + main._aux_models.clear() + main._model_devices.clear() + main._models_dir = None + # _detect_gpu 가 CUDA 장치를 찾은 상태를 흉내 낸다. + main._gpu_available = True + main._gpu_choice = GPU_FLOAT16 + fake_module = types.ModuleType("faster_whisper") + fake_module.WhisperModel = FakeWhisperModel # type: ignore[attr-defined] + self._patch = mock.patch.dict(sys.modules, {"faster_whisper": fake_module}) + self._patch.start() + # lifespan(_detect_gpu)을 돌리지 않도록 컨텍스트 매니저 없이 쓴다. + self.client = TestClient(main.app) + + def tearDown(self) -> None: + self._patch.stop() + main._model = None + main._model_id = None + main._aux_models.clear() + main._model_devices.clear() + + def _transcribe(self, model_id: str = "") -> "object": + return self.client.post( + "/transcribe", + files={"audio": ("a.pcm", _pcm(3.0), "application/octet-stream")}, + data={"language": "ko", "vad_filter": "false", "model_id": model_id}, + ) + + def test_working_gpu_stays_on_gpu(self) -> None: + res = self.client.post("/load", json={"model_id": "small"}) + self.assertEqual(res.status_code, 200) + self.assertEqual(main._model.device, "cuda") # type: ignore[union-attr] + self.assertEqual(self._transcribe().status_code, 200) + self.assertTrue(self.client.get("/health").json()["gpu"]) + + def test_missing_cublas_detected_at_load_uses_cpu(self) -> None: + _FakeGpuEnv.cublas_missing = True + res = self.client.post("/load", json={"model_id": "medium"}) + self.assertEqual(res.status_code, 200) + self.assertEqual(main._model.device, "cpu") # type: ignore[union-attr] + self.assertEqual(main._model.compute_type, "int8") # type: ignore[union-attr] + + res = self._transcribe() + self.assertEqual(res.status_code, 200) + self.assertEqual(res.json()["text"], "안녕하세요") + + health = self.client.get("/health").json() + self.assertFalse(health["gpu"]) + self.assertEqual(health["device"], "cpu") + + def test_later_models_go_straight_to_cpu_after_fallback(self) -> None: + _FakeGpuEnv.cublas_missing = True + self.client.post("/load", json={"model_id": "medium"}) + FakeWhisperModel.created = [] + res = self.client.post("/load", json={"model_id": "small", "slot": "aux"}) + self.assertEqual(res.status_code, 200) + self.assertEqual([m.device for m in FakeWhisperModel.created], ["cpu"]) + + def test_pascal_float16_error_at_load_uses_cpu(self) -> None: + _FakeGpuEnv.float16_unsupported = True + res = self.client.post("/load", json={"model_id": "medium"}) + self.assertEqual(res.status_code, 200) + self.assertEqual(main._model.device, "cpu") # type: ignore[union-attr] + + def test_gpu_error_during_transcribe_reloads_primary_on_cpu_and_retries(self) -> None: + _FakeGpuEnv.cublas_missing = True + _FakeGpuEnv.fail_only_long_audio = True + self.client.post("/load", json={"model_id": "medium"}) + self.assertEqual(main._model.device, "cuda") # type: ignore[union-attr] + + res = self._transcribe() + self.assertEqual(res.status_code, 200) + self.assertEqual(res.json()["text"], "안녕하세요") + self.assertEqual(main._model.device, "cpu") # type: ignore[union-attr] + self.assertEqual(main._model_id, "medium") + self.assertFalse(main._gpu_available) + + def test_gpu_error_during_transcribe_reloads_aux_on_cpu(self) -> None: + self.client.post("/load", json={"model_id": "medium"}) + self.client.post("/load", json={"model_id": "small", "slot": "aux"}) + _FakeGpuEnv.cublas_missing = True + _FakeGpuEnv.fail_only_long_audio = True + + res = self._transcribe(model_id="small") + self.assertEqual(res.status_code, 200) + self.assertEqual(main._aux_models["small"].device, "cpu") + # 기본 모델 자리는 건드리지 않는다. + self.assertEqual(main._model.device, "cuda") # type: ignore[union-attr] + + def test_cpu_model_errors_are_not_retried(self) -> None: + main._gpu_available = False + main._gpu_choice = None + self.client.post("/load", json={"model_id": "medium"}) + FakeWhisperModel.created = [] + with mock.patch.object( + FakeWhisperModel, "transcribe", side_effect=RuntimeError("CUDA weird") + ): + res = self._transcribe() + self.assertEqual(res.status_code, 500) + self.assertEqual(FakeWhisperModel.created, []) + + +class DetectGpuTest(unittest.TestCase): + def _detect(self, count: int, supported: set[str]) -> None: + fake = types.ModuleType("ctranslate2") + fake.get_cuda_device_count = lambda: count # type: ignore[attr-defined] + fake.get_supported_compute_types = lambda _device: supported # type: ignore[attr-defined] + with mock.patch.dict(sys.modules, {"ctranslate2": fake}): + main._detect_gpu() + + def test_pascal_gets_int8_float32(self) -> None: + self._detect(1, {"float32", "int8", "int8_float32"}) + self.assertTrue(main._gpu_available) + self.assertEqual(main._gpu_choice, DeviceChoice("cuda", "int8_float32")) + + def test_no_supported_types_means_cpu(self) -> None: + self._detect(1, set()) + self.assertFalse(main._gpu_available) + self.assertIsNone(main._gpu_choice) + + +if __name__ == "__main__": + unittest.main()