feat: 批量流水线按 GPU 资源调度(非 GPU 阶段并行、GPU 阶段互斥)
此前组内阶段是串行的(先全部提音、再全部转写、再全部翻译),LLM 走线上端点时 翻译阶段不占显存、GPU 全程空转——实测占整轮挂钟约 40%(19.6W / 272MiB)。 - `_run_job` 改为按组启动在途流水线:每个视频独立推进自己的阶段,最多 `WOV_BATCH_PIPELINE_WORKERS`(默认 4)个阶段在途。 - 派发只看资源:`stage_gpu_need_mb` 为 0 的阶段(提音、线上翻译、ASS)立刻派发, 可与其它视频的转写并行;需要 GPU 的阶段由 `GpuGate` 互斥准入,并按"阶段索引 最小者优先"派发,组内仍是先跑完全部转写再进翻译——本机 Ollama 模型每组只 加载一次,不需要按"是否云端"写分支。 - 同一阶段只在途一份(派发即标记 running),单视频异常不带走整组;暂停沿用 run 级 paused.flag,暂停后不再派发新阶段。 - 测试:远端翻译与其它视频转写重叠、本机端点下全部转写先于翻译且翻译互斥、 提音与转写重叠,以及既有分组/暂停/失败隔离用例。
This commit is contained in:
@@ -14,6 +14,7 @@ from pathlib import Path
|
||||
import pytest
|
||||
|
||||
from wov_app import registry
|
||||
from wov_app.resources import GpuGate
|
||||
from wov_app.batch import (
|
||||
KEEP_MODEL_FLAG,
|
||||
MARKER_NAME,
|
||||
@@ -561,7 +562,7 @@ def test_worker_processes_recovered_zombie_job(tmp_path: Path, monkeypatch) -> N
|
||||
|
||||
|
||||
def test_worker_runs_grouped_stage_pipeline(tmp_path: Path, monkeypatch) -> None:
|
||||
"""分块流水线:组内按节点顺序跑完全部视频,而不是每个视频跑完整链路。"""
|
||||
"""分组流水线:组内每个视频独立推进阶段,阶段顺序仍按 DAG,组间分批。"""
|
||||
# 数据:3 个视频 + 三节点链路,分组大小 2(前两个一组、第三个一组)。
|
||||
folder = tmp_path / "videos"
|
||||
for name in ("a.mp4", "b.mp4", "c.mp4"):
|
||||
@@ -576,13 +577,12 @@ def test_worker_runs_grouped_stage_pipeline(tmp_path: Path, monkeypatch) -> None
|
||||
job_id = create_job(db, str(folder), "wf")
|
||||
BatchWorker(db, interval_seconds=999)._process_job(db.get_batch_job(job_id))
|
||||
|
||||
# 验证结果:组1 三个阶段各跑 a/b,再轮到组2 的 c。
|
||||
assert trace == [
|
||||
("prep", "a.mp4"), ("prep", "b.mp4"),
|
||||
("translate", "a.mp4"), ("translate", "b.mp4"),
|
||||
("post", "a.mp4"), ("post", "b.mp4"),
|
||||
("prep", "c.mp4"), ("translate", "c.mp4"), ("post", "c.mp4"),
|
||||
]
|
||||
# 验证结果:每个视频的三阶段按 DAG 顺序各跑一次;组 1(a/b)全部跑完才轮到组 2(c)。
|
||||
for name in ("a.mp4", "b.mp4", "c.mp4"):
|
||||
assert [tag for tag, video in trace if video == name] == ["prep", "translate", "post"]
|
||||
first_group = [i for i, entry in enumerate(trace) if entry[1] in ("a.mp4", "b.mp4")]
|
||||
second_group = [i for i, entry in enumerate(trace) if entry[1] == "c.mp4"]
|
||||
assert max(first_group) < min(second_group)
|
||||
# 三个视频都完成且产物按约定名落到视频旁。
|
||||
assert db.get_batch_job(job_id)["status"] == "COMPLETED"
|
||||
assert sorted(p.name for p in folder.glob("*.srt")) == ["a.CN.srt", "b.CN.srt", "c.CN.srt"]
|
||||
@@ -616,6 +616,155 @@ def test_worker_releases_local_llm_once_per_group(tmp_path: Path, monkeypatch) -
|
||||
assert not list((tmp_path / "storage").rglob(KEEP_MODEL_FLAG))
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _timed_nodes(durations: dict[str, float]):
|
||||
"""把三类节点换成"记录起止时间"的假节点,用于断言阶段之间的并行/互斥。
|
||||
|
||||
节点类型决定显存需求(`faster-whisper` 需 GPU、`llm-translate` 看端点、
|
||||
`echo` 不需要),因此可以直接观察资源门控的调度结果。
|
||||
"""
|
||||
import time as _time
|
||||
|
||||
registry.register_all()
|
||||
events: list[tuple[str, str, str, float]] = []
|
||||
|
||||
def handler(request: InvokeRequest) -> InvokeResponse:
|
||||
tag = str(request.params.get("node_tag"))
|
||||
# 首阶段可能从 audio_uri 拿到视频路径,后续阶段从 file_uri 拿到上游产物。
|
||||
source = str(request.inputs.get("file_uri") or request.inputs.get("audio_uri") or "")
|
||||
if source and Path(source).suffix.lower() in VIDEO_EXTENSIONS:
|
||||
video_name = Path(source).name
|
||||
else:
|
||||
video_name = Path(source).read_text(encoding="utf-8").strip() if source else ""
|
||||
events.append((tag, video_name, "start", _time.monotonic()))
|
||||
# 时长可按 "阶段:视频" 细化(用于制造"A 已转写、B 还在提音"的时序)。
|
||||
_time.sleep(durations.get(f"{tag}:{video_name}", durations.get(tag, 0.05)))
|
||||
output_dir = Path(request.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
output = output_dir / "payload.srt"
|
||||
output.write_text(video_name, encoding="utf-8")
|
||||
events.append((tag, video_name, "end", _time.monotonic()))
|
||||
return InvokeResponse(status="completed", outputs={"file_uri": str(output)})
|
||||
|
||||
for node_type in ("echo", "llm-translate", "faster-whisper"):
|
||||
registry.register(registry.get_node(node_type), handler)
|
||||
yield events
|
||||
|
||||
|
||||
def _overlaps(events, first: tuple[str, str], second: tuple[str, str]) -> bool:
|
||||
"""两个 (阶段, 视频) 执行区间是否重叠(用于断言并行/串行)。"""
|
||||
def window(key):
|
||||
start = next(e[3] for e in events if (e[0], e[1], e[2]) == (*key, "start"))
|
||||
end = next(e[3] for e in events if (e[0], e[1], e[2]) == (*key, "end"))
|
||||
return start, end
|
||||
|
||||
a_start, a_end = window(first)
|
||||
b_start, b_end = window(second)
|
||||
return a_start < b_end and b_start < a_end
|
||||
|
||||
|
||||
def _asr_translate_db(tmp_path: Path) -> Database:
|
||||
"""两阶段链路 asr(faster-whisper) → translate(llm) 的临时库。"""
|
||||
db = Database(tmp_path / "wov.db")
|
||||
db.upsert_workflow({
|
||||
"id": "wf", "name": "转写翻译", "description": "", "published": 1, "latest_version": 1,
|
||||
})
|
||||
definition = {
|
||||
"name": "转写翻译", "version": 1,
|
||||
"nodes": [
|
||||
{"id": "asr", "node_type": "faster-whisper", "params": {"node_tag": "asr"},
|
||||
"inputs": {"audio_uri": "input.video_uri"}},
|
||||
{"id": "translate", "node_type": "llm-translate", "params": {"node_tag": "translate"},
|
||||
"inputs": {"file_uri": "asr.file_uri"}},
|
||||
],
|
||||
"edges": [{"from": "asr", "to": "translate"}],
|
||||
"entry_inputs": {"video_uri": "file"},
|
||||
"final_outputs": {"cn_srt": "translate.file_uri"},
|
||||
}
|
||||
db.create_workflow_version("wf", 1, WorkflowDefinition.from_dict(definition).to_dict())
|
||||
return db
|
||||
|
||||
|
||||
def test_remote_translate_overlaps_other_video_transcription(tmp_path: Path, monkeypatch) -> None:
|
||||
"""线上翻译不占 GPU:A 在翻译时,B 的转写可以同时跑(GPU 不空转)。"""
|
||||
# 数据:2 个视频 + asr→translate 链路;翻译端点在远端,显存充裕。
|
||||
folder = tmp_path / "videos"
|
||||
for name in ("a.mp4", "b.mp4"):
|
||||
_make_video(folder / name)
|
||||
db = _asr_translate_db(tmp_path)
|
||||
monkeypatch.setattr("wov_app.batch.BATCH_WORK_ROOT", tmp_path / "storage" / "batch")
|
||||
monkeypatch.setenv("LLM_API_BASE", "https://api.siliconflow.cn/v1/chat/completions")
|
||||
monkeypatch.setattr("wov_app.batch._make_gate", lambda: GpuGate(probe=lambda: 24000))
|
||||
|
||||
# 测试过程
|
||||
with _timed_nodes({"asr": 0.4, "translate": 0.6}) as events:
|
||||
job_id = create_job(db, str(folder), "wf")
|
||||
BatchWorker(db, interval_seconds=999)._process_job(db.get_batch_job(job_id))
|
||||
|
||||
# 验证结果:A 的翻译与 B 的转写时间区间重叠。
|
||||
assert _overlaps(events, ("translate", "a.mp4"), ("asr", "b.mp4"))
|
||||
|
||||
|
||||
def test_local_llm_endpoint_serializes_and_runs_all_asr_first(tmp_path: Path, monkeypatch) -> None:
|
||||
"""本机 LLM 端点占显存:GPU 阶段互斥,且组内先跑完全部转写再翻译(模型只载一次)。"""
|
||||
# 数据:2 个视频 + asr→translate 链路;翻译端点在本机(占整卡显存)。
|
||||
folder = tmp_path / "videos"
|
||||
for name in ("a.mp4", "b.mp4"):
|
||||
_make_video(folder / name)
|
||||
db = _asr_translate_db(tmp_path)
|
||||
monkeypatch.setattr("wov_app.batch.BATCH_WORK_ROOT", tmp_path / "storage" / "batch")
|
||||
monkeypatch.setenv("LLM_API_BASE", "http://localhost:11434/v1/chat/completions")
|
||||
monkeypatch.setattr("wov_app.batch._make_gate", lambda: GpuGate(probe=lambda: 24000))
|
||||
# 本用例只验调度顺序,不必真去卸载本机模型。
|
||||
monkeypatch.setattr("wov_app.batch.release_local_model", lambda model=None: None)
|
||||
|
||||
# 测试过程
|
||||
with _timed_nodes({"asr": 0.3, "translate": 0.3}) as events:
|
||||
job_id = create_job(db, str(folder), "wf")
|
||||
BatchWorker(db, interval_seconds=999)._process_job(db.get_batch_job(job_id))
|
||||
|
||||
# 验证结果:全部转写先于任何翻译;翻译之间互斥不重叠。
|
||||
last_asr = max(e[3] for e in events if e[0] == "asr" and e[2] == "end")
|
||||
first_translate = min(e[3] for e in events if e[0] == "translate" and e[2] == "start")
|
||||
assert last_asr < first_translate
|
||||
assert _overlaps(events, ("translate", "a.mp4"), ("translate", "b.mp4")) is False
|
||||
|
||||
|
||||
def test_extract_overlaps_transcription_of_other_video(tmp_path: Path, monkeypatch) -> None:
|
||||
"""提音(不需要 GPU)与其它视频的转写并行——组内不再"先全部提音再转写"。"""
|
||||
# 数据:2 个视频 + extract(echo) → asr(faster-whisper) 链路。
|
||||
folder = tmp_path / "videos"
|
||||
for name in ("a.mp4", "b.mp4"):
|
||||
_make_video(folder / name)
|
||||
db = Database(tmp_path / "wov.db")
|
||||
db.upsert_workflow({
|
||||
"id": "wf", "name": "提音转写", "description": "", "published": 1, "latest_version": 1,
|
||||
})
|
||||
definition = {
|
||||
"name": "提音转写", "version": 1,
|
||||
"nodes": [
|
||||
{"id": "extract", "node_type": "echo", "params": {"node_tag": "extract"},
|
||||
"inputs": {"file_uri": "input.video_uri"}},
|
||||
{"id": "asr", "node_type": "faster-whisper", "params": {"node_tag": "asr"},
|
||||
"inputs": {"audio_uri": "extract.file_uri"}},
|
||||
],
|
||||
"edges": [{"from": "extract", "to": "asr"}],
|
||||
"entry_inputs": {"video_uri": "file"},
|
||||
"final_outputs": {"cn_srt": "asr.file_uri"},
|
||||
}
|
||||
db.create_workflow_version("wf", 1, WorkflowDefinition.from_dict(definition).to_dict())
|
||||
monkeypatch.setattr("wov_app.batch.BATCH_WORK_ROOT", tmp_path / "storage" / "batch")
|
||||
monkeypatch.setattr("wov_app.batch._make_gate", lambda: GpuGate(probe=lambda: 24000))
|
||||
|
||||
# 测试过程
|
||||
with _timed_nodes({"extract:a.mp4": 0.05, "extract:b.mp4": 0.8, "asr": 0.4}) as events:
|
||||
job_id = create_job(db, str(folder), "wf")
|
||||
BatchWorker(db, interval_seconds=999)._process_job(db.get_batch_job(job_id))
|
||||
|
||||
# 验证结果:A 的转写与 B 的提音重叠。
|
||||
assert _overlaps(events, ("asr", "a.mp4"), ("extract", "b.mp4"))
|
||||
|
||||
|
||||
def test_worker_defers_video_failed_in_earlier_stage(tmp_path: Path, monkeypatch) -> None:
|
||||
"""上一阶段失败的视频不在后续阶段重跑(避免 LLM 已常驻时重跑 ASR 抢显存)。"""
|
||||
# 数据:2 个视频(同一组)+ 三节点链路,prep 阶段让 a 失败。
|
||||
|
||||
Reference in New Issue
Block a user