sgl-project--sglang/python/sglang/multimodal_gen/test/unit/test_server_args.py

import json
import os
import sys
import tempfile
import unittest
from contextlib import contextmanager
from unittest.mock import patch

from sglang.multimodal_gen.configs.models.fsdp import (
    is_module_list_entry,
    is_module_list_entry_in,
    is_zimage_layer,
)
from sglang.multimodal_gen.configs.pipeline_configs.base import (
    ModelTaskType,
    PipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import FastHunyuanConfig
from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
    LTX2PipelineConfig,
    LTX23PipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig
from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
    QwenImagePipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.sana_wm import (
    SanaWMPipelineConfig,
    SanaWMRealtimeConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.wan import (
    FastWan2_2_TI2V_5B_Config,
    TurboWanT2V480PConfig,
    Wan2_2_I2V_A14B_Config,
    Wan2_2_T2V_A14B_Config,
    WanI2V480PConfig,
    WanI2V720PConfig,
    WanT2V480PConfig,
    WanT2V720PConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.zimage import ZImagePipelineConfig
from sglang.multimodal_gen.registry import _get_config_info
from sglang.multimodal_gen.runtime.models.dits.qwen_image import (
    QwenImageTransformer2DModel,
)
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.utils import FlexibleArgumentParser


@contextmanager
def _mock_cuda_platform(
    *,
    memory_gb: int = 80,
    available_memory_gb: int | dict[int, int] | None = None,
):
    def get_available_gpu_memory(device_id=0, **_kwargs):
        if isinstance(available_memory_gb, dict):
            return available_memory_gb[device_id]
        if available_memory_gb is not None:
            return available_memory_gb
        return memory_gb

    with (
        patch(
            "sglang.multimodal_gen.runtime.platforms.current_platform.is_cpu",
            return_value=False,
        ),
        patch(
            "sglang.multimodal_gen.runtime.platforms.current_platform.is_mps",
            return_value=False,
        ),
        patch(
            "sglang.multimodal_gen.runtime.platforms.current_platform.is_cuda",
            return_value=True,
        ),
        patch(
            "sglang.multimodal_gen.runtime.platforms.current_platform.get_device_total_memory",
            return_value=memory_gb * 1024**3,
        ),
        patch(
            "sglang.multimodal_gen.runtime.platforms.current_platform.get_available_gpu_memory",
            side_effect=get_available_gpu_memory,
        ),
        patch(
            "sglang.multimodal_gen.runtime.platforms.current_platform.enable_dit_layerwise_offload_for_wan_by_default",
            return_value=True,
        ),
    ):
        yield


def _from_dict_without_model_resolution(
    kwargs, pipeline_config: PipelineConfig | None = None
):
    pipeline_config = pipeline_config or QwenImagePipelineConfig()
    with (
        patch.object(PipelineConfig, "from_kwargs", return_value=pipeline_config),
        _mock_cuda_platform(),
    ):
        return ServerArgs.from_dict(kwargs)


class TestServerArgsPathExpansion(unittest.TestCase):
    def _from_dict_without_model_resolution(self, kwargs):
        return _from_dict_without_model_resolution(kwargs)

    def test_tilde_model_path_is_expanded(self):
        args = self._from_dict_without_model_resolution(
            {"model_path": "~/fake/local/model"}
        )
        expected = os.path.expanduser("~/fake/local/model")
        self.assertEqual(args.model_path, expected)
        self.assertFalse(args.model_path.startswith("~"))

    def test_absolute_path_is_unchanged(self):
        args = self._from_dict_without_model_resolution(
            {"model_path": "/data/my-model"}
        )
        self.assertEqual(args.model_path, "/data/my-model")

    def test_component_paths_are_expanded_before_pipeline_resolution(self):
        args = self._from_dict_without_model_resolution(
            {
                "model_path": "/data/my-model",
                "component_paths": {"vae": "~/fake/local/vae"},
            }
        )

        self.assertEqual(
            args.component_paths["vae"], os.path.expanduser("~/fake/local/vae")
        )

    def test_component_attention_backends_are_normalized(self):
        args = self._from_dict_without_model_resolution(
            {
                "model_path": "/data/my-model",
                "component_attention_backends": "text-encoder=torch_sdpa,transformer=fa3",
            }
        )

        self.assertEqual(
            args.component_attention_backends,
            {"text_encoder": "torch_sdpa", "transformer": "fa"},
        )

    def test_component_attention_backend_lookup(self):
        args = self._from_dict_without_model_resolution(
            {
                "model_path": "/data/my-model",
                "component_attention_backends": {"text_encoder": "torch_sdpa"},
            }
        )

        backend, matched_key = args.resolve_component_attention_backend(
            "text_encoder", "transformer"
        )

        self.assertEqual(backend.name, "TORCH_SDPA")
        self.assertEqual(matched_key, "text_encoder")

    def test_invalid_component_attention_backend_raises(self):
        with self.assertRaises(ValueError):
            self._from_dict_without_model_resolution(
                {
                    "model_path": "/data/my-model",
                    "component_attention_backends": {"text_encoder": "bad_backend"},
                }
            )
        with self.assertRaises(ValueError):
            self._from_dict_without_model_resolution(
                {
                    "model_path": "/data/my-model",
                    "component_attention_backends": "text_encoder",
                }
            )

    def test_dynamic_component_attention_backend_cli_args(self):
        parser = FlexibleArgumentParser()
        ServerArgs.add_cli_args(parser)
        argv = [
            "--model-path",
            "/fake",
            "--component-attention-backends.text-encoder",
            "torch_sdpa",
        ]

        with (
            patch.object(sys, "argv", ["sglang"] + argv),
            patch.object(
                PipelineConfig, "from_kwargs", return_value=QwenImagePipelineConfig()
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.is_cpu",
                return_value=False,
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.is_mps",
                return_value=False,
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.is_cuda",
                return_value=True,
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.get_device_total_memory",
                return_value=80 * 1024**3,
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.get_available_gpu_memory",
                return_value=80,
            ),
        ):
            args, unknown_args = parser.parse_known_args(argv)
            server_args = ServerArgs.from_cli_args(args, unknown_args)

        self.assertEqual(
            server_args.component_attention_backends, {"text_encoder": "torch_sdpa"}
        )

    def test_layerwise_offload_components_imply_layerwise(self):
        args = self._from_dict_without_model_resolution(
            {
                "model_path": "/data/my-model",
                "performance_mode": "manual",
            }
        )
        args.layerwise_offload_components = ["text_encoder", "transformer"]
        args._adjust_layerwise_offload_components()

        self.assertTrue(args.layerwise_offload_components)
        self.assertEqual(
            args.layerwise_offload_components, ["text_encoder", "transformer"]
        )

    def test_dit_layerwise_offload_selects_dit_group(self):
        args = self._from_dict_without_model_resolution(
            {
                "model_path": "/data/my-model",
                "performance_mode": "manual",
                "dit_layerwise_offload": True,
            }
        )

        self.assertTrue(args.layerwise_offload_components)
        self.assertEqual(args.layerwise_offload_components, ["dit"])

    def test_dit_layerwise_offload_from_kwargs(self):
        with patch.object(
            PipelineConfig, "from_kwargs", return_value=QwenImagePipelineConfig()
        ):
            args = ServerArgs.from_kwargs(
                model_path="/data/my-model",
                performance_mode="manual",
                dit_layerwise_offload=True,
            )

        self.assertTrue(args.layerwise_offload_components)
        self.assertEqual(args.layerwise_offload_components, ["dit"])

    def test_layerwise_offload_components_normalize_commas(self):
        args = self._from_dict_without_model_resolution(
            {
                "model_path": "/data/my-model",
                "performance_mode": "manual",
            }
        )
        args.layerwise_offload_components = ["text-encoder,transformer"]
        args._adjust_layerwise_offload_components()

        self.assertEqual(
            args.layerwise_offload_components, ["text_encoder", "transformer"]
        )

    def test_layerwise_offload_components_normalize_default_group(self):
        args = self._from_dict_without_model_resolution(
            {
                "model_path": "/data/my-model",
                "performance_mode": "manual",
            }
        )
        args.layerwise_offload_components = ["default", "text_encoder"]
        args._adjust_layerwise_offload_components()

        self.assertEqual(
            args.layerwise_offload_components,
            ["text_encoder", "image_encoder", "vae"],
        )

    def test_dit_layerwise_offload_cli_arg(self):
        parser = FlexibleArgumentParser()
        ServerArgs.add_cli_args(parser)
        argv = [
            "--model-path",
            "/fake",
            "--performance-mode",
            "manual",
            "--dit-layerwise-offload",
            "true",
        ]

        with patch.object(sys, "argv", ["sglang"] + argv):
            args, unknown_args = parser.parse_known_args(argv)
            with patch.object(
                PipelineConfig, "from_kwargs", return_value=QwenImagePipelineConfig()
            ):
                server_args = ServerArgs.from_cli_args(args, unknown_args)

        self.assertTrue(server_args.layerwise_offload_components)
        self.assertEqual(server_args.layerwise_offload_components, ["dit"])

    def test_layerwise_offload_components_cli_args(self):
        parser = FlexibleArgumentParser()
        ServerArgs.add_cli_args(parser)
        argv = [
            "--model-path",
            "/fake",
            "--performance-mode",
            "manual",
            "--layerwise-offload-components",
            "transformer",
            "text_encoder",
        ]

        with patch.object(sys, "argv", ["sglang"] + argv):
            args, unknown_args = parser.parse_known_args(argv)
            with patch.object(
                PipelineConfig, "from_kwargs", return_value=QwenImagePipelineConfig()
            ):
                server_args = ServerArgs.from_cli_args(args, unknown_args)

        self.assertEqual(
            server_args.layerwise_offload_components, ["transformer", "text_encoder"]
        )

    def test_serve_cli_preserves_config_and_dynamic_unknown_args(self):
        from sglang.multimodal_gen.runtime.entrypoints.cli.serve import (
            add_multimodal_gen_serve_args,
        )

        with tempfile.NamedTemporaryFile("w", suffix=".json") as config_file:
            json.dump({"model_path": "/from/config", "num_gpus": 2}, config_file)
            config_file.flush()
            parser = FlexibleArgumentParser()
            add_multimodal_gen_serve_args(parser)
            argv = [
                "--config",
                config_file.name,
                "--model-path",
                "/from/cli",
                "--vae-path",
                "/custom/vae",
                "--component-attention-backends.transformer",
                "fa3",
            ]

            with patch.object(sys, "argv", ["sglang", "serve"] + argv):
                args, unknown_args = parser.parse_known_args(argv)
                with (
                    patch.object(
                        PipelineConfig,
                        "from_kwargs",
                        return_value=QwenImagePipelineConfig(),
                    ),
                    patch(
                        "sglang.multimodal_gen.registry.get_model_info",
                        return_value=None,
                    ),
                    patch(
                        "sglang.multimodal_gen.runtime.platforms.current_platform.get_device_total_memory",
                        return_value=80 * 1024**3,
                    ),
                    patch(
                        "sglang.multimodal_gen.runtime.platforms.current_platform.get_available_gpu_memory",
                        return_value=80,
                    ),
                ):
                    server_args = ServerArgs.from_cli_args(args, unknown_args)

        self.assertEqual("/from/cli", server_args.model_path)
        self.assertEqual(2, server_args.num_gpus)
        self.assertEqual("/custom/vae", server_args.component_paths["vae"])
        self.assertEqual(
            {"transformer": "fa"},
            server_args.component_attention_backends,
        )

    def test_serve_cli_defaults_warmup_on(self):
        from sglang.multimodal_gen.runtime.entrypoints.cli.serve import (
            add_multimodal_gen_serve_args,
            execute_serve_cmd,
        )

        parser = FlexibleArgumentParser()
        add_multimodal_gen_serve_args(parser)
        argv = [
            "--model-path",
            "/fake",
        ]

        with (
            patch.object(sys, "argv", ["sglang", "serve"] + argv),
            patch.object(
                PipelineConfig, "from_kwargs", return_value=QwenImagePipelineConfig()
            ),
            patch(
                "sglang.multimodal_gen.runtime.entrypoints.cli.serve.dispatch_launch"
            ) as dispatch_launch,
        ):
            args, unknown_args = parser.parse_known_args(argv)
            execute_serve_cmd(args, unknown_args)

        server_args = dispatch_launch.call_args.args[0]
        self.assertTrue(server_args.warmup)
        self.assertTrue(server_args.server_warmup)
        self.assertFalse(server_args.is_arg_explicitly_set("warmup"))
        self.assertFalse(server_args.is_arg_explicitly_set("server_warmup"))

    def test_serve_cli_preserves_explicit_warmup_false(self):
        from sglang.multimodal_gen.runtime.entrypoints.cli.serve import (
            add_multimodal_gen_serve_args,
            execute_serve_cmd,
        )

        parser = FlexibleArgumentParser()
        add_multimodal_gen_serve_args(parser)
        argv = [
            "--model-path",
            "/fake",
            "--warmup",
            "false",
        ]

        with (
            patch.object(sys, "argv", ["sglang", "serve"] + argv),
            patch.object(
                PipelineConfig, "from_kwargs", return_value=QwenImagePipelineConfig()
            ),
            patch(
                "sglang.multimodal_gen.runtime.entrypoints.cli.serve.dispatch_launch"
            ) as dispatch_launch,
        ):
            args, unknown_args = parser.parse_known_args(argv)
            execute_serve_cmd(args, unknown_args)

        server_args = dispatch_launch.call_args.args[0]
        self.assertFalse(server_args.warmup)
        self.assertFalse(server_args.server_warmup)
        self.assertTrue(server_args.is_arg_explicitly_set("warmup"))

    def test_serve_cli_preserves_config_warmup_false(self):
        from sglang.multimodal_gen.runtime.entrypoints.cli.serve import (
            add_multimodal_gen_serve_args,
            execute_serve_cmd,
        )

        with tempfile.NamedTemporaryFile("w", suffix=".json") as config_file:
            json.dump({"model_path": "/fake", "warmup": False}, config_file)
            config_file.flush()

            parser = FlexibleArgumentParser()
            add_multimodal_gen_serve_args(parser)
            argv = [
                "--config",
                config_file.name,
            ]

            with (
                patch.object(sys, "argv", ["sglang", "serve"] + argv),
                patch.object(
                    PipelineConfig,
                    "from_kwargs",
                    return_value=QwenImagePipelineConfig(),
                ),
                patch(
                    "sglang.multimodal_gen.runtime.entrypoints.cli.serve.dispatch_launch"
                ) as dispatch_launch,
            ):
                args, unknown_args = parser.parse_known_args(argv)
                execute_serve_cmd(args, unknown_args)

        server_args = dispatch_launch.call_args.args[0]
        self.assertFalse(server_args.warmup)
        self.assertFalse(server_args.server_warmup)
        self.assertTrue(server_args.is_arg_explicitly_set("warmup"))

    def test_disagg_role_disables_server_warmup(self):
        with patch.object(
            PipelineConfig, "from_kwargs", return_value=QwenImagePipelineConfig()
        ):
            server_args = ServerArgs.from_dict(
                {
                    "model_path": "/fake",
                    "warmup": True,
                    "server_warmup": True,
                    "disagg_role": "server",
                }
            )

        self.assertTrue(server_args.warmup)
        self.assertFalse(server_args.server_warmup)


class TestWarmupModeNormalization(unittest.TestCase):
    """`_adjust_warmup` resolves the canonical warmup_mode and its derived booleans."""

    def _resolve(
        self,
        *,
        warmup_mode=None,
        warmup=False,
        server_warmup=False,
        warmup_resolutions=None,
        enable_torch_compile=False,
        disagg_role=None,
        explicit=(),
    ):
        from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType

        sa = ServerArgs.__new__(ServerArgs)
        sa.warmup_mode = warmup_mode
        sa.warmup = warmup
        sa.server_warmup = server_warmup
        sa.warmup_resolutions = warmup_resolutions
        sa.enable_torch_compile = enable_torch_compile
        sa.disagg_role = RoleType.MONOLITHIC if disagg_role is None else disagg_role
        sa._explicit_arg_names = set(explicit)
        sa._adjust_warmup()
        return sa

    def test_explicit_mode_off_disables_all(self):
        sa = self._resolve(warmup_mode="off", explicit=("warmup_mode",))
        self.assertEqual(sa.warmup_mode, "off")
        self.assertFalse(sa.warmup)
        self.assertFalse(sa.server_warmup)

    def test_explicit_mode_request(self):
        sa = self._resolve(warmup_mode="request", explicit=("warmup_mode",))
        self.assertEqual(sa.warmup_mode, "request")
        self.assertTrue(sa.warmup)
        self.assertFalse(sa.server_warmup)

    def test_explicit_mode_server(self):
        sa = self._resolve(warmup_mode="server", explicit=("warmup_mode",))
        self.assertEqual(sa.warmup_mode, "server")
        self.assertTrue(sa.warmup)
        self.assertTrue(sa.server_warmup)

    def test_explicit_mode_overrides_explicit_legacy(self):
        sa = self._resolve(
            warmup_mode="request",
            warmup=True,
            server_warmup=True,
            explicit=("warmup_mode", "warmup", "server_warmup"),
        )
        self.assertEqual(sa.warmup_mode, "request")
        self.assertTrue(sa.warmup)
        self.assertFalse(sa.server_warmup)

    def test_explicit_legacy_false_beats_defaulted_mode(self):
        # serve defaults warmup_mode="server" (not explicit); `--warmup false` wins.
        sa = self._resolve(
            warmup_mode="server",
            warmup=False,
            server_warmup=False,
            explicit=("warmup",),
        )
        self.assertEqual(sa.warmup_mode, "off")
        self.assertFalse(sa.warmup)
        self.assertFalse(sa.server_warmup)

    def test_defaulted_mode_applies_without_legacy_flags(self):
        # bare `sglang serve`: warmup_mode="server" defaulted, no legacy override.
        sa = self._resolve(warmup_mode="server")
        self.assertEqual(sa.warmup_mode, "server")
        self.assertTrue(sa.warmup)
        self.assertTrue(sa.server_warmup)

    def test_legacy_only_maps_to_request(self):
        sa = self._resolve(warmup_mode=None, warmup=True, explicit=("warmup",))
        self.assertEqual(sa.warmup_mode, "request")
        self.assertTrue(sa.warmup)
        self.assertFalse(sa.server_warmup)

    def test_resolutions_force_warmup_on(self):
        sa = self._resolve(
            warmup_mode="off",
            warmup_resolutions=["512x512"],
            explicit=("warmup_mode",),
        )
        self.assertTrue(sa.warmup)
        self.assertFalse(sa.server_warmup)
        self.assertEqual(sa.warmup_mode, "request")

    def test_torch_compile_defaults_to_server_warmup(self):
        sa = self._resolve(enable_torch_compile=True)

        self.assertEqual(sa.warmup_mode, "server")
        self.assertTrue(sa.warmup)
        self.assertTrue(sa.server_warmup)

    def test_legacy_warmup_on_uses_defaulted_server_mode(self):
        # `serve --warmup` (legacy ON, mode defaulted to "server" but not
        # explicit) must resolve to server-based warmup, not silently downgrade
        # to request mode.
        sa = self._resolve(warmup_mode="server", warmup=True, explicit=("warmup",))

        self.assertEqual(sa.warmup_mode, "server")
        self.assertTrue(sa.warmup)
        self.assertTrue(sa.server_warmup)

    def test_torch_compile_respects_explicit_warmup_off(self):
        sa = self._resolve(
            warmup_mode="off",
            enable_torch_compile=True,
            explicit=("warmup_mode",),
        )
        self.assertEqual(sa.warmup_mode, "off")
        self.assertFalse(sa.warmup)
        self.assertFalse(sa.server_warmup)

    def test_torch_compile_uses_server_warmup_for_explicit_resolutions(self):
        sa = self._resolve(
            warmup_resolutions=["1024x1024"],
            enable_torch_compile=True,
            explicit=("warmup_resolutions",),
        )
        self.assertEqual(sa.warmup_mode, "server")
        self.assertTrue(sa.warmup)
        self.assertTrue(sa.server_warmup)

    def test_legacy_warmup_with_resolutions_runs_server_warmup(self):
        # Dead-zone regression: `serve --warmup --warmup-resolutions X` must run
        # server-based (synthetic) warmup, not end up with no warmup at all
        # (request-based warmup bails out when warmup_resolutions is set).
        sa = self._resolve(
            warmup_mode="server",
            warmup=True,
            warmup_resolutions=["1024x1024"],
            explicit=("warmup",),
        )
        self.assertTrue(sa.warmup)
        self.assertTrue(sa.server_warmup)
        self.assertEqual(sa.warmup_mode, "server")

    def test_disagg_role_disables_server_warmup(self):
        from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType

        sa = self._resolve(
            warmup_mode="server",
            disagg_role=RoleType.DENOISER,
            explicit=("warmup_mode",),
        )
        self.assertTrue(sa.warmup)
        self.assertFalse(sa.server_warmup)
        self.assertEqual(sa.warmup_mode, "request")

    def test_torch_compile_server_warmup_disabled_for_disagg_role(self):
        from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType

        sa = self._resolve(enable_torch_compile=True, disagg_role=RoleType.DENOISER)
        self.assertEqual(sa.warmup_mode, "request")
        self.assertTrue(sa.warmup)
        self.assertFalse(sa.server_warmup)

    def test_invalid_mode_raises(self):
        with self.assertRaises(ValueError):
            self._resolve(warmup_mode="bogus", explicit=("warmup_mode",))


class TestWarmupImageIsModelValid(unittest.TestCase):
    """The server-warmup placeholder image must be large enough for real pipelines."""

    def test_minimum_warmup_image_is_at_least_64px(self):
        import base64
        import struct

        from sglang.multimodal_gen.runtime.server_warmup import (
            MINIMUM_PICTURE_BASE64_FOR_WARMUP,
        )

        payload = MINIMUM_PICTURE_BASE64_FOR_WARMUP.split(",", 1)[-1]
        raw = base64.b64decode(payload)
        self.assertEqual(raw[:8], b"\x89PNG\r\n\x1a\n")
        # IHDR width/height are the two big-endian uint32 after the chunk header.
        width, height = struct.unpack(">II", raw[16:24])
        self.assertGreaterEqual(width, 64)
        self.assertGreaterEqual(height, 64)


class TestOffloadDefaults(unittest.TestCase):
    def _from_dict_with_pipeline_config(
        self,
        pipeline_config,
        *,
        memory_gb=80,
        available_memory_gb=None,
        kwargs=None,
    ):
        def get_available_gpu_memory(device_id=0, **_kwargs):
            if isinstance(available_memory_gb, dict):
                return available_memory_gb[device_id]
            if available_memory_gb is not None:
                return available_memory_gb
            return memory_gb

        with (
            patch.object(PipelineConfig, "from_kwargs", return_value=pipeline_config),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.is_cpu",
                return_value=False,
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.is_mps",
                return_value=False,
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.is_cuda",
                return_value=True,
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.enable_dit_layerwise_offload_for_wan_by_default",
                return_value=True,
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.get_device_total_memory",
                return_value=memory_gb * 1024**3,
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.get_available_gpu_memory",
                side_effect=get_available_gpu_memory,
            ),
        ):
            return ServerArgs.from_dict({"model_path": "/fake", **(kwargs or {})})

    def _from_dict_with_task_type(
        self,
        task_type,
        *,
        memory_gb=80,
        kwargs=None,
    ):
        pipeline_config = PipelineConfig()
        pipeline_config.task_type = task_type
        with (
            patch.object(PipelineConfig, "from_kwargs", return_value=pipeline_config),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.is_cpu",
                return_value=False,
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.is_cuda",
                return_value=True,
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.get_device_total_memory",
                return_value=memory_gb * 1024**3,
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.get_available_gpu_memory",
                return_value=memory_gb,
            ),
        ):
            return ServerArgs.from_dict({"model_path": "/fake", **(kwargs or {})})

    def test_vae_cpu_offload_defaults_false_for_video_generation(self):
        args = self._from_dict_with_task_type(ModelTaskType.T2V)

        self.assertFalse(args.vae_cpu_offload)

    def test_vae_cpu_offload_defaults_false_on_low_memory_gpu(self):
        args = self._from_dict_with_task_type(
            ModelTaskType.T2V,
            memory_gb=16,
            kwargs={"performance_mode": "memory"},
        )

        self.assertFalse(args.vae_cpu_offload)
        self.assertTrue(args.dit_cpu_offload)
        self.assertFalse(args.text_encoder_cpu_offload)
        self.assertFalse(args.image_encoder_cpu_offload)
        self.assertEqual(
            args.layerwise_offload_components,
            ["text_encoder", "image_encoder", "vae"],
        )

    def test_explicit_vae_cpu_offload_true_is_preserved_by_default_layerwise(
        self,
    ):
        args = self._from_dict_with_task_type(
            ModelTaskType.T2V,
            kwargs={"vae_cpu_offload": True},
        )

        self.assertTrue(args.vae_cpu_offload)
        self.assertEqual(
            args.layerwise_offload_components, ["text_encoder", "image_encoder"]
        )

    def test_explicit_component_resident_is_preserved_by_default_layerwise(self):
        args = self._from_dict_with_task_type(
            ModelTaskType.T2V,
            kwargs={"text_encoder_cpu_offload": False},
        )

        self.assertFalse(args.text_encoder_cpu_offload)
        self.assertEqual(args.layerwise_offload_components, ["image_encoder", "vae"])

    def test_layerwise_components_disable_matching_non_dit_cpu_offloads(self):
        args = self._from_dict_with_task_type(
            ModelTaskType.T2V,
            memory_gb=16,
            kwargs={
                "performance_mode": "manual",
                "dit_cpu_offload": True,
                "text_encoder_cpu_offload": True,
                "image_encoder_cpu_offload": True,
                "vae_cpu_offload": True,
            },
        )
        args.layerwise_offload_components = [
            "text_encoder",
            "image_encoder",
            "video_dit",
            "vae",
        ]
        args._adjust_layerwise_offload_components()

        self.assertTrue(args.layerwise_offload_components)
        # dit_cpu_offload is complementary to DiT layerwise offload (keeps
        # weights off-device during load), so it must be preserved here.
        self.assertTrue(args.dit_cpu_offload)
        self.assertFalse(args.text_encoder_cpu_offload)
        self.assertFalse(args.image_encoder_cpu_offload)
        self.assertFalse(args.vae_cpu_offload)

    def test_dit_layerwise_offload_preserves_dit_cpu_offload(self):
        """Combining --dit-cpu-offload with --dit-layerwise-offload must keep both on.

        dit_cpu_offload controls initial residency (host memory), while
        dit_layerwise_offload only swaps layers on/off device at inference.
        Force-disabling dit_cpu_offload here would push the full DiT to GPU at
        load time and OOM low-VRAM cards.
        """
        args = self._from_dict_with_task_type(
            ModelTaskType.T2I,
            memory_gb=32,
            kwargs={
                "dit_cpu_offload": True,
                "dit_layerwise_offload": True,
            },
        )

        self.assertTrue(args.dit_cpu_offload)
        self.assertTrue(args.dit_layerwise_offload)
        self.assertEqual(args.layerwise_offload_components, ["dit"])

    def test_pipeline_configs_declare_auto_tune_hints(self):
        qwen_deployment = QwenImagePipelineConfig().get_model_deployment_config()
        wan_deployment = WanT2V480PConfig().get_model_deployment_config()
        mova_deployment = MOVAPipelineConfig().get_model_deployment_config()
        zimage_deployment = ZImagePipelineConfig().get_model_deployment_config()
        ltx_deployment = LTX2PipelineConfig().get_model_deployment_config()
        ltx23_config = LTX23PipelineConfig()
        sana_wm_deployment = SanaWMPipelineConfig().get_model_deployment_config()

        self.assertIsNone(qwen_deployment.fsdp_auto_min_available_memory_gb)
        self.assertFalse(qwen_deployment.auto_dit_layerwise_offload)

        self.assertIsNone(wan_deployment.fsdp_auto_min_available_memory_gb)
        self.assertTrue(wan_deployment.auto_dit_layerwise_offload)

        self.assertIsNone(mova_deployment.fsdp_auto_min_available_memory_gb)
        self.assertTrue(mova_deployment.auto_dit_layerwise_offload)

        self.assertEqual(zimage_deployment.fsdp_auto_min_available_memory_gb, 40)
        self.assertTrue(zimage_deployment.fsdp_auto_requires_cfg)
        self.assertFalse(zimage_deployment.auto_dit_layerwise_offload)

        self.assertEqual(ltx_deployment.keep_resident_min_available_gb, 70)
        self.assertEqual(ltx_deployment.keep_resident_components, ("dit",))
        self.assertEqual(
            ltx_deployment.auto_cfg_parallel_degree_by_num_gpus, ((4, 1), (8, 1))
        )
        self.assertEqual(ltx_deployment.get_auto_cfg_parallel_degree(4), 1)
        self.assertEqual(ltx_deployment.get_auto_cfg_parallel_degree(8), 1)
        self.assertEqual(ltx_deployment.get_auto_cfg_parallel_degree(2), 2)
        self.assertFalse(
            LTX2PipelineConfig().dit_config.arch_config.enable_packed_qkv_input_a2a
        )
        self.assertFalse(
            ltx23_config.dit_config.arch_config.enable_packed_qkv_input_a2a
        )

        self.assertEqual(sana_wm_deployment.fsdp_auto_min_available_memory_gb, 60)
        self.assertTrue(sana_wm_deployment.auto_dit_layerwise_offload)

        # fasthunyuan no longer pins 150gb -- falls back to the global video default
        fast_hunyuan_deployment = FastHunyuanConfig().get_model_deployment_config()
        self.assertIsNone(fast_hunyuan_deployment.keep_resident_min_available_gb)
        self.assertEqual(fast_hunyuan_deployment.keep_resident_components, ("vae",))

        # default keeps only vae resident (encoders are large, dit owned by FSDP)
        self.assertEqual(qwen_deployment.keep_resident_components, ("vae",))
        self.assertIsNone(qwen_deployment.keep_resident_min_available_gb)

    def test_auto_multi_gpu_sana_wm_prefers_fsdp_and_cfg_parallel(self):
        args = self._from_dict_with_pipeline_config(
            SanaWMPipelineConfig(),
            kwargs={
                "model_path": "Efficient-Large-Model/SANA-WM_bidirectional",
                "num_gpus": 2,
                "performance_mode": "auto",
            },
        )

        self.assertTrue(args.use_fsdp_inference)
        self.assertTrue(args.enable_cfg_parallel)

    def test_cache_dit_rejects_explicit_fsdp(self):
        with patch.dict(os.environ, {"SGLANG_CACHE_DIT_ENABLED": "true"}):
            with self.assertRaisesRegex(ValueError, "FSDP inference"):
                self._from_dict_with_pipeline_config(
                    SanaWMPipelineConfig(),
                    kwargs={
                        "model_path": "Efficient-Large-Model/SANA-WM_bidirectional",
                        "num_gpus": 2,
                        "use_fsdp_inference": True,
                    },
                )

    def test_cache_dit_auto_disables_implicit_fsdp(self):
        with patch.dict(os.environ, {"SGLANG_CACHE_DIT_ENABLED": "true"}):
            args = self._from_dict_with_pipeline_config(
                SanaWMPipelineConfig(),
                kwargs={
                    "model_path": "Efficient-Large-Model/SANA-WM_bidirectional",
                    "num_gpus": 2,
                    "performance_mode": "auto",
                },
            )

        self.assertFalse(args.use_fsdp_inference)
        self.assertTrue(args.enable_cfg_parallel)

    def test_auto_multi_gpu_sana_wm_realtime_disables_cfg_parallel(self):
        args = self._from_dict_with_pipeline_config(
            SanaWMRealtimeConfig(),
            kwargs={
                "model_path": "Efficient-Large-Model/SANA-WM_streaming",
                "num_gpus": 2,
                "performance_mode": "auto",
            },
        )

        self.assertFalse(args.use_fsdp_inference)
        self.assertFalse(args.enable_cfg_parallel)

    def test_auto_ltx23_large_gpu_counts_prefer_sp_over_cfg_parallel(self):
        for num_gpus in (4, 8):
            with self.subTest(num_gpus=num_gpus):
                args = self._from_dict_with_pipeline_config(
                    LTX2PipelineConfig(),
                    kwargs={
                        "model_path": "Lightricks/LTX-2.3",
                        "num_gpus": num_gpus,
                        "performance_mode": "auto",
                    },
                )

                self.assertFalse(args.enable_cfg_parallel)
                self.assertEqual(args.cfg_parallel_degree, 1)
                self.assertEqual(args.sp_degree, num_gpus)
                self.assertEqual(args.ulysses_degree, num_gpus)
                self.assertEqual(args.ring_degree, 1)

    def test_manual_mode_preserves_unset_performance_args(self):
        args = self._from_dict_with_pipeline_config(
            QwenImagePipelineConfig(),
            kwargs={
                "model_path": "Qwen/Qwen-Image",
                "num_gpus": 2,
                "performance_mode": "manual",
            },
        )

        self.assertEqual(args.performance_mode, "manual")
        self.assertIsNone(args.use_fsdp_inference)
        self.assertIsNone(args.dit_cpu_offload)
        self.assertIsNone(args.dit_layerwise_offload)
        self.assertIsNone(args.layerwise_offload_components)
        self.assertIsNone(args.text_encoder_cpu_offload)
        self.assertIsNone(args.image_encoder_cpu_offload)
        self.assertFalse(args.enable_cfg_parallel)

    def test_default_auto_keeps_image_vae_resident_when_memory_allows(self):
        args = self._from_dict_with_pipeline_config(
            QwenImagePipelineConfig(),
            kwargs={"model_path": "Qwen/Qwen-Image"},
        )

        self.assertEqual(args.performance_mode, "auto")
        self.assertFalse(args.use_fsdp_inference)
        # 80gb > image threshold (45gb): only vae kept resident, encoders stay
        # offloaded layerwise, dit unchanged
        self.assertTrue(args.dit_cpu_offload)
        self.assertEqual(
            args.layerwise_offload_components,
            ["text_encoder", "image_encoder"],
        )
        self.assertFalse(args.vae_cpu_offload)

    def test_auto_image_offloads_aux_below_resident_threshold(self):
        # 40gb < image threshold (45gb): aux incl. vae still offloaded to save vram
        args = self._from_dict_with_pipeline_config(
            QwenImagePipelineConfig(),
            memory_gb=40,
            kwargs={"model_path": "Qwen/Qwen-Image"},
        )

        self.assertEqual(args.performance_mode, "auto")
        self.assertTrue(args.dit_cpu_offload)
        self.assertEqual(
            args.layerwise_offload_components,
            ["text_encoder", "image_encoder", "vae"],
        )

    def test_auto_ltx_original_replaces_component_cpu_offload(
        self,
    ):
        args = self._from_dict_with_pipeline_config(
            LTX2PipelineConfig(),
            available_memory_gb=76,
            kwargs={
                "model_path": "Lightricks/LTX-2.3",
                "pipeline_class_name": "LTX2TwoStageHQPipeline",
                "performance_mode": "auto",
            },
        )

        self.assertEqual(args.ltx2_two_stage_device_mode, "original")
        self.assertFalse(args.dit_cpu_offload)
        self.assertTrue(args.layerwise_offload_components)
        self.assertFalse(args.text_encoder_cpu_offload)
        self.assertFalse(args.image_encoder_cpu_offload)
        self.assertEqual(
            args.layerwise_offload_components,
            ["text_encoder", "image_encoder", "vae"],
        )

    def test_auto_wan_layerwise_offload_is_enabled_without_fsdp(self):
        args = self._from_dict_with_pipeline_config(
            WanT2V480PConfig(),
            kwargs={"performance_mode": "auto"},
        )

        self.assertTrue(args.layerwise_offload_components)
        self.assertFalse(args.use_fsdp_inference)
        self.assertTrue(args.dit_cpu_offload)
        self.assertFalse(args.text_encoder_cpu_offload)
        self.assertFalse(args.image_encoder_cpu_offload)
        self.assertEqual(
            args.layerwise_offload_components,
            ["text_encoder", "image_encoder", "vae"],
        )

    def test_auto_wan2_2_a14b_layerwise_offload_adds_dit(self):
        for pipeline_config, model_path in (
            (Wan2_2_T2V_A14B_Config(), "Wan-AI/Wan2.2-T2V-A14B-Diffusers"),
            (Wan2_2_I2V_A14B_Config(), "Wan-AI/Wan2.2-I2V-A14B-Diffusers"),
        ):
            with self.subTest(pipeline_config=pipeline_config.__class__.__name__):
                args = self._from_dict_with_pipeline_config(
                    pipeline_config,
                    kwargs={
                        "model_path": model_path,
                        "performance_mode": "auto",
                    },
                )

                self.assertTrue(args.layerwise_offload_components)
                self.assertFalse(args.use_fsdp_inference)
                # dit_cpu_offload is complementary to DiT layerwise offload:
                # layerwise only moves layers on/off device at runtime, while
                # dit_cpu_offload keeps the initial weights on host memory.
                self.assertTrue(args.dit_cpu_offload)
                self.assertFalse(args.text_encoder_cpu_offload)
                self.assertFalse(args.image_encoder_cpu_offload)
                self.assertEqual(args.dit_offload_prefetch_size, 2)
                self.assertEqual(
                    args.layerwise_offload_components,
                    ["dit", "text_encoder", "image_encoder", "vae"],
                )

    def test_auto_wan2_1_14b_layerwise_offload_uses_non_dit_default(self):
        for pipeline_config, model_path in (
            (WanT2V720PConfig(), "Wan-AI/Wan2.1-T2V-14B-Diffusers"),
            (WanI2V480PConfig(), "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"),
            (WanI2V720PConfig(), "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers"),
        ):
            with self.subTest(pipeline_config=pipeline_config.__class__.__name__):
                args = self._from_dict_with_pipeline_config(
                    pipeline_config,
                    kwargs={
                        "model_path": model_path,
                        "performance_mode": "auto",
                    },
                )

                self.assertTrue(args.layerwise_offload_components)
                self.assertTrue(args.dit_cpu_offload)
                self.assertEqual(args.dit_offload_prefetch_size, 0.0)
                self.assertEqual(
                    args.layerwise_offload_components,
                    ["text_encoder", "image_encoder", "vae"],
                )

    def test_memory_wan_layerwise_offload_is_enabled_without_fsdp(self):
        args = self._from_dict_with_pipeline_config(
            WanT2V480PConfig(),
            kwargs={"performance_mode": "memory"},
        )

        self.assertTrue(args.layerwise_offload_components)
        self.assertFalse(args.use_fsdp_inference)
        self.assertTrue(args.dit_cpu_offload)
        self.assertFalse(args.text_encoder_cpu_offload)
        self.assertFalse(args.image_encoder_cpu_offload)
        self.assertEqual(
            args.layerwise_offload_components,
            ["dit", "text_encoder", "image_encoder", "vae"],
        )

    def test_auto_wan_layerwise_offload_does_not_disable_explicit_fsdp(self):
        args = self._from_dict_with_pipeline_config(
            WanT2V480PConfig(),
            kwargs={
                "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
                "num_gpus": 2,
                "performance_mode": "auto",
                "use_fsdp_inference": True,
            },
        )

        self.assertEqual(
            args.layerwise_offload_components,
            ["text_encoder", "image_encoder", "vae"],
        )
        self.assertTrue(args.use_fsdp_inference)

    def test_auto_wan_layerwise_offload_preserves_explicit_dit_cpu_offload(self):
        args = self._from_dict_with_pipeline_config(
            WanT2V480PConfig(),
            kwargs={
                "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
                "performance_mode": "auto",
                "dit_cpu_offload": True,
            },
        )

        self.assertTrue(args.dit_cpu_offload)
        self.assertEqual(
            args.layerwise_offload_components,
            ["text_encoder", "image_encoder", "vae"],
        )

    def test_auto_mova_layerwise_offload_does_not_implicitly_add_dit(self):
        args = self._from_dict_with_pipeline_config(
            MOVAPipelineConfig(),
            kwargs={
                "model_path": "OpenMOSS-Team/MOVA-360p",
                "performance_mode": "auto",
            },
        )

        self.assertTrue(args.dit_cpu_offload)
        self.assertEqual(
            args.layerwise_offload_components,
            ["text_encoder", "image_encoder", "vae"],
        )

    def test_auto_fastwan_layerwise_offload_does_not_implicitly_add_dit(self):
        args = self._from_dict_with_pipeline_config(
            FastWan2_2_TI2V_5B_Config(),
            kwargs={
                "model_path": "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
                "performance_mode": "auto",
            },
        )

        self.assertTrue(args.dit_cpu_offload)
        self.assertEqual(
            args.layerwise_offload_components,
            ["text_encoder", "image_encoder", "vae"],
        )

    def test_auto_turbo_wan_layerwise_offload_does_not_implicitly_add_dit(self):
        args = self._from_dict_with_pipeline_config(
            TurboWanT2V480PConfig(),
            kwargs={
                "model_path": "IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers",
                "performance_mode": "auto",
            },
        )

        self.assertTrue(args.dit_cpu_offload)
        self.assertEqual(
            args.layerwise_offload_components,
            ["text_encoder", "image_encoder", "vae"],
        )

    def test_explicit_fastwan_dit_layerwise_still_selects_dit_group(self):
        args = self._from_dict_with_pipeline_config(
            FastWan2_2_TI2V_5B_Config(),
            kwargs={
                "model_path": "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
                "dit_layerwise_offload": True,
            },
        )

        # dit_cpu_offload defaults to True from _adjust_offload and is now
        # preserved alongside DiT layerwise offload (the two are complementary).
        self.assertTrue(args.dit_cpu_offload)
        self.assertEqual(args.layerwise_offload_components, ["dit"])

    def test_auto_multi_gpu_wan_uses_layerwise_offload_without_cfg(self):
        with patch.object(ServerArgs, "_model_default_uses_cfg", return_value=False):
            args = self._from_dict_with_pipeline_config(
                WanT2V480PConfig(),
                kwargs={
                    "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
                    "num_gpus": 2,
                    "performance_mode": "auto",
                },
            )

        self.assertFalse(args.use_fsdp_inference)
        self.assertFalse(args.enable_cfg_parallel)
        self.assertTrue(args.dit_cpu_offload)
        self.assertTrue(args.layerwise_offload_components)
        self.assertFalse(args.text_encoder_cpu_offload)
        self.assertFalse(args.image_encoder_cpu_offload)
        self.assertEqual(
            args.layerwise_offload_components,
            ["text_encoder", "image_encoder", "vae"],
        )

    def test_explicit_multi_gpu_dit_layerwise_only_selects_dit_group(self):
        args = self._from_dict_with_pipeline_config(
            MOVAPipelineConfig(),
            kwargs={
                "model_path": "OpenMOSS-Team/MOVA-360p",
                "num_gpus": 2,
                "dit_layerwise_offload": True,
            },
        )

        self.assertFalse(args.use_fsdp_inference)
        self.assertTrue(args.dit_cpu_offload)
        self.assertTrue(args.layerwise_offload_components)
        self.assertTrue(args.text_encoder_cpu_offload)
        self.assertTrue(args.image_encoder_cpu_offload)
        self.assertEqual(args.layerwise_offload_components, ["dit"])

    def test_auto_multi_gpu_ltx_replaces_component_cpu_offload_with_resident_dit(self):
        args = self._from_dict_with_pipeline_config(
            LTX2PipelineConfig(),
            available_memory_gb=76,
            kwargs={
                "model_path": "Lightricks/LTX-2",
                "num_gpus": 2,
                "pipeline_class_name": "LTX2TwoStagePipeline",
            },
        )

        self.assertFalse(args.use_fsdp_inference)
        self.assertFalse(args.dit_cpu_offload)
        self.assertTrue(args.layerwise_offload_components)
        self.assertFalse(args.text_encoder_cpu_offload)
        self.assertFalse(args.image_encoder_cpu_offload)
        self.assertEqual(
            args.layerwise_offload_components,
            ["text_encoder", "image_encoder", "vae"],
        )

    def test_auto_high_memory_ltx23_resident_keeps_aux_components_resident(self):
        args = self._from_dict_with_pipeline_config(
            LTX2PipelineConfig(),
            memory_gb=140,
            available_memory_gb=134,
            kwargs={
                "model_path": "Lightricks/LTX-2.3",
                "num_gpus": 2,
                "pipeline_class_name": "LTX2TwoStagePipeline",
            },
        )

        self.assertEqual(args.ltx2_two_stage_device_mode, "resident")
        self.assertFalse(args.use_fsdp_inference)
        self.assertFalse(args.dit_cpu_offload)
        self.assertFalse(args.text_encoder_cpu_offload)
        self.assertFalse(args.image_encoder_cpu_offload)
        self.assertFalse(args.vae_cpu_offload)
        self.assertIsNone(args.layerwise_offload_components)

    def test_auto_high_memory_ltx23_original_keeps_default_layerwise_components(self):
        args = self._from_dict_with_pipeline_config(
            LTX2PipelineConfig(),
            memory_gb=140,
            available_memory_gb=134,
            kwargs={
                "model_path": "Lightricks/LTX-2.3",
                "num_gpus": 2,
                "pipeline_class_name": "LTX2TwoStagePipeline",
                "ltx2_two_stage_device_mode": "original",
            },
        )

        self.assertEqual(
            args.layerwise_offload_components,
            ["text_encoder", "image_encoder", "vae"],
        )

    def test_ltx23_snapshot_device_mode_is_deprecated_alias_for_original(self):
        args = self._from_dict_with_pipeline_config(
            LTX2PipelineConfig(),
            memory_gb=140,
            available_memory_gb=134,
            kwargs={
                "model_path": "Lightricks/LTX-2.3",
                "num_gpus": 2,
                "pipeline_class_name": "LTX2TwoStagePipeline",
                "ltx2_two_stage_device_mode": "snapshot",
            },
        )

        self.assertEqual(args.ltx2_two_stage_device_mode, "original")
        self.assertEqual(
            args.layerwise_offload_components,
            ["text_encoder", "image_encoder", "vae"],
        )

    def test_explicit_layerwise_components_preserved_in_ltx23_resident(self):
        args = self._from_dict_with_pipeline_config(
            LTX2PipelineConfig(),
            memory_gb=140,
            available_memory_gb=134,
            kwargs={
                "model_path": "Lightricks/LTX-2.3",
                "num_gpus": 2,
                "pipeline_class_name": "LTX2TwoStagePipeline",
                "layerwise_offload_components": ["text_encoder"],
            },
        )

        self.assertEqual(args.ltx2_two_stage_device_mode, "resident")
        self.assertEqual(args.layerwise_offload_components, ["text_encoder"])

    def test_auto_multi_gpu_qwen_keeps_vae_resident_with_cfg(self):
        args = self._from_dict_with_pipeline_config(
            QwenImagePipelineConfig(),
            kwargs={
                "model_path": "Qwen/Qwen-Image",
                "num_gpus": 2,
                "performance_mode": "auto",
            },
        )

        self.assertFalse(args.use_fsdp_inference)
        self.assertTrue(args.enable_cfg_parallel)
        # 80gb > image threshold (45gb): only vae resident, encoders offloaded;
        # cfg/dit unchanged
        self.assertTrue(args.dit_cpu_offload)
        self.assertEqual(
            args.layerwise_offload_components,
            ["text_encoder", "image_encoder"],
        )
        self.assertFalse(args.vae_cpu_offload)

    def test_auto_multi_gpu_zimage_base_prefers_fsdp(self):
        args = self._from_dict_with_pipeline_config(
            ZImagePipelineConfig(),
            kwargs={
                "model_path": "Tongyi-MAI/Z-Image",
                "num_gpus": 2,
                "performance_mode": "auto",
            },
        )

        self.assertTrue(args.use_fsdp_inference)
        self.assertTrue(args.enable_cfg_parallel)

    def test_auto_multi_gpu_zimage_turbo_skips_fsdp(self):
        args = self._from_dict_with_pipeline_config(
            ZImagePipelineConfig(),
            kwargs={
                "model_path": "Tongyi-MAI/Z-Image-Turbo",
                "num_gpus": 2,
                "performance_mode": "auto",
            },
        )

        self.assertFalse(args.use_fsdp_inference)
        self.assertFalse(args.enable_cfg_parallel)

    def test_auto_multi_gpu_qwen_preserves_explicit_fsdp_false(self):
        args = self._from_dict_with_pipeline_config(
            QwenImagePipelineConfig(),
            kwargs={
                "model_path": "Qwen/Qwen-Image",
                "num_gpus": 2,
                "performance_mode": "auto",
                "use_fsdp_inference": False,
            },
        )

        self.assertFalse(args.use_fsdp_inference)
        self.assertTrue(args.enable_cfg_parallel)
        self.assertTrue(args.dit_cpu_offload)
        self.assertFalse(args.vae_cpu_offload)
        # explicit use_fsdp_inference skips the residency pass, but the layerwise
        # filter still drops vae (kept resident); encoders stay offloaded
        self.assertEqual(
            args.layerwise_offload_components,
            ["text_encoder", "image_encoder"],
        )

    def test_auto_multi_gpu_qwen_skips_fsdp_when_available_memory_is_low(self):
        args = self._from_dict_with_pipeline_config(
            QwenImagePipelineConfig(),
            memory_gb=50,
            kwargs={
                "model_path": "Qwen/Qwen-Image",
                "num_gpus": 2,
                "performance_mode": "auto",
            },
        )

        self.assertFalse(args.use_fsdp_inference)
        self.assertTrue(args.enable_cfg_parallel)
        # 50gb still > image threshold (45gb): vae resident, encoders offloaded;
        # fsdp skipped (qwen does not opt into auto fsdp)
        self.assertTrue(args.dit_cpu_offload)
        self.assertEqual(
            args.layerwise_offload_components,
            ["text_encoder", "image_encoder"],
        )
        self.assertFalse(args.vae_cpu_offload)

    def test_auto_multi_gpu_qwen_uses_selected_gpu_min_available_memory(self):
        args = self._from_dict_with_pipeline_config(
            QwenImagePipelineConfig(),
            available_memory_gb={1: 50, 2: 80},
            kwargs={
                "model_path": "Qwen/Qwen-Image",
                "base_gpu_id": 1,
                "num_gpus": 2,
                "performance_mode": "auto",
            },
        )

        self.assertFalse(args.use_fsdp_inference)
        self.assertTrue(args.enable_cfg_parallel)

    def test_auto_multi_gpu_qwen_keeps_vae_resident_with_headroom(self):
        args = self._from_dict_with_pipeline_config(
            QwenImagePipelineConfig(),
            available_memory_gb={1: 72, 2: 80},
            kwargs={
                "model_path": "Qwen/Qwen-Image",
                "base_gpu_id": 1,
                "num_gpus": 2,
                "performance_mode": "auto",
            },
        )

        self.assertFalse(args.use_fsdp_inference)
        self.assertTrue(args.enable_cfg_parallel)
        # min available across selected gpus is 72gb > image threshold (45gb):
        # vae resident, encoders offloaded
        self.assertTrue(args.dit_cpu_offload)
        self.assertEqual(
            args.layerwise_offload_components,
            ["text_encoder", "image_encoder"],
        )
        self.assertFalse(args.vae_cpu_offload)

    def test_speed_mode_single_gpu_disables_offload(self):
        args = self._from_dict_with_pipeline_config(
            QwenImagePipelineConfig(),
            kwargs={
                "model_path": "Qwen/Qwen-Image",
                "performance_mode": "speed",
            },
        )

        self.assertEqual(args.performance_mode, "speed")
        self.assertFalse(args.use_fsdp_inference)
        self.assertFalse(args.dit_cpu_offload)
        self.assertFalse(args.layerwise_offload_components)
        self.assertFalse(args.text_encoder_cpu_offload)
        self.assertFalse(args.image_encoder_cpu_offload)

    def test_speed_mode_preserves_explicit_offload(self):
        args = self._from_dict_with_pipeline_config(
            QwenImagePipelineConfig(),
            kwargs={
                "model_path": "Qwen/Qwen-Image",
                "performance_mode": "speed",
                "dit_cpu_offload": True,
            },
        )

        self.assertEqual(args.performance_mode, "speed")
        self.assertTrue(args.dit_cpu_offload)
        self.assertFalse(args.text_encoder_cpu_offload)
        self.assertFalse(args.image_encoder_cpu_offload)

    def test_speed_mode_enables_torch_compile_by_default(self):
        args = self._from_dict_with_pipeline_config(
            QwenImagePipelineConfig(),
            kwargs={
                "model_path": "Qwen/Qwen-Image",
                "performance_mode": "speed",
            },
        )

        self.assertTrue(args.enable_torch_compile)

    def test_speed_mode_preserves_explicit_torch_compile_off(self):
        args = self._from_dict_with_pipeline_config(
            QwenImagePipelineConfig(),
            kwargs={
                "model_path": "Qwen/Qwen-Image",
                "performance_mode": "speed",
                "enable_torch_compile": False,
            },
        )

        self.assertFalse(args.enable_torch_compile)

    def test_auto_mode_leaves_torch_compile_off(self):
        args = self._from_dict_with_pipeline_config(
            QwenImagePipelineConfig(),
            kwargs={
                "model_path": "Qwen/Qwen-Image",
                "performance_mode": "auto",
            },
        )

        self.assertFalse(args.enable_torch_compile)

    def test_memory_mode_wan_uses_layerwise_offload(self):
        args = self._from_dict_with_pipeline_config(
            WanT2V480PConfig(),
            kwargs={
                "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
                "performance_mode": "memory",
            },
        )

        self.assertFalse(args.use_fsdp_inference)
        self.assertTrue(args.layerwise_offload_components)
        self.assertTrue(args.dit_cpu_offload)
        self.assertFalse(args.text_encoder_cpu_offload)
        self.assertFalse(args.image_encoder_cpu_offload)
        self.assertEqual(
            args.layerwise_offload_components,
            ["dit", "text_encoder", "image_encoder", "vae"],
        )

    def test_memory_mode_preserves_explicit_fsdp(self):
        args = self._from_dict_with_pipeline_config(
            WanT2V480PConfig(),
            kwargs={
                "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
                "num_gpus": 2,
                "performance_mode": "memory",
                "use_fsdp_inference": True,
            },
        )

        self.assertTrue(args.use_fsdp_inference)
        self.assertEqual(
            args.layerwise_offload_components,
            ["text_encoder", "image_encoder", "vae"],
        )
        self.assertFalse(args.dit_cpu_offload)

    def test_invalid_performance_mode_raises(self):
        with self.assertRaises(ValueError):
            self._from_dict_with_pipeline_config(
                QwenImagePipelineConfig(),
                kwargs={"performance_mode": "turbo"},
            )

    def test_cfg_parallel_cli_can_be_disabled_explicitly(self):
        parser = FlexibleArgumentParser()
        ServerArgs.add_cli_args(parser)
        argv = [
            "--model-path",
            "Qwen/Qwen-Image",
            "--num-gpus",
            "2",
            "--performance-mode",
            "auto",
            "--enable-cfg-parallel",
            "false",
        ]

        with (
            patch.object(sys, "argv", ["sglang"] + argv),
            patch.object(
                PipelineConfig, "from_kwargs", return_value=QwenImagePipelineConfig()
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.is_cpu",
                return_value=False,
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.is_mps",
                return_value=False,
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.is_cuda",
                return_value=True,
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.get_device_total_memory",
                return_value=80 * 1024**3,
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.get_available_gpu_memory",
                return_value=80,
            ),
        ):
            args, unknown_args = parser.parse_known_args(argv)
            server_args = ServerArgs.from_cli_args(args, unknown_args)

        self.assertFalse(server_args.use_fsdp_inference)
        self.assertFalse(server_args.enable_cfg_parallel)

    def test_ltx23_snapshot_device_mode_cli_alias_is_accepted(self):
        parser = FlexibleArgumentParser()
        ServerArgs.add_cli_args(parser)
        argv = [
            "--model-path",
            "Lightricks/LTX-2.3",
            "--pipeline-class-name",
            "LTX2TwoStagePipeline",
            "--ltx2-two-stage-device-mode",
            "snapshot",
        ]

        with (
            patch.object(sys, "argv", ["sglang"] + argv),
            patch.object(
                PipelineConfig, "from_kwargs", return_value=LTX2PipelineConfig()
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.is_cpu",
                return_value=False,
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.is_mps",
                return_value=False,
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.is_cuda",
                return_value=True,
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.get_device_total_memory",
                return_value=140 * 1024**3,
            ),
            patch(
                "sglang.multimodal_gen.runtime.platforms.current_platform.get_available_gpu_memory",
                return_value=134,
            ),
        ):
            args, unknown_args = parser.parse_known_args(argv)
            server_args = ServerArgs.from_cli_args(args, unknown_args)

        self.assertEqual(server_args.ltx2_two_stage_device_mode, "original")


class TestFSDPShardConditions(unittest.TestCase):
    def test_helpers_match_only_direct_block_entries(self):
        self.assertTrue(
            is_module_list_entry("transformer_blocks.0", "transformer_blocks")
        )
        self.assertFalse(
            is_module_list_entry("transformer_blocks.0.ff.net.0", "transformer_blocks")
        )
        self.assertTrue(
            is_module_list_entry_in(
                "single_transformer_blocks.12",
                ("transformer_blocks", "single_transformer_blocks"),
            )
        )
        self.assertFalse(
            is_module_list_entry_in(
                "single_transformer_blocks.12.attn.to_out.0",
                ("transformer_blocks", "single_transformer_blocks"),
            )
        )

    def test_qwen_dit_has_fsdp_shard_condition(self):
        conditions = QwenImageTransformer2DModel._fsdp_shard_conditions

        self.assertTrue(conditions)
        self.assertTrue(conditions[0]("transformer_blocks.0", None))
        self.assertFalse(conditions[0]("transformer_blocks.0.attn", None))
        self.assertFalse(conditions[0]("transformer_blocks.0.ff.net.0", None))

    def test_zimage_condition_keeps_inner_numbered_modules(self):
        self.assertTrue(is_zimage_layer("layers.0.mlp.0", None))
        self.assertTrue(is_zimage_layer("noise_refiner.0.attention.to_out.0", None))
        self.assertFalse(is_zimage_layer("transformer_blocks.0", None))


class TestModelIdResolution(unittest.TestCase):
    def setUp(self):
        _get_config_info.cache_clear()

    def test_model_id_overrides_arbitrary_local_path(self):
        # a local path whose directory name does not match any HF repo name;
        # --model-id tells the engine which config to use
        info = _get_config_info("/data/my-custom-qwen", model_id="Qwen-Image")
        self.assertIsNotNone(info)

        self.assertIs(info.pipeline_config_cls, QwenImagePipelineConfig)

    def test_model_id_works_after_tilde_expansion(self):
        # simulate the full flow: user passes ~/..., engine expands and resolves
        expanded = os.path.expanduser("~/.cache/huggingface/hub/bbb/snapshots/ccc")
        _get_config_info.cache_clear()
        info = _get_config_info(expanded, model_id="Qwen-Image")
        self.assertIsNotNone(info)

    def test_hf_cache_snapshot_path_resolves_registered_nvfp4_model(self):
        path = (
            "/root/.cache/huggingface/hub/"
            "models--black-forest-labs--FLUX.2-dev-NVFP4/"
            "snapshots/142b87e70bc3006937b7093d89ff287b5f59f071"
        )
        info = _get_config_info(path)
        self.assertIsNotNone(info)

    def test_sana_wm_model_path_resolves_registry(self):
        info = _get_config_info("Efficient-Large-Model/SANA-WM_bidirectional")
        self.assertIs(info.pipeline_config_cls, SanaWMPipelineConfig)

    def test_model_id_unknown_falls_back_without_crash(self):
        # unrecognized model_id: should warn and fall back to path-based detection
        # with an unresolvable path, expect RuntimeError from the detector step
        with self.assertRaises((RuntimeError, Exception)):
            _get_config_info("/data/no-such-model", model_id="NonExistentModelXYZ")


class TestPerRoleParallelism(unittest.TestCase):
    """Test per-role parallelism args and get_role_parallelism helper."""

    def _from_dict(self, kwargs):
        return _from_dict_without_model_resolution(kwargs)

    def test_defaults_are_none(self):
        args = self._from_dict({"model_path": "/fake"})
        from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType

        for role in [RoleType.ENCODER, RoleType.DENOISER, RoleType.DECODER]:
            par = args.get_role_parallelism(role)
            self.assertIsNone(par["tp_size"])
            self.assertIsNone(par["sp_degree"])
            self.assertIsNone(par["ulysses_degree"])
            self.assertIsNone(par["ring_degree"])

    def test_encoder_overrides(self):
        args = self._from_dict({"model_path": "/fake", "encoder_tp": 2})
        from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType

        par = args.get_role_parallelism(RoleType.ENCODER)
        self.assertEqual(par["tp_size"], 2)
        self.assertIsNone(par["sp_degree"])
        self.assertIsNone(par["ulysses_degree"])
        self.assertIsNone(par["ring_degree"])

    def test_denoiser_overrides(self):
        args = self._from_dict(
            {
                "model_path": "/fake",
                "denoiser_tp": 1,
                "denoiser_sp": 8,
                "denoiser_ulysses": 4,
                "denoiser_ring": 2,
            }
        )
        from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType

        par = args.get_role_parallelism(RoleType.DENOISER)
        self.assertEqual(par["tp_size"], 1)
        self.assertEqual(par["sp_degree"], 8)
        self.assertEqual(par["ulysses_degree"], 4)
        self.assertEqual(par["ring_degree"], 2)

    def test_decoder_overrides(self):
        args = self._from_dict({"model_path": "/fake", "decoder_sp": 2})
        from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType

        par = args.get_role_parallelism(RoleType.DECODER)
        self.assertIsNone(par["tp_size"])
        self.assertEqual(par["sp_degree"], 2)
        self.assertIsNone(par["ulysses_degree"])
        self.assertIsNone(par["ring_degree"])

    def test_decoder_tp_is_alias_of_decoder_sp(self):
        args = self._from_dict({"model_path": "/fake", "decoder_tp": 2})
        from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType

        self.assertEqual(args.decoder_sp, 2)
        par = args.get_role_parallelism(RoleType.DECODER)
        self.assertIsNone(par["tp_size"])
        self.assertEqual(par["sp_degree"], 2)

    def test_conflicting_decoder_tp_and_decoder_sp_raise(self):
        with self.assertRaisesRegex(ValueError, "decoder_tp is deprecated"):
            self._from_dict(
                {
                    "model_path": "/fake",
                    "decoder_tp": 2,
                    "decoder_sp": 4,
                }
            )

    def test_monolithic_returns_all_none(self):
        args = self._from_dict({"model_path": "/fake", "encoder_tp": 2})
        from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType

        par = args.get_role_parallelism(RoleType.MONOLITHIC)
        self.assertIsNone(par["tp_size"])
        self.assertIsNone(par["sp_degree"])

    def test_mixed_roles_independent(self):
        """Per-role args don't interfere with each other."""
        args = self._from_dict(
            {
                "model_path": "/fake",
                "encoder_tp": 1,
                "denoiser_tp": 2,
                "decoder_sp": 4,
            }
        )
        from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType

        self.assertEqual(args.get_role_parallelism(RoleType.ENCODER)["tp_size"], 1)
        self.assertEqual(args.get_role_parallelism(RoleType.DENOISER)["tp_size"], 2)
        self.assertEqual(args.get_role_parallelism(RoleType.DECODER)["sp_degree"], 4)

    def test_disagg_args_import_path_matches_server_args_package(self):
        from sglang.multimodal_gen.runtime.disaggregation import disagg_args
        from sglang.multimodal_gen.runtime.server_args.disagg import (
            DisaggServerArgsMixin,
        )

        self.assertIs(disagg_args.DisaggArgsMixin, DisaggServerArgsMixin)
        self.assertIs(
            disagg_args.DISAGG_RESULT_PORT_OFFSETS,
            DisaggServerArgsMixin.DISAGG_RESULT_PORT_OFFSETS,
        )

    def test_gpu_ids_normalize_lists_and_commas(self):
        args = self._from_dict({"model_path": "/fake", "gpu_ids": ["0,1", "6", "7 8"]})

        self.assertEqual(args.gpu_ids, [0, 1, 6, 7, 8])

    def test_gpu_ids_reject_duplicates(self):
        with self.assertRaisesRegex(ValueError, "duplicate GPU ids"):
            self._from_dict({"model_path": "/fake", "gpu_ids": ["0,1", "1"]})

    def test_pool_endpoints_use_role_and_scheduler_ports(self):
        args = self._from_dict(
            {
                "model_path": "/fake",
                "disagg_role": "denoiser",
                "disagg_server_addr": "tcp://127.0.0.1:30000",
                "scheduler_port": 5600,
                "host": "0.0.0.0",
                "disagg_p2p_hostname": "10.0.0.7",
            }
        )

        self.assertEqual(args.derive_pool_result_endpoint(), "tcp://127.0.0.1:30002")
        self.assertEqual(
            args.derive_pool_work_endpoint(),
            f"tcp://0.0.0.0:{args.scheduler_port}",
        )
        self.assertEqual(
            args.derive_pool_control_endpoint(),
            f"tcp://0.0.0.0:{args.scheduler_port + 1}",
        )
        self.assertEqual(
            args.derive_pool_control_advertised_endpoint(),
            f"tcp://10.0.0.7:{args.scheduler_port + 1}",
        )

    def test_pool_result_endpoint_validates_addr_and_role(self):
        args = self._from_dict({"model_path": "/fake", "disagg_server_addr": "bad"})
        with self.assertRaisesRegex(ValueError, "disagg_server_addr must be"):
            args.derive_pool_result_endpoint()

        args = self._from_dict(
            {"model_path": "/fake", "disagg_server_addr": "127.0.0.1:30000"}
        )
        with self.assertRaisesRegex(ValueError, "only defined for encoder"):
            args.derive_pool_result_endpoint()

    def test_cli_args_parsed(self):
        """Per-role parallelism args are parsed from CLI."""
        parser = FlexibleArgumentParser()
        ServerArgs.add_cli_args(parser)
        argv = [
            "--model-path",
            "/fake",
            "--denoiser-tp",
            "2",
            "--denoiser-sp",
            "4",
            "--denoiser-ulysses",
            "2",
            "--denoiser-ring",
            "2",
            "--encoder-tp",
            "1",
            "--decoder-sp",
            "8",
        ]
        args, unknown = parser.parse_known_args(argv)
        self.assertEqual(args.denoiser_tp, 2)
        self.assertEqual(args.denoiser_sp, 4)
        self.assertEqual(args.denoiser_ulysses, 2)
        self.assertEqual(args.denoiser_ring, 2)
        self.assertEqual(args.encoder_tp, 1)
        self.assertEqual(args.decoder_sp, 8)
        self.assertIsNone(args.decoder_tp)


class TestPipelineResolutionCliOverride(unittest.TestCase):
    def setUp(self):
        _get_config_info.cache_clear()

    def test_resolution_flag_overrides_qwen_image_layered_pipeline_config(self):
        parser = FlexibleArgumentParser()
        ServerArgs.add_cli_args(parser)
        argv = [
            "--model-path",
            "Qwen/Qwen-Image-Layered",
            "--resolution",
            "768",
        ]

        with (
            patch.object(sys, "argv", ["sglang"] + argv),
            _mock_cuda_platform(),
        ):
            args, unknown_args = parser.parse_known_args(argv)
            server_args = ServerArgs.from_cli_args(args, unknown_args)

        self.assertEqual(server_args.pipeline_config.resolution, 768)

    def test_disable_autocast_is_preserved_after_pipeline_config_resolution(self):
        parser = FlexibleArgumentParser()
        ServerArgs.add_cli_args(parser)
        argv = [
            "--model-path",
            "Qwen/Qwen-Image-Layered",
            "--disable-autocast",
            "true",
        ]

        with (
            patch.object(sys, "argv", ["sglang"] + argv),
            _mock_cuda_platform(),
        ):
            args, unknown_args = parser.parse_known_args(argv)
            server_args = ServerArgs.from_cli_args(args, unknown_args)

        self.assertTrue(server_args.pipeline_config.disable_autocast)
        self.assertTrue(server_args.disable_autocast)


class TestDisaggTimeoutArgs(unittest.TestCase):
    def test_disagg_defaults_match_reviewed_values(self):
        args = _from_dict_without_model_resolution({"model_path": "/fake"})
        self.assertEqual(args.disagg_max_slots_per_instance, 8)
        self.assertEqual(args.disagg_downstream_wait_timeout, 1800)
        self.assertEqual(args.disagg_timeout, 3600)

    def test_downstream_wait_timeout_cli_arg_is_parsed(self):
        parser = FlexibleArgumentParser()
        ServerArgs.add_cli_args(parser)
        argv = [
            "--model-path",
            "/fake",
            "--disagg-downstream-wait-timeout",
            "45",
        ]

        args, _unknown = parser.parse_known_args(argv)
        self.assertEqual(args.disagg_downstream_wait_timeout, 45)

    def test_disagg_timeout_help_uses_current_defaults(self):
        parser = FlexibleArgumentParser()
        ServerArgs.add_cli_args(parser)
        help_text = parser.format_help()

        self.assertIn("Default: 3600.", help_text)
        self.assertIn("Default: 1800.", help_text)

    def test_disagg_role_alias_cli_arg_is_accepted(self):
        parser = FlexibleArgumentParser()
        ServerArgs.add_cli_args(parser)
        args, _unknown = parser.parse_known_args(
            ["--model-path", "/fake", "--disagg-role", "denoising"]
        )

        self.assertEqual(args.disagg_role, "denoising")

    def test_disagg_role_alias_normalizes_to_denoiser(self):
        from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType

        args = _from_dict_without_model_resolution(
            {"model_path": "/fake", "disagg_role": "denoising"}
        )

        self.assertEqual(args.disagg_role, RoleType.DENOISER)


class TestDisaggTransferBackendArgs(unittest.TestCase):
    def test_transfer_backend_defaults_to_auto(self):
        args = _from_dict_without_model_resolution({"model_path": "/fake"})
        self.assertEqual(args.disagg_transfer_backend, "auto")

    def test_transfer_backend_cli_arg_is_parsed(self):
        parser = FlexibleArgumentParser()
        ServerArgs.add_cli_args(parser)
        argv = [
            "--model-path",
            "/fake",
            "--disagg-transfer-backend",
            "mock",
        ]

        args, _unknown = parser.parse_known_args(argv)
        self.assertEqual(args.disagg_transfer_backend, "mock")


if __name__ == "__main__":
    unittest.main()