feat: whisper 分块间暂停检查与默认 float16 计算类型

- 分块转写在每个分块前检查 run 根目录 paused.flag,批量/任务暂停时
  分块粒度内中止(默认 60s 一块),当前块执行完才停
- 默认 compute_type 由 auto 改为 float16,测试断言同步更新
This commit is contained in:
2026-08-23 16:25:15 +08:00
parent 3b5bd42b60
commit 3aa05bbf76
2 changed files with 42 additions and 3 deletions
+10 -1
View File
@@ -22,6 +22,9 @@ from wov_sdk.models import InvokeRequest, InvokeResponse
# 转写进度日志:输出到主进程控制台,长音频分块时可见每块进度。 # 转写进度日志:输出到主进程控制台,长音频分块时可见每块进度。
logger = get_logger("whisper") logger = get_logger("whisper")
# 暂停信号文件名:位于 run 根目录(<storage>/runs/<run_id>/paused.flag),
# 与 subtitle-ocr 节点约定一致;批量暂停时由暂停接口写入,分块间检查即中止。
PAUSE_FLAG = "paused.flag"
def _is_windows() -> bool: def _is_windows() -> bool:
"""判断当前是否为 Windows,供测试单独注入覆盖。""" """判断当前是否为 Windows,供测试单独注入覆盖。"""
@@ -216,7 +219,7 @@ def invoke(request: InvokeRequest) -> InvokeResponse:
model_path = resolve_model_path(request.params) model_path = resolve_model_path(request.params)
device = str(request.params.get("device") or os.getenv("WHISPER_DEVICE", "auto")) device = str(request.params.get("device") or os.getenv("WHISPER_DEVICE", "auto"))
# auto 让 faster-whisper 根据硬件自动选择 float16/int8 等计算类型。 # auto 让 faster-whisper 根据硬件自动选择 float16/int8 等计算类型。
compute_type = str(request.params.get("compute_type") or "auto") compute_type = str(request.params.get("compute_type") or "float16")
model = WhisperModel( model = WhisperModel(
model_path, model_path,
device=device, device=device,
@@ -241,6 +244,12 @@ def invoke(request: InvokeRequest) -> InvokeResponse:
srt_number = 1 srt_number = 1
transcribe_started = time.monotonic() transcribe_started = time.monotonic()
for chunk_index, chunk in enumerate(chunks, start=1): for chunk_index, chunk in enumerate(chunks, start=1):
# 暂停检查:批量暂停时在 run 根目录写 paused.flagwhisper 在分块
# 之间检查该信号(默认 60s 一块,暂停粒度不超过一块),检测到即抛
# 异常,由 invoke 转 failed、调度器保持任务 PAUSED,继续时整个节点
# 重新转写(whisper 没有节点级断点存档,产物只在结束时一次性写出)。
if (Path(request.output_dir).parent.parent / PAUSE_FLAG).exists():
raise RuntimeError(f"whisper 被暂停(run {request.run_id}")
chunk_started = time.monotonic() chunk_started = time.monotonic()
segments, _info = model.transcribe( segments, _info = model.transcribe(
str(chunk), str(chunk),
+32 -2
View File
@@ -309,7 +309,7 @@ def test_format_timestamp() -> None:
def test_whisper_success(tmp_path, monkeypatch) -> None: def test_whisper_success(tmp_path, monkeypatch) -> None:
"""验证成功转写会生成 SRT 并默认使用 auto 设备/计算类型。""" """验证成功转写会生成 SRT 并默认使用 auto 设备、float16 计算类型。"""
FakeWhisperModel.instances.clear() FakeWhisperModel.instances.clear()
_install_fake_whisper(monkeypatch) _install_fake_whisper(monkeypatch)
_make_wav(tmp_path / "audio.wav", 5) _make_wav(tmp_path / "audio.wav", 5)
@@ -320,9 +320,39 @@ def test_whisper_success(tmp_path, monkeypatch) -> None:
assert "01:00:00,500 --> 01:00:02,250" in content assert "01:00:00,500 --> 01:00:02,250" in content
_, kwargs = FakeWhisperModel.instances[-1] _, kwargs = FakeWhisperModel.instances[-1]
assert kwargs["device"] == "auto" assert kwargs["device"] == "auto"
assert kwargs["compute_type"] == "auto" assert kwargs["compute_type"] == "float16"
def test_whisper_pause_flag_stops_between_chunks(tmp_path, monkeypatch) -> None:
"""验证 paused.flag 存在时 whisper 在分块边界中止(批量暂停机制)。
run 根目录(output_dir 的上上级)的 paused.flag 由暂停接口写入,whisper
在每个分块转写前检查;检测到即抛异常,invoke 统一转 failed 响应,调度器
捕获后保持任务 PAUSED。
"""
_install_fake_whisper(monkeypatch)
_make_wav(tmp_path / "audio.wav", 5)
# 暂停信号位于 run 根目录:output_dir 为 runs/run_p/steps/whisper
# 其上上级即 runs/run_p(与调度器/OCR 的目录约定一致)。
(tmp_path / "runs" / "run_p").mkdir(parents=True)
(tmp_path / "runs" / "run_p" / "paused.flag").write_text("", encoding="utf-8")
# 注入两个分块路径,确保进入分块循环并执行至少一次暂停检查。
monkeypatch.setattr(
"nodes.whisper._split_audio",
lambda *args, **kwargs: [tmp_path / "chunk_1.wav", tmp_path / "chunk_2.wav"],
)
response = whisper_invoke(
_whisper_request(
tmp_path,
params={"language": "ja", "chunk_seconds": 60},
output_dir=str(tmp_path / "runs" / "run_p" / "steps" / "whisper"),
)
)
assert response.status == "failed"
assert "被暂停" in response.error
# 暂停时不会写出 SRT 产物。
assert not (tmp_path / "runs" / "run_p" / "steps" / "whisper" / "transcript.srt").exists()
def test_load_cuda_libraries_linux(monkeypatch) -> None: def test_load_cuda_libraries_linux(monkeypatch) -> None:
"""验证 Linux 下进程内预加载 nvidia 动态库,含失败跳过分支。""" """验证 Linux 下进程内预加载 nvidia 动态库,含失败跳过分支。"""
import ctypes import ctypes