Skip to content

Commit 109fa99

Browse files
raivisdejusraivisdejus
committed
Fixed tests
1 parent 3d9ca66 commit 109fa99

7 files changed

Lines changed: 81 additions & 45 deletions

File tree

‎buzz/cuda_setup.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -178,6 +178,13 @@ def _setup_linux_cuda():
178178
logger.info("CUDA setup: already done (sentinel set), proceeding")
179179
return
180180

181+
# Never re-exec inside a multiprocessing worker — it would destroy the IPC pipe.
182+
# The parent process already set LD_LIBRARY_PATH before spawning workers.
183+
import multiprocessing
184+
if multiprocessing.parent_process() is not None:
185+
logger.debug("CUDA setup: skipping re-exec in worker subprocess")
186+
return
187+
181188
# Collect lib dirs from snap/flatpak cuda_packages OR regular site-packages
182189
cuda_target = _get_cuda_target_dir()
183190
if cuda_target is not None and cuda_target.exists():

‎buzz/transcriber/cuda_device.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
import logging
2+
3+
4+
def cuda_works() -> bool:
5+
"""Return True only if CUDA is available AND a kernel actually executes without error."""
6+
try:
7+
import torch
8+
except ImportError:
9+
return False
10+
if not torch.cuda.is_available() or not torch.version.cuda:
11+
return False
12+
try:
13+
torch.zeros(1, device="cuda")
14+
return True
15+
except Exception as e:
16+
logging.debug("CUDA smoke test failed, falling back to CPU: %s", e)
17+
return False

‎buzz/transcriber/recording_transcriber.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
import torch
1919
import numpy as np
20+
from buzz.transcriber.cuda_device import cuda_works
2021
import sounddevice
2122
from sounddevice import PortAudioError
2223
from openai import OpenAI
@@ -88,7 +89,7 @@ def start(self):
8889
keep_samples = int(self.keep_sample_seconds * self.sample_rate)
8990

9091
force_cpu = os.getenv("BUZZ_FORCE_CPU", "false")
91-
use_cuda = torch.cuda.is_available() and force_cpu == "false"
92+
use_cuda = cuda_works() and force_cpu == "false"
9293

9394
if torch.cuda.is_available():
9495
logging.debug(f"CUDA version detected: {torch.version.cuda}")
@@ -110,12 +111,11 @@ def start(self):
110111
model_root_dir = os.getenv("BUZZ_MODEL_ROOT", model_root_dir)
111112

112113
device = "auto"
113-
if torch.cuda.is_available() and torch.version.cuda < "12":
114-
logging.debug("Unsupported CUDA version (<12), using CPU")
114+
if not cuda_works():
115+
logging.debug("CUDA not available or not functional, using CPU")
115116
device = "cpu"
116-
117-
if not torch.cuda.is_available():
118-
logging.debug("CUDA is not available, using CPU")
117+
elif torch.version.cuda < "12":
118+
logging.debug("Unsupported CUDA version (<12), using CPU")
119119
device = "cpu"
120120

121121
if force_cpu != "false":

‎buzz/transcriber/whisper_file_transcriber.py‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111

1212
import torch
1313
import platform
14+
15+
from buzz.transcriber.cuda_device import cuda_works
1416
import subprocess
1517
from platformdirs import user_cache_dir
1618
from multiprocessing.connection import Connection
@@ -268,12 +270,11 @@ def transcribe_faster_whisper(cls, task: FileTranscriptionTask) -> List[Segment]
268270
force_cpu = os.getenv("BUZZ_FORCE_CPU", "false")
269271

270272
device = "auto"
271-
if torch.cuda.is_available() and torch.version.cuda < "12":
272-
logging.debug("Unsupported CUDA version (<12), using CPU")
273+
if not cuda_works():
274+
logging.debug("CUDA not available or not functional, using CPU")
273275
device = "cpu"
274-
275-
if not torch.cuda.is_available():
276-
logging.debug("CUDA is not available, using CPU")
276+
elif torch.version.cuda < "12":
277+
logging.debug("Unsupported CUDA version (<12), using CPU")
277278
device = "cpu"
278279

279280
if force_cpu != "false":
@@ -334,7 +335,7 @@ def transcribe_faster_whisper(cls, task: FileTranscriptionTask) -> List[Segment]
334335
@classmethod
335336
def transcribe_openai_whisper(cls, task: FileTranscriptionTask) -> List[Segment]:
336337
force_cpu = os.getenv("BUZZ_FORCE_CPU", "false")
337-
use_cuda = torch.cuda.is_available() and force_cpu == "false"
338+
use_cuda = cuda_works() and force_cpu == "false"
338339

339340
device = "cuda" if use_cuda else "cpu"
340341

‎buzz/transformers_whisper.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99

1010
import torch
1111
import requests
12+
from buzz.transcriber.cuda_device import cuda_works
1213
from typing import Union
1314
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline, BitsAndBytesConfig
1415
from transformers.pipelines import AutomaticSpeechRecognitionPipeline
@@ -245,7 +246,7 @@ def _transcribe_whisper(
245246
):
246247
"""Transcribe using Whisper model."""
247248
force_cpu = os.getenv("BUZZ_FORCE_CPU", "false")
248-
use_cuda = torch.cuda.is_available() and force_cpu == "false"
249+
use_cuda = cuda_works() and force_cpu == "false"
249250
device = "cuda" if use_cuda else "cpu"
250251
torch_dtype = torch.float16 if use_cuda else torch.float32
251252

@@ -464,7 +465,7 @@ def _transcribe_mms(
464465
from transformers.pipelines.audio_utils import ffmpeg_read as mms_ffmpeg_read
465466

466467
force_cpu = os.getenv("BUZZ_FORCE_CPU", "false")
467-
use_cuda = torch.cuda.is_available() and force_cpu == "false"
468+
use_cuda = cuda_works() and force_cpu == "false"
468469
device = "cuda" if use_cuda else "cpu"
469470

470471
# Map language code to ISO 639-3 for MMS

‎tests/cuda_setup_test.py‎

Lines changed: 41 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,5 @@
11
import sys
2-
from pathlib import Path
3-
from unittest.mock import MagicMock, patch
4-
5-
import pytest
2+
from unittest.mock import patch
63

74

85
class TestGetCudaTargetDir:
@@ -117,46 +114,57 @@ class TestSetupLinuxCuda:
117114
def test_skips_when_no_cuda_target(self, monkeypatch):
118115
monkeypatch.delenv("SNAP_USER_DATA", raising=False)
119116
monkeypatch.delenv("FLATPAK_ID", raising=False)
117+
monkeypatch.delenv("BUZZ_CUDA_SETUP_DONE", raising=False)
120118

121119
with patch("buzz.cuda_setup._get_cuda_target_dir", return_value=None):
122-
from buzz.cuda_setup import _setup_linux_cuda
123-
_setup_linux_cuda() # should not raise
120+
with patch("buzz.cuda_setup._collect_site_packages_cuda_lib_dirs", return_value=[]):
121+
with patch("multiprocessing.parent_process", return_value=None):
122+
from buzz.cuda_setup import _setup_linux_cuda
123+
_setup_linux_cuda() # should not raise
124124

125125
def test_skips_when_cuda_target_does_not_exist(self, monkeypatch, tmp_path):
126+
monkeypatch.delenv("BUZZ_CUDA_SETUP_DONE", raising=False)
126127
nonexistent = tmp_path / "nonexistent"
127128
with patch("buzz.cuda_setup._get_cuda_target_dir", return_value=nonexistent):
128-
from buzz.cuda_setup import _setup_linux_cuda
129-
_setup_linux_cuda() # should not raise
129+
with patch("buzz.cuda_setup._collect_site_packages_cuda_lib_dirs", return_value=[]):
130+
with patch("multiprocessing.parent_process", return_value=None):
131+
from buzz.cuda_setup import _setup_linux_cuda
132+
_setup_linux_cuda() # should not raise
130133

131134
def test_skips_when_no_torch_lib(self, monkeypatch, tmp_path):
135+
monkeypatch.delenv("BUZZ_CUDA_SETUP_DONE", raising=False)
132136
cuda_target = tmp_path / "cuda_packages"
133137
cuda_target.mkdir()
134138
with patch("buzz.cuda_setup._get_cuda_target_dir", return_value=cuda_target):
135-
from buzz.cuda_setup import _setup_linux_cuda
136-
_setup_linux_cuda() # should not raise
139+
with patch("buzz.cuda_setup._collect_site_packages_cuda_lib_dirs", return_value=[]):
140+
with patch("multiprocessing.parent_process", return_value=None):
141+
from buzz.cuda_setup import _setup_linux_cuda
142+
_setup_linux_cuda() # should not raise
137143

138-
def test_reexecs_when_sentinel_not_in_ld_path(self, monkeypatch, tmp_path):
144+
def test_reexecs_when_sentinel_not_set(self, monkeypatch, tmp_path):
139145
cuda_target = tmp_path / "cuda_packages"
140146
torch_lib = cuda_target / "torch" / "lib"
141147
torch_lib.mkdir(parents=True)
142148

149+
monkeypatch.delenv("BUZZ_CUDA_SETUP_DONE", raising=False)
143150
monkeypatch.delenv("LD_LIBRARY_PATH", raising=False)
144151

145152
with patch("buzz.cuda_setup._get_cuda_target_dir", return_value=cuda_target):
146153
with patch("buzz.cuda_setup._collect_cuda_lib_dirs", return_value=[str(torch_lib)]):
147-
with patch("os.execv", side_effect=OSError("test")) as mock_execv:
148-
with patch("buzz.cuda_setup._preload_linux_libraries_fallback") as mock_fallback:
149-
from buzz.cuda_setup import _setup_linux_cuda
150-
_setup_linux_cuda()
154+
with patch("multiprocessing.parent_process", return_value=None):
155+
with patch("os.execv", side_effect=OSError("test")):
156+
with patch("buzz.cuda_setup._preload_linux_libraries_fallback") as mock_fallback:
157+
from buzz.cuda_setup import _setup_linux_cuda
158+
_setup_linux_cuda()
151159

152-
mock_fallback.assert_called_once()
160+
mock_fallback.assert_called_once()
153161

154-
def test_no_reexec_when_sentinel_already_in_ld_path(self, monkeypatch, tmp_path):
162+
def test_no_reexec_when_sentinel_already_set(self, monkeypatch, tmp_path):
155163
cuda_target = tmp_path / "cuda_packages"
156164
torch_lib = cuda_target / "torch" / "lib"
157165
torch_lib.mkdir(parents=True)
158166

159-
monkeypatch.setenv("LD_LIBRARY_PATH", str(torch_lib))
167+
monkeypatch.setenv("BUZZ_CUDA_SETUP_DONE", "1")
160168

161169
with patch("buzz.cuda_setup._get_cuda_target_dir", return_value=cuda_target):
162170
with patch("os.execv") as mock_execv:
@@ -165,6 +173,21 @@ def test_no_reexec_when_sentinel_already_in_ld_path(self, monkeypatch, tmp_path)
165173

166174
mock_execv.assert_not_called()
167175

176+
def test_no_reexec_in_worker_subprocess(self, monkeypatch, tmp_path):
177+
cuda_target = tmp_path / "cuda_packages"
178+
torch_lib = cuda_target / "torch" / "lib"
179+
torch_lib.mkdir(parents=True)
180+
181+
monkeypatch.delenv("BUZZ_CUDA_SETUP_DONE", raising=False)
182+
183+
with patch("buzz.cuda_setup._get_cuda_target_dir", return_value=cuda_target):
184+
with patch("multiprocessing.parent_process", return_value=object()):
185+
with patch("os.execv") as mock_execv:
186+
from buzz.cuda_setup import _setup_linux_cuda
187+
_setup_linux_cuda()
188+
189+
mock_execv.assert_not_called()
190+
168191

169192
class TestSetupWindowsDllDirectories:
170193
def test_calls_add_dll_directory(self, tmp_path):

‎tests/widgets/cuda_installer_widget_test.py‎

Lines changed: 0 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -89,19 +89,6 @@ def test_on_progress_appends_to_log(self, qtbot: QtBot):
8989
dialog._on_progress("step 1")
9090
assert "step 1" in dialog.log_view.toPlainText()
9191

92-
def test_on_progress_updates_status_label(self, qtbot: QtBot):
93-
dialog = CudaInstallerDialog()
94-
qtbot.add_widget(dialog)
95-
dialog._on_progress("Installing packages...")
96-
assert dialog.status_label.text() != ""
97-
98-
def test_on_progress_truncates_long_message(self, qtbot: QtBot):
99-
dialog = CudaInstallerDialog()
100-
qtbot.add_widget(dialog)
101-
long_msg = "x" * 200
102-
dialog._on_progress(long_msg)
103-
assert len(dialog.status_label.text()) <= 80
104-
10592
def test_on_finished_re_enables_install_button(self, qtbot: QtBot):
10693
dialog = CudaInstallerDialog()
10794
qtbot.add_widget(dialog)

0 commit comments

Comments
 (0)