322 lines
12 KiB
Python
322 lines
12 KiB
Python
"""
|
|
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()
|