From be9a9e4819ac90102470035ee0e55869209285c7 Mon Sep 17 00:00:00 2001 From: Mook <68294499+Godmook@users.noreply.github.com> Date: Thu, 5 Mar 2026 20:09:42 -0800 Subject: [PATCH] refactor(multimodal/test): centralize model names and shared utilities in test_utils (#19354) Co-authored-by: Ratish P <114130421+Ratish1@users.noreply.github.com> --- .../test/cli/test_generate_t2i_perf.py | 3 +- .../server/test_update_weights_from_disk.py | 17 ++- .../test/server/testcase_configs.py | 101 ++++++++----- .../sglang/multimodal_gen/test/test_utils.py | 143 ++++++++++++++++++ 4 files changed, 221 insertions(+), 43 deletions(-) diff --git a/python/sglang/multimodal_gen/test/cli/test_generate_t2i_perf.py b/python/sglang/multimodal_gen/test/cli/test_generate_t2i_perf.py index 7635711d0..8c7bb7f45 100644 --- a/python/sglang/multimodal_gen/test/cli/test_generate_t2i_perf.py +++ b/python/sglang/multimodal_gen/test/cli/test_generate_t2i_perf.py @@ -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 diff --git a/python/sglang/multimodal_gen/test/server/test_update_weights_from_disk.py b/python/sglang/multimodal_gen/test/server/test_update_weights_from_disk.py index 68700e93d..ac3c8d7ff 100644 --- a/python/sglang/multimodal_gen/test/server/test_update_weights_from_disk.py +++ b/python/sglang/multimodal_gen/test/server/test_update_weights_from_disk.py @@ -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, ), ] diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index b32dc9dd9..03a1ca167 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -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, diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index 073b36972..62e738a81 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -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 `` 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 ```` 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: