@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 91 lines 3.59 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Sanity tests for the small registry / CLI-arg / config additions made
4 by the Cohere v2 chat API:
5
6 * ``vllm/renderers/registry.py``: new ``"cohere"`` renderer entry.
7 * ``vllm/tokenizers/registry.py``: new ``"cohere"`` tokenizer entry
8 (aliased to the cached HF tokenizer).
9 * ``vllm/entrypoints/openai/cli_args.py``: new
10 ``cohere_is_reasoning_model`` field.
11 * ``vllm/config/model.py``: ``"cohere"`` added to ``TokenizerMode``
12 Literal.
13 """
14
15 import dataclasses
16 import typing
17
18 import pytest
19
20 from vllm.config.model import TokenizerMode
21 from vllm.entrypoints.launchers.cli_args import BaseFrontendArgs, make_arg_parser
22 from vllm.renderers.registry import RENDERER_REGISTRY
23 from vllm.tokenizers.registry import TokenizerRegistry
24 from vllm.utils.argparse_utils import FlexibleArgumentParser
25
26
27 class TestRendererRegistry:
28 def test_cohere_renderer_registered(self):
29 # The registry resolves to the importable ``CohereRenderer`` class.
30 cls = RENDERER_REGISTRY.load_renderer_cls("cohere")
31 assert cls.__name__ == "CohereRenderer"
32 # Sanity: class lives in the cohere renderer module.
33 assert cls.__module__ == "vllm.renderers.cohere"
34
35
36 class TestTokenizerRegistry:
37 def test_cohere_aliased_to_cached_hf_tokenizer(self):
38 # ``cohere`` mode uses the standard HF tokenizer; only the
39 # renderer stage is replaced. This test guards against accidental
40 # divergence.
41 cls = TokenizerRegistry.load_tokenizer_cls("cohere")
42 assert cls.__name__ == "CachedHfTokenizer"
43
44
45 class TestTokenizerModeLiteral:
46 def test_cohere_is_a_valid_tokenizer_mode(self):
47 # The Literal must enumerate ``"cohere"`` so engine arg parsing
48 # accepts ``--tokenizer-mode cohere``.
49 modes = typing.get_args(TokenizerMode)
50 assert "cohere" in modes
51
52
53 # ----------------------------------------------------------------------
54 # ``--cohere-is-reasoning-model`` CLI flag
55 # ----------------------------------------------------------------------
56
57
58 class TestCohereCliArg:
59 """Verifies the new ``--cohere-is-reasoning-model`` flag end-to-end
60 through ``make_arg_parser`` — mirrors the pattern used in
61 :mod:`tests.entrypoints.openai.test_cli_args`.
62 """
63
64 def test_default_value_on_dataclass_is_true(self):
65 fields = {f.name: f for f in dataclasses.fields(BaseFrontendArgs)}
66 assert "cohere_is_reasoning_model" in fields
67 field = fields["cohere_is_reasoning_model"]
68 assert field.default is True
69 assert field.type is bool
70
71 @pytest.fixture
72 def serve_parser(self) -> FlexibleArgumentParser:
73 parser = FlexibleArgumentParser()
74 return make_arg_parser(parser)
75
76 def test_default_via_argparse_is_true(self, serve_parser: FlexibleArgumentParser):
77 # No flag supplied → dataclass default (True) wins.
78 args = serve_parser.parse_args(["--model", "m"])
79 assert args.cohere_is_reasoning_model is True
80
81 def test_explicit_false_via_argparse(self, serve_parser: FlexibleArgumentParser):
82 # Boolean dataclass fields are wired up as ``--flag value`` /
83 # ``--no-flag`` pairs by FlexibleArgumentParser.
84 args = serve_parser.parse_args(
85 ["--model", "m", "--no-cohere-is-reasoning-model"]
86 )
87 assert args.cohere_is_reasoning_model is False
88
89 def test_explicit_true_via_argparse(self, serve_parser: FlexibleArgumentParser):
90 args = serve_parser.parse_args(["--model", "m", "--cohere-is-reasoning-model"])
91 assert args.cohere_is_reasoning_model is True