diff --git a/tensorrt_llm/commands/mooncake.py b/tensorrt_llm/commands/mooncake.py index e7d6acaf5c27..5d1a7be53f4a 100644 --- a/tensorrt_llm/commands/mooncake.py +++ b/tensorrt_llm/commands/mooncake.py @@ -33,15 +33,17 @@ import tensorrt_llm.usage as usage from tensorrt_llm.commands import _telemetry as _command_telemetry +from tensorrt_llm.commands._serve_stability import stability_option from tensorrt_llm.logger import logger from tensorrt_llm.usage.config import UsageContext #: Both commands sit in the telemetry-aware `trtllm-serve` group, so they offer #: the same opt-out the group documents. -_telemetry_option = click.option( +_telemetry_option = stability_option( "--telemetry/--no-telemetry", default=True, help="Enable or disable anonymous usage telemetry collection.", + status="beta", ) @@ -88,29 +90,32 @@ def _signal_handoff(): @click.command("mooncake_master") -@click.option( +@stability_option( "--rpc_port", type=int, default=50051, show_default=True, help="Port the store clients reach the master on.", + status="prototype", ) -@click.option( +@stability_option( "--metrics_port", type=int, default=9004, show_default=True, help="Prometheus port. Pool occupancy and eviction are read " "from here or from the master's log.", + status="prototype", ) -@click.option( +@stability_option( "--eviction_ratio", type=float, default=0.05, show_default=True, help="Fraction of the pool freed per eviction pass.", + status="prototype", ) -@click.option( +@stability_option( "--pool_file", type=str, default=None, @@ -119,8 +124,9 @@ def _signal_handoff(): "file://, which is how they reach a master whose host " "the scheduler chose. Removed on exit so a stale address is " "never dialed. One is written to --run_dir regardless.", + status="prototype", ) -@click.option( +@stability_option( "--protocol", type=str, default="rdma", @@ -128,30 +134,34 @@ def _signal_handoff(): help="Transport every participant will use: 'rdma' or 'tcp'. " "Recorded in the manifest so the servers need not restate it. " "TCP is for bring-up only; it invalidates performance conclusions.", + status="prototype", ) -@click.option( +@stability_option( "--metadata_server", type=str, default="P2PHANDSHAKE", show_default=True, help="Mooncake metadata service, recorded in the manifest. " "P2PHANDSHAKE keeps a separate metadata process out of the deployment.", + status="prototype", ) -@click.option( +@stability_option( "--namespace", type=str, default="trtllm", show_default=True, help="Default key namespace for this pool, recorded in the manifest. A server may override it.", + status="prototype", ) -@click.option( +@stability_option( "--run_dir", type=str, default=None, help="Where to keep the master's log and a copy of the " "manifest. Defaults to a temporary directory.", + status="prototype", ) -@click.option( +@stability_option( "--binary", type=str, default="mooncake_master", @@ -159,20 +169,23 @@ def _signal_handoff(): help="The master executable to run, looked up on PATH. It " "ships with the Mooncake runtime that " "docker/common/install_mooncake.sh installs.", + status="prototype", ) -@click.option( +@stability_option( "--timeout", type=float, default=60.0, show_default=True, help="Seconds to wait for the master to accept connections before giving up on it.", + status="prototype", ) -@click.option( +@stability_option( "--heartbeat_seconds", type=int, default=300, show_default=True, help="Interval between liveness lines. 0 disables them.", + status="prototype", ) @_telemetry_option def mooncake_master( @@ -245,19 +258,21 @@ def mooncake_master( @click.command("mooncake_pool_report") -@click.option( +@stability_option( "--run_dir", type=str, required=True, help="Directory holding the pool manifest and the segments/ " "records each participating rank wrote.", + status="prototype", ) -@click.option( +@stability_option( "--master_log", type=str, default=None, help="The master's log, for the block-placement table. " "Defaults to mooncake_master.log in --run_dir.", + status="prototype", ) @_telemetry_option def mooncake_pool_report(run_dir: str, master_log: Optional[str], telemetry: bool): diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 184bd8539e20..9470bf3e6295 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -1817,63 +1817,80 @@ def serve_encoder( @click.command("embeddings") @click.argument("model", type=str) -@click.option("--host", - type=str, - default="localhost", - help="Hostname of the server.") -@click.option("--port", type=int, default=8000, help="Port of the server.") -@click.option('--log_level', - type=click.Choice(severity_map.keys()), - default='info', - help="The logging level.") -@click.option("--max_batch_size", - type=int, - default=_LLM_ARGS_FIELDS["max_batch_size"].default, - help="Maximum batch size coalesced into a single encode() call.") -@click.option( +@stability_option("--host", + type=str, + default="localhost", + help="Hostname of the server.", + status="beta") +@stability_option("--port", + type=int, + default=8000, + help="Port of the server.", + status="beta") +@stability_option('--log_level', + type=click.Choice(severity_map.keys()), + default='info', + help="The logging level.", + status="beta") +@stability_option( + "--max_batch_size", + type=int, + default=_LLM_ARGS_FIELDS["max_batch_size"].default, + help="Maximum batch size coalesced into a single encode() call.", + status="beta") +@stability_option( "--max_num_tokens", type=int, default=8192, - help="Maximum number of batched input tokens in each encode() call.") -@click.option( + help="Maximum number of batched input tokens in each encode() call.", + status="beta") +@stability_option( "--max_queue_delay", type=click.FloatRange(min=0.0), default=0.005, help="Dynamic-batching hold window in seconds: how long an incoming request " "waits for others to join its batch before being dispatched (mirrors Triton's " - "max_queue_delay_microseconds).") -@click.option( + "max_queue_delay_microseconds).", + status="prototype") +@stability_option( "--max_queue_size", type=click.IntRange(min=1), default=2048, help="Maximum number of in-flight queued requests; further requests are " - "rejected with HTTP 429 (mirrors Triton's max_queue_size).") -@click.option("--trust_remote_code", - is_flag=True, - default=False, - help="Flag for HF transformers.") -@click.option( + "rejected with HTTP 429 (mirrors Triton's max_queue_size).", + status="prototype") +@stability_option("--trust_remote_code", + is_flag=True, + default=False, + help="Flag for HF transformers.", + status="beta") +@stability_option( "--config", "--extra_llm_api_options", "extra_llm_api_options", type=str, default=None, help="Path to a YAML configuration file. Explicit CLI flags take precedence " - "over values in this file.") -@click.option("--hf_revision", - "--revision", - "revision", - type=str, - default=None, - help="The revision to use for the HuggingFace model " - "(branch name, tag name, or commit id).") -@click.option("--metadata_server_config_file", - type=str, - default=None, - help="Path to metadata server config file") -@click.option("--telemetry/--no-telemetry", - default=True, - help="Enable or disable anonymous usage telemetry collection.") + "over values in this file.", + status="prototype") +@stability_option("--hf_revision", + "--revision", + "revision", + type=str, + default=None, + help="The revision to use for the HuggingFace model " + "(branch name, tag name, or commit id).", + status="beta") +@stability_option("--metadata_server_config_file", + type=str, + default=None, + help="Path to metadata server config file", + status="prototype") +@stability_option( + "--telemetry/--no-telemetry", + default=True, + help="Enable or disable anonymous usage telemetry collection.", + status="beta") def serve_embedding( model: str, host: str, diff --git a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml index 8c06a6d0ea57..fb16a92f82d1 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml @@ -3,10 +3,13 @@ # Stability reference for the `trtllm-serve` CLI surface. # -# Every Click option on the `serve`, `disaggregated`, `disaggregated_mpi_worker`, -# and `mm_embedding_serve` subcommands MUST be listed here with an explicit -# stability status. The CI test (`tests/unittest/api_stability/test_serve_cli.py`) -# diffs the live Click group against this file and fails on any divergence. +# Every Click option of every subcommand registered on `trtllm-serve` MUST be +# listed here with an explicit stability status. The gated subcommand list is +# derived from `main.commands` in tensorrt_llm/commands/serve.py, so a newly +# registered subcommand is gated automatically and needs a section here. The +# CI test (`tests/unittest/api_stability/test_serve_cli.py`) diffs the live +# Click group against this file and fails on any divergence, including a +# subcommand with no section and a section for an unregistered subcommand. # # Allowed statuses (forward transitions only): # prototype -> beta -> stable -> deprecated (-> removed after grace period) @@ -781,3 +784,257 @@ commands: is_flag: true flags: - "--trust_remote_code" + embeddings: + options: + extra_llm_api_options: + type: str + default: null + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--config" + - "--extra_llm_api_options" + host: + type: str + default: localhost + status: beta + required: false + multiple: false + is_flag: false + flags: + - "--host" + log_level: + type: Choice(['debug', 'error', 'info', 'internal_error', 'trace', 'verbose', 'warning']) + default: info + status: beta + required: false + multiple: false + is_flag: false + flags: + - "--log_level" + max_batch_size: + type: int + default: 2048 + status: beta + required: false + multiple: false + is_flag: false + flags: + - "--max_batch_size" + max_num_tokens: + type: int + default: 8192 + status: beta + required: false + multiple: false + is_flag: false + flags: + - "--max_num_tokens" + max_queue_delay: + type: float + default: 0.005 + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--max_queue_delay" + max_queue_size: + type: int + default: 2048 + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--max_queue_size" + metadata_server_config_file: + type: str + default: null + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--metadata_server_config_file" + port: + type: int + default: 8000 + status: beta + required: false + multiple: false + is_flag: false + flags: + - "--port" + revision: + type: str + default: null + status: beta + required: false + multiple: false + is_flag: false + flags: + - "--hf_revision" + - "--revision" + telemetry: + type: bool + default: true + status: beta + required: false + multiple: false + is_flag: true + flags: + - "--no-telemetry" + - "--telemetry" + trust_remote_code: + type: bool + default: false + status: beta + required: false + multiple: false + is_flag: true + flags: + - "--trust_remote_code" + mooncake_master: + options: + binary: + type: str + default: mooncake_master + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--binary" + eviction_ratio: + type: float + default: 0.05 + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--eviction_ratio" + heartbeat_seconds: + type: int + default: 300 + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--heartbeat_seconds" + metadata_server: + type: str + default: P2PHANDSHAKE + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--metadata_server" + metrics_port: + type: int + default: 9004 + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--metrics_port" + namespace: + type: str + default: trtllm + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--namespace" + pool_file: + type: str + default: null + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--pool_file" + protocol: + type: str + default: rdma + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--protocol" + rpc_port: + type: int + default: 50051 + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--rpc_port" + run_dir: + type: str + default: null + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--run_dir" + telemetry: + type: bool + default: true + status: beta + required: false + multiple: false + is_flag: true + flags: + - "--no-telemetry" + - "--telemetry" + timeout: + type: float + default: 60.0 + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--timeout" + mooncake_pool_report: + options: + master_log: + type: str + default: null + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--master_log" + run_dir: + type: str + default: null + status: prototype + required: true + multiple: false + is_flag: false + flags: + - "--run_dir" + telemetry: + type: bool + default: true + status: beta + required: false + multiple: false + is_flag: true + flags: + - "--no-telemetry" + - "--telemetry" diff --git a/tests/unittest/api_stability/test_serve_cli.py b/tests/unittest/api_stability/test_serve_cli.py index 95015ed200ec..6fac7ff82a0e 100644 --- a/tests/unittest/api_stability/test_serve_cli.py +++ b/tests/unittest/api_stability/test_serve_cli.py @@ -61,19 +61,52 @@ REFERENCE_PATH = pathlib.Path(__file__).parent / "references" / "trtllm_serve_cli.yaml" -# Subcommands we gate. Names must match the keys in `main.commands` in -# tensorrt_llm/commands/serve.py. -GATED_SUBCOMMANDS = ( - "serve", - "disaggregated", - "disaggregated_mpi_worker", - "mm_embedding_serve", -) - - -def _load_reference() -> dict[str, Any]: - with open(REFERENCE_PATH) as f: - return yaml.safe_load(f) +# Every subcommand registered on `trtllm-serve` is gated. The list is derived +# from `main.commands` (in registration order) so that a newly registered +# subcommand is gated automatically: its options must carry stability tags and +# the reference YAML must gain a section for it, or this file's tests fail. +GATED_SUBCOMMANDS = tuple(serve_cmd.main.commands) + +_YAML_MERGE_TAG = "tag:yaml.org,2002:merge" + + +class _UniqueKeyLoader(yaml.SafeLoader): + """SafeLoader that rejects a repeated mapping key instead of keeping its last value. + + With plain ``yaml.safe_load`` a second ``serve:`` section or a repeated + option entry silently replaces the first, so the gate would compare + against whichever copy came last. + """ + + def __init__(self, stream: str) -> None: + super().__init__(stream) + self._checked_mappings: set[yaml.MappingNode] = set() + + def flatten_mapping(self, node: yaml.MappingNode) -> None: + # SafeLoader flattens merge sources before constructing them. Check each + # mapping's own keys once, so an explicit override of a merged key stays legal. + if node in self._checked_mappings: + return + self._checked_mappings.add(node) + seen: set[object] = set() + for key_node, _ in node.value: + # Merge keys and (unhashable) non-scalar keys are left to SafeLoader. + if not isinstance(key_node, yaml.ScalarNode) or key_node.tag == _YAML_MERGE_TAG: + continue + key = self.construct_object(key_node) + if key in seen: + raise yaml.constructor.ConstructorError( + "while constructing a mapping", + node.start_mark, + f"found duplicate key {key!r}", + key_node.start_mark, + ) + seen.add(key) + super().flatten_mapping(node) + + +def _load_reference(path: pathlib.Path = REFERENCE_PATH) -> dict[str, Any]: + return yaml.load(path.read_text(encoding="utf-8"), Loader=_UniqueKeyLoader) def _live_surface() -> dict[str, dict[str, dict[str, Any]]]: @@ -102,6 +135,40 @@ def reference(self) -> dict[str, Any]: def live(self) -> dict[str, dict[str, dict[str, Any]]]: return _live_surface() + def test_reference_lists_exactly_the_registered_subcommands(self, reference): + """The YAML must have exactly one section per registered subcommand. + + `GATED_SUBCOMMANDS` is derived from `main.commands`, so a new + subcommand is gated automatically. This check keeps the YAML in step + in both directions: a registered subcommand without a section would + otherwise ship ungated options, and a section for a subcommand that + is no longer registered would never be read again. + """ + registered = set(serve_cmd.main.commands) + listed = set(reference.get("commands", {}) or {}) + errors: list[str] = [] + missing = sorted(registered - listed) + if missing: + errors.append( + f"trtllm-serve subcommand(s) {missing} have no section in " + "tests/unittest/api_stability/references/trtllm_serve_cli.yaml. " + "Add `commands: : options:` with one entry per option " + "(status: prototype | beta | stable | deprecated) and declare " + "each option with stability_option(..., status=...). The API " + "review committee approves changes to this file." + ) + extra = sorted(listed - registered) + if extra: + errors.append( + "tests/unittest/api_stability/references/trtllm_serve_cli.yaml " + f"lists subcommand(s) {extra} that trtllm-serve no longer " + "registers. If one was removed or renamed, was it deprecated " + "for at least one minor release first? Remove the section in " + "the same PR; the API review committee approves." + ) + if errors: + pytest.fail("\n".join(errors)) + def test_every_option_has_a_stability_tag(self, live): """Every CLI option must carry an explicit stability tag. @@ -257,3 +324,41 @@ def test_no_status_demotion(self, reference, live): if demotions: pytest.fail("Forbidden stability demotions:\n" + "\n".join(demotions)) + + +@pytest.mark.parametrize( + "text, duplicate", + [ + pytest.param( + "commands:\n serve:\n options: {}\n serve:\n options: {}\n", + "serve", + id="subcommand-section", + ), + pytest.param( + "commands:\n serve:\n options:\n" + " port: {status: stable}\n port: {status: beta}\n", + "port", + id="option-entry", + ), + ], +) +def test_reference_with_a_duplicate_key_fails_to_load(tmp_path, text, duplicate): + """A repeated section or option fails to load instead of collapsing into one.""" + path = tmp_path / "trtllm_serve_cli.yaml" + path.write_text(text, encoding="utf-8") + with pytest.raises(yaml.constructor.ConstructorError, match=f"duplicate key '{duplicate}'"): + _load_reference(path) + + +def test_reference_may_override_a_merged_key(tmp_path): + """Keys pulled in through a YAML merge (``<<``) are not duplicates of the mapping's own.""" + path = tmp_path / "trtllm_serve_cli.yaml" + path.write_text( + "base: &base {status: beta, type: int}\n" + "derived: &derived {<<: *base, status: stable}\n" + "nested: {<<: *derived}\n", + encoding="utf-8", + ) + reference = _load_reference(path) + assert reference["derived"] == {"status": "stable", "type": "int"} + assert reference["nested"] == {"status": "stable", "type": "int"}