Dymas/whisper-api-server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix thread safety, implement prompt, fix resource leaks, add history rotation

Browse Source 
- Add threading.Lock around asr_pipeline calls (thread-safe inference on shared GPU)
- Implement prompt parameter: propagate through process_file → transcribe → generate_kwargs
- Fix prompt_ids device mismatch (.to(self.device) for CUDA)
- Replace mkdtemp() with create_temp_file() in sources.py (no more leaked /tmp dirs)
- Add max_history_days (default 30) with automatic cleanup of old history directories
- Add intentional comments for SSRF and /config decisions (internal service)
- Lower normalize_audio log level from INFO to DEBUG

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Serge Zaigraeff
2026-03-31 00:43:06 +03:00
co-authored by Claude Sonnet 4.6
parent 2138651474
commit 1e778cae86
7 changed files with 171 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/core/transcriber.py
+74 -31
View File
@@ -7,6 +7,7 @@ OpenAI для транскрибации аудиофайлов в текст.
"""
import time
import threading
import traceback
from typing import Dict, Tuple, Union
import logging
@@ -63,8 +64,10 @@ class WhisperTranscriber:
self.return_timestamps = config["return_timestamps"]
self.temperature = config["temperature"]
# Оптимальный тип для тензоров
self.torch_dtype = torch.bfloat16
# Lock для потокобезопасного доступа к модели —
# Waitress обслуживает запросы в нескольких потоках,
# а HuggingFace pipeline не является thread-safe
self._inference_lock = threading.Lock()
# Создаем объект для обработки аудио
self.audio_processor = AudioProcessor(config)
@@ -72,6 +75,9 @@ class WhisperTranscriber:
# Определяем устройство для вычислений
self.device = self._get_device()
# Оптимальный тип для тензоров (зависит от устройства)
self.torch_dtype = self._get_torch_dtype()
# Загружаем модель при инициализации
self._load_model()
@@ -88,26 +94,40 @@ class WhisperTranscriber:
# Проверяем, что device_id является целым числом
if not isinstance(device_id, int):
logger.warning(f"device_id должен быть целым числом, получено: {device_id}. Используем значение по умолчанию 0")
logger.warning("device_id должен быть целым числом, получено: %s. Используем значение по умолчанию 0", device_id)
device_id = 0
# Проверяем, доступен ли запрошенный GPU
device_count = torch.cuda.device_count()
if device_id >= device_count:
logger.warning(f"Запрошенный GPU с индексом {device_id} недоступен. Доступно GPU: {device_count}. Используем GPU с индексом 0")
logger.warning("Запрошенный GPU с индексом %s недоступен. Доступно GPU: %s. Используем GPU с индексом 0", device_id, device_count)
device_id = 0
logger.info(f"Используется CUDA GPU с индексом {device_id} для вычислений")
logger.info("Используется CUDA GPU с индексом %s для вычислений", device_id)
return torch.device(f"cuda:{device_id}")
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
logger.info("Используется MPS (Apple Silicon) для вычислений")
# Обходное решение для MPS
# Обходное решение для MPS: PyTorch проверяет is_initialized()
# при создании тензоров на MPS-устройстве, что вызывает ошибку
# в однопроцессном режиме.
# TODO: Удалить после обновления до PyTorch >= 2.5
setattr(torch.distributed, "is_initialized", lambda: False)
return torch.device("mps")
else:
logger.info("Используется CPU для вычислений")
return torch.device("cpu")
def _get_torch_dtype(self) -> torch.dtype:
"""
Определяет оптимальный тип данных для тензоров в зависимости от устройства.
Returns:
torch.float32 для CPU, torch.bfloat16 для GPU/MPS.
"""
if self.device.type == "cpu":
return torch.float32
return torch.bfloat16
def _load_model(self) -> None:
"""
Загрузка модели и процессора.
@@ -115,7 +135,7 @@ class WhisperTranscriber:
Raises:
Exception: Если не удалось загрузить модель.
"""
logger.info(f"Загрузка модели из {self.model_path}")
logger.info("Загрузка модели из %s", self.model_path)
model_kwargs = dict(
torch_dtype=self.torch_dtype,
@@ -132,7 +152,7 @@ class WhisperTranscriber:
if self.device.type == "cuda":
logger.info("Используется Flash Attention 2")
except Exception as e:
logger.warning(f"Не удалось загрузить модель с Flash Attention: {e}")
logger.warning("Не удалось загрузить модель с Flash Attention: %s", e)
model_kwargs.pop("attn_implementation", None)
self.model = WhisperForConditionalGeneration.from_pretrained(
self.model_path, **model_kwargs
@@ -154,13 +174,18 @@ class WhisperTranscriber:
logger.info("Модель успешно загружена и готова к использованию")
def transcribe(self, audio_path: str, return_timestamps: bool = None) -> Union[str, Dict]:
def transcribe(self, audio_path: str, return_timestamps: bool = None,
language: str = None, temperature: float = None,
prompt: str = None) -> Union[str, Dict]:
"""
Транскрибация аудиофайла.
Args:
audio_path: Путь к обработанному аудиофайлу.
return_timestamps: Флаг возврата временных меток. Если None — берётся из конфига.
language: Язык распознавания. Если None — берётся из конфига.
temperature: Параметр температуры для генерации. Если None — берётся из конфига.
prompt: Текстовая подсказка для модели (имена, термины, контекст).
Returns:
В зависимости от параметра return_timestamps:
@@ -169,28 +194,39 @@ class WhisperTranscriber:
"""
if return_timestamps is None:
return_timestamps = self.return_timestamps
if language is None:
language = self.language
if temperature is None:
temperature = self.temperature
logger.info(f"Начало транскрибации файла: {audio_path}")
logger.info("Начало транскрибации файла: %s", audio_path)
try:
# Загрузка аудио в формате numpy array
audio_array, sampling_rate = load_audio(audio_path, sr=16000)
# Транскрибация с корректным форматом данных
result = self.asr_pipeline(
{"raw": audio_array, "sampling_rate": sampling_rate},
generate_kwargs={
"language": self.language,
"max_new_tokens": self.max_new_tokens,
"temperature": self.temperature
},
return_timestamps=return_timestamps
)
generate_kwargs = {
"language": language,
"max_new_tokens": self.max_new_tokens,
"temperature": temperature,
}
if prompt:
generate_kwargs["prompt_ids"] = self.processor.get_prompt_ids(
prompt, return_tensors="pt"
).to(self.device)
with self._inference_lock:
result = self.asr_pipeline(
{"raw": audio_array, "sampling_rate": sampling_rate},
generate_kwargs=generate_kwargs,
return_timestamps=return_timestamps,
)
# Если временные метки не запрошены, возвращаем только текст
if not return_timestamps:
transcribed_text = result.get("text", "")
logger.info(f"Транскрибация завершена: получено {len(transcribed_text)} символов текста")
logger.info("Транскрибация завершена: получено %s символов текста", len(transcribed_text))
return transcribed_text
# Если временные метки запрошены, обрабатываем и форматируем результат
@@ -224,7 +260,7 @@ class WhisperTranscriber:
else:
logger.warning("Временные метки запрошены, но не найдены в результате транскрибации")
logger.info(f"Транскрибация с временными метками завершена: получено {len(segments)} сегментов")
logger.info("Транскрибация с временными метками завершена: получено %s сегментов", len(segments))
# Возвращаем словарь с сегментами и полным текстом
return {
@@ -233,18 +269,23 @@ class WhisperTranscriber:
}
except Exception as e:
logger.error(f"Ошибка в процессе транскрибации аудиофайла '{audio_path}': {str(e)}")
logger.error(f"Тип исключения: {type(e).__name__}")
logger.error(f"Traceback: {traceback.format_exc()}")
logger.error("Ошибка в процессе транскрибации аудиофайла '%s': %s", audio_path, e)
logger.error("Тип исключения: %s", type(e).__name__)
logger.error("Traceback: %s", traceback.format_exc())
raise
def process_file(self, input_path: str, return_timestamps: bool = None) -> Union[str, Dict]:
def process_file(self, input_path: str, return_timestamps: bool = None,
language: str = None, temperature: float = None,
prompt: str = None) -> Union[str, Dict]:
"""
Полный процесс обработки и транскрибации аудиофайла.
Args:
input_path: Путь к исходному аудиофайлу.
return_timestamps: Флаг возврата временных меток. Если None — берётся из конфига.
language: Язык распознавания. Если None — берётся из конфига.
temperature: Параметр температуры для генерации. Если None — берётся из конфига.
prompt: Текстовая подсказка для модели (имена, термины, контекст).
Returns:
В зависимости от параметра return_timestamps:
@@ -252,7 +293,7 @@ class WhisperTranscriber:
- Если return_timestamps=True: словарь с ключами "segments" и "text"
"""
start_time = time.time()
logger.info(f"Начало обработки файла: {input_path}")
logger.info("Начало обработки файла: %s", input_path)
temp_files = []
@@ -261,18 +302,20 @@ class WhisperTranscriber:
processed_path, temp_files = self.audio_processor.process_audio(input_path)
# Транскрибация
result = self.transcribe(processed_path, return_timestamps=return_timestamps)
result = self.transcribe(processed_path, return_timestamps=return_timestamps,
language=language, temperature=temperature,
prompt=prompt)
elapsed_time = time.time() - start_time
logger.info(f"Обработка и транскрибация завершены за {elapsed_time:.2f} секунд")
logger.info("Обработка и транскрибация завершены за %.2f секунд", elapsed_time)
return result
except Exception as e:
elapsed_time = time.time() - start_time
logger.error(f"Ошибка при обработке файла '{input_path}' через {elapsed_time:.2f} секунд: {str(e)}")
logger.error(f"Тип исключения: {type(e).__name__}")
logger.error(f"Traceback: {traceback.format_exc()}")
logger.error("Ошибка при обработке файла '%s' через %.2f секунд: %s", input_path, elapsed_time, e)
logger.error("Тип исключения: %s", type(e).__name__)
logger.error("Traceback: %s", traceback.format_exc())
raise
finally:
app/core/transcription_service.py
+12 -6
View File
@@ -50,11 +50,15 @@ class TranscriptionService:
try:
duration = get_audio_duration(file_path)
except Exception as e:
logger.error(f"Ошибка при определении длительности файла: {e}")
logger.error("Ошибка при определении длительности файла: %s", e)
return {"error": f"Не удалось определить длительность аудиофайла: {e}"}, 500
start_time = time.time()
result = self.transcriber.process_file(file_path, return_timestamps=return_timestamps)
result = self.transcriber.process_file(
file_path, return_timestamps=return_timestamps,
language=language, temperature=temperature,
prompt=prompt
)
processing_time = time.time() - start_time
# Формируем ответ
@@ -63,7 +67,7 @@ class TranscriptionService:
"segments": result.get("segments", []),
"text": result.get("text", ""),
"processing_time": processing_time,
"response_size_bytes": len(json.dumps(result, ensure_ascii=False).encode('utf-8')),
"response_size_bytes": 0,
"duration_seconds": duration,
"model": os.path.basename(self.config["model_path"])
}
@@ -71,15 +75,17 @@ class TranscriptionService:
response = {
"text": result,
"processing_time": processing_time,
"response_size_bytes": len(json.dumps(result, ensure_ascii=False).encode('utf-8')),
"response_size_bytes": 0,
"duration_seconds": duration,
"model": os.path.basename(self.config["model_path"])
}
response["response_size_bytes"] = len(json.dumps(response, ensure_ascii=False).encode('utf-8'))
save_history(response, filename, self.config)
return response, 200
except Exception as e:
logger.error(f"Ошибка при транскрибации файла '{filename}': {str(e)}")
logger.error(f"Traceback: {traceback.format_exc()}")
logger.error("Ошибка при транскрибации файла '%s': %s", filename, e)
logger.error("Traceback: %s", traceback.format_exc())
return {"error": str(e)}, 500