Commit 5dab65566a98860d981734b788b38102c33cf2aa

Authored by yibo
1 parent 854911f3

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>
Showing 1 changed file with 98 additions and 3 deletions
raylight/scripts/server.py
@@ -12,6 +12,9 @@ RAY_INIT_CFG 是全局常量,任何请求都不应该改它。 @@ -12,6 +12,9 @@ RAY_INIT_CFG 是全局常量,任何请求都不应该改它。
12 import asyncio 12 import asyncio
13 import contextlib 13 import contextlib
14 import json 14 import json
  15 +import hashlib
  16 +import io
  17 +import math
15 import os 18 import os
16 import queue 19 import queue
17 import re 20 import re
@@ -30,6 +33,7 @@ import websockets @@ -30,6 +33,7 @@ import websockets
30 from fastapi import FastAPI, HTTPException 33 from fastapi import FastAPI, HTTPException
31 from fastapi.responses import FileResponse 34 from fastapi.responses import FileResponse
32 from pydantic import BaseModel, Field, model_validator 35 from pydantic import BaseModel, Field, model_validator
  36 +from PIL import Image, ImageOps
33 37
34 COMFY_HOST = os.environ.get("COMFY_HOST", "127.0.0.1") 38 COMFY_HOST = os.environ.get("COMFY_HOST", "127.0.0.1")
35 COMFY_PORT = os.environ.get("COMFY_PORT", "8188") 39 COMFY_PORT = os.environ.get("COMFY_PORT", "8188")
@@ -130,7 +134,8 @@ class GenReq(BaseModel): @@ -130,7 +134,8 @@ class GenReq(BaseModel):
130 negative_prompt: Optional[str] = None # 同上,暂不生效,见 notes.md 134 negative_prompt: Optional[str] = None # 同上,暂不生效,见 notes.md
131 first_frame_url: Optional[str] = None 135 first_frame_url: Optional[str] = None
132 last_frame_url: Optional[str] = None 136 last_frame_url: Optional[str] = None
133 - reference_image_urls: list[str] = [] 137 + reference_image_urls: list[str] = Field(default_factory=list, max_length=9)
  138 + ref_image_short_edge: Optional[int] = Field(None, ge=64, le=2048, multiple_of=32)
134 # 2026-08-23:MiniMaxH3ReferenceToVideo 节点本身支持参考视频/独立音频参考(ComfyUI 139 # 2026-08-23:MiniMaxH3ReferenceToVideo 节点本身支持参考视频/独立音频参考(ComfyUI
135 # /object_info 里的 ref_videos/ref_audios),之前只接了 ref_images 这一种,见 notes.md。 140 # /object_info 里的 ref_videos/ref_audios),之前只接了 ref_images 这一种,见 notes.md。
136 reference_video_urls: list[str] = Field(default_factory=list, max_length=3) # 最多3段, 141 reference_video_urls: list[str] = Field(default_factory=list, max_length=3) # 最多3段,
@@ -168,6 +173,7 @@ class Job(BaseModel): @@ -168,6 +173,7 @@ class Job(BaseModel):
168 backend: str = BACKEND_URL 173 backend: str = BACKEND_URL
169 queue_seq: int = 0 # 2026-09-03: 排队顺序+插队用,见 _queue_position/prioritize_video 174 queue_seq: int = 0 # 2026-09-03: 排队顺序+插队用,见 _queue_position/prioritize_video
170 restart_recoveries: int = 0 # 2026-09-09: 见 _load_jobs 里 status=='running' 分支的注释 175 restart_recoveries: int = 0 # 2026-09-09: 见 _load_jobs 里 status=='running' 分支的注释
  176 + peak_vram_gb: Optional[float] = None # 2026-10-31: 任务运行期间四卡里单卡显存占用的峰值(GB),见 _sample_peak_vram
171 177
172 178
173 JOBS: dict[str, Job] = {} 179 JOBS: dict[str, Job] = {}
@@ -336,6 +342,18 @@ def build_fl2va_graph(req: GenReq) -&gt; dict: @@ -336,6 +342,18 @@ def build_fl2va_graph(req: GenReq) -&gt; dict:
336 return g 342 return g
337 343
338 344
  345 +def reference_geometry(req: GenReq, width: int, height: int) -> dict:
  346 + """The explicit resize and subsequent node 'match' sizing share one calculation."""
  347 + scale = min(1.0, req.ref_image_short_edge / min(width, height)) if req.ref_image_short_edge else 1.0
  348 + rw = max(32, round(width * scale / 32) * 32) if req.ref_image_short_edge else width
  349 + rh = max(32, round(height * scale / 32) * 32) if req.ref_image_short_edge else height
  350 + match = min(1.0, math.sqrt((req.width * req.height) / (rw * rh)))
  351 + cw = max(32, round(rw * match / 32) * 32)
  352 + ch = max(32, round(rh * match / 32) * 32)
  353 + return {"source_width": width, "source_height": height, "resize_width": rw, "resize_height": rh,
  354 + "conditioning_width": cw, "conditioning_height": ch, "reference_tokens": (cw // 32) * (ch // 32)}
  355 +
  356 +
339 def build_ref2va_graph(req: GenReq) -> dict: 357 def build_ref2va_graph(req: GenReq) -> dict:
340 """多参考图生视频 —— MiniMaxH3ReferenceToVideo,最多9张参考图(协议规格上限)。""" 358 """多参考图生视频 —— MiniMaxH3ReferenceToVideo,最多9张参考图(协议规格上限)。"""
341 g = _static_nodes() 359 g = _static_nodes()
@@ -348,10 +366,16 @@ def build_ref2va_graph(req: GenReq) -&gt; dict: @@ -348,10 +366,16 @@ def build_ref2va_graph(req: GenReq) -&gt; dict:
348 "audio_vae": ["s_vae_a_gpu", 0], "width": req.width, "height": req.height, 366 "audio_vae": ["s_vae_a_gpu", 0], "width": req.width, "height": req.height,
349 "length": req.frames, "ref_image_size": "match", 367 "length": req.frames, "ref_image_size": "match",
350 } 368 }
351 - for i, url in enumerate(req.reference_image_urls[:9]): 369 + for i, url in enumerate(req.reference_image_urls):
352 fname = _fetch_to_input(url) 370 fname = _fetch_to_input(url)
353 nid = f"n_ref{i}" 371 nid = f"n_ref{i}"
354 g[nid] = {"class_type": "LoadImage", "inputs": {"image": fname}} 372 g[nid] = {"class_type": "LoadImage", "inputs": {"image": fname}}
  373 + if req.ref_image_short_edge is not None:
  374 + with Image.open(INPUT_DIR / fname) as image:
  375 + geometry = reference_geometry(req, *ImageOps.exif_transpose(image).size)
  376 + resized = f"{nid}_resize"
  377 + g[resized] = {"class_type": "ImageScale", "inputs": {"image": [nid, 0], "upscale_method": "lanczos", "width": geometry["resize_width"], "height": geometry["resize_height"], "crop": "disabled"}}
  378 + nid = resized
355 conditioning_inputs[f"ref_images.ref_image_{i}"] = [nid, 0] 379 conditioning_inputs[f"ref_images.ref_image_{i}"] = [nid, 0]
356 for i, url in enumerate(req.reference_video_urls[:3]): 380 for i, url in enumerate(req.reference_video_urls[:3]):
357 fname = _fetch_to_input(url, kind="video") 381 fname = _fetch_to_input(url, kind="video")
@@ -508,6 +532,19 @@ async def _run_job_async(j: Job): @@ -508,6 +532,19 @@ async def _run_job_async(j: Job):
508 j.size_bytes = fpath.stat().st_size 532 j.size_bytes = fpath.stat().st_size
509 533
510 534
  535 +
  536 +def _sample_peak_vram(j: Job, stop: threading.Event, interval: float = 3.0):
  537 + # 2026-10-31:console 详情页的"峰值显存"一直是"—",后端根本没这个字段。运行期间每隔几秒读一次
  538 + # nvidia-smi,记四卡里单卡占用的最大值。注意是整张卡的占用(含别的进程),不是只算本任务;
  539 + # 采样间隔内的瞬时尖峰会漏掉,所以是"下限估计"——要精确得上 torch.cuda.max_memory_allocated。
  540 + while not stop.is_set():
  541 + used = [g["memory_used_mb"] for g in _gpu_status()]
  542 + if used:
  543 + peak = round(max(used) / 1024, 1)
  544 + if j.peak_vram_gb is None or peak > j.peak_vram_gb:
  545 + j.peak_vram_gb = peak
  546 + stop.wait(interval)
  547 +
511 def worker(): 548 def worker():
512 STATE["model"] = "ready" # ComfyUI 自己管加载,这里没有独立的"模型加载"阶段 549 STATE["model"] = "ready" # ComfyUI 自己管加载,这里没有独立的"模型加载"阶段
513 STATE["loaded_at"] = time.time() 550 STATE["loaded_at"] = time.time()
@@ -526,6 +563,9 @@ def worker(): @@ -526,6 +563,9 @@ def worker():
526 j.status = "running" 563 j.status = "running"
527 j.started_at = time.time() 564 j.started_at = time.time()
528 _persist(j) 565 _persist(j)
  566 + j.peak_vram_gb = None # 重跑/重启恢复的任务不沿用上一轮的峰值
  567 + stop_vram = threading.Event()
  568 + threading.Thread(target=_sample_peak_vram, args=(j, stop_vram), daemon=True).start()
529 try: 569 try:
530 loop.run_until_complete(_run_job_async(j)) 570 loop.run_until_complete(_run_job_async(j))
531 j.status = "done" 571 j.status = "done"
@@ -534,6 +574,7 @@ def worker(): @@ -534,6 +574,7 @@ def worker():
534 j.error = f"{type(e).__name__}: {e}" 574 j.error = f"{type(e).__name__}: {e}"
535 print(f"[job {jid}] failed: {j.error}") 575 print(f"[job {jid}] failed: {j.error}")
536 finally: 576 finally:
  577 + stop_vram.set()
537 j.finished_at = time.time() 578 j.finished_at = time.time()
538 _persist(j) 579 _persist(j)
539 Q.task_done() 580 Q.task_done()
@@ -766,9 +807,63 @@ def list_videos(limit: int = 50): @@ -766,9 +807,63 @@ def list_videos(limit: int = 50):
766 "size_bytes": j.size_bytes, "error": j.error} for j in jobs] 807 "size_bytes": j.size_bytes, "error": j.error} for j in jobs]
767 808
768 809
  810 +@app.post("/v1/videos/validate")
  811 +def validate_video(payload: dict):
  812 + """Resolve actual images and reserve future-frame geometry; never enqueue or write media."""
  813 + try:
  814 + req = GenReq.model_validate(payload.get("request"))
  815 + except ValueError:
  816 + raise HTTPException(422, "invalid generation request")
  817 + reserved = payload.get("reserved_references", [])
  818 + if not isinstance(reserved, list) or len(req.reference_image_urls) + len(reserved) > 9:
  819 + raise HTTPException(422, "at most nine actual and reserved reference images")
  820 + ceiling = RESOLUTION_CEILINGS.get((req.width, req.height))
  821 + if ceiling is None or req.width % 32 or req.height % 32 or req.frames % 17 != 5:
  822 + raise HTTPException(422, "unsupported resolution or frame grid")
  823 + if req.first_frame_url or req.last_frame_url or req.reference_video_urls or req.reference_audio_urls:
  824 + raise HTTPException(422, "this validation contract supports image-reference requests only")
  825 + if req.turbo or req.cfg_scale is not None or req.negative_prompt is not None:
  826 + raise HTTPException(422, "requested option is not implemented by this deployment")
  827 + references = []
  828 + for url in req.reference_image_urls:
  829 + if urllib.parse.urlsplit(url).scheme not in ("http", "https"):
  830 + raise HTTPException(422, "invalid reference URL scheme")
  831 + try:
  832 + with urllib.request.urlopen(url, timeout=20) as response:
  833 + data = response.read(10 * 1024 * 1024 + 1)
  834 + if len(data) > 10 * 1024 * 1024:
  835 + raise ValueError("reference too large")
  836 + with Image.open(io.BytesIO(data)) as image:
  837 + geometry = reference_geometry(req, *ImageOps.exif_transpose(image).size)
  838 + except Exception:
  839 + raise HTTPException(422, "reference could not be read as a bounded image")
  840 + references.append({"kind": "resolved_image", "sha256": hashlib.sha256(data).hexdigest(), **geometry})
  841 + for entry in reserved:
  842 + 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")):
  843 + raise HTTPException(422, "invalid reserved frame geometry")
  844 + references.append({"kind": "reserved_runtime_frame", **reference_geometry(req, entry["width"], entry["height"])})
  845 + video_tokens = (2 + ((req.frames - 5) // 17) * 5) * (req.width // 32) * (req.height // 32)
  846 + reference_tokens = sum(item["reference_tokens"] for item in references)
  847 + prompt_reserve = len(req.prompt.encode("utf-8"))
  848 + envelope = (2 + ((ceiling - 5) // 17) * 5) * (req.width // 32) * (req.height // 32)
  849 + if video_tokens + reference_tokens + prompt_reserve > envelope:
  850 + raise HTTPException(422, "combined image/video/prompt reserve exceeds deployed resolution envelope")
  851 + normalized = req.model_dump()
  852 + return {"schema": "h3-request-validation/v1", "validated": True, "observed_at": time.time(),
  853 + "normalized_request_sha256": hashlib.sha256(json.dumps(normalized, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode()).hexdigest(),
  854 + "parameters": {k: normalized[k] for k in ("width", "height", "frames", "steps", "turbo", "seed", "ref_image_short_edge")},
  855 + "references": references, "budget": {"video_tokens": video_tokens, "reference_tokens": reference_tokens,
  856 + "prompt_utf8_reserve": prompt_reserve, "resolution_envelope": envelope,
  857 + "basis": "deployed_resolution_ceiling_equivalent_video_tokens_with_reference_and_prompt_reserve"},
  858 + "runtime_dependencies_pending": bool(reserved), "inference_success_guaranteed": False,
  859 + "job_created": False, "media_written": False}
  860 +
  861 +
769 @app.get("/v1/resolutions") 862 @app.get("/v1/resolutions")
770 def list_resolutions(): 863 def list_resolutions():
771 """实测过显存/画质边界的分辨率档位——RESOLUTION_CEILINGS 的只读镜像,给 h3-console 864 """实测过显存/画质边界的分辨率档位——RESOLUTION_CEILINGS 的只读镜像,给 h3-console
772 的 /v1/models 之类的外部查询接口用,不要在别处重复硬编码这份表。""" 865 的 /v1/models 之类的外部查询接口用,不要在别处重复硬编码这份表。"""
773 return {"resolutions": [{"width": w, "height": h, "max_frames": max_frames} 866 return {"resolutions": [{"width": w, "height": h, "max_frames": max_frames}
774 - for (w, h), max_frames in RESOLUTION_CEILINGS.items()]} 867 + for (w, h), max_frames in RESOLUTION_CEILINGS.items()],
  868 + "request_schema": GenReq.model_json_schema(),
  869 + "request_validation": {"path": "/v1/videos/validate", "schema": "h3-request-validation/v1", "reference_sizing": "explicit_short_edge_then_node_match"}}