raylight: 任务记录运行期间单卡显存峰值（peak_vram_gb），并合入服务器上的参考图缩放改动

- Job 新增 peak_vram_gb；worker 运行任务时起采样线程，每 3s 读 nvidia-smi 取四卡最大占用 - 合并服务器上先于仓库的改动（reference_image_urls 上限 9、ref_image_short_edge、reference_geometry） - 实测 1024x576/120 帧文生：峰值 11.1 GB，console 详情页可见 Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

raylight: 任务记录运行期间单卡显存峰值（peak_vram_gb），并合入服务器上的参考图缩放改动
- Job 新增 peak_vram_gb；worker 运行任务时起采样线程，每 3s 读 nvidia-smi 取四卡最大占用 - 合并服务器上先于仓库的改动（reference_image_urls 上限 9、ref_image_short_edge、reference_geometry） - 实测 1024x576/120 帧文生：峰值 11.1 GB，console 详情页可见 Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
yibo
1 parent 854911f3
Showing 1 changed file with 98 additions and 3 deletions
raylight/scripts/server.py
@@ -12,6 +12,9 @@ RAY_INIT_CFG 是全局常量，任何请求都不应该改它。
 import asyncio
 import contextlib
 import json
+import hashlib
+import io
+import math
 import os
 import queue
 import re
@@ -30,6 +33,7 @@ import websockets
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import FileResponse
 from pydantic import BaseModel, Field, model_validator
+from PIL import Image, ImageOps
  
 COMFY_HOST = os.environ.get("COMFY_HOST", "127.0.0.1")
 COMFY_PORT = os.environ.get("COMFY_PORT", "8188")
@@ -130,7 +134,8 @@ class GenReq(BaseModel):
     negative_prompt: Optional[str] = None   # 同上，暂不生效，见 notes.md
     first_frame_url: Optional[str] = None
     last_frame_url: Optional[str] = None
-    reference_image_urls: list[str] = []
+    reference_image_urls: list[str] = Field(default_factory=list, max_length=9)
+    ref_image_short_edge: Optional[int] = Field(None, ge=64, le=2048, multiple_of=32)
     # 2026-08-23：MiniMaxH3ReferenceToVideo 节点本身支持参考视频/独立音频参考（ComfyUI
     # /object_info 里的 ref_videos/ref_audios），之前只接了 ref_images 这一种，见 notes.md。
     reference_video_urls: list[str] = Field(default_factory=list, max_length=3)  # 最多3段，
@@ -168,6 +173,7 @@ class Job(BaseModel):
     backend: str = BACKEND_URL
     queue_seq: int = 0  # 2026-09-03: 排队顺序+插队用，见 _queue_position/prioritize_video
     restart_recoveries: int = 0  # 2026-09-09: 见 _load_jobs 里 status=='running' 分支的注释
+    peak_vram_gb: Optional[float] = None  # 2026-10-31: 任务运行期间四卡里单卡显存占用的峰值（GB），见 _sample_peak_vram
  
  
 JOBS: dict[str, Job] = {}
@@ -336,6 +342,18 @@ def build_fl2va_graph(req: GenReq) -&gt; dict:
     return g
  
  
+def reference_geometry(req: GenReq, width: int, height: int) -> dict:
+    """The explicit resize and subsequent node 'match' sizing share one calculation."""
+    scale = min(1.0, req.ref_image_short_edge / min(width, height)) if req.ref_image_short_edge else 1.0
+    rw = max(32, round(width * scale / 32) * 32) if req.ref_image_short_edge else width
+    rh = max(32, round(height * scale / 32) * 32) if req.ref_image_short_edge else height
+    match = min(1.0, math.sqrt((req.width * req.height) / (rw * rh)))
+    cw = max(32, round(rw * match / 32) * 32)
+    ch = max(32, round(rh * match / 32) * 32)
+    return {"source_width": width, "source_height": height, "resize_width": rw, "resize_height": rh,
+            "conditioning_width": cw, "conditioning_height": ch, "reference_tokens": (cw // 32) * (ch // 32)}
+
+
 def build_ref2va_graph(req: GenReq) -> dict:
     """多参考图生视频 —— MiniMaxH3ReferenceToVideo，最多9张参考图（协议规格上限）。"""
     g = _static_nodes()
@@ -348,10 +366,16 @@ def build_ref2va_graph(req: GenReq) -&gt; dict:
         "audio_vae": ["s_vae_a_gpu", 0], "width": req.width, "height": req.height,
         "length": req.frames, "ref_image_size": "match",
     }
-    for i, url in enumerate(req.reference_image_urls[:9]):
+    for i, url in enumerate(req.reference_image_urls):
         fname = _fetch_to_input(url)
         nid = f"n_ref{i}"
         g[nid] = {"class_type": "LoadImage", "inputs": {"image": fname}}
+        if req.ref_image_short_edge is not None:
+            with Image.open(INPUT_DIR / fname) as image:
+                geometry = reference_geometry(req, *ImageOps.exif_transpose(image).size)
+            resized = f"{nid}_resize"
+            g[resized] = {"class_type": "ImageScale", "inputs": {"image": [nid, 0], "upscale_method": "lanczos", "width": geometry["resize_width"], "height": geometry["resize_height"], "crop": "disabled"}}
+            nid = resized
         conditioning_inputs[f"ref_images.ref_image_{i}"] = [nid, 0]
     for i, url in enumerate(req.reference_video_urls[:3]):
         fname = _fetch_to_input(url, kind="video")
@@ -508,6 +532,19 @@ async def _run_job_async(j: Job):
     j.size_bytes = fpath.stat().st_size
  
  
+
+def _sample_peak_vram(j: Job, stop: threading.Event, interval: float = 3.0):
+    # 2026-10-31：console 详情页的"峰值显存"一直是"—"，后端根本没这个字段。运行期间每隔几秒读一次
+    # nvidia-smi，记四卡里单卡占用的最大值。注意是整张卡的占用（含别的进程），不是只算本任务；
+    # 采样间隔内的瞬时尖峰会漏掉，所以是"下限估计"——要精确得上 torch.cuda.max_memory_allocated。
+    while not stop.is_set():
+        used = [g["memory_used_mb"] for g in _gpu_status()]
+        if used:
+            peak = round(max(used) / 1024, 1)
+            if j.peak_vram_gb is None or peak > j.peak_vram_gb:
+                j.peak_vram_gb = peak
+        stop.wait(interval)
+
 def worker():
     STATE["model"] = "ready"  # ComfyUI 自己管加载，这里没有独立的"模型加载"阶段
     STATE["loaded_at"] = time.time()
@@ -526,6 +563,9 @@ def worker():
         j.status = "running"
         j.started_at = time.time()
         _persist(j)
+        j.peak_vram_gb = None  # 重跑/重启恢复的任务不沿用上一轮的峰值
+        stop_vram = threading.Event()
+        threading.Thread(target=_sample_peak_vram, args=(j, stop_vram), daemon=True).start()
         try:
             loop.run_until_complete(_run_job_async(j))
             j.status = "done"
@@ -534,6 +574,7 @@ def worker():
             j.error = f"{type(e).__name__}: {e}"
             print(f"[job {jid}] failed: {j.error}")
         finally:
+            stop_vram.set()
             j.finished_at = time.time()
             _persist(j)
             Q.task_done()
@@ -766,9 +807,63 @@ def list_videos(limit: int = 50):
              "size_bytes": j.size_bytes, "error": j.error} for j in jobs]
  
  
+@app.post("/v1/videos/validate")
+def validate_video(payload: dict):
+    """Resolve actual images and reserve future-frame geometry; never enqueue or write media."""
+    try:
+        req = GenReq.model_validate(payload.get("request"))
+    except ValueError:
+        raise HTTPException(422, "invalid generation request")
+    reserved = payload.get("reserved_references", [])
+    if not isinstance(reserved, list) or len(req.reference_image_urls) + len(reserved) > 9:
+        raise HTTPException(422, "at most nine actual and reserved reference images")
+    ceiling = RESOLUTION_CEILINGS.get((req.width, req.height))
+    if ceiling is None or req.width % 32 or req.height % 32 or req.frames % 17 != 5:
+        raise HTTPException(422, "unsupported resolution or frame grid")
+    if req.first_frame_url or req.last_frame_url or req.reference_video_urls or req.reference_audio_urls:
+        raise HTTPException(422, "this validation contract supports image-reference requests only")
+    if req.turbo or req.cfg_scale is not None or req.negative_prompt is not None:
+        raise HTTPException(422, "requested option is not implemented by this deployment")
+    references = []
+    for url in req.reference_image_urls:
+        if urllib.parse.urlsplit(url).scheme not in ("http", "https"):
+            raise HTTPException(422, "invalid reference URL scheme")
+        try:
+            with urllib.request.urlopen(url, timeout=20) as response:
+                data = response.read(10 * 1024 * 1024 + 1)
+            if len(data) > 10 * 1024 * 1024:
+                raise ValueError("reference too large")
+            with Image.open(io.BytesIO(data)) as image:
+                geometry = reference_geometry(req, *ImageOps.exif_transpose(image).size)
+        except Exception:
+            raise HTTPException(422, "reference could not be read as a bounded image")
+        references.append({"kind": "resolved_image", "sha256": hashlib.sha256(data).hexdigest(), **geometry})
+    for entry in reserved:
+        if not isinstance(entry, dict) or any(type(entry.get(k)) is not int or not 32 <= entry[k] <= 2048 for k in ("width", "height")):
+            raise HTTPException(422, "invalid reserved frame geometry")
+        references.append({"kind": "reserved_runtime_frame", **reference_geometry(req, entry["width"], entry["height"])})
+    video_tokens = (2 + ((req.frames - 5) // 17) * 5) * (req.width // 32) * (req.height // 32)
+    reference_tokens = sum(item["reference_tokens"] for item in references)
+    prompt_reserve = len(req.prompt.encode("utf-8"))
+    envelope = (2 + ((ceiling - 5) // 17) * 5) * (req.width // 32) * (req.height // 32)
+    if video_tokens + reference_tokens + prompt_reserve > envelope:
+        raise HTTPException(422, "combined image/video/prompt reserve exceeds deployed resolution envelope")
+    normalized = req.model_dump()
+    return {"schema": "h3-request-validation/v1", "validated": True, "observed_at": time.time(),
+            "normalized_request_sha256": hashlib.sha256(json.dumps(normalized, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode()).hexdigest(),
+            "parameters": {k: normalized[k] for k in ("width", "height", "frames", "steps", "turbo", "seed", "ref_image_short_edge")},
+            "references": references, "budget": {"video_tokens": video_tokens, "reference_tokens": reference_tokens,
+            "prompt_utf8_reserve": prompt_reserve, "resolution_envelope": envelope,
+            "basis": "deployed_resolution_ceiling_equivalent_video_tokens_with_reference_and_prompt_reserve"},
+            "runtime_dependencies_pending": bool(reserved), "inference_success_guaranteed": False,
+            "job_created": False, "media_written": False}
+
+
 @app.get("/v1/resolutions")
 def list_resolutions():
     """实测过显存/画质边界的分辨率档位——RESOLUTION_CEILINGS 的只读镜像，给 h3-console
     的 /v1/models 之类的外部查询接口用，不要在别处重复硬编码这份表。"""
     return {"resolutions": [{"width": w, "height": h, "max_frames": max_frames}
-                             for (w, h), max_frames in RESOLUTION_CEILINGS.items()]}
+                             for (w, h), max_frames in RESOLUTION_CEILINGS.items()],
+            "request_schema": GenReq.model_json_schema(),
+            "request_validation": {"path": "/v1/videos/validate", "schema": "h3-request-validation/v1", "reference_sizing": "explicit_short_edge_then_node_match"}}