feat: whisper 分块间暂停检查与默认 float16 计算类型
- 分块转写在每个分块前检查 run 根目录 paused.flag,批量/任务暂停时 分块粒度内中止(默认 60s 一块),当前块执行完才停 - 默认 compute_type 由 auto 改为 float16,测试断言同步更新
This commit is contained in:
+10
-1
@@ -22,6 +22,9 @@ from wov_sdk.models import InvokeRequest, InvokeResponse
|
||||
# 转写进度日志:输出到主进程控制台,长音频分块时可见每块进度。
|
||||
logger = get_logger("whisper")
|
||||
|
||||
# 暂停信号文件名:位于 run 根目录(<storage>/runs/<run_id>/paused.flag),
|
||||
# 与 subtitle-ocr 节点约定一致;批量暂停时由暂停接口写入,分块间检查即中止。
|
||||
PAUSE_FLAG = "paused.flag"
|
||||
|
||||
def _is_windows() -> bool:
|
||||
"""判断当前是否为 Windows,供测试单独注入覆盖。"""
|
||||
@@ -216,7 +219,7 @@ def invoke(request: InvokeRequest) -> InvokeResponse:
|
||||
model_path = resolve_model_path(request.params)
|
||||
device = str(request.params.get("device") or os.getenv("WHISPER_DEVICE", "auto"))
|
||||
# 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_path,
|
||||
device=device,
|
||||
@@ -241,6 +244,12 @@ def invoke(request: InvokeRequest) -> InvokeResponse:
|
||||
srt_number = 1
|
||||
transcribe_started = time.monotonic()
|
||||
for chunk_index, chunk in enumerate(chunks, start=1):
|
||||
# 暂停检查:批量暂停时在 run 根目录写 paused.flag,whisper 在分块
|
||||
# 之间检查该信号(默认 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()
|
||||
segments, _info = model.transcribe(
|
||||
str(chunk),
|
||||
|
||||
+32
-2
@@ -309,7 +309,7 @@ def test_format_timestamp() -> None:
|
||||
|
||||
|
||||
def test_whisper_success(tmp_path, monkeypatch) -> None:
|
||||
"""验证成功转写会生成 SRT 并默认使用 auto 设备/计算类型。"""
|
||||
"""验证成功转写会生成 SRT 并默认使用 auto 设备、float16 计算类型。"""
|
||||
FakeWhisperModel.instances.clear()
|
||||
_install_fake_whisper(monkeypatch)
|
||||
_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
|
||||
_, kwargs = FakeWhisperModel.instances[-1]
|
||||
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:
|
||||
"""验证 Linux 下进程内预加载 nvidia 动态库,含失败跳过分支。"""
|
||||
import ctypes
|
||||
|
||||
Reference in New Issue
Block a user