|
2 | 2 | import json |
3 | 3 | import os |
4 | 4 | import subprocess |
| 5 | +import sys |
5 | 6 |
|
6 | 7 | import pytest |
7 | 8 |
|
|
12 | 13 | FileTranscriptionTask, |
13 | 14 | FileTranscriptionOptions, |
14 | 15 | ) |
| 16 | +from buzz.transcriber import whisper_cpp |
15 | 17 | from buzz.transcriber.whisper_cpp import WhisperCpp, _is_cli_safe_path |
16 | 18 | from tests.audio import test_audio_path, test_multibyte_utf8_audio_path |
17 | 19 | from tests.model_loader import get_model_path |
18 | 20 |
|
19 | 21 |
|
| 22 | +def _whisper_env_source() -> dict: |
| 23 | + # On Windows _whisper_env() copies the app_env snapshot, not os.environ |
| 24 | + return whisper_cpp.app_env if sys.platform == "win32" else os.environ |
| 25 | + |
| 26 | + |
20 | 27 | class TestWhisperCpp: |
21 | 28 | def test_transcribe(self): |
22 | 29 | transcription_options = TranscriptionOptions( |
@@ -617,3 +624,69 @@ def test_parse_word_level_timings_filters_aborted_superseded_fragments(self): |
617 | 624 | segments = WhisperCpp._parse_word_level_timings(mock_result, language="en", vad_enabled=False) |
618 | 625 | words = [s.text for s in segments] |
619 | 626 | assert words == ["Hello", "world", "Goodbye", "all"] |
| 627 | + |
| 628 | + def test_whisper_env_selects_discrete_gpu(self, monkeypatch): |
| 629 | + monkeypatch.delitem(_whisper_env_source(), "GGML_VK_VISIBLE_DEVICES", raising=False) |
| 630 | + with patch("buzz.transcriber.whisper_cpp.DISCRETE_GPU_INDEX", 1): |
| 631 | + assert WhisperCpp._whisper_env()["GGML_VK_VISIBLE_DEVICES"] == "1" |
| 632 | + |
| 633 | + def test_whisper_env_keeps_user_device_choice(self, monkeypatch): |
| 634 | + monkeypatch.setitem(_whisper_env_source(), "GGML_VK_VISIBLE_DEVICES", "0") |
| 635 | + with patch("buzz.transcriber.whisper_cpp.DISCRETE_GPU_INDEX", 1): |
| 636 | + assert WhisperCpp._whisper_env()["GGML_VK_VISIBLE_DEVICES"] == "0" |
| 637 | + |
| 638 | + def test_detect_language_uses_discrete_gpu(self, monkeypatch): |
| 639 | + monkeypatch.delitem(_whisper_env_source(), "GGML_VK_VISIBLE_DEVICES", raising=False) |
| 640 | + completed = MagicMock(stderr="auto-detected language: en (p = 0.9)\n", stdout="") |
| 641 | + with patch( |
| 642 | + "buzz.transcriber.whisper_cpp.subprocess.run", return_value=completed |
| 643 | + ) as run_mock, patch( |
| 644 | + "buzz.transcriber.whisper_cpp.get_whisper_cli_path", |
| 645 | + return_value="/fake/whisper-cli", |
| 646 | + ), patch("buzz.transcriber.whisper_cpp.DISCRETE_GPU_INDEX", 1): |
| 647 | + WhisperCpp.detect_language("/fake/audio.wav", "/fake/model.bin") |
| 648 | + |
| 649 | + assert run_mock.call_args.kwargs["env"]["GGML_VK_VISIBLE_DEVICES"] == "1" |
| 650 | + |
| 651 | + def test_whisper_env_without_discrete_gpu(self, monkeypatch): |
| 652 | + monkeypatch.delitem(_whisper_env_source(), "GGML_VK_VISIBLE_DEVICES", raising=False) |
| 653 | + with patch("buzz.transcriber.whisper_cpp.DISCRETE_GPU_INDEX", None): |
| 654 | + assert "GGML_VK_VISIBLE_DEVICES" not in WhisperCpp._whisper_env() |
| 655 | + |
| 656 | + def test_vulkan_instance_uses_loader_api_version(self): |
| 657 | + # ggml-vulkan creates its instance at the loader's highest API version, and the |
| 658 | + # device order can differ between versions, so the probe must use the same one |
| 659 | + vk = MagicMock() |
| 660 | + vk.vkEnumerateInstanceVersion.return_value = 4206592 # 1.3.0 |
| 661 | + |
| 662 | + instance = whisper_cpp._create_vulkan_instance(vk) |
| 663 | + |
| 664 | + vk.VkApplicationInfo.assert_called_once_with(apiVersion=4206592) |
| 665 | + vk.VkInstanceCreateInfo.assert_called_once_with( |
| 666 | + pApplicationInfo=vk.VkApplicationInfo.return_value |
| 667 | + ) |
| 668 | + vk.vkCreateInstance.assert_called_once_with(vk.VkInstanceCreateInfo.return_value, None) |
| 669 | + assert instance is vk.vkCreateInstance.return_value |
| 670 | + |
| 671 | + @staticmethod |
| 672 | + def _fake_vulkan(device_types): |
| 673 | + vk = MagicMock() |
| 674 | + vk.VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU = 1 |
| 675 | + vk.VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU = 2 |
| 676 | + vk.VK_PHYSICAL_DEVICE_TYPE_CPU = 4 |
| 677 | + devices = [MagicMock(deviceType=t, deviceName=f"gpu{i}") for i, t in enumerate(device_types)] |
| 678 | + vk.vkEnumeratePhysicalDevices.return_value = list(range(len(devices))) |
| 679 | + vk.vkGetPhysicalDeviceProperties.side_effect = lambda index: devices[index] |
| 680 | + return vk |
| 681 | + |
| 682 | + def test_find_discrete_gpu_index_after_integrated(self): |
| 683 | + vk = self._fake_vulkan([1, 2, 4]) |
| 684 | + assert whisper_cpp._find_discrete_gpu_index(vk, MagicMock()) == 1 |
| 685 | + |
| 686 | + def test_find_discrete_gpu_index_returns_first_discrete(self): |
| 687 | + vk = self._fake_vulkan([2, 1, 2]) |
| 688 | + assert whisper_cpp._find_discrete_gpu_index(vk, MagicMock()) == 0 |
| 689 | + |
| 690 | + def test_find_discrete_gpu_index_without_discrete(self): |
| 691 | + vk = self._fake_vulkan([1, 4]) |
| 692 | + assert whisper_cpp._find_discrete_gpu_index(vk, MagicMock()) is None |
0 commit comments