fix(ceus): 降低长视频内存峰值,快照开跑即写,build 核对 CEUS 运行与交付包

- 内存:prepare_segmentation 改为流式解码,只留分屏和配准的抽样帧,绿框逐帧在两半上找;
  分割那一半由 SegmentationInput.target_frames() 到用时再解码,segment_video 边读边写 jpg;
  MedSAM2 init_state 的读帧换成直接在存放设备上建张量(不改 MedSAM2 源码,结果逐位相同);
  传播结果写进预分配数组。帧数最多的一段视频常驻内存峰值 11.2 GiB → 2.3 GiB,分辨率最大的
  一段原代码被 OOM 杀掉、现为 4.2 GiB;两段的掩码和全部产物与改前逐字节相同,33 段视频的
  第 1–2 步结果与改前一致,29 例配准与 review.csv 一致。
- 快照:ceus_quant 开跑时写 snapshot/(git、ceus.yaml、review.csv、uv.lock、segment.json),
  配置和核对表从快照副本读;收尾由新函数 runs.finish_snapshot 只新建 finished.json。
  build 与 ceus_align_review 同样改用 finish_snapshot 记结束时间。
- build:summary.json 加 video_sha256;装 CEUS 结果时核对快照里的 delivery_dir、逐例核对
  视频 SHA-256 和推注时间,不同就整次停下;新字段做类型检查,不进 tic.json。

来源:主会话任务(问题 1–3 及补充的 delivery_dir 比对)。

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
2026-10-16 09:01:44 -04:00
parent e165b32b9d
commit 1de88a5d0a
23 changed files with 561 additions and 202 deletions
+32 -22
View File
@@ -7,6 +7,8 @@
第 3 步(MedSAM2 分割)的结果按内容键缓存在 `<work_root>/segment_cache/<键>/masks.npz`,键相同就
直接读缓存、不调显卡。某一例在第 1–6 步失败(`CeusCaseError`)记进 `failed.csv`,接着跑下一例;
其余异常整次运行停下,留下 `.partial` 目录。终端只打印计数,不打印视频号。
快照在开跑时写进 `.partial/snapshot/`(代码、配置和核对表原文、权重与 MedSAM2 版本),配置和核对表
随后从快照里的副本读;收尾只补 `finished.json`(结束时间、显存峰值、分割缓存命中数)。
"""
import argparse
@@ -44,7 +46,7 @@ from cadx.ceus import (
tic_plot_markers,
)
from cadx.delivery import CaseRecord, read_delivery
from cadx.runs import make_build_id, read_git_state, write_snapshot
from cadx.runs import finish_snapshot, make_build_id, read_git_state, write_snapshot
from cadx.segment import SegmentConfig, cuda_peak_memory_bytes, medsam2_commit, segment_video
REPO_ROOT = Path(__file__).resolve().parents[1]
@@ -241,7 +243,7 @@ def _masks(
if final.exists():
with np.load(final / "masks.npz", allow_pickle=False) as data:
return data["masks"], True
masks = segment_video(prepared.target_frames, prepared.prompt_masks, segment_cfg)
masks = segment_video(prepared.target_frames(), prepared.prompt_masks, segment_cfg)
partial = cache_root / f"{key}.partial"
# 上次中途停下留下的 .partial 会让这里报错:work/ 可以删掉重算,删掉它再跑。
partial.mkdir(parents=True, exist_ok=False)
@@ -283,6 +285,7 @@ def _run_case(
registration = prepared.registration
summary = {
"video_id": case.video_id,
"video_sha256": video_sha256,
"segment_cache_key": key,
"registration": {
"contrast_side": prepared.contrast.side,
@@ -323,10 +326,31 @@ def main(argv: Sequence[str]) -> None:
raise FileExistsError(f"运行目录已存在,不覆盖:{path}")
partial.mkdir(parents=True, exist_ok=False)
review = read_review(cfg.paths.align_review)
delivery = read_delivery(cfg.paths.delivery_dir, None)
checkpoint_sha256 = file_sha256(cfg.segment.checkpoint)
commit = medsam2_commit(cfg.segment.medsam2_dir)
# 开跑时就写快照:运行期间改了 ceus.yaml 或 review.csv,快照里仍是这次用的那份。
snapshot_dir = partial / "snapshot"
write_snapshot(
snapshot_dir=snapshot_dir,
build_id=run_id,
git=git,
started_at=started_at,
uv_lock=REPO_ROOT / "uv.lock",
configs={"ceus.yaml": config_path, "align_review.csv": cfg.paths.align_review},
)
segment_record = {
"checkpoint_sha256": checkpoint_sha256,
"medsam2_commit": commit,
"device": cfg.segment.device,
}
(snapshot_dir / "segment.json").write_text(
json.dumps(segment_record, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
)
# 配置和核对表读快照里的副本,与快照记的是同一份字节。
written = load_config(snapshot_dir / "ceus.yaml")
cfg = resolve_paths(written)
review = read_review(snapshot_dir / "align_review.csv")
delivery = read_delivery(cfg.paths.delivery_dir, None)
# 缓存键用配置文件里的写法(路径不随仓库位置变)。
def key_of(video_sha256: str) -> str:
@@ -352,24 +376,10 @@ def main(argv: Sequence[str]) -> None:
writer.writerow(_FAILED_HEADER)
writer.writerows(failures)
peak = cuda_peak_memory_bytes(cfg.segment.device)
write_snapshot(
snapshot_dir=partial / "snapshot",
build_id=run_id,
git=git,
started_at=started_at,
finished_at=datetime.now().astimezone(),
uv_lock=REPO_ROOT / "uv.lock",
configs={"ceus.yaml": config_path, "align_review.csv": cfg.paths.align_review},
)
segment_record = {
"checkpoint_sha256": checkpoint_sha256,
"medsam2_commit": commit,
"device": cfg.segment.device,
"cuda_peak_memory_bytes": peak,
"segment_cache_hits": cache_hits,
}
(partial / "snapshot" / "segment.json").write_text(
json.dumps(segment_record, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
finish_snapshot(
snapshot_dir,
datetime.now().astimezone(),
{"cuda_peak_memory_bytes": peak, "segment_cache_hits": cache_hits},
)
# POSIX 的 rename 会悄悄覆盖空目录,所以改名前再查一次。
if run_dir.exists():