chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

refactor(multimodal/test): centralize model names and shared utilities in test_utils (#19354)

Browse Source 
Co-authored-by: Ratish P <114130421+Ratish1@users.noreply.github.com>
This commit is contained in:
Mook
2026-03-05 20:09:42 -08:00
committed by GitHub
parent 51e5dc845a
commit be9a9e4819
4 changed files with 221 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/multimodal_gen/test/cli/test_generate_t2i_perf.py
View File

@@ -5,12 +5,13 @@ import unittest
from sglang.multimodal_gen.configs.sample.sampling_params import DataType
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.test.cli.test_generate_common import CLIBase
from sglang.multimodal_gen.test.test_utils import DEFAULT_FLUX_1_DEV_MODEL_NAME_FOR_TEST
logger = init_logger(__name__)
class TestFlux_T2V(CLIBase):
model_path = "black-forest-labs/FLUX.1-dev"
model_path = DEFAULT_FLUX_1_DEV_MODEL_NAME_FOR_TEST
extra_args = []
data_type: DataType = DataType.IMAGE

17
python/sglang/multimodal_gen/test/server/test_update_weights_from_disk.py
View File

@@ -159,7 +159,14 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.test.server.test_server_utils import (
ServerManager,
)
from sglang.multimodal_gen.test.test_utils import get_dynamic_server_port, is_in_ci
from sglang.multimodal_gen.test.test_utils import (
DEFAULT_FLUX_2_KLEIN_4B_MODEL_NAME_FOR_TEST,
DEFAULT_FLUX_2_KLEIN_BASE_4B_MODEL_NAME_FOR_TEST,
DEFAULT_QWEN_IMAGE_2512_MODEL_NAME_FOR_TEST,
DEFAULT_QWEN_IMAGE_MODEL_NAME_FOR_TEST,
get_dynamic_server_port,
is_in_ci,
)
logger = init_logger(__name__)
@@ -175,12 +182,12 @@ _DIFFERING_MODULES: list[str] = [_TRANSFORMER_MODULE, _VAE_MODULE]
_ALL_MODEL_PAIRS: list[tuple[str, str]] = [
(
"black-forest-labs/FLUX.2-klein-base-4B",
"black-forest-labs/FLUX.2-klein-4B",
DEFAULT_FLUX_2_KLEIN_BASE_4B_MODEL_NAME_FOR_TEST,
DEFAULT_FLUX_2_KLEIN_4B_MODEL_NAME_FOR_TEST,
),
(
"Qwen/Qwen-Image",
"Qwen/Qwen-Image-2512",
DEFAULT_QWEN_IMAGE_MODEL_NAME_FOR_TEST,
DEFAULT_QWEN_IMAGE_2512_MODEL_NAME_FOR_TEST,
),
]

101
python/sglang/multimodal_gen/test/server/testcase_configs.py
View File

@@ -27,8 +27,24 @@ from typing import Sequence
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord
DEFAULT_SMALL_MODEL = "Tongyi-MAI/Z-Image-Turbo"
from sglang.multimodal_gen.test.test_utils import (
DEFAULT_FLUX_1_DEV_MODEL_NAME_FOR_TEST,
DEFAULT_FLUX_2_DEV_MODEL_NAME_FOR_TEST,
DEFAULT_FLUX_2_KLEIN_4B_MODEL_NAME_FOR_TEST,
DEFAULT_QWEN_IMAGE_EDIT_2509_MODEL_NAME_FOR_TEST,
DEFAULT_QWEN_IMAGE_EDIT_2511_MODEL_NAME_FOR_TEST,
DEFAULT_QWEN_IMAGE_EDIT_MODEL_NAME_FOR_TEST,
DEFAULT_QWEN_IMAGE_LAYERED_MODEL_NAME_FOR_TEST,
DEFAULT_QWEN_IMAGE_MODEL_NAME_FOR_TEST,
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
DEFAULT_WAN_2_1_I2V_14B_480P_MODEL_NAME_FOR_TEST,
DEFAULT_WAN_2_1_I2V_14B_720P_MODEL_NAME_FOR_TEST,
DEFAULT_WAN_2_1_T2V_1_3B_MODEL_NAME_FOR_TEST,
DEFAULT_WAN_2_1_T2V_14B_MODEL_NAME_FOR_TEST,
DEFAULT_WAN_2_2_I2V_A14B_MODEL_NAME_FOR_TEST,
DEFAULT_WAN_2_2_T2V_A14B_MODEL_NAME_FOR_TEST,
DEFAULT_WAN_2_2_TI2V_5B_MODEL_NAME_FOR_TEST,
)
@dataclass
@@ -301,8 +317,6 @@ class PerformanceSummary:
)
SMALL_T2I_MODEL = "Tongyi-MAI/Z-Image-Turbo"
T2I_sampling_params = DiffusionSamplingParams(
prompt="Doraemon is eating dorayaki",
output_size="1024x1024",
@@ -362,7 +376,7 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
DiffusionTestCase(
"qwen_image_t2i",
DiffusionServerArgs(
model_path="Qwen/Qwen-Image",
model_path=DEFAULT_QWEN_IMAGE_MODEL_NAME_FOR_TEST,
modality="image",
),
T2I_sampling_params,
@@ -370,7 +384,7 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
DiffusionTestCase(
"qwen_image_t2i_cache_dit_enabled",
DiffusionServerArgs(
model_path="Qwen/Qwen-Image",
model_path=DEFAULT_QWEN_IMAGE_MODEL_NAME_FOR_TEST,
modality="image",
enable_cache_dit=True,
),
@@ -379,7 +393,7 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
DiffusionTestCase(
"flux_image_t2i",
DiffusionServerArgs(
model_path="black-forest-labs/FLUX.1-dev", modality="image"
model_path=DEFAULT_FLUX_1_DEV_MODEL_NAME_FOR_TEST, modality="image"
),
T2I_sampling_params,
),
@@ -396,14 +410,14 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
DiffusionTestCase(
"flux_2_image_t2i",
DiffusionServerArgs(
model_path="black-forest-labs/FLUX.2-dev", modality="image"
model_path=DEFAULT_FLUX_2_DEV_MODEL_NAME_FOR_TEST, modality="image"
),
T2I_sampling_params,
),
DiffusionTestCase(
"flux_2_klein_image_t2i",
DiffusionServerArgs(
model_path="black-forest-labs/FLUX.2-klein-4B",
model_path=DEFAULT_FLUX_2_KLEIN_4B_MODEL_NAME_FOR_TEST,
modality="image",
),
T2I_sampling_params,
@@ -414,7 +428,7 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
DiffusionTestCase(
"layerwise_offload",
DiffusionServerArgs(
model_path=SMALL_T2I_MODEL,
model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
modality="image",
dit_layerwise_offload=True,
dit_offload_prefetch_size=2,
@@ -423,13 +437,15 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
),
DiffusionTestCase(
"zimage_image_t2i",
DiffusionServerArgs(model_path="Tongyi-MAI/Z-Image-Turbo", modality="image"),
DiffusionServerArgs(
model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST, modality="image"
),
T2I_sampling_params,
),
DiffusionTestCase(
"zimage_image_t2i_fp8",
DiffusionServerArgs(
model_path="Tongyi-MAI/Z-Image-Turbo",
model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
modality="image",
extras=["--transformer-path MickJ/Z-Image-Turbo-fp8"],
),
@@ -439,7 +455,7 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
DiffusionTestCase(
"zimage_image_t2i_multi_lora",
DiffusionServerArgs(
model_path="Tongyi-MAI/Z-Image-Turbo",
model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
modality="image",
lora_path="reverentelusarca/elusarca-anime-style-lora-z-image-turbo",
second_lora_path="tarn59/pixel_art_style_lora_z_image_turbo",
@@ -449,22 +465,33 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
# === Text and Image to Image (TI2I) ===
DiffusionTestCase(
"qwen_image_edit_ti2i",
DiffusionServerArgs(model_path="Qwen/Qwen-Image-Edit", modality="image"),
DiffusionServerArgs(
model_path=DEFAULT_QWEN_IMAGE_EDIT_MODEL_NAME_FOR_TEST, modality="image"
),
TI2I_sampling_params,
),
DiffusionTestCase(
"qwen_image_edit_2509_ti2i",
DiffusionServerArgs(model_path="Qwen/Qwen-Image-Edit-2509", modality="image"),
DiffusionServerArgs(
model_path=DEFAULT_QWEN_IMAGE_EDIT_2509_MODEL_NAME_FOR_TEST,
modality="image",
),
MULTI_IMAGE_TI2I_sampling_params,
),
DiffusionTestCase(
"qwen_image_edit_2511_ti2i",
DiffusionServerArgs(model_path="Qwen/Qwen-Image-Edit-2511", modality="image"),
DiffusionServerArgs(
model_path=DEFAULT_QWEN_IMAGE_EDIT_2511_MODEL_NAME_FOR_TEST,
modality="image",
),
TI2I_sampling_params,
),
DiffusionTestCase(
"qwen_image_layered_i2i",
DiffusionServerArgs(model_path="Qwen/Qwen-Image-Layered", modality="image"),
DiffusionServerArgs(
model_path=DEFAULT_QWEN_IMAGE_LAYERED_MODEL_NAME_FOR_TEST,
modality="image",
),
MULTI_FRAME_I2I_sampling_params,
),
]
@@ -479,7 +506,7 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
DiffusionTestCase(
"wan2_1_t2v_1.3b",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
model_path=DEFAULT_WAN_2_1_T2V_1_3B_MODEL_NAME_FOR_TEST,
modality="video",
custom_validator="video",
),
@@ -490,7 +517,7 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
DiffusionTestCase(
"wan2_1_t2v_1.3b_text_encoder_cpu_offload",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
model_path=DEFAULT_WAN_2_1_T2V_1_3B_MODEL_NAME_FOR_TEST,
modality="video",
custom_validator="video",
text_encoder_cpu_offload=True,
@@ -503,7 +530,7 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
DiffusionTestCase(
"wan2_1_t2v_1.3b_teacache_enabled",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
model_path=DEFAULT_WAN_2_1_T2V_1_3B_MODEL_NAME_FOR_TEST,
modality="video",
custom_validator="video",
),
@@ -533,7 +560,7 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
DiffusionTestCase(
"wan2_1_t2v_1_3b_lora_1gpu",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
model_path=DEFAULT_WAN_2_1_T2V_1_3B_MODEL_NAME_FOR_TEST,
modality="video",
custom_validator="video",
num_gpus=1,
@@ -558,14 +585,14 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
DiffusionTestCase(
"flux_2_ti2i",
DiffusionServerArgs(
model_path="black-forest-labs/FLUX.2-dev", modality="image"
model_path=DEFAULT_FLUX_2_DEV_MODEL_NAME_FOR_TEST, modality="image"
),
TI2I_sampling_params,
),
DiffusionTestCase(
"flux_2_t2i_customized_vae_path",
DiffusionServerArgs(
model_path="black-forest-labs/FLUX.2-dev",
model_path=DEFAULT_FLUX_2_DEV_MODEL_NAME_FOR_TEST,
modality="image",
extras=["--vae-path=fal/FLUX.2-Tiny-AutoEncoder"],
),
@@ -587,7 +614,7 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
DiffusionTestCase(
"wan2_2_ti2v_5b",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.2-TI2V-5B-Diffusers",
model_path=DEFAULT_WAN_2_2_TI2V_5B_MODEL_NAME_FOR_TEST,
modality="video",
custom_validator="video",
),
@@ -636,7 +663,7 @@ TWO_GPU_CASES_A = [
DiffusionTestCase(
"wan2_2_i2v_a14b_2gpu",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
model_path=DEFAULT_WAN_2_2_I2V_A14B_MODEL_NAME_FOR_TEST,
modality="video",
custom_validator="video",
),
@@ -645,7 +672,7 @@ TWO_GPU_CASES_A = [
DiffusionTestCase(
"wan2_2_t2v_a14b_2gpu",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
model_path=DEFAULT_WAN_2_2_T2V_A14B_MODEL_NAME_FOR_TEST,
modality="video",
custom_validator="video",
num_gpus=2,
@@ -658,7 +685,7 @@ TWO_GPU_CASES_A = [
DiffusionTestCase(
"wan2_2_t2v_a14b_lora_2gpu",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
model_path=DEFAULT_WAN_2_2_T2V_A14B_MODEL_NAME_FOR_TEST,
modality="video",
custom_validator="video",
num_gpus=2,
@@ -671,7 +698,7 @@ TWO_GPU_CASES_A = [
DiffusionTestCase(
"wan2_1_t2v_14b_2gpu",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.1-T2V-14B-Diffusers",
model_path=DEFAULT_WAN_2_1_T2V_14B_MODEL_NAME_FOR_TEST,
modality="video",
num_gpus=2,
custom_validator="video",
@@ -684,7 +711,7 @@ TWO_GPU_CASES_A = [
DiffusionTestCase(
"wan2_1_t2v_1.3b_cfg_parallel",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
model_path=DEFAULT_WAN_2_1_T2V_1_3B_MODEL_NAME_FOR_TEST,
modality="video",
custom_validator="video",
num_gpus=2,
@@ -697,7 +724,7 @@ TWO_GPU_CASES_A = [
DiffusionTestCase(
"fsdp-inference",
DiffusionServerArgs(
model_path=DEFAULT_SMALL_MODEL,
model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
modality="image",
num_gpus=2,
extras=["--use-fsdp-inference"],
@@ -710,7 +737,7 @@ TWO_GPU_CASES_B = [
DiffusionTestCase(
"wan2_1_i2v_14b_480P_2gpu",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
model_path=DEFAULT_WAN_2_1_I2V_14B_480P_MODEL_NAME_FOR_TEST,
modality="video",
custom_validator="video",
num_gpus=2,
@@ -721,7 +748,7 @@ TWO_GPU_CASES_B = [
DiffusionTestCase(
"wan2_1_i2v_14b_lora_2gpu",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.1-I2V-14B-720P-Diffusers",
model_path=DEFAULT_WAN_2_1_I2V_14B_720P_MODEL_NAME_FOR_TEST,
modality="video",
custom_validator="video",
num_gpus=2,
@@ -732,7 +759,7 @@ TWO_GPU_CASES_B = [
DiffusionTestCase(
"wan2_1_i2v_14b_720P_2gpu",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.1-I2V-14B-720P-Diffusers",
model_path=DEFAULT_WAN_2_1_I2V_14B_720P_MODEL_NAME_FOR_TEST,
modality="video",
custom_validator="video",
num_gpus=2,
@@ -742,7 +769,7 @@ TWO_GPU_CASES_B = [
DiffusionTestCase(
"qwen_image_t2i_2_gpus",
DiffusionServerArgs(
model_path="Qwen/Qwen-Image",
model_path=DEFAULT_QWEN_IMAGE_MODEL_NAME_FOR_TEST,
modality="image",
num_gpus=2,
# test ring attn
@@ -754,7 +781,7 @@ TWO_GPU_CASES_B = [
DiffusionTestCase(
"zimage_image_t2i_2_gpus",
DiffusionServerArgs(
model_path="Tongyi-MAI/Z-Image-Turbo",
model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
modality="image",
num_gpus=2,
ulysses_degree=2,
@@ -764,7 +791,7 @@ TWO_GPU_CASES_B = [
DiffusionTestCase(
"flux_image_t2i_2_gpus",
DiffusionServerArgs(
model_path="black-forest-labs/FLUX.1-dev",
model_path=DEFAULT_FLUX_1_DEV_MODEL_NAME_FOR_TEST,
modality="image",
num_gpus=2,
),
@@ -773,7 +800,7 @@ TWO_GPU_CASES_B = [
DiffusionTestCase(
"flux_2_image_t2i_2_gpus",
DiffusionServerArgs(
model_path="black-forest-labs/FLUX.2-dev",
model_path=DEFAULT_FLUX_2_DEV_MODEL_NAME_FOR_TEST,
modality="image",
num_gpus=2,
tp_size=2,

143
python/sglang/multimodal_gen/test/test_utils.py
View File

@@ -4,11 +4,14 @@ import io
import json
import os
import socket
import subprocess
import tempfile
import time
from pathlib import Path
from urllib.parse import urljoin
import cv2
import httpx
import numpy as np
from PIL import Image
@@ -21,6 +24,45 @@ from sglang.multimodal_gen.runtime.utils.perf_logger import (
logger = init_logger(__name__)
# ---------------------------------------------------------------------------
# Common model IDs for diffusion tests
#
# Centralised here so every test file references the same constants instead
# of scattering hard-coded strings. When adding a new model that will be
# reused across tests, define it here.
# ---------------------------------------------------------------------------
DEFAULT_SMALL_MODEL_NAME_FOR_TEST = "Tongyi-MAI/Z-Image-Turbo"
# Qwen image generation models
DEFAULT_QWEN_IMAGE_MODEL_NAME_FOR_TEST = "Qwen/Qwen-Image"
DEFAULT_QWEN_IMAGE_2512_MODEL_NAME_FOR_TEST = "Qwen/Qwen-Image-2512"
DEFAULT_QWEN_IMAGE_EDIT_MODEL_NAME_FOR_TEST = "Qwen/Qwen-Image-Edit"
DEFAULT_QWEN_IMAGE_EDIT_2509_MODEL_NAME_FOR_TEST = "Qwen/Qwen-Image-Edit-2509"
DEFAULT_QWEN_IMAGE_EDIT_2511_MODEL_NAME_FOR_TEST = "Qwen/Qwen-Image-Edit-2511"
DEFAULT_QWEN_IMAGE_LAYERED_MODEL_NAME_FOR_TEST = "Qwen/Qwen-Image-Layered"
# FLUX image generation models
DEFAULT_FLUX_1_DEV_MODEL_NAME_FOR_TEST = "black-forest-labs/FLUX.1-dev"
DEFAULT_FLUX_2_DEV_MODEL_NAME_FOR_TEST = "black-forest-labs/FLUX.2-dev"
DEFAULT_FLUX_2_KLEIN_4B_MODEL_NAME_FOR_TEST = "black-forest-labs/FLUX.2-klein-4B"
DEFAULT_FLUX_2_KLEIN_BASE_4B_MODEL_NAME_FOR_TEST = (
"black-forest-labs/FLUX.2-klein-base-4B"
)
# Wan video generation models
DEFAULT_WAN_2_1_T2V_1_3B_MODEL_NAME_FOR_TEST = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
DEFAULT_WAN_2_1_T2V_14B_MODEL_NAME_FOR_TEST = "Wan-AI/Wan2.1-T2V-14B-Diffusers"
DEFAULT_WAN_2_1_I2V_14B_480P_MODEL_NAME_FOR_TEST = (
"Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
)
DEFAULT_WAN_2_1_I2V_14B_720P_MODEL_NAME_FOR_TEST = (
"Wan-AI/Wan2.1-I2V-14B-720P-Diffusers"
)
DEFAULT_WAN_2_2_TI2V_5B_MODEL_NAME_FOR_TEST = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
DEFAULT_WAN_2_2_T2V_A14B_MODEL_NAME_FOR_TEST = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
DEFAULT_WAN_2_2_I2V_A14B_MODEL_NAME_FOR_TEST = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
def print_value_formatted(description: str, value: int | float | str):
"""Helper function to print a metric value formatted."""
@@ -84,6 +126,107 @@ def get_dynamic_server_port() -> int:
return base_port + 1000
def find_free_port(host: str = "127.0.0.1") -> int:
"""Bind to port 0 and let the OS assign an available port."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind((host, 0))
return s.getsockname()[1]
def wait_for_server_health(
base_url: str,
path: str = "/health",
timeout: float = 180.0,
interval: float = 1.0,
) -> None:
"""Poll ``GET <base_url><path>`` until it returns HTTP 200."""
deadline = time.time() + timeout
last_err: httpx.RequestError | None = None
last_status: int | None = None
while time.time() < deadline:
try:
r = httpx.get(urljoin(base_url, path), timeout=5.0)
last_status = r.status_code
if r.status_code == 200:
return
except httpx.RequestError as e:
last_err = e
time.sleep(interval)
raise TimeoutError(
f"Server at {urljoin(base_url, path)} not healthy after {timeout}s. "
f"{last_status=} {last_err=}"
)
def post_json(
base_url: str,
path: str,
payload: dict,
timeout: float = 300.0,
) -> httpx.Response:
"""POST JSON to ``<base_url><path>`` and return the response."""
return httpx.post(urljoin(base_url, path), json=payload, timeout=timeout)
# ---------------------------------------------------------------------------
# GPU memory helpers (nvidia-smi)
# ---------------------------------------------------------------------------
def query_gpu_mem_used_mib(gpu_index: int = 0, required: bool = False) -> int | None:
"""Return GPU memory usage in MiB via ``nvidia-smi``, or *None* on failure.
When *required* is ``True`` the function raises instead of returning ``None``.
"""
try:
out = subprocess.check_output(
[
"nvidia-smi",
f"--id={gpu_index}",
"--query-gpu=memory.used",
"--format=csv,noheader,nounits",
],
text=True,
).strip()
return int(out.splitlines()[0].strip())
except Exception as e:
logger.warning(f"nvidia-smi memory query failed: {type(e).__name__}: {e}")
assert not required, (
"nvidia-smi memory query is unavailable; "
"cannot enforce GPU memory assertions."
)
return None
def require_gpu_mem_query(gpu_index: int = 0) -> int:
"""Same as :func:`query_gpu_mem_used_mib` but asserts availability.
Raises ``AssertionError`` when ``nvidia-smi`` is unavailable instead of
returning ``None``, so callers can rely on a valid ``int`` result.
"""
mem = query_gpu_mem_used_mib(gpu_index, required=True)
assert mem is not None
return mem
def assert_gpu_mem_changed(
label: str,
before_mib: int,
after_mib: int,
min_delta_mib: int,
) -> None:
"""Assert that GPU memory changed by at least *min_delta_mib* MiB."""
delta = abs(after_mib - before_mib)
logger.debug(
f"[MEM] {label}: before={before_mib} MiB after={after_mib} MiB |delta|={delta} MiB"
)
assert delta >= min_delta_mib, (
f"GPU memory change too small for '{label}': "
f"|after-before|={delta} MiB < {min_delta_mib} MiB "
f"(before={before_mib} MiB, after={after_mib} MiB)"
)
def is_mp4(data: bytes) -> bool:
"""Check if data represents a valid MP4 file by magic bytes."""
if len(data) < 8: