lightseekorg--tokenspeed/test/runtime/test_cli_config_compat.py

"""Tests for vLLM-style CLI configuration arguments.

Verifies that vLLM-style CLI args are correctly parsed and mapped
to TokenSpeed's internal ServerArgs configuration.
"""

import os
import sys

# CI Registration (parsed via AST, runtime no-op)
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from ci_system.ci_register import register_cuda_ci

register_cuda_ci(est_time=10, suite="runtime-1gpu")

import argparse
import contextlib
import io
import unittest
from types import SimpleNamespace
from unittest.mock import patch

from tokenspeed.runtime.utils.server_args import ServerArgs


class TestCLIConfigCompat(unittest.TestCase):
    """Test that vLLM-style CLI arguments map to TokenSpeed config."""

    def _parse_args(self, argv: list[str]) -> argparse.Namespace:
        parser = argparse.ArgumentParser()
        ServerArgs.add_cli_args(parser)
        return parser.parse_args(argv)

    def _from_cli_args_no_init(self, args: argparse.Namespace) -> ServerArgs:
        with patch.object(ServerArgs, "__post_init__"):
            return ServerArgs.from_cli_args(args)

    def _parallelism_snapshot(self, argv: list[str]) -> tuple[int, ...]:
        args = self._parse_args(argv)
        sa = self._from_cli_args_no_init(args)
        sa.resolve_basic_defaults()
        sa.resolve_parallelism()
        mapping = sa.mapping
        return (
            mapping.world_size,
            mapping.attn.tp_size,
            mapping.attn.cp_size,
            mapping.attn.dp_size,
            mapping.dense.tp_size,
            mapping.dense.dp_size,
            mapping.moe.tp_size,
            mapping.moe.ep_size,
            mapping.moe.dp_size,
        )

    # ---- Positional model arg ----

    def test_positional_model_arg(self):
        args = self._parse_args(["deepseek-ai/DeepSeek-V3"])
        self.assertEqual(args.model_path, "deepseek-ai/DeepSeek-V3")
        self.assertIsNone(args.model)

    def test_model_flag(self):
        args = self._parse_args(["--model", "deepseek-ai/DeepSeek-V3"])
        self.assertIsNone(args.model_path)
        self.assertEqual(args.model, "deepseek-ai/DeepSeek-V3")

    def test_positional_model_resolved_in_from_cli_args(self):
        args = self._parse_args(["deepseek-ai/DeepSeek-V3"])
        sa = self._from_cli_args_no_init(args)
        self.assertEqual(sa.model, "deepseek-ai/DeepSeek-V3")

    def test_both_positional_and_model_raises(self):
        args = self._parse_args(["deepseek-ai/DeepSeek-V3", "--model", "other/model"])
        with self.assertRaises(ValueError):
            self._from_cli_args_no_init(args)

    def test_no_model_raises(self):
        args = self._parse_args([])
        with self.assertRaises(ValueError):
            self._from_cli_args_no_init(args)

    # ---- Tensor parallel size ----

    def test_tensor_parallel_size_maps_to_attn_tp_size(self):
        args = self._parse_args(
            ["--model", "test/model", "--tensor-parallel-size", "8"]
        )
        sa = self._from_cli_args_no_init(args)
        self.assertEqual(sa.attn_tp_size, 8)

    def test_tp_long_alias(self):
        args = self._parse_args(["--model", "test/model", "--tp", "4"])
        sa = self._from_cli_args_no_init(args)
        self.assertEqual(sa.attn_tp_size, 4)

    def test_tensor_parallel_aliases_match_explicit_attn_moe_tp(self):
        explicit = self._parallelism_snapshot(
            [
                "--model",
                "nvidia/Kimi-K2.5-NVFP4",
                "--attn-tp-size",
                "4",
                "--moe-tp-size",
                "4",
            ]
        )
        tensor_parallel_size = self._parallelism_snapshot(
            [
                "--model",
                "nvidia/Kimi-K2.5-NVFP4",
                "--tensor-parallel-size",
                "4",
            ]
        )
        tp = self._parallelism_snapshot(
            [
                "--model",
                "nvidia/Kimi-K2.5-NVFP4",
                "--tp",
                "4",
            ]
        )

        self.assertEqual(tensor_parallel_size, explicit)
        self.assertEqual(tp, explicit)

    def test_tensor_parallel_size_conflicts_with_attn_tp_size(self):
        args = self._parse_args(
            [
                "--model",
                "test/model",
                "--tensor-parallel-size",
                "8",
                "--attn-tp-size",
                "4",
            ]
        )
        with self.assertRaises(ValueError):
            self._from_cli_args_no_init(args)

    # ---- Enable expert parallel ----

    def test_enable_expert_parallel_flag(self):
        args = self._parse_args(["--model", "test/model", "--enable-expert-parallel"])
        sa = self._from_cli_args_no_init(args)
        self.assertTrue(sa.enable_expert_parallel)

    def test_enable_expert_parallel_default_false(self):
        args = self._parse_args(["--model", "test/model"])
        sa = self._from_cli_args_no_init(args)
        self.assertFalse(sa.enable_expert_parallel)

    # ---- vLLM config names ----

    def test_tokenizer_arg(self):
        args = self._parse_args(
            ["--model", "test/model", "--tokenizer", "my/tokenizer"]
        )
        self.assertEqual(args.tokenizer, "my/tokenizer")

    def test_max_model_len_arg(self):
        args = self._parse_args(["--model", "test/model", "--max-model-len", "4096"])
        self.assertEqual(args.max_model_len, 4096)

    def test_gpu_memory_utilization_arg(self):
        args = self._parse_args(
            ["--model", "test/model", "--gpu-memory-utilization", "0.9"]
        )
        self.assertEqual(args.gpu_memory_utilization, 0.9)

    def test_seed_arg(self):
        args = self._parse_args(["--model", "test/model", "--seed", "42"])
        self.assertEqual(args.seed, 42)

    def test_max_num_seqs_arg(self):
        args = self._parse_args(["--model", "test/model", "--max-num-seqs", "256"])
        self.assertEqual(args.max_num_seqs, 256)

    def test_dp_sampling_backend_arg_removed(self):
        with self.assertRaises(SystemExit):
            self._parse_args(
                ["--model", "test/model", "--dp-sampling-backend", "onesided"]
            )

    def test_max_prefill_tokens_arg(self):
        args = self._parse_args(
            ["--model", "test/model", "--max-prefill-tokens", "4096"]
        )
        self.assertEqual(args.max_prefill_tokens, 4096)

    def test_chunked_prefill_size_arg(self):
        args = self._parse_args(
            ["--model", "test/model", "--chunked-prefill-size", "2048"]
        )
        self.assertEqual(args.chunked_prefill_size, 2048)

    def test_prefill_token_defaults(self):
        args = self._parse_args(["--model", "test/model"])
        self.assertEqual(args.max_prefill_tokens, 8192)
        self.assertIsNone(args.chunked_prefill_size)
        self.assertFalse(args.enable_mixed_batch)

        sa = self._from_cli_args_no_init(args)
        sa.mapping = SimpleNamespace(world_size=1)
        platform = SimpleNamespace(is_amd=False, is_nvidia=False)
        with patch(
            "tokenspeed.runtime.utils.server_args.current_platform",
            return_value=platform,
        ):
            sa.resolve_memory_and_scheduling()

        self.assertEqual(sa.max_prefill_tokens, 8192)
        self.assertEqual(sa.chunked_prefill_size, 8192)
        self.assertFalse(sa.enable_mixed_batch)

    def test_mixed_batch_can_be_enabled(self):
        args = self._parse_args(["--model", "test/model", "--enable-mixed-batch"])
        self.assertTrue(args.enable_mixed_batch)

    def test_distributed_timeout_seconds_arg(self):
        args = self._parse_args(
            ["--model", "test/model", "--distributed-timeout-seconds", "600"]
        )
        self.assertEqual(args.distributed_timeout_seconds, 600)

    def test_enforce_eager_arg(self):
        args = self._parse_args(["--model", "test/model", "--enforce-eager"])
        self.assertTrue(args.enforce_eager)

    def test_cudagraph_capture_size_arg(self):
        args = self._parse_args(
            ["--model", "test/model", "--max-cudagraph-capture-size", "32"]
        )
        self.assertEqual(args.max_cudagraph_capture_size, 32)

    def test_cudagraph_capture_sizes_arg(self):
        args = self._parse_args(
            [
                "--model",
                "test/model",
                "--cudagraph-capture-sizes",
                "1",
                "2",
                "4",
            ]
        )
        self.assertEqual(args.cudagraph_capture_sizes, [1, 2, 4])

    def test_block_size_arg(self):
        args = self._parse_args(["--model", "test/model", "--block-size", "128"])
        self.assertEqual(args.block_size, 128)

    def test_moe_backend_arg(self):
        args = self._parse_args(["--model", "test/model", "--moe-backend", "triton"])
        self.assertEqual(args.moe_backend, "triton")

    def test_sampling_backend_arg(self):
        for backend in (
            "greedy",
            "flashinfer",
            "flashinfer_full",
            "triton",
            "triton_full",
        ):
            args = self._parse_args(
                ["--model", "test/model", "--sampling-backend", backend]
            )
            self.assertEqual(args.sampling_backend, backend)

    def test_all2all_backend_arg(self):
        args = self._parse_args(
            ["--model", "test/model", "--all2all-backend", "deepep"]
        )
        self.assertEqual(args.all2all_backend, "deepep")

    def test_recipe_all2all_backend_alias_arg(self):
        args = self._parse_args(
            [
                "--model",
                "test/model",
                "--all2all-backend",
                "flashinfer_nvlink_one_sided",
            ]
        )
        self.assertEqual(args.all2all_backend, "flashinfer_nvlink_one_sided")

    def test_recipe_moe_backend_alias_arg(self):
        args = self._parse_args(
            ["--model", "test/model", "--moe-backend", "deep_gemm_mega_moe"]
        )
        self.assertEqual(args.moe_backend, "deep_gemm_mega_moe")

    def test_kv_cache_dtype_fp8_alias_arg(self):
        args = self._parse_args(["--model", "test/model", "--kv-cache-dtype", "fp8"])
        sa = self._from_cli_args_no_init(args)
        sa.resolve_basic_defaults()
        self.assertEqual(sa.kv_cache_dtype, "fp8_e4m3")

    def test_tokenizer_mode_deepseek_v4_arg(self):
        args = self._parse_args(
            ["--model", "test/model", "--tokenizer-mode", "deepseek_v4"]
        )
        self.assertEqual(args.tokenizer_mode, "deepseek_v4")

    def test_hf_overrides_arg(self):
        args = self._parse_args(
            ["--model", "test/model", "--hf-overrides", '{"rope_scaling": null}']
        )
        self.assertEqual(args.hf_overrides, '{"rope_scaling": null}')

    def test_enable_log_requests_arg(self):
        args = self._parse_args(["--model", "test/model", "--enable-log-requests"])
        self.assertTrue(args.enable_log_requests)

    def test_no_enable_log_requests_arg(self):
        args = self._parse_args(["--model", "test/model", "--no-enable-log-requests"])
        self.assertFalse(args.enable_log_requests)

    def test_no_trust_remote_code_arg(self):
        args = self._parse_args(["--model", "test/model", "--no-trust-remote-code"])
        self.assertFalse(args.trust_remote_code)

    def test_enable_prefix_caching_arg(self):
        args = self._parse_args(["--model", "test/model", "--enable-prefix-caching"])
        self.assertTrue(args.enable_prefix_caching)

    def test_no_enable_prefix_caching_arg(self):
        args = self._parse_args(["--model", "test/model", "--no-enable-prefix-caching"])
        self.assertFalse(args.enable_prefix_caching)

    def test_kv_events_config_arg(self):
        config = (
            '{"publisher":"zmq","endpoint":"tcp://*:5557",'
            '"topic":"kv-events","enable_kv_cache_events":true}'
        )
        args = self._parse_args(["--model", "test/model", "--kv-events-config", config])
        sa = self._from_cli_args_no_init(args)
        self.assertEqual(sa.kv_events_config, config)

    def test_speculative_draft_quantization_defaults_to_unquant(self):
        args = self._parse_args(["--model", "test/model", "--quantization", "nvfp4"])
        self.assertEqual(args.speculative_draft_model_quantization, "unquant")

        sa = self._from_cli_args_no_init(args)
        sa.resolve_speculative_decoding()
        self.assertIsNone(sa.speculative_draft_model_quantization)

    def test_mxfp4_quantization_arg(self):
        args = self._parse_args(["--model", "test/model", "--quantization", "mxfp4"])
        self.assertEqual(args.quantization, "mxfp4")

    def test_dotted_attention_config_args(self):
        args = self._parse_args(
            [
                "--model",
                "test/model",
                "--attention_config.use_fp4_indexer_cache=True",
                "--attention-config.use_trtllm_ragged_deepseek_prefill=True",
            ]
        )
        self.assertTrue(args.attention_use_fp4_indexer_cache)
        self.assertTrue(args.use_trtllm_ragged_deepseek_prefill)

    def test_vllm_recipe_speculative_config_arg(self):
        args = self._parse_args(
            [
                "--model",
                "test/model",
                "--speculative-config",
                '{"method": "mtp", "model": "draft/model", "num_speculative_tokens": 3}',
            ]
        )
        sa = self._from_cli_args_no_init(args)
        sa.resolve_basic_defaults()
        self.assertEqual(sa.speculative_algorithm, "MTP")
        self.assertEqual(sa.speculative_draft_model_path, "draft/model")
        self.assertEqual(sa.speculative_num_steps, 3)
        self.assertEqual(sa.speculative_num_draft_tokens, 4)

    def test_speculative_config_matches_explicit_eagle3_args(self):
        draft_model = "lightseekorg/kimi-k2.5-eagle3-mla"

        config_args = self._parse_args(
            [
                "--model",
                "test/model",
                "--speculative-config",
                (
                    f'{{"model":"{draft_model}",'
                    '"method":"eagle3",'
                    '"num_speculative_tokens":1}'
                ),
            ]
        )
        explicit_args = self._parse_args(
            [
                "--model",
                "test/model",
                "--speculative-algorithm",
                "EAGLE3",
                "--speculative-draft-model-path",
                draft_model,
                "--speculative-num-steps",
                "1",
            ]
        )

        config_server_args = self._from_cli_args_no_init(config_args)
        explicit_server_args = self._from_cli_args_no_init(explicit_args)
        config_server_args.resolve_basic_defaults()
        explicit_server_args.resolve_basic_defaults()

        self.assertEqual(
            config_server_args.speculative_algorithm,
            explicit_server_args.speculative_algorithm,
        )
        self.assertEqual(
            config_server_args.speculative_draft_model_path,
            explicit_server_args.speculative_draft_model_path,
        )
        self.assertEqual(
            config_server_args.speculative_num_steps,
            explicit_server_args.speculative_num_steps,
        )
        self.assertEqual(
            config_server_args.speculative_num_draft_tokens,
            explicit_server_args.speculative_num_draft_tokens,
        )

    def test_speculative_config_matches_explicit_mtp_args(self):
        target_model = "nvidia/Qwen3.5-397B-A17B-NVFP4"

        config_args = self._parse_args(
            [
                "--model",
                target_model,
                "--speculative-config",
                '{"method":"mtp","num_speculative_tokens":3}',
            ]
        )
        explicit_args = self._parse_args(
            [
                "--model",
                target_model,
                "--speculative-algorithm",
                "MTP",
                "--speculative-num-steps",
                "3",
            ]
        )

        config_server_args = self._from_cli_args_no_init(config_args)
        explicit_server_args = self._from_cli_args_no_init(explicit_args)
        config_server_args.resolve_basic_defaults()
        explicit_server_args.resolve_basic_defaults()
        config_server_args.resolve_speculative_decoding()
        explicit_server_args.resolve_speculative_decoding()

        self.assertEqual(
            config_server_args.speculative_algorithm,
            explicit_server_args.speculative_algorithm,
        )
        self.assertEqual(
            config_server_args.speculative_draft_model_path,
            explicit_server_args.speculative_draft_model_path,
        )
        self.assertEqual(
            config_server_args.speculative_draft_model_path,
            target_model,
        )
        self.assertTrue(explicit_server_args.draft_model_path_use_base)
        self.assertEqual(
            config_server_args.speculative_num_steps,
            explicit_server_args.speculative_num_steps,
        )
        self.assertEqual(
            config_server_args.speculative_num_draft_tokens,
            explicit_server_args.speculative_num_draft_tokens,
        )

    def test_speculative_config_must_be_json_object(self):
        args = self._parse_args(["--model", "test/model", "--speculative-config", "[]"])
        sa = self._from_cli_args_no_init(args)
        with self.assertRaisesRegex(
            ValueError, "--speculative-config must be a JSON object"
        ):
            sa.resolve_basic_defaults()

    def test_speculative_defaults(self):
        args = self._parse_args(["--model", "test/model"])
        sa = self._from_cli_args_no_init(args)
        sa.resolve_basic_defaults()
        self.assertEqual(sa.speculative_num_steps, 3)
        self.assertEqual(sa.speculative_eagle_topk, 1)
        self.assertEqual(sa.speculative_num_draft_tokens, 4)

    def test_speculative_draft_tokens_default_to_steps_plus_one(self):
        args = self._parse_args(
            ["--model", "test/model", "--speculative-num-steps", "1"]
        )
        sa = self._from_cli_args_no_init(args)
        sa.resolve_basic_defaults()
        self.assertEqual(sa.speculative_num_steps, 1)
        self.assertEqual(sa.speculative_num_draft_tokens, 2)

    def test_speculative_eagle_topk_cli_rejects_non_1(self):
        # Only chain spec (topk=1) is wired end-to-end; the CLI choices
        # set is the gate, so non-1 values must fail at parse time.
        with self.assertRaises(SystemExit):
            self._parse_args(["--model", "test/model", "--speculative-eagle-topk", "4"])

    def test_speculative_eagle_topk_runtime_rejects_non_1_when_spec_on(self):
        # ServerArgs can be built programmatically (e.g. by smg_grpc_servicer),
        # bypassing argparse — keep the resolve-time defensive check covered.
        args = self._parse_args(
            [
                "--model",
                "test/model",
                "--speculative-algorithm",
                "EAGLE3",
            ]
        )
        sa = self._from_cli_args_no_init(args)
        sa.speculative_eagle_topk = 4
        sa.resolve_basic_defaults()
        with self.assertRaisesRegex(ValueError, "speculative_eagle_topk"):
            sa.resolve_speculative_decoding()

    def test_dp_sampling_is_opt_in(self):
        args = self._parse_args(["--model", "test/model"])
        sa = self._from_cli_args_no_init(args)
        self.assertFalse(sa.dp_sampling)
        self.assertIsNone(sa.dp_sampling_min_bs)

        args = self._parse_args(["--model", "test/model", "--dp-sampling"])
        sa = self._from_cli_args_no_init(args)
        self.assertTrue(sa.dp_sampling)

    def test_dp_sampling_min_bs_arg(self):
        args = self._parse_args(["--model", "test/model", "--dp-sampling-min-bs", "16"])
        sa = self._from_cli_args_no_init(args)
        self.assertEqual(sa.dp_sampling_min_bs, 16)

    # ---- Full server command example ----

    def test_full_server_command(self):
        """Test a full server command example:
        tokenspeed serve deepseek-ai/DeepSeek-V3.1 \\
          --enable-expert-parallel \\
          --tensor-parallel-size 8 \\
          --served-model-name ds31
        """
        args = self._parse_args(
            [
                "deepseek-ai/DeepSeek-V3.1",
                "--enable-expert-parallel",
                "--tensor-parallel-size",
                "8",
                "--served-model-name",
                "ds31",
            ]
        )
        sa = self._from_cli_args_no_init(args)

        self.assertEqual(sa.model, "deepseek-ai/DeepSeek-V3.1")
        self.assertEqual(sa.attn_tp_size, 8)
        self.assertTrue(sa.enable_expert_parallel)
        self.assertEqual(sa.served_model_name, "ds31")

    def test_data_parallel_size_arg(self):
        args = self._parse_args(["--model", "test/model", "--data-parallel-size", "2"])
        sa = self._from_cli_args_no_init(args)
        self.assertEqual(sa.data_parallel_size, 2)

    def test_help_uses_expected_metavars(self):
        parser = argparse.ArgumentParser()
        ServerArgs.add_cli_args(parser)

        with contextlib.redirect_stdout(io.StringIO()) as stdout:
            with self.assertRaises(SystemExit):
                parser.parse_args(["--help"])

        help_text = stdout.getvalue()
        self.assertIn("--max-num-seqs MAX_NUM_SEQS", help_text)
        self.assertIn("--max-prefill-tokens MAX_PREFILL_TOKENS", help_text)
        self.assertIn("--chunked-prefill-size CHUNKED_PREFILL_SIZE", help_text)
        self.assertIn("--gpu-memory-utilization GPU_MEMORY_UTILIZATION", help_text)
        self.assertIn(
            "--distributed-timeout-seconds DISTRIBUTED_TIMEOUT_SECONDS", help_text
        )
        self.assertIn("--all2all-backend ALL2ALL_BACKEND", help_text)
        self.assertIn("--hf-overrides HF_OVERRIDES", help_text)
        self.assertNotIn("MAX_RUNNING_REQUESTS", help_text)
        self.assertNotIn("MEM_FRACTION_STATIC", help_text)
        self.assertNotIn("DIST_TIMEOUT", help_text)
        self.assertNotIn("MOE_A2A_BACKEND", help_text)
        self.assertNotIn("JSON_MODEL_OVERRIDE_ARGS", help_text)


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