diff --git a/README.md b/README.md index f36926afc608..f47898bf1a24 100644 --- a/README.md +++ b/README.md @@ -325,8 +325,9 @@ deployment, and not linked to users. The data we collect includes: - Feature summary flags (LoRA, speculative decoding, prefix caching, CUDA graphs, chunked context, data parallelism) - Disaggregated serving metadata (role and deployment ID) - Process-local LLM lifecycle counts (initialization attempts, successful and active instances, peak concurrency, and initialization failures) +- VisualGen pipeline identity, capability, parallelism, feature settings, and process-local lifecycle counts - Bounded terminal status for instrumented processes (known exit code or signal, termination category, lifecycle phase, component, and reporting source) -- Selected LLM API configuration values: parallelism, dtype, KV cache, scheduler, CUDA graph, and compile settings +- Selected LLM and VisualGen configuration values: parallelism, dtype, caching, attention, CUDA graph, and compile settings - Capture diagnostics for that payload: a schema checksum (for provenance), the count of captured fields, and whether any free-form value was skipped Telemetry is automatically disabled in CI and test environments. @@ -337,7 +338,7 @@ To disable telemetry data collection, use any of the following methods: - **Environment variable**: Set `TRTLLM_NO_USAGE_STATS=1`, `DO_NOT_TRACK=1`, or `TELEMETRY_DISABLED=true` - **File-based**: Create the file `~/.config/trtllm/do_not_track` -- **Python API**: Pass `TelemetryConfig(disabled=True)` to `LLM()` +- **Python API**: Pass `TelemetryConfig(disabled=True)` to `LLM()`, or pass `VisualGenArgs(telemetry_config=TelemetryConfig(disabled=True))` to `VisualGen()` - **CLI flag**: Use `--no-telemetry` on `trtllm-serve`, `trtllm-bench`, or `trtllm-eval` CLI terminal telemetry may report that an argument or YAML configuration was diff --git a/docs/source/_ext/llmapi_config_telemetry.py b/docs/source/_ext/llmapi_config_telemetry.py index 0b0bf9e63427..be15be822bf0 100644 --- a/docs/source/_ext/llmapi_config_telemetry.py +++ b/docs/source/_ext/llmapi_config_telemetry.py @@ -33,7 +33,7 @@ [telemetry schema reference](source:tensorrt_llm/usage/schemas/README.md) for the wire schema. -**No PII or free-form fields are captured.** LLM API configuration capture is +**No PII or free-form fields are captured.** Runtime configuration capture is *type-driven*: fields whose type is categorical (`Literal`/`Enum`/`bool`) or numeric (`int`/`float`), plus safe collections of those, are captured automatically. Free-form `str`/`Any`/`Path`/`dict`/`Callable` are never captured @@ -44,7 +44,7 @@ If the manifest check fails, run `python3 scripts/generate_llm_args_golden_manifest.py`, then commit `tensorrt_llm/usage/llm_args_golden_manifest.json`; new fields require telemetry/privacy CODEOWNER approval. -## LLM API Configuration Fields +## Runtime Configuration Fields A field can still be absent from a specific payload when its parent config is unset or when the safety sanitizer rejects the runtime value. @@ -76,16 +76,22 @@ def _table(rows: list[dict]) -> str: def generate_telemetry_reference(repo_root: Path | str, output_path: Path | str) -> None: repo_root = Path(repo_root) golden = json.loads((repo_root / _GOLDEN_REL).read_text()) - rows = golden.get("TorchLlmArgs", []) - content = [ - _REFERENCE_PREAMBLE, - "### `TorchLlmArgs`", - "", - f"{len(rows)} captured fields.", - "", - _table(rows), - "", - ] + preferred = ("TorchLlmArgs", "VisualGenArgs") + model_names = [name for name in preferred if name in golden] + model_names.extend(sorted(set(golden) - set(model_names))) + content = [_REFERENCE_PREAMBLE] + for model_name in model_names: + rows = golden[model_name] + content.extend( + [ + f"### `{model_name}`", + "", + f"{len(rows)} captured fields.", + "", + _table(rows), + "", + ] + ) output = Path(output_path) output.parent.mkdir(parents=True, exist_ok=True) output.write_text("\n".join(content)) diff --git a/docs/source/developer-guide/telemetry.md b/docs/source/developer-guide/telemetry.md index 27ff48294b6e..82c30b982459 100644 --- a/docs/source/developer-guide/telemetry.md +++ b/docs/source/developer-guide/telemetry.md @@ -10,7 +10,7 @@ for the user-facing collection and opt-out overview, and the [telemetry schema reference](source:tensorrt_llm/usage/schemas/README.md) for the wire schema. -**No PII or free-form fields are captured.** LLM API configuration capture is +**No PII or free-form fields are captured.** Runtime configuration capture is *type-driven*: fields whose type is categorical (`Literal`/`Enum`/`bool`) or numeric (`int`/`float`), plus safe collections of those, are captured automatically. Free-form `str`/`Any`/`Path`/`dict`/`Callable` are never captured @@ -21,14 +21,14 @@ below; the runtime can capture nothing absent from this list. If the manifest check fails, run `python3 scripts/generate_llm_args_golden_manifest.py`, then commit `tensorrt_llm/usage/llm_args_golden_manifest.json`; new fields require telemetry/privacy CODEOWNER approval. -## LLM API Configuration Fields +## Runtime Configuration Fields A field can still be absent from a specific payload when its parent config is unset or when the safety sanitizer rejects the runtime value. ### `TorchLlmArgs` -294 captured fields. +298 captured fields. | Captured key | Annotation | Kind | Converter | Allowed values | |--------------|------------|------|-----------|----------------| @@ -51,12 +51,14 @@ unset or when the safety sanitizer rejects the runtime value. | `batch_wait_timeout_iters` | `` | `value` | | | | `batch_wait_timeout_ms` | `` | `value` | | | | `cache_transceiver_config.backend` | `Optional[Literal['DEFAULT', 'UCX', 'NIXL', 'MOONCAKE', 'MPI']]` | `categorical` | | `DEFAULT`, `UCX`, `NIXL`, `MOONCAKE`, `MPI` | +| `cache_transceiver_config.enable_pipelined_transfer` | `` | `value` | | | | `cache_transceiver_config.kv_cache_bounce_size_mb` | `` | `value` | | | | `cache_transceiver_config.kv_transfer_poll_interval_ms` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` | | | | `cache_transceiver_config.kv_transfer_sender_future_timeout_ms` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` | | | | `cache_transceiver_config.kv_transfer_timeout_ms` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` | | | | `cache_transceiver_config.max_tokens_in_buffer` | `Optional[int]` | `value` | | | | `cache_transceiver_config.transceiver_runtime` | `Optional[Literal['CPP', 'PYTHON', 'auto']]` | `categorical` | | `CPP`, `PYTHON`, `auto` | +| `checkpoint_io_policy` | `Literal['auto', 'native', 'rank_striped_read_ahead']` | `categorical` | | `auto`, `native`, `rank_striped_read_ahead` | | `context_parallel_size` | `` | `value` | | | | `cp_config.cp_type` | `` | `categorical` | | `ULYSSES`, `RING`, `HELIX` | | `cp_config.fifo_version` | `Optional[int]` | `value` | | | @@ -112,11 +114,12 @@ unset or when the safety sanitizer rejects the runtime value. | `gpus_per_node` | `Optional[int]` | `value` | | | | `guided_decoding_backend` | `Optional[Literal['xgrammar', 'llguidance']]` | `categorical` | | `xgrammar`, `llguidance` | | `iter_stats_max_iterations` | `Optional[int]` | `value` | | | -| `kv_cache_compression_config.algorithm` | `Literal['triattention']` | `categorical` | | `triattention` | +| `kv_cache_compression_config.algorithm` | `Literal['triattention']` | `categorical` | allowlist | `quantization_for_cold_page`, `triattention` | | `kv_cache_compression_config.beta` | `` | `value` | | | | `kv_cache_compression_config.budget` | `` | `value` | | | | `kv_cache_compression_config.eviction_mode` | `Literal['union', 'per_head', 'per_layer_perhead']` | `categorical` | | `union`, `per_head`, `per_layer_perhead` | | `kv_cache_compression_config.normalize_scores` | `` | `value` | | | +| `kv_cache_compression_config.quant` | `Literal['nvfp4']` | `categorical` | | `nvfp4` | | `kv_cache_config.attention_dp_events_gather_period_ms` | `` | `value` | | | | `kv_cache_config.avg_seq_len` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` | | | | `kv_cache_config.block_reuse_config.max_num_turns` | `` | `value` | | | @@ -125,7 +128,7 @@ unset or when the safety sanitizer rejects the runtime value. | `kv_cache_config.cross_kv_cache_fraction` | `Optional[float]` | `value` | | | | `kv_cache_config.disk_cache_size` | `Optional[Annotated[int, Ge(ge=0)]]` | `value` | | | | `kv_cache_config.disk_prefetch_num_reqs` | `` | `value` | | | -| `kv_cache_config.dtype` | `` | `categorical` | allowlist | `auto`, `float16`, `bfloat16`, `float32`, `fp8`, `nvfp4` | +| `kv_cache_config.dtype` | `` | `categorical` | allowlist | `auto`, `float16`, `bfloat16`, `float32`, `fp8`, `fp8_ds_mla`, `nvfp4` | | `kv_cache_config.enable_block_reuse` | `` | `value` | | | | `kv_cache_config.enable_kv_pool_rebalance` | `` | `value` | | | | `kv_cache_config.enable_partial_reuse` | `` | `value` | | | @@ -154,10 +157,12 @@ unset or when the safety sanitizer rejects the runtime value. | `layer_wise_benchmarks_config.calibration_layer_indices` | `Optional[List[int]]` | `value` | | | | `layer_wise_benchmarks_config.calibration_mode` | `Literal['NONE', 'MARK', 'COLLECT']` | `categorical` | | `NONE`, `MARK`, `COLLECT` | | `load_format` | `Union[str, tensorrt_llm.llmapi.llm_args.LoadFormat]` | `categorical` | allowlist | `auto`, `dummy`, `vision_only`, `gms` | +| `lora_config.cuda_graph_specialize_lora` | `` | `value` | | | | `lora_config.lora_ckpt_source` | `Literal['hf', 'nemo']` | `categorical` | | `hf`, `nemo` | | `lora_config.max_cpu_loras` | `Optional[int]` | `value` | | | | `lora_config.max_lora_rank` | `` | `value` | | | | `lora_config.max_loras` | `Optional[int]` | `value` | | | +| `lora_config.overlap_lora_and_base` | `` | `value` | | | | `lora_config.swap_gate_up_proj_lora_b_weight` | `` | `value` | | | | `max_batch_size` | `Optional[int]` | `value` | | | | `max_beam_width` | `Optional[int]` | `value` | | | @@ -167,7 +172,7 @@ unset or when the safety sanitizer rejects the runtime value. | `max_stats_len` | `` | `value` | | | | `mm_encoder_only` | `` | `value` | | | | `moe_cluster_parallel_size` | `Optional[int]` | `value` | | | -| `moe_config.backend` | `Literal['AUTO', 'CUTLASS', 'CUTEDSL', 'WIDEEP', 'TRTLLM', 'DEEPGEMM', 'DENSEGEMM', 'VANILLA', 'TRITON', 'MARLIN', 'MEGAMOE_DEEPGEMM', 'MEGAMOE_CUTEDSL']` | `categorical` | | `AUTO`, `CUTLASS`, `CUTEDSL`, `WIDEEP`, `TRTLLM`, `DEEPGEMM`, `DENSEGEMM`, `VANILLA`, `TRITON`, `MARLIN`, `MEGAMOE_DEEPGEMM`, `MEGAMOE_CUTEDSL` | +| `moe_config.backend` | `Literal['AUTO', 'CUTLASS', 'CUTEDSL', 'TRTLLM', 'DEEPGEMM', 'DENSEGEMM', 'VANILLA', 'TRITON', 'MARLIN', 'MEGAMOE_DEEPGEMM', 'MEGAMOE_CUTEDSL']` | `categorical` | | `AUTO`, `CUTLASS`, `CUTEDSL`, `TRTLLM`, `DEEPGEMM`, `DENSEGEMM`, `VANILLA`, `TRITON`, `MARLIN`, `MEGAMOE_DEEPGEMM`, `MEGAMOE_CUTEDSL` | | `moe_config.disable_finalize_fusion` | `` | `value` | | | | `moe_config.max_num_tokens` | `Optional[int]` | `value` | | | | `moe_config.use_low_precision_moe_combine` | `` | `value` | | | @@ -274,6 +279,7 @@ unset or when the safety sanitizer rejects the runtime value. | `speculative_config.eagle3_one_model` | `Optional[bool]` | `value` | | | | `speculative_config.eagle_choices` | `Optional[List[List[int]]]` | `value` | | | | `speculative_config.enable_global_pool` | `` | `value` | | | +| `speculative_config.enable_penalty` | `` | `value` | | | | `speculative_config.end_thinking_phase_token` | `` | `value` | | | | `speculative_config.global_pool_size` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` | | | | `speculative_config.greedy_sampling` | `Optional[bool]` | `value` | | | @@ -309,7 +315,7 @@ unset or when the safety sanitizer rejects the runtime value. | `speculative_config.write_interval` | `` | `value` | | | | `stream_interval` | `` | `value` | | | | `telemetry_config.disabled` | `` | `value` | | | -| `telemetry_config.usage_context` | `` | `categorical` | | `unknown`, `llm_class`, `cli_serve`, `cli_bench`, `cli_eval`, `disaggregated` | +| `telemetry_config.usage_context` | `` | `categorical` | | `unknown`, `llm_class`, `visual_gen_class`, `cli_serve`, `cli_bench`, `cli_eval`, `disaggregated` | | `tensor_parallel_size` | `` | `value` | | | | `tokenizer_mode` | `Literal['auto', 'slow']` | `categorical` | | `auto`, `slow` | | `torch_compile_config.capture_num_tokens` | `Optional[List[Annotated[int, Gt(gt=0)]]]` | `value` | | | @@ -323,3 +329,70 @@ unset or when the safety sanitizer rejects the runtime value. | `use_cute_dsl_bf16_gemm` | `` | `value` | | | | `use_cute_dsl_blockscaling_bmm` | `` | `value` | | | | `use_cute_dsl_blockscaling_mm` | `` | `value` | | | +| `use_fine_grained_sync` | `` | `value` | | | + +### `VisualGenArgs` + +59 captured fields. + +| Captured key | Annotation | Kind | Converter | Allowed values | +|--------------|------------|------|-----------|----------------| +| `attention_config.backend` | `Literal['VANILLA', 'TRTLLM', 'FA4', 'CUTEDSL']` | `categorical` | | `VANILLA`, `TRTLLM`, `FA4`, `CUTEDSL` | +| `attention_config.quant_attention_config.k_block_size` | `` | `value` | | | +| `attention_config.quant_attention_config.q_block_size` | `` | `value` | | | +| `attention_config.quant_attention_config.qk_dtype` | `Literal['bf16', 'int8', 'fp8', 'mxfp8', 'nvfp4']` | `categorical` | | `bf16`, `int8`, `fp8`, `mxfp8`, `nvfp4` | +| `attention_config.quant_attention_config.v_block_size` | `` | `value` | | | +| `attention_config.quant_attention_config.v_dtype` | `Literal['fp8']` | `categorical` | | `fp8` | +| `attention_config.sparse_attention_config.algorithm` | `Literal['skip_softmax']` | `categorical` | | `skip_softmax`, `vsa` | +| `attention_config.sparse_attention_config.disabled_until_timestep` | `Optional[float]` | `value` | | | +| `attention_config.sparse_attention_config.target_sparsity` | `Optional[float]` | `value` | | | +| `attention_config.sparse_attention_config.threshold_scale_factor` | `Optional[float]` | `value` | | | +| `attention_config.sparse_attention_config.vsa_sparsity` | `` | `value` | | | +| `cache_config.Bn_compute_blocks` | `` | `value` | | | +| `cache_config.Fn_compute_blocks` | `` | `value` | | | +| `cache_config.cache_backend` | `Literal['cache_dit']` | `categorical` | | `cache_dit`, `teacache` | +| `cache_config.coefficients` | `Optional[List[float]]` | `value` | | | +| `cache_config.coefficients_2` | `Optional[List[float]]` | `value` | | | +| `cache_config.enable_separate_cfg` | `Optional[bool]` | `value` | | | +| `cache_config.enable_taylorseer` | `` | `value` | | | +| `cache_config.force_refresh_step_hint` | `Optional[int]` | `value` | | | +| `cache_config.force_refresh_step_policy` | `Literal['once', 'repeat']` | `categorical` | | `once`, `repeat` | +| `cache_config.max_cached_steps` | `` | `value` | | | +| `cache_config.max_continuous_cached_steps` | `` | `value` | | | +| `cache_config.max_warmup_steps` | `` | `value` | | | +| `cache_config.residual_diff_threshold` | `` | `value` | | | +| `cache_config.scm_steps_policy` | `Literal['dynamic', 'static']` | `categorical` | | `dynamic`, `static` | +| `cache_config.taylorseer_order` | `` | `value` | | | +| `cache_config.teacache_thresh` | `` | `value` | | | +| `cache_config.use_ret_steps` | `` | `value` | | | +| `compilation_config.num_frames` | `Optional[List[int]]` | `value` | | | +| `compilation_config.resolutions` | `Optional[List[Tuple[int, int]]]` | `value` | | | +| `compilation_config.skip_warmup` | `` | `value` | | | +| `cpu_offload_config.enable` | `` | `value` | | | +| `cpu_offload_config.pin_memory` | `` | `value` | | | +| `cuda_graph_config.enable` | `` | `value` | | | +| `enable_layerwise_nvtx_marker` | `` | `value` | | | +| `parallel_config.async_ulysses` | `` | `value` | | | +| `parallel_config.attn2d_size` | `Tuple[Annotated[int, FieldInfo(annotation=NoneType, required=True, metadata=[Ge(ge=1)])], Annotated[int, FieldInfo(annotation=NoneType, required=True, metadata=[Ge(ge=1)])]]` | `value` | | | +| `parallel_config.cfg_size` | `` | `value` | | | +| `parallel_config.parallel_vae_size` | `` | `value` | | | +| `parallel_config.parallel_vae_split_dim` | `Literal['width', 'height']` | `categorical` | | `width`, `height` | +| `parallel_config.ring_size` | `` | `value` | | | +| `parallel_config.tp_size` | `` | `value` | | | +| `parallel_config.ulysses_size` | `` | `value` | | | +| `quant_config.clamp_val` | `Optional[List[float]]` | `value` | | | +| `quant_config.group_size` | `Optional[int]` | `value` | | | +| `quant_config.has_zero_point` | `` | `value` | | | +| `quant_config.kv_cache_quant_algo` | `Optional[tensorrt_llm.quantization.mode.QuantAlgo]` | `categorical` | | `W8A16`, `W4A16`, `W4A16_AWQ`, `W4A8_AWQ`, `W8A16_GPTQ`, `W4A16_GPTQ`, `W8A8_SQ_PER_CHANNEL`, `W8A8_SQ_PER_TENSOR_PLUGIN`, `W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN`, `W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN`, `W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN`, `W4A8_QSERVE_PER_GROUP`, `W4A8_QSERVE_PER_CHANNEL`, `FP8`, `FP8_PER_CHANNEL_PER_TOKEN`, `FP8_BLOCK_SCALES`, `INT8`, `MIXED_PRECISION`, `NVFP4`, `W4A8_NVFP4_FP8`, `W4A8_MXFP4_FP8`, `W4A8_MXFP4_MXFP8`, `W4A16_MXFP4`, `MXFP8`, `W4A16_NVFP4`, `NVFP4_AWQ`, `NVFP4_ARC`, `NO_QUANT` | +| `quant_config.mamba_ssm_philox_rounds` | `` | `value` | | | +| `quant_config.mamba_ssm_stochastic_rounding` | `` | `value` | | | +| `quant_config.pre_quant_scale` | `` | `value` | | | +| `quant_config.quant_algo` | `Optional[tensorrt_llm.quantization.mode.QuantAlgo]` | `categorical` | | `W8A16`, `W4A16`, `W4A16_AWQ`, `W4A8_AWQ`, `W8A16_GPTQ`, `W4A16_GPTQ`, `W8A8_SQ_PER_CHANNEL`, `W8A8_SQ_PER_TENSOR_PLUGIN`, `W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN`, `W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN`, `W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN`, `W4A8_QSERVE_PER_GROUP`, `W4A8_QSERVE_PER_CHANNEL`, `FP8`, `FP8_PER_CHANNEL_PER_TOKEN`, `FP8_BLOCK_SCALES`, `INT8`, `MIXED_PRECISION`, `NVFP4`, `W4A8_NVFP4_FP8`, `W4A8_MXFP4_FP8`, `W4A8_MXFP4_MXFP8`, `W4A16_MXFP4`, `MXFP8`, `W4A16_NVFP4`, `NVFP4_AWQ`, `NVFP4_ARC`, `NO_QUANT` | +| `quant_config.smoothquant_val` | `` | `value` | | | +| `quant_config.use_meta_recipe` | `` | `value` | | | +| `runtime_lora_config.fuse_qkv` | `` | `value` | | | +| `runtime_lora_config.scale` | `` | `value` | | | +| `runtime_lora_config.strict` | `` | `value` | | | +| `torch_compile_config.enable` | `` | `value` | | | +| `torch_compile_config.enable_autotune` | `` | `value` | | | +| `torch_compile_config.enable_fullgraph` | `` | `value` | | | diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py index 04efe4b0efbd..dbf699d51afe 100644 --- a/tensorrt_llm/_torch/visual_gen/executor.py +++ b/tensorrt_llm/_torch/visual_gen/executor.py @@ -18,6 +18,7 @@ from tensorrt_llm._torch.shared_tensor import SharedTensorContainer from tensorrt_llm._torch.visual_gen.output import PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader +from tensorrt_llm._torch.visual_gen.pipeline_registry import PIPELINE_REGISTRY from tensorrt_llm.executor.ipc import ZeroMqQueue from tensorrt_llm.llmapi.utils import configure_cpu_affinity from tensorrt_llm.logger import logger @@ -36,6 +37,116 @@ # Default cap on the size of the iteration-stats snapshot buffer used by the # /metrics endpoint. Mirrors the LLM ``iter_stats_max_iterations`` default. _DEFAULT_ITER_STATS_MAX = 1000 +_TELEMETRY_TRANSFORMER_COMPONENTS = frozenset(("transformer", "transformer_2")) + + +def _visual_gen_launch_metadata( + external_launch: Optional[Tuple[int, int, int, str, int]], +) -> Tuple[str, int]: + """Return a bounded launch mode and best-known deployment node count.""" + if external_launch is None: + return "local_spawn", 1 + + if "SLURM_PROCID" in os.environ: + launch_mode = "slurm" + node_count_value = os.environ.get("SLURM_NNODES") + else: + launch_mode = "torchrun" + node_count_value = os.environ.get("GROUP_WORLD_SIZE") + + try: + node_count = int(node_count_value) if node_count_value is not None else 0 + except ValueError: + node_count = 0 + if node_count <= 0: + try: + local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", "0")) + world_size = external_launch[2] + node_count = (world_size + local_world_size - 1) // local_world_size + except (TypeError, ValueError, ZeroDivisionError): + node_count = 0 + return launch_mode, max(node_count, 0) + + +def _quant_algo_name(quant_config: Any) -> str: + """Return a bounded enum label from a resolved component config.""" + quant_algo = getattr(quant_config, "quant_algo", None) + if quant_algo is None: + return "" + name = getattr(quant_algo, "name", None) + return name if isinstance(name, str) else "" + + +def _visual_gen_telemetry_metadata(pipeline: Any, args: VisualGenArgs) -> Dict[str, Any]: + """Build bounded READY metadata from the successfully loaded pipeline.""" + defaults: Dict[str, Any] = { + "model_id": "other", + "pipeline_class_name": "unknown", + "resolved_pipeline_class": "unknown", + "modality": "unknown", + "quantization_algo": "", + "dynamic_weight_quant": False, + "quantized_components": [], + } + try: + pipeline_class_name = getattr(pipeline, "_telemetry_pipeline_class_name", "unknown") + entry = PIPELINE_REGISTRY.get(pipeline_class_name) + if entry is None or not entry.telemetry_safe: + pipeline_class_name = "unknown" + entry = None + + resolved_entry = next( + ( + candidate + for candidate in PIPELINE_REGISTRY.values() + if candidate.telemetry_safe and type(pipeline) is candidate.pipeline_cls + ), + None, + ) + resolved_pipeline_class = ( + type(pipeline).__name__ if resolved_entry is not None else "unknown" + ) + + model_id = "other" + if any( + candidate.telemetry_safe and args.model in candidate.hf_ids + for candidate in PIPELINE_REGISTRY.values() + ): + model_id = args.model + + pipeline_config = pipeline.pipeline_config + quantized_components = [] + quantization_algos = [] + for component in pipeline.transformer_components: + if component not in _TELEMETRY_TRANSFORMER_COMPONENTS: + continue + component_quant = pipeline_config.get_quant_config(component) + quantization_algo = _quant_algo_name(component_quant) + if quantization_algo: + quantized_components.append(component) + if quantization_algo not in quantization_algos: + quantization_algos.append(quantization_algo) + + if len(quantization_algos) == 1: + quantization_algo = quantization_algos[0] + elif quantization_algos: + quantization_algo = "mixed" + else: + quantization_algo = "" + + defaults.update( + model_id=model_id, + pipeline_class_name=pipeline_class_name, + resolved_pipeline_class=resolved_pipeline_class, + modality=entry.modality if entry is not None else "unknown", + quantization_algo=quantization_algo, + dynamic_weight_quant=bool(pipeline_config.dynamic_weight_quant), + quantized_components=quantized_components, + ) + except Exception: + # Telemetry metadata is optional and must never disrupt worker startup. + pass + return defaults class _IterationStatsTracker: @@ -441,6 +552,9 @@ def _load_pipeline(self): "extra_param_specs": self.pipeline.extra_param_specs, "supports_image_edit": self.pipeline.supports_image_edit, "ref_slot_specs": self.pipeline.ref_slot_specs, + "telemetry_metadata": _visual_gen_telemetry_metadata( + self.pipeline, self.visual_gen_args + ), }, ) ) @@ -756,6 +870,7 @@ def __init__( # --- Detect external launcher (torchrun / srun) --- ext = _detect_external_launch() + self.launch_mode, self.node_count = _visual_gen_launch_metadata(ext) if ext is None: # Single-node: coordinator spawns all workers locally @@ -824,6 +939,7 @@ def __init__( self.extra_param_specs: Dict = {} self.supports_image_edit: bool = False self.ref_slot_specs: Dict = {} + self.telemetry_metadata: Dict[str, Any] = {} # --- Launch workers --- self.worker_processes = [] @@ -1182,6 +1298,9 @@ async def _wait_ready_async(self): self.extra_param_specs = payload.get("extra_param_specs", {}) self.supports_image_edit = bool(payload.get("supports_image_edit", False)) self.ref_slot_specs = payload.get("ref_slot_specs", {}) + telemetry_metadata = payload.get("telemetry_metadata") + if isinstance(telemetry_metadata, dict): + self.telemetry_metadata = telemetry_metadata elapsed = time.time() - start_time logger.info(f"DiffusionClient: Workers ready ({elapsed:.1f}s)") return diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 69a4434710a8..ab4a96c40acd 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -273,6 +273,8 @@ def _load_reference_image(data: bytes): "nvidia/Cosmos3-Edge", ], doc="Cosmos3 Omnimodal world models.", + modality="mixed", + telemetry_safe=True, ) class Cosmos3OmniMoTPipeline(BasePipeline): def __init__(self, pipeline_config): diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py index 5881614c872b..014223dccbf6 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py @@ -43,6 +43,8 @@ "FluxPipeline", hf_ids=["black-forest-labs/FLUX.1-dev"], doc="Black Forest Labs FLUX.1 family (text-to-image).", + modality="image", + telemetry_safe=True, ) class FluxPipeline(BasePipeline): """FLUX.1 Text-to-Image Pipeline. diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py index 724695babc08..4e536b729136 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py @@ -97,6 +97,8 @@ def format_input(prompts: List[str], system_message: str) -> List[List[dict]]: "Flux2Pipeline", hf_ids=["black-forest-labs/FLUX.2-dev"], doc="Black Forest Labs FLUX.2 family (text-to-image and reference-image generation).", + modality="image", + telemetry_safe=True, ) class Flux2Pipeline(BasePipeline): """FLUX.2 text-to-image and reference-image pipeline. diff --git a/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py b/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py index 0b4d161ef3d5..64eff5b9f3b4 100644 --- a/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py +++ b/tensorrt_llm/_torch/visual_gen/models/glm_image/pipeline_glm_image.py @@ -128,6 +128,8 @@ def calculate_shift( "GlmImagePipeline", hf_ids=["zai-org/GLM-Image"], doc="GlmImage family (text-to-image).", + modality="image", + telemetry_safe=True, ) class GlmImagePipeline(BasePipeline): # ------------------------------------------------------------------ diff --git a/tensorrt_llm/_torch/visual_gen/models/hunyuan_video1_5/pipeline_hunyuan_video1_5.py b/tensorrt_llm/_torch/visual_gen/models/hunyuan_video1_5/pipeline_hunyuan_video1_5.py index 015dc4452f50..7c10144df42a 100644 --- a/tensorrt_llm/_torch/visual_gen/models/hunyuan_video1_5/pipeline_hunyuan_video1_5.py +++ b/tensorrt_llm/_torch/visual_gen/models/hunyuan_video1_5/pipeline_hunyuan_video1_5.py @@ -143,6 +143,8 @@ def retrieve_timesteps( "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v_distilled", ], doc="Tencent HunyuanVideo 1.5 family (text-to-video).", + modality="video", + telemetry_safe=True, ) class HunyuanVideo15Pipeline(BasePipeline): def __init__(self, pipeline_config: DiffusionPipelineConfig): diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py index de6411754878..d6da5b724bf2 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py @@ -616,6 +616,8 @@ def _load_component_weights( "``distilled_lora_path`` are not ``None``, either set by the " "user or auto-discovered from the checkpoint." ), + modality="video_audio", + telemetry_safe=True, ) class LTX2Pipeline(BasePipeline): """Pipeline for text-to-video generation with audio using LTX2 model. diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py index d7223d91b413..ef69dd35f1db 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py @@ -1012,7 +1012,7 @@ def get_graph_key(self, *args, **kwargs): # class-name dispatch from the registry, so its own discovery surface # entry would duplicate the canonical ``LTX2Pipeline`` entry above and # make ``supported_models()`` / ``pipeline_config()`` ambiguous. -@register_pipeline("LTX2TwoStagesPipeline") +@register_pipeline("LTX2TwoStagesPipeline", modality="video_audio", telemetry_safe=True) class LTX2TwoStagesPipeline(LTX2Pipeline): """Lightricks LTX-Video two-stage text-to-video with audio. diff --git a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py index b7d3a69015cf..828ee916709d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py +++ b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py @@ -78,7 +78,12 @@ def _calculate_shift( } -@register_pipeline("QwenImagePipeline") +@register_pipeline( + "QwenImagePipeline", + doc="Qwen-Image text-to-image family.", + modality="image", + telemetry_safe=True, +) class QwenImagePipeline(BasePipeline): """Qwen-Image text-to-image pipeline. diff --git a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py index 7cf64dabb121..75286d2635e0 100644 --- a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py +++ b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py @@ -76,6 +76,8 @@ def _retrieve_latents( "QwenImageEditPlusPipeline", hf_ids=["Qwen/Qwen-Image-Edit-2511"], doc="Qwen-Image-Edit image editing pipeline.", + modality="image", + telemetry_safe=True, ) class QwenImageEditPlusPipeline(QwenImagePipeline): """Qwen-Image-Edit pipeline using the existing Qwen-Image transformer. diff --git a/tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.py b/tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.py index 547e4f960811..0f8f061cd347 100644 --- a/tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.py +++ b/tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.py @@ -135,6 +135,8 @@ def _retrieve_latents( "Use VisualGen.default_params or VisualGen.extra_param_specs to set request knobs " "such as extra_params['layers'] and extra_params['resolution']." ), + modality="layered_image", + telemetry_safe=True, ) class QwenImageLayeredPipeline(BasePipeline): """Qwen-Image-Layered image decomposition pipeline.""" diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_fastwan.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_fastwan.py index f3f31167ce04..063bad31c899 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_fastwan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_fastwan.py @@ -40,6 +40,8 @@ "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers", ], doc="FastWan 2.2 distilled (DMD) — 3-step Wan 2.2 TI2V-5B text-to-video.", + modality="video", + telemetry_safe=True, ) class WanDMDPipeline(WanPipeline): """Wan 2.2 TI2V-5B with the DMD 3-step sampling loop. diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py index fbf9684736f1..65354c56a09a 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -121,6 +121,8 @@ "nvidia/Wan2.2-T2V-A14B-Diffusers-NVFP4", ], doc="Wan 2.1 & 2.2 text-to-video family.", + modality="video", + telemetry_safe=True, ) class WanPipeline(BasePipeline): def __init__(self, pipeline_config): diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py index 332fb7dc530c..a0c85b93f152 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py @@ -103,6 +103,8 @@ "Wan-AI/Wan2.2-I2V-A14B-Diffusers", ], doc="Wan 2.1 & 2.2 image-to-video family.", + modality="video", + telemetry_safe=True, ) class WanImageToVideoPipeline(BasePipeline): def __init__(self, pipeline_config): diff --git a/tensorrt_llm/_torch/visual_gen/pipeline_registry.py b/tensorrt_llm/_torch/visual_gen/pipeline_registry.py index 751f7aec8da0..8e3a7fc1d6a9 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline_registry.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline_registry.py @@ -7,7 +7,7 @@ All pipelines (Wan, Flux, Flux2, LTX2, QwenImage) register via @register_pipeline decorator. The registry value is a private ``_PipelineEntry`` dataclass that carries -the pipeline class plus three pieces of per-family metadata: +the pipeline class plus five pieces of per-family metadata: * ``hf_ids`` — canonical HuggingFace model IDs that dispatch to this pipeline. Powers ``VisualGen.supported_models()`` and @@ -17,6 +17,8 @@ * ``defaults`` — default per-family ``pipeline_config`` knobs (schema-by-example for the strict-validated dict). * ``doc`` — short human-readable description for discovery tooling. + * ``modality`` — bounded pipeline capability used by telemetry and discovery. + * ``telemetry_safe`` — explicit opt-in for transmitting public registry metadata. The dataclass and the registry itself are deliberately private — users go through ``VisualGenArgs(model=...)``, ``VisualGen.supported_models()``, @@ -30,7 +32,7 @@ import os from dataclasses import dataclass, field from enum import Enum -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type from tensorrt_llm.logger import logger @@ -61,6 +63,16 @@ class PipelineComponent(str, Enum): VISION_LANGUAGE_ENCODER = "vision_language_encoder" +VisualGenModality = Literal[ + "image", + "video", + "video_audio", + "layered_image", + "mixed", + "unknown", +] + + @dataclass class _PipelineEntry: """Private per-pipeline-family metadata stored in PIPELINE_REGISTRY.""" @@ -69,6 +81,8 @@ class _PipelineEntry: hf_ids: List[str] = field(default_factory=list) defaults: Dict[str, Any] = field(default_factory=dict) doc: str = "" + modality: VisualGenModality = "unknown" + telemetry_safe: bool = False # Keyed by Diffusers ``_class_name`` (from model_index.json). ~3-5 entries @@ -83,6 +97,8 @@ def register_pipeline( hf_ids: Optional[List[str]] = None, defaults: Optional[Dict[str, Any]] = None, doc: str = "", + modality: VisualGenModality = "unknown", + telemetry_safe: bool = False, ): """Register a pipeline class with optional per-family metadata. @@ -96,6 +112,8 @@ class WanPipeline(BasePipeline): hf_ids=["Lightricks/LTX-Video"], defaults={"text_encoder_path": ""}, doc="Lightricks LTX-Video family.", + modality="video_audio", + telemetry_safe=True, ) class LTX2Pipeline(BasePipeline): ... @@ -114,6 +132,8 @@ def decorator(cls: Type["BasePipeline"]) -> Type["BasePipeline"]: hf_ids=list(hf_ids or []), defaults=dict(defaults or {}), doc=doc, + modality=modality, + telemetry_safe=telemetry_safe, ) logger.debug(f"Registered pipeline: {name} -> {cls.__name__}") return cls @@ -142,7 +162,8 @@ def from_config( f"Checkpoint: {checkpoint_dir}" ) - pipeline_class = PIPELINE_REGISTRY[class_name].pipeline_cls + entry = PIPELINE_REGISTRY[class_name] + pipeline_class = entry.pipeline_cls # Let the pipeline class upgrade itself to a specialised variant # (e.g. LTX2Pipeline → LTX2TwoStagesPipeline) based on config. @@ -150,8 +171,15 @@ def from_config( logger.info(f"AutoPipeline: Creating {pipeline_class.__name__} from {checkpoint_dir}") - # Instantiate pipeline with DiffusionPipelineConfig - return pipeline_class(config) + # Preserve the allowlisted registry key separately from the resolved + # runtime class (a pipeline may select a specialized variant). + pipeline = pipeline_class(config) + if entry.telemetry_safe: + try: + pipeline._telemetry_pipeline_class_name = class_name + except Exception as exc: + logger.debug(f"Could not attach pipeline telemetry metadata: {exc}") + return pipeline @staticmethod def _detect_from_checkpoint(checkpoint_dir: str) -> str: diff --git a/tensorrt_llm/bench/benchmark/visual_gen.py b/tensorrt_llm/bench/benchmark/visual_gen.py index bc29527ad865..265e0586427a 100644 --- a/tensorrt_llm/bench/benchmark/visual_gen.py +++ b/tensorrt_llm/bench/benchmark/visual_gen.py @@ -197,24 +197,57 @@ def visual_gen_command( """Benchmark VisualGen (image/video generation) models offline.""" import yaml + from tensorrt_llm.commands import _telemetry as _command_telemetry + from tensorrt_llm.usage import apply_usage_session_config + from tensorrt_llm.usage import config as _telemetry_config from tensorrt_llm.visual_gen import VisualGen, VisualGenParams from tensorrt_llm.visual_gen.args import VisualGenArgs - if prompt is None and prompt_file is None: - raise click.UsageError("Either --prompt or --prompt_file must be specified.") - if prompt is not None and prompt_file is not None: - raise click.UsageError("--prompt and --prompt_file are mutually exclusive.") - model = bench_env.model model_path = str(bench_env.checkpoint_path or model) + apply_usage_session_config( + bench_env.telemetry_config, + default_usage_context=_telemetry_config.UsageContext.CLI_BENCH.value, + component="visual_gen", + lifecycle_phase="config_validation", + ) # Build VisualGenArgs (same pattern as trtllm-serve _serve_visual_gen) extra_args: dict = {} if visual_gen_args is not None: with open(visual_gen_args, "r") as f: extra_args = yaml.safe_load(f) or {} + if not isinstance(extra_args, dict): + raise ValueError("VisualGenArgs YAML must contain a mapping at the document root") + + ctx = click.get_current_context() + explicit_cli_telemetry = ( + ctx.parent is not None + and ctx.parent.get_parameter_source("telemetry") is click.core.ParameterSource.COMMANDLINE + ) + _command_telemetry.apply_raw_config_telemetry_opt_out( + extra_args, + usage_context=_telemetry_config.UsageContext.CLI_BENCH, + component="visual_gen", + explicit_cli_telemetry=explicit_cli_telemetry, + ) + + if prompt is None and prompt_file is None: + raise click.UsageError("Either --prompt or --prompt_file must be specified.") + if prompt is not None and prompt_file is not None: + raise click.UsageError("--prompt and --prompt_file are mutually exclusive.") - visual_gen_args = VisualGenArgs(**extra_args) if extra_args else None + visual_gen_args = VisualGenArgs(**extra_args) + telemetry_config = visual_gen_args.telemetry_config + cli_telemetry_disabled = ( + bench_env.telemetry_config.disabled if bench_env.telemetry_config is not None else False + ) + if explicit_cli_telemetry or cli_telemetry_disabled: + telemetry_config = telemetry_config.model_copy(update={"disabled": cli_telemetry_disabled}) + telemetry_config = telemetry_config.model_copy( + update={"usage_context": _telemetry_config.UsageContext.CLI_BENCH} + ) + visual_gen_args = visual_gen_args.model_copy(update={"telemetry_config": telemetry_config}) n_workers = visual_gen_args.parallel_config.n_workers if visual_gen_args is not None else 1 parallel_config = extra_args.get("parallel_config", {}) diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index e798b82dcb9b..43cffa15aede 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -114,6 +114,27 @@ def _apply_effective_telemetry_config(llm_args: dict, ) +def _apply_visual_gen_telemetry_config( + visual_gen_args: "VisualGenArgs", *, telemetry: bool, + explicit_cli_telemetry: bool) -> "VisualGenArgs": + """Apply the authoritative CLI ingress and opt-out to VisualGen args.""" + telemetry_config = visual_gen_args.telemetry_config + if explicit_cli_telemetry or not telemetry: + telemetry_config = telemetry_config.model_copy( + update={"disabled": not telemetry}) + telemetry_config = telemetry_config.model_copy( + update={"usage_context": _telemetry_config.UsageContext.CLI_SERVE}) + visual_gen_args = visual_gen_args.model_copy( + update={"telemetry_config": telemetry_config}) + apply_usage_session_config( + telemetry_config, + default_usage_context=_telemetry_config.UsageContext.CLI_SERVE.value, + component="server", + lifecycle_phase="config_validation", + ) + return visual_gen_args + + def _pop_bool_config_option(config: dict[str, Any], key: str) -> bool: return validate_config_bool(config.pop(key, False), key) @@ -1549,8 +1570,26 @@ def _serve_llm(): def _serve_visual_gen(): from tensorrt_llm.visual_gen.args import VisualGenArgs - parsed_visual_gen_args = (VisualGenArgs.from_yaml(visual_gen_args) - if visual_gen_args is not None else None) + raw_visual_gen_args = {} + if visual_gen_args is not None: + with open(visual_gen_args, "r") as config_file: + raw_visual_gen_args = yaml.safe_load(config_file) or {} + if not isinstance(raw_visual_gen_args, dict): + raise ValueError( + "VisualGenArgs YAML must contain a mapping at the document root" + ) + _command_telemetry.apply_raw_config_telemetry_opt_out( + raw_visual_gen_args, + usage_context=_telemetry_config.UsageContext.CLI_SERVE, + component="server", + explicit_cli_telemetry="telemetry" in explicit_cli_keys, + ) + parsed_visual_gen_args = VisualGenArgs(**raw_visual_gen_args) + parsed_visual_gen_args = _apply_visual_gen_telemetry_config( + parsed_visual_gen_args, + telemetry=telemetry, + explicit_cli_telemetry="telemetry" in explicit_cli_keys, + ) metadata_server_cfg = parse_metadata_server_config_file( metadata_server_config_file) diff --git a/tensorrt_llm/usage/__init__.py b/tensorrt_llm/usage/__init__.py index 6c0c77b1ac16..28689fb2f7ab 100644 --- a/tensorrt_llm/usage/__init__.py +++ b/tensorrt_llm/usage/__init__.py @@ -22,7 +22,7 @@ - Set environment variable TELEMETRY_DISABLED=true or TELEMETRY_DISABLED=1 - Set environment variable DO_NOT_TRACK=1 - Create file ~/.config/trtllm/do_not_track - - Pass TelemetryConfig(disabled=True) to LLM() or --no-telemetry via CLI + - Pass TelemetryConfig(disabled=True) to LLM()/VisualGen() or --no-telemetry via CLI - Automatically disabled in CI/test environments (override with TRTLLM_USAGE_FORCE_ENABLED=1) """ @@ -40,6 +40,10 @@ record_llm_initialization_failure = _usage_lib.record_llm_initialization_failure record_llm_initialized = _usage_lib.record_llm_initialized record_llm_shutdown = _usage_lib.record_llm_shutdown +record_visual_gen_initialization_attempt = _usage_lib.record_visual_gen_initialization_attempt +record_visual_gen_initialization_failure = _usage_lib.record_visual_gen_initialization_failure +record_visual_gen_initialized = _usage_lib.record_visual_gen_initialized +record_visual_gen_shutdown = _usage_lib.record_visual_gen_shutdown record_observed_signal = _usage_lib.record_observed_signal record_termination_observation = _usage_lib.record_termination_observation set_lifecycle_phase = _usage_lib.set_lifecycle_phase @@ -59,6 +63,10 @@ "record_llm_initialization_failure", "record_llm_initialized", "record_llm_shutdown", + "record_visual_gen_initialization_attempt", + "record_visual_gen_initialization_failure", + "record_visual_gen_initialized", + "record_visual_gen_shutdown", "record_observed_signal", "record_termination_observation", "report_exit", diff --git a/tensorrt_llm/usage/config.py b/tensorrt_llm/usage/config.py index c3b38b777160..6efce6985ddf 100644 --- a/tensorrt_llm/usage/config.py +++ b/tensorrt_llm/usage/config.py @@ -42,6 +42,7 @@ class UsageContext(str, Enum): UNKNOWN = "unknown" LLM_CLASS = "llm_class" + VISUAL_GEN_CLASS = "visual_gen_class" CLI_SERVE = "cli_serve" CLI_BENCH = "cli_bench" CLI_EVAL = "cli_eval" @@ -95,5 +96,5 @@ class TelemetryConfig(_StrictUsageBaseModel): default=UsageContext.UNKNOWN, description="Identifies how TRT-LLM was invoked (CLI command vs Python API). " "Set automatically by CLI commands; defaults to UNKNOWN (promoted to " - "LLM_CLASS by BaseLLM.__init__ for direct Python API usage).", + "LLM_CLASS or VISUAL_GEN_CLASS for direct Python API usage).", ) diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index d4a688885151..c89d1a20e935 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -2219,6 +2219,7 @@ "allowed_values": [ "unknown", "llm_class", + "visual_gen_class", "cli_serve", "cli_bench", "cli_eval", @@ -2330,5 +2331,506 @@ "kind": "value", "path": "use_fine_grained_sync" } + ], + "VisualGenArgs": [ + { + "allowed_values": [ + "VANILLA", + "TRTLLM", + "FA4", + "CUTEDSL" + ], + "annotation": "Literal['VANILLA', 'TRTLLM', 'FA4', 'CUTEDSL']", + "converter": "", + "kind": "categorical", + "path": "attention_config.backend" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "attention_config.quant_attention_config.k_block_size" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "attention_config.quant_attention_config.q_block_size" + }, + { + "allowed_values": [ + "bf16", + "int8", + "fp8", + "mxfp8", + "nvfp4" + ], + "annotation": "Literal['bf16', 'int8', 'fp8', 'mxfp8', 'nvfp4']", + "converter": "", + "kind": "categorical", + "path": "attention_config.quant_attention_config.qk_dtype" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "attention_config.quant_attention_config.v_block_size" + }, + { + "allowed_values": [ + "fp8" + ], + "annotation": "Literal['fp8']", + "converter": "", + "kind": "categorical", + "path": "attention_config.quant_attention_config.v_dtype" + }, + { + "allowed_values": [ + "skip_softmax", + "vsa" + ], + "annotation": "Literal['skip_softmax']", + "converter": "", + "kind": "categorical", + "path": "attention_config.sparse_attention_config.algorithm" + }, + { + "allowed_values": [], + "annotation": "Optional[float]", + "converter": "", + "kind": "value", + "path": "attention_config.sparse_attention_config.disabled_until_timestep" + }, + { + "allowed_values": [], + "annotation": "Optional[float]", + "converter": "", + "kind": "value", + "path": "attention_config.sparse_attention_config.target_sparsity" + }, + { + "allowed_values": [], + "annotation": "Optional[float]", + "converter": "", + "kind": "value", + "path": "attention_config.sparse_attention_config.threshold_scale_factor" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "attention_config.sparse_attention_config.vsa_sparsity" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "cache_config.Bn_compute_blocks" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "cache_config.Fn_compute_blocks" + }, + { + "allowed_values": [ + "cache_dit", + "teacache" + ], + "annotation": "Literal['cache_dit']", + "converter": "", + "kind": "categorical", + "path": "cache_config.cache_backend" + }, + { + "allowed_values": [], + "annotation": "Optional[List[float]]", + "converter": "", + "kind": "value", + "path": "cache_config.coefficients" + }, + { + "allowed_values": [], + "annotation": "Optional[List[float]]", + "converter": "", + "kind": "value", + "path": "cache_config.coefficients_2" + }, + { + "allowed_values": [], + "annotation": "Optional[bool]", + "converter": "", + "kind": "value", + "path": "cache_config.enable_separate_cfg" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "cache_config.enable_taylorseer" + }, + { + "allowed_values": [], + "annotation": "Optional[int]", + "converter": "", + "kind": "value", + "path": "cache_config.force_refresh_step_hint" + }, + { + "allowed_values": [ + "once", + "repeat" + ], + "annotation": "Literal['once', 'repeat']", + "converter": "", + "kind": "categorical", + "path": "cache_config.force_refresh_step_policy" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "cache_config.max_cached_steps" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "cache_config.max_continuous_cached_steps" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "cache_config.max_warmup_steps" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "cache_config.residual_diff_threshold" + }, + { + "allowed_values": [ + "dynamic", + "static" + ], + "annotation": "Literal['dynamic', 'static']", + "converter": "", + "kind": "categorical", + "path": "cache_config.scm_steps_policy" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "cache_config.taylorseer_order" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "cache_config.teacache_thresh" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "cache_config.use_ret_steps" + }, + { + "allowed_values": [], + "annotation": "Optional[List[int]]", + "converter": "", + "kind": "value", + "path": "compilation_config.num_frames" + }, + { + "allowed_values": [], + "annotation": "Optional[List[Tuple[int, int]]]", + "converter": "", + "kind": "value", + "path": "compilation_config.resolutions" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "compilation_config.skip_warmup" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "cpu_offload_config.enable" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "cpu_offload_config.pin_memory" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "cuda_graph_config.enable" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "enable_layerwise_nvtx_marker" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "parallel_config.async_ulysses" + }, + { + "allowed_values": [], + "annotation": "Tuple[Annotated[int, FieldInfo(annotation=NoneType, required=True, metadata=[Ge(ge=1)])], Annotated[int, FieldInfo(annotation=NoneType, required=True, metadata=[Ge(ge=1)])]]", + "converter": "", + "kind": "value", + "path": "parallel_config.attn2d_size" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "parallel_config.cfg_size" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "parallel_config.parallel_vae_size" + }, + { + "allowed_values": [ + "width", + "height" + ], + "annotation": "Literal['width', 'height']", + "converter": "", + "kind": "categorical", + "path": "parallel_config.parallel_vae_split_dim" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "parallel_config.ring_size" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "parallel_config.tp_size" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "parallel_config.ulysses_size" + }, + { + "allowed_values": [], + "annotation": "Optional[List[float]]", + "converter": "", + "kind": "value", + "path": "quant_config.clamp_val" + }, + { + "allowed_values": [], + "annotation": "Optional[int]", + "converter": "", + "kind": "value", + "path": "quant_config.group_size" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "quant_config.has_zero_point" + }, + { + "allowed_values": [ + "W8A16", + "W4A16", + "W4A16_AWQ", + "W4A8_AWQ", + "W8A16_GPTQ", + "W4A16_GPTQ", + "W8A8_SQ_PER_CHANNEL", + "W8A8_SQ_PER_TENSOR_PLUGIN", + "W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN", + "W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN", + "W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN", + "W4A8_QSERVE_PER_GROUP", + "W4A8_QSERVE_PER_CHANNEL", + "FP8", + "FP8_PER_CHANNEL_PER_TOKEN", + "FP8_BLOCK_SCALES", + "INT8", + "MIXED_PRECISION", + "NVFP4", + "W4A8_NVFP4_FP8", + "W4A8_MXFP4_FP8", + "W4A8_MXFP4_MXFP8", + "W4A16_MXFP4", + "MXFP8", + "W4A16_NVFP4", + "NVFP4_AWQ", + "NVFP4_ARC", + "NO_QUANT" + ], + "annotation": "Optional[tensorrt_llm.quantization.mode.QuantAlgo]", + "converter": "", + "kind": "categorical", + "path": "quant_config.kv_cache_quant_algo" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "quant_config.mamba_ssm_philox_rounds" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "quant_config.mamba_ssm_stochastic_rounding" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "quant_config.pre_quant_scale" + }, + { + "allowed_values": [ + "W8A16", + "W4A16", + "W4A16_AWQ", + "W4A8_AWQ", + "W8A16_GPTQ", + "W4A16_GPTQ", + "W8A8_SQ_PER_CHANNEL", + "W8A8_SQ_PER_TENSOR_PLUGIN", + "W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN", + "W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN", + "W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN", + "W4A8_QSERVE_PER_GROUP", + "W4A8_QSERVE_PER_CHANNEL", + "FP8", + "FP8_PER_CHANNEL_PER_TOKEN", + "FP8_BLOCK_SCALES", + "INT8", + "MIXED_PRECISION", + "NVFP4", + "W4A8_NVFP4_FP8", + "W4A8_MXFP4_FP8", + "W4A8_MXFP4_MXFP8", + "W4A16_MXFP4", + "MXFP8", + "W4A16_NVFP4", + "NVFP4_AWQ", + "NVFP4_ARC", + "NO_QUANT" + ], + "annotation": "Optional[tensorrt_llm.quantization.mode.QuantAlgo]", + "converter": "", + "kind": "categorical", + "path": "quant_config.quant_algo" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "quant_config.smoothquant_val" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "quant_config.use_meta_recipe" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "runtime_lora_config.fuse_qkv" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "runtime_lora_config.scale" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "runtime_lora_config.strict" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "torch_compile_config.enable" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "torch_compile_config.enable_autotune" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "torch_compile_config.enable_fullgraph" + } ] } diff --git a/tensorrt_llm/usage/llmapi_config.py b/tensorrt_llm/usage/llmapi_config.py index 25c661af8d61..58eb840a2274 100644 --- a/tensorrt_llm/usage/llmapi_config.py +++ b/tensorrt_llm/usage/llmapi_config.py @@ -12,7 +12,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""LLM API configuration capture for usage telemetry.""" +"""Validated runtime-configuration capture for usage telemetry.""" from __future__ import annotations @@ -34,8 +34,9 @@ FIELD_POLICY_VERSION = "2" API_CONTRACT_VERSION = "0.2.0" CAPTURE_SOURCE = "effective_validated_llm_args" +VISUAL_GEN_CAPTURE_SOURCE = "effective_validated_visual_gen_args" -# Cap total serialized bytes of llmApiConfigJson. The wire field is unbounded and +# Cap each serialized runtime-config payload. The wire fields are unbounded and # the reporter is fail-silent, so an oversized payload is dropped whole by the # endpoint; truncate and flag instead. Conservative bound until the endpoint limit # is confirmed. @@ -388,9 +389,11 @@ def manifest_rows(model_cls: type[BaseModel]) -> list[dict[str, Any]]: def golden_manifest() -> dict[str, list[dict[str, Any]]]: from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + from tensorrt_llm.visual_gen.args import VisualGenArgs return { "TorchLlmArgs": manifest_rows(TorchLlmArgs), + "VisualGenArgs": manifest_rows(VisualGenArgs), } @@ -568,7 +571,7 @@ def _truncate_to_budget(values: dict[str, Any]) -> tuple[dict[str, Any], str]: return kept, _canonical_json(kept) -def _failure_meta(args_class: str = "") -> dict[str, Any]: +def _failure_meta(args_class: str = "", source: str = CAPTURE_SOURCE) -> dict[str, Any]: """Metadata for capture failure. One shape used by collector and reporter.""" return { "api_contract_version": API_CONTRACT_VERSION, @@ -583,32 +586,51 @@ def _failure_meta(args_class: str = "") -> dict[str, Any]: "payload_truncated": False, "schema_digest": "", "sequence_truncated": False, - "source": CAPTURE_SOURCE, + "source": source, "unsafe_excluded": False, } -def _failure_llm_api_config_payloads(args_class: str = "") -> tuple[str, str]: +def _failure_config_payloads(args_class: str = "", source: str = CAPTURE_SOURCE) -> tuple[str, str]: """Return empty config plus canonical failure metadata JSON.""" - return "{}", _canonical_json(_failure_meta(args_class=args_class)) + return "{}", _canonical_json(_failure_meta(args_class=args_class, source=source)) -def collect_llm_api_config_payloads(llm_args: Any) -> tuple[str, str]: - """Return sanitized LLM API config and capture metadata JSON strings. +def _failure_llm_api_config_payloads(args_class: str = "") -> tuple[str, str]: + """Return empty LLM config plus canonical failure metadata JSON.""" + return _failure_config_payloads(args_class=args_class) + + +def _failure_visual_gen_config_payloads(args_class: str = "") -> tuple[str, str]: + """Return empty VisualGen config plus canonical failure metadata JSON.""" + return _failure_config_payloads( + args_class=args_class, + source=VISUAL_GEN_CAPTURE_SOURCE, + ) + + +def _collect_config_payloads( + config: Any, + *, + source: str, + manifest_cls: type[BaseModel] | None = None, +) -> tuple[str, str]: + """Return sanitized Pydantic config and capture metadata JSON strings. Manifest-driven: capture exactly the keys build_capture_manifest lists for - this class, so the runtime can never emit a key absent from the committed - golden (runtime_keys subset of manifest_keys, by construction). + the selected manifest class. Callers may pin that class to prevent runtime + subclasses from expanding the approved field set. """ try: - if not _is_pydantic_model(llm_args): - return _failure_llm_api_config_payloads() + if not _is_pydantic_model(config): + args_class = manifest_cls.__name__ if manifest_cls is not None else "" + return _failure_config_payloads(args_class=args_class, source=source) - cls = llm_args.__class__ + cls = manifest_cls or config.__class__ entries = build_capture_manifest(cls) state = _CaptureState() for entry in entries: - present, value = _resolve_path(llm_args, entry.path) + present, value = _resolve_path(config, entry.path) if not present: continue is_safe, sanitized = _sanitize_value(value, entry.annotation, entry.metadata, state) @@ -637,7 +659,7 @@ def collect_llm_api_config_payloads(llm_args: Any) -> tuple[str, str]: "payload_truncated": state.payload_truncated, "schema_digest": _schema_digest(cls), "sequence_truncated": state.sequence_truncated, - "source": CAPTURE_SOURCE, + "source": source, "unsafe_excluded": state.unsafe_excluded, } return config_json, _canonical_json(metadata) @@ -645,5 +667,27 @@ def collect_llm_api_config_payloads(llm_args: Any) -> tuple[str, str]: # Stay fail-silent only for the sanitizer/walk error family we expect. # Unexpected exceptions propagate to the daemon-thread guard in # usage_lib so genuine collector bugs are not silently masked. - args_class = type(llm_args).__name__ if llm_args is not None else "" - return _failure_llm_api_config_payloads(args_class=args_class) + args_class = ( + manifest_cls.__name__ + if manifest_cls is not None + else type(config).__name__ + if config is not None + else "" + ) + return _failure_config_payloads(args_class=args_class, source=source) + + +def collect_llm_api_config_payloads(llm_args: Any) -> tuple[str, str]: + """Return sanitized LLM API config and capture metadata JSON strings.""" + return _collect_config_payloads(llm_args, source=CAPTURE_SOURCE) + + +def collect_visual_gen_config_payloads(visual_gen_args: Any) -> tuple[str, str]: + """Return sanitized VisualGen config and capture metadata JSON strings.""" + from tensorrt_llm.visual_gen.args import VisualGenArgs + + return _collect_config_payloads( + visual_gen_args, + source=VISUAL_GEN_CAPTURE_SOURCE, + manifest_cls=VisualGenArgs, + ) diff --git a/tensorrt_llm/usage/schema.py b/tensorrt_llm/usage/schema.py index 3e50bef55c5b..d111926c1a89 100644 --- a/tensorrt_llm/usage/schema.py +++ b/tensorrt_llm/usage/schema.py @@ -39,7 +39,7 @@ CLIENT_ID = "616561816355034" EVENT_PROTOCOL = "1.6" -EVENT_SCHEMA_VER = "0.7" +EVENT_SCHEMA_VER = "0.8" EVENT_SYS_VER = "trtllm-telemetry/1.0" CLIENT_TYPE = "Native" CLIENT_VARIANT = "Release" @@ -100,6 +100,37 @@ class _LlmCounterSnapshot(BaseModel): model_config = {"populate_by_name": True} +class _VisualGenCounterSnapshot(BaseModel): + """Process-local aggregate VisualGen lifecycle counters.""" + + visual_gen_initialization_attempts: int = Field( + default=0, + ge=0, + le=_UINT32_MAX, + alias="visualGenInitializationAttempts", + ) + visual_gen_instances_created: int = Field( + default=0, + ge=0, + le=_UINT32_MAX, + alias="visualGenInstancesCreated", + ) + active_visual_gen_instances: int = Field( + default=0, + ge=0, + le=_UINT32_MAX, + alias="activeVisualGenInstances", + ) + visual_gen_initialization_failures: int = Field( + default=0, + ge=0, + le=_UINT32_MAX, + alias="visualGenInitializationFailures", + ) + + model_config = {"populate_by_name": True} + + class TrtllmInitialReport(_LlmCounterSnapshot): """TRT-LLM initial report event parameters. @@ -232,6 +263,98 @@ class TrtllmHeartbeat(_LlmCounterSnapshot): model_config = {"populate_by_name": True} +VisualGenModality = Literal[ + "image", + "video", + "video_audio", + "layered_image", + "mixed", + "unknown", +] +VisualGenLaunchMode = Literal["local_spawn", "torchrun", "slurm", "unknown"] + + +class TrtllmVisualGenInitialReport(_VisualGenCounterSnapshot): + """VisualGen initial report emitted after the pipeline is ready.""" + + trtllm_version: str = Field(default="", max_length=_SHORT_STR, alias="trtllmVersion") + platform_info: str = Field(default="", max_length=_LONG_STR, alias="platform") + python_version: str = Field(default="", max_length=_SHORT_STR, alias="pythonVersion") + cpu_architecture: str = Field(default="", max_length=_SHORT_STR, alias="cpuArchitecture") + cpu_count: int = Field(default=0, ge=0, le=_UINT32_MAX, alias="cpuCount") + gpu_count: int = Field(default=0, ge=0, le=_UINT32_MAX, alias="gpuCount") + gpu_name: str = Field(default="", max_length=_LONG_STR, alias="gpuName") + gpu_memory_mb: int = Field(default=0, ge=0, le=_UINT32_MAX, alias="gpuMemoryMB") + cuda_version: str = Field(default="", max_length=_SHORT_STR, alias="cudaVersion") + + runtime_kind: Literal["visual_gen"] = Field(default="visual_gen", alias="runtimeKind") + ingress_point: str = Field(default="", max_length=_SHORT_STR, alias="ingressPoint") + model_id: str = Field(default="other", max_length=_LONG_STR, alias="modelId") + pipeline_class_name: str = Field( + default="unknown", max_length=_LONG_STR, alias="pipelineClassName" + ) + resolved_pipeline_class: str = Field( + default="unknown", max_length=_LONG_STR, alias="resolvedPipelineClass" + ) + modality: VisualGenModality = Field(default="unknown", alias="modality") + launch_mode: VisualGenLaunchMode = Field(default="unknown", alias="launchMode") + node_count: int = Field(default=0, ge=0, le=_UINT32_MAX, alias="nodeCount") + n_workers: int = Field(default=1, ge=0, le=_UINT32_MAX, alias="nWorkers") + + quantization_algo: str = Field(default="", max_length=_SHORT_STR, alias="quantizationAlgo") + dynamic_weight_quant: bool = Field(default=False, alias="dynamicWeightQuant") + quantized_components_json: str = Field(default="[]", alias="quantizedComponentsJson") + + cfg_size: int = Field(default=1, ge=0, le=_UINT32_MAX, alias="cfgSize") + ulysses_size: int = Field(default=1, ge=0, le=_UINT32_MAX, alias="ulyssesSize") + async_ulysses: bool = Field(default=False, alias="asyncUlysses") + ring_size: int = Field(default=1, ge=0, le=_UINT32_MAX, alias="ringSize") + attn2d_row_size: int = Field(default=1, ge=0, le=_UINT32_MAX, alias="attn2dRowSize") + attn2d_col_size: int = Field(default=1, ge=0, le=_UINT32_MAX, alias="attn2dColSize") + tensor_parallel_size: int = Field(default=1, ge=0, le=_UINT32_MAX, alias="tensorParallelSize") + parallel_vae_size: int = Field(default=1, ge=0, le=_UINT32_MAX, alias="parallelVaeSize") + parallel_vae_split_dim: Literal["width", "height"] = Field( + default="width", alias="parallelVaeSplitDim" + ) + + parallel_vae: bool = Field(default=False, alias="parallelVae") + step_caching: bool = Field(default=False, alias="stepCaching") + cache_backend: Literal["teacache", "cache_dit", "none"] = Field( + default="none", alias="cacheBackend" + ) + sparse_attention: bool = Field(default=False, alias="sparseAttention") + sparse_attention_algorithm: Literal["skip_softmax", "vsa", "none"] = Field( + default="none", alias="sparseAttentionAlgorithm" + ) + sparse_attention_sparsity: float = Field( + default=0.0, ge=0.0, le=1.0, alias="sparseAttentionSparsity" + ) + quantized_attention: bool = Field(default=False, alias="quantizedAttention") + attention_backend: Literal["VANILLA", "TRTLLM", "FA4", "CUTEDSL"] = Field( + default="VANILLA", alias="attentionBackend" + ) + cuda_graphs: bool = Field(default=False, alias="cudaGraphs") + torch_compile: bool = Field(default=False, alias="torchCompile") + quantized_weights: bool = Field(default=False, alias="quantizedWeights") + + visual_gen_config_json: str = Field(default="{}", alias="visualGenConfigJson") + visual_gen_config_meta_json: str = Field(default="{}", alias="visualGenConfigMetaJson") + + model_config = {"populate_by_name": True} + + +class TrtllmVisualGenHeartbeat(_VisualGenCounterSnapshot): + """Periodic VisualGen liveness and lifecycle-counter snapshot.""" + + seq: int = Field(..., ge=0, le=_UINT32_MAX) + runtime_kind: Literal["visual_gen"] = Field(default="visual_gen", alias="runtimeKind") + ingress_point: str = Field(default="", max_length=_SHORT_STR, alias="ingressPoint") + n_workers: int = Field(default=1, ge=0, le=_UINT32_MAX, alias="nWorkers") + gpu_count: int = Field(default=0, ge=0, le=_UINT32_MAX, alias="gpuCount") + + model_config = {"populate_by_name": True} + + TerminationKind = Literal[ "clean", "exception", @@ -249,6 +372,7 @@ class TrtllmHeartbeat(_LlmCounterSnapshot): ] TerminationComponent = Literal[ "llm", + "visual_gen", "server", "engine_worker", "disagg_worker", @@ -257,7 +381,10 @@ class TrtllmHeartbeat(_LlmCounterSnapshot): ReportingSource = Literal["self", "supervisor", "executor_proxy"] -class TrtllmExitReport(_LlmCounterSnapshot): +RuntimeKind = Literal["llm", "visual_gen", "mixed", "unknown"] + + +class TrtllmExitReport(_LlmCounterSnapshot, _VisualGenCounterSnapshot): """TRT-LLM terminal event parameters. Sent at most once for a telemetry session when TRT-LLM or a surviving @@ -322,6 +449,11 @@ class TrtllmExitReport(_LlmCounterSnapshot): alias="deploymentId", description="Identifier used to correlate a disaggregated deployment.", ) + runtime_kind: RuntimeKind = Field( + default="unknown", + alias="runtimeKind", + description="Runtime families observed in this process.", + ) model_config = {"populate_by_name": True} @@ -427,7 +559,13 @@ def get_iso_timestamp(dt: Optional[datetime] = None) -> str: def build_gxt_payload( - event: Union[TrtllmInitialReport, TrtllmHeartbeat, TrtllmExitReport], + event: Union[ + TrtllmInitialReport, + TrtllmHeartbeat, + TrtllmVisualGenInitialReport, + TrtllmVisualGenHeartbeat, + TrtllmExitReport, + ], *, session_id: str, trtllm_version: str, @@ -448,6 +586,10 @@ def build_gxt_payload( event_name = "trtllm_initial_report" elif isinstance(event, TrtllmHeartbeat): event_name = "trtllm_heartbeat" + elif isinstance(event, TrtllmVisualGenInitialReport): + event_name = "trtllm_visual_gen_initial_report" + elif isinstance(event, TrtllmVisualGenHeartbeat): + event_name = "trtllm_visual_gen_heartbeat" elif isinstance(event, TrtllmExitReport): event_name = "trtllm_exit_report" else: diff --git a/tensorrt_llm/usage/schemas/README.md b/tensorrt_llm/usage/schemas/README.md index 19c16068ab80..94c113591cfb 100644 --- a/tensorrt_llm/usage/schemas/README.md +++ b/tensorrt_llm/usage/schemas/README.md @@ -1,12 +1,12 @@ # TRT-LLM Telemetry Schema Reference -Schema version: **0.7** | Client ID: `616561816355034` | Protocol: GXT Event Protocol v1.6 +Schema version: **0.8** | Client ID: `616561816355034` | Protocol: GXT Event Protocol v1.6 ## Overview TRT-LLM collects anonymous, session-level deployment telemetry to understand -how the library is used in production (GPU types, parallelism configs, model -architectures). No PII, model weights, prompts, outputs, model paths, tokenizer +how the library is used in production (GPU types, LLM and VisualGen parallelism, +and bounded model or pipeline categories). No PII, model weights, prompts, outputs, model paths, tokenizer paths, or raw free-form configuration strings are collected. **Opt-out** (any one of these disables telemetry): @@ -31,7 +31,7 @@ these top-level fields in Kibana alongside the event parameters. | `clientType` | string | Always `"Native"`. | | `clientVer` | string | TRT-LLM version, e.g. `"1.3.0rc9"`. | | `eventProtocol` | string | Always `"1.6"`. | -| `eventSchemaVer` | string | Schema version, currently `"0.7"`. | +| `eventSchemaVer` | string | Schema version, currently `"0.8"`. | | `eventSysVer` | string | Always `"trtllm-telemetry/1.0"`. | | `sessionId` | string | Unique hex UUID per telemetry session. Use this to correlate initial, heartbeat, and terminal events. | | `sentTs` | string | ISO 8601 UTC timestamp of when the payload was sent. | @@ -94,7 +94,7 @@ fails earlier can send a terminal report without an initial report. |-------|------|-------------|---------| | `ingressPoint` | ShortString | How TRT-LLM was invoked. See [Ingress point values](#ingress-point-values). | `"cli_serve"` | | `featuresJson` | string | Legacy JSON-serialized summary of feature flags. See [featuresJson keys](#featuresjson-keys). | `'{"lora":false,...}'` | -| `llmApiConfigJson` | string | JSON-serialized sanitized, type-driven effective LLM API configuration. See [LLM API config capture](#llm-api-config-capture). | `'{"tensor_parallel_size":2,...}'` | +| `llmApiConfigJson` | string | JSON-serialized sanitized, type-driven effective LLM API configuration. See [runtime config capture](#runtime-config-capture). | `'{"tensor_parallel_size":2,...}'` | | `llmApiConfigMetaJson` | string | JSON-serialized metadata for LLM API configuration capture. | `'{"capture_succeeded":true,...}'` | | `disaggRole` | ShortString | Disaggregated serving role. Empty if not disaggregated. | `""`, `"context"`, `"generation"`, `"coordinator"`, `"server_coordinator"`, `"ctx0"`, `"gen0"` | | `deploymentId` | ShortString | Shared ID across disaggregated workers. Empty if not disaggregated. | `""`, `"dep-abc123"` | @@ -127,6 +127,37 @@ heartbeats per session. Every heartbeat also contains the five aggregate LLM lifecycle counters above. +### `trtllm_visual_gen_initial_report` + +Sent once after the first successful VisualGen pipeline initialization. It +contains the system and GPU fields listed above plus bounded, effective runtime +metadata received from the ready worker. + +| Field group | Fields | Description | +|-------------|--------|-------------| +| Identity | `runtimeKind`, `ingressPoint`, `modelId`, `pipelineClassName`, `resolvedPipelineClass`, `modality` | Runtime and explicitly allowlisted built-in pipeline identity. Custom and unrecognized model IDs and classes use `other` or `unknown`; local paths are never sent. `modality` is pipeline capability, not a request's input type. | +| Deployment | `launchMode`, `nodeCount`, `nWorkers`, `gpuCount` | Bounded launch topology. `gpuCount` is local process visibility; `nWorkers` is total VisualGen workers. | +| Quantization | `quantizationAlgo`, `dynamicWeightQuant`, `quantizedComponentsJson`, `quantizedWeights` | Quantization resolved after loading. Components are restricted to `transformer` and `transformer_2`. | +| Parallelism | `cfgSize`, `ulyssesSize`, `asyncUlysses`, `ringSize`, `attn2dRowSize`, `attn2dColSize`, `tensorParallelSize`, `parallelVaeSize`, `parallelVaeSplitDim` | Effective VisualGen parallelism settings. Attention 2D topology is represented by separate row and column sizes. | +| Features | `parallelVae`, `stepCaching`, `cacheBackend`, `sparseAttention`, `sparseAttentionAlgorithm`, `sparseAttentionSparsity`, `quantizedAttention`, `attentionBackend`, `cudaGraphs`, `torchCompile` | Bounded feature flags and categories. | +| Safe config | `visualGenConfigJson`, `visualGenConfigMetaJson` | Type-driven, sanitized `VisualGenArgs` fields and capture diagnostics. | + +The event also contains these process-local VisualGen counters: + +| Field | Type | Description | +|-------|------|-------------| +| `visualGenInitializationAttempts` | PositiveInt | Number of entries into `VisualGen.__init__`. | +| `visualGenInstancesCreated` | PositiveInt | Number of VisualGen objects initialized successfully. | +| `activeVisualGenInstances` | PositiveInt | Successfully initialized objects not yet shut down. | +| `visualGenInitializationFailures` | PositiveInt | Initialization attempts that raised a handled Python exception. | + +### `trtllm_visual_gen_heartbeat` + +Sent periodically on the same cadence as the LLM heartbeat. It contains `seq`, +`runtimeKind`, `ingressPoint`, `nWorkers`, `gpuCount`, and the four VisualGen +lifecycle counters above. The heartbeat sequence provides approximate session +duration; this version does not send an exact duration field. + ### `trtllm_exit_report` Sent at most once when TRT-LLM or a surviving observer can classify the session @@ -139,14 +170,16 @@ outcome. Missing terminal events remain unknown; they are not confirmed crashes. | `signalNumber` | PositiveInt | Signal number, or `0` when not applicable or unknown. | | `terminationKind` | enum | `clean`, `exception`, `signal`, `worker_failure`, `timeout`, or `unknown`. | | `lifecyclePhase` | enum | Last known phase reached before termination: `cli_parsing`, `config_validation`, `model_initialization`, `serving`, or `unknown`. | -| `component` | enum | `llm`, `server`, `engine_worker`, `disagg_worker`, or `unknown`. | +| `component` | enum | `llm`, `visual_gen`, `server`, `engine_worker`, `disagg_worker`, or `unknown`. | | `reportingSource` | enum | `self`, `supervisor`, or `executor_proxy`. | +| `runtimeKind` | enum | Runtime families observed in the process: `llm`, `visual_gen`, `mixed`, or `unknown`. | | `ingressPoint` | ShortString | Entry point copied onto the terminal event so terminal-only early failures remain attributable. | | `disaggRole` | ShortString | Disaggregated role (`context`, `generation`, `coordinator`, `server_coordinator`, or compatible legacy `ctx0`/`gen0`), or empty when unavailable/not applicable. | | `deploymentId` | ShortString | Optional shared disaggregated deployment ID. | Every terminal report also contains the five aggregate LLM lifecycle counters -above. Delivery is best-effort and waits no more than 0.5 seconds; the local +and four VisualGen lifecycle counters above. Delivery is best-effort and waits +no more than 0.5 seconds; the local terminal lock permits at most one delivery attempt per process session. When a surviving parent observes a subprocess return code such as `-9`, it is @@ -169,6 +202,7 @@ before it can send a terminal report. | ShortString | string | 0–128 characters | | LongString | string | 0–256 characters | | PositiveInt | integer | 0–4,294,967,295 | +| Fraction | number | 0.0–1.0 | ## Ingress Point Values @@ -180,6 +214,7 @@ The `ingressPoint` field identifies which TRT-LLM entry point started the sessio | `"cli_bench"` | Started via `trtllm-bench` CLI | | `"cli_eval"` | Started via evaluation CLI | | `"llm_class"` | Started via `LLM()` Python API directly | +| `"visual_gen_class"` | Started via `VisualGen()` Python API directly | | `"disaggregated"` | Started as a disaggregated coordinator or fleet worker | | `"unknown"` | Entry point not identified | @@ -200,10 +235,11 @@ flags such as LoRA/speculative decoding have explicit safe config fields. | `chunked_context` | bool | `false` | Chunked prefill enabled (`enable_chunked_prefill=True`). | | `data_parallel_size` | int | `1` | Data parallel degree. `1` = no data parallelism. Derived from `tp_size` when attention DP is enabled. | -## LLM API Config Capture +## Runtime Config Capture -The `llmApiConfigJson` field is a JSON-serialized dict containing a type-driven -subset of the validated, effective LLM API configuration. Capture is +The `llmApiConfigJson` and `visualGenConfigJson` fields are JSON-serialized +dicts containing type-driven subsets of the validated, effective LLM or +VisualGen configuration. Capture is **type-driven**: a field is captured automatically when its type is categorical (`Literal`/`Enum`/`bool`) or numeric (`int`/`float`), or a safe collection of those. Free-form `str`/`Any`/`Path`/`dict`/`Callable` are not captured unless the @@ -215,15 +251,15 @@ field is a `Literal[...]` or uses an explicit `allowlist` converter. Paths, tokenizer locations, dicts, objects, callables, raw `Any` values, non-finite floats (`nan`/`inf`), and unsafe or heterogeneous sequences are excluded. Captured sequences are capped at a fixed length and any clipping is reported in -`llmApiConfigMetaJson`. Exclusion is fail-closed: the value is omitted instead -of being serialized, and `llmApiConfigMetaJson` reports whether any resolved field -was excluded as unsafe. +the corresponding metadata field. Exclusion is fail-closed: the value is +omitted instead of being serialized, and the metadata reports whether any +resolved field was excluded as unsafe. The table below is a non-exhaustive set of examples for readers building dashboards. The exhaustive source of truth is `tensorrt_llm/usage/llm_args_golden_manifest.json` (regenerated from `build_capture_manifest`), after the safety sanitizer has excluded unsafe values. -Use `llmApiConfigMetaJson` digests and field counts to track the exact capture +Use the config metadata digests and field counts to track the exact capture manifest for a given release. The rendered documentation generates the exhaustive field table at docs build time under **Developer Guide > Telemetry**. @@ -250,7 +286,7 @@ exhaustive field table at docs build time under **Developer Guide > Telemetry**. | `sparse_attention_config.algorithm` | Sparse attention algorithm discriminator; arm-specific knobs appear under `sparse_attention_config.*`. | | `reasoning_parser` | Reasoning parser selection, captured through an allowlist mirroring the `ReasoningParserFactory` registry. | -`llmApiConfigMetaJson` describes the capture process itself. It includes +The matching config metadata field describes the capture process itself. It includes contract/version fields, schema and manifest digests, source args class, field counts (`capturable_field_count`, `captured_field_count`, `excluded_field_count`), capture success, unsafe-exclusion status, a `sequence_truncated` flag set when any captured @@ -276,7 +312,7 @@ over time. Checklist for adding a telemetry field: -1. **`tensorrt_llm/usage/schema.py`** — Add field to `TrtllmInitialReport` (or `TrtllmHeartbeat`) Pydantic model with alias. +1. **`tensorrt_llm/usage/schema.py`** — Add the field to the appropriate event Pydantic model with an alias. 2. **`tensorrt_llm/usage/schemas/trtllm_usage_event_schema.json`** — Add to `properties` and `required` array. 3. **`tensorrt_llm/usage/usage_lib.py`** — Populate the field in `_background_reporter()` and add extraction logic in `_extract_trtllm_config()` or `_collect_gpu_info()` as appropriate. 4. **`tests/unittest/usage/test_schema.py`** — Update test fixtures and expected field sets. @@ -285,7 +321,7 @@ Checklist for adding a telemetry field: 7. **SMS schema upload** — Upload the updated JSON schema to the NvTelemetry Schema Management Service and toggle "on stage" / "on prod". 8. **Update this README** — Add the field to the appropriate table above. -Checklist for adding an LLM API config capture field inside `llmApiConfigJson`: +Checklist for adding a runtime config capture field inside `llmApiConfigJson` or `visualGenConfigJson`: 1. **Add the field with its natural type.** If it is categorical (`Literal`/`Enum`/`bool`) or numeric (`int`/`float`) — or a safe collection of @@ -317,10 +353,17 @@ Checklist for adding an LLM API config capture field inside `llmApiConfigJson`: important enough for dashboard users to know by name. Dashboard note: payloads carry `capture_version` and `field_policy_version` in -`llmApiConfigMetaJson`. During release adoption, v1 (opt-in) and v2 (type-driven) +their config metadata. During release adoption, v1 (opt-in) and v2 (type-driven) payloads coexist in the same index — **bucket by these before aggregating** `captured_field_count` or any `llmApiConfigJson.`. +### VisualGen V1 Boundaries + +- No prompts, input media, outputs, request-level parameters, latency, queue, or throughput data are collected. +- Pipeline and quantization identity are reported only after a worker becomes ready; earlier failures have lifecycle and terminal counters but no resolved pipeline event. +- Dedicated flags without an authoritative effective runtime value, including two-stage, guardrails, prefetch, and request-level audio or layered-image behavior, are deferred. `resolvedPipelineClass` can still identify a built-in two-stage pipeline. +- A process has one reporter thread. With multiple runtimes or VisualGen objects, the first runtime supplies static initial and heartbeat metadata while lifecycle counters aggregate all objects. If LLM and VisualGen are both created, the terminal event reports `runtimeKind=mixed` and both counter sets. + ### Conventions - Use **camelCase** aliases for JSON wire format (Pydantic `alias=`). diff --git a/tensorrt_llm/usage/schemas/trtllm_usage_event_schema.json b/tensorrt_llm/usage/schemas/trtllm_usage_event_schema.json index 81ef2e9d4f74..ebd11a32feb0 100644 --- a/tensorrt_llm/usage/schemas/trtllm_usage_event_schema.json +++ b/tensorrt_llm/usage/schemas/trtllm_usage_event_schema.json @@ -6,18 +6,24 @@ { "$ref": "#/definitions/events/trtllm_heartbeat" }, + { + "$ref": "#/definitions/events/trtllm_visual_gen_initial_report" + }, + { + "$ref": "#/definitions/events/trtllm_visual_gen_heartbeat" + }, { "$ref": "#/definitions/events/trtllm_exit_report" } ], "$schema": "http://json-schema.org/draft-07/schema#", "schemaMeta": { - "schemaVersion": "0.7", + "schemaVersion": "0.8", "clientId": "616561816355034", "clientName": "TrtllmTelemetry", "definitionVersion": "2.0" }, - "description": "TensorRT-LLM usage telemetry events. Collects anonymous session-level deployment data (GPU type, parallelism config, model architecture class) with opt-out via TRTLLM_NO_USAGE_STATS=1. Auto-disabled in CI/test environments.", + "description": "TensorRT-LLM usage telemetry events. Collects anonymous session-level LLM and VisualGen deployment data with opt-out via TRTLLM_NO_USAGE_STATS=1. Auto-disabled in CI/test environments.", "definitions": { "types": { "ShortString": { @@ -34,6 +40,11 @@ "type": "integer", "minimum": 0, "maximum": 4294967295 + }, + "Fraction": { + "type": "number", + "minimum": 0.0, + "maximum": 1.0 } }, "events": { @@ -117,7 +128,7 @@ }, "llmApiConfigJson": { "type": "string", - "description": "JSON-serialized sanitized, opt-in subset of effective LLM API configuration" + "description": "JSON-serialized sanitized, type-driven subset of effective LLM API configuration" }, "llmApiConfigMetaJson": { "type": "string", @@ -230,6 +241,303 @@ "llmInitializationFailures" ] }, + "trtllm_visual_gen_initial_report": { + "eventMeta": { + "service": "telemetry", + "gdpr": { + "category": "functional", + "description": "TRT-LLM VisualGen initial session report with bounded pipeline, topology, feature, and sanitized configuration details. No prompts, outputs, model paths, or request data are collected." + } + }, + "additionalProperties": false, + "type": "object", + "properties": { + "trtllmVersion": { + "$ref": "#/definitions/types/ShortString" + }, + "platform": { + "$ref": "#/definitions/types/LongString" + }, + "pythonVersion": { + "$ref": "#/definitions/types/ShortString" + }, + "cpuArchitecture": { + "$ref": "#/definitions/types/ShortString" + }, + "cpuCount": { + "$ref": "#/definitions/types/PositiveInt" + }, + "gpuCount": { + "$ref": "#/definitions/types/PositiveInt" + }, + "gpuName": { + "$ref": "#/definitions/types/LongString" + }, + "gpuMemoryMB": { + "$ref": "#/definitions/types/PositiveInt" + }, + "cudaVersion": { + "$ref": "#/definitions/types/ShortString" + }, + "runtimeKind": { + "type": "string", + "enum": [ + "visual_gen" + ] + }, + "ingressPoint": { + "$ref": "#/definitions/types/ShortString" + }, + "modelId": { + "$ref": "#/definitions/types/LongString" + }, + "pipelineClassName": { + "$ref": "#/definitions/types/LongString" + }, + "resolvedPipelineClass": { + "$ref": "#/definitions/types/LongString" + }, + "modality": { + "type": "string", + "enum": [ + "image", + "video", + "video_audio", + "layered_image", + "mixed", + "unknown" + ] + }, + "launchMode": { + "type": "string", + "enum": [ + "local_spawn", + "torchrun", + "slurm", + "unknown" + ] + }, + "nodeCount": { + "$ref": "#/definitions/types/PositiveInt" + }, + "nWorkers": { + "$ref": "#/definitions/types/PositiveInt" + }, + "quantizationAlgo": { + "$ref": "#/definitions/types/ShortString" + }, + "dynamicWeightQuant": { + "type": "boolean" + }, + "quantizedComponentsJson": { + "type": "string", + "description": "JSON-serialized allowlisted VisualGen component names with enabled weight quantization" + }, + "cfgSize": { + "$ref": "#/definitions/types/PositiveInt" + }, + "ulyssesSize": { + "$ref": "#/definitions/types/PositiveInt" + }, + "asyncUlysses": { + "type": "boolean" + }, + "ringSize": { + "$ref": "#/definitions/types/PositiveInt" + }, + "attn2dRowSize": { + "$ref": "#/definitions/types/PositiveInt" + }, + "attn2dColSize": { + "$ref": "#/definitions/types/PositiveInt" + }, + "tensorParallelSize": { + "$ref": "#/definitions/types/PositiveInt" + }, + "parallelVaeSize": { + "$ref": "#/definitions/types/PositiveInt" + }, + "parallelVaeSplitDim": { + "type": "string", + "enum": [ + "width", + "height" + ] + }, + "parallelVae": { + "type": "boolean" + }, + "stepCaching": { + "type": "boolean" + }, + "cacheBackend": { + "type": "string", + "enum": [ + "teacache", + "cache_dit", + "none" + ] + }, + "sparseAttention": { + "type": "boolean" + }, + "sparseAttentionAlgorithm": { + "type": "string", + "enum": [ + "skip_softmax", + "vsa", + "none" + ] + }, + "sparseAttentionSparsity": { + "$ref": "#/definitions/types/Fraction" + }, + "quantizedAttention": { + "type": "boolean" + }, + "attentionBackend": { + "type": "string", + "enum": [ + "VANILLA", + "TRTLLM", + "FA4", + "CUTEDSL" + ] + }, + "cudaGraphs": { + "type": "boolean" + }, + "torchCompile": { + "type": "boolean" + }, + "quantizedWeights": { + "type": "boolean" + }, + "visualGenConfigJson": { + "type": "string", + "description": "JSON-serialized sanitized, type-driven subset of effective VisualGen configuration" + }, + "visualGenConfigMetaJson": { + "type": "string", + "description": "JSON-serialized metadata for VisualGen configuration capture" + }, + "visualGenInitializationAttempts": { + "$ref": "#/definitions/types/PositiveInt" + }, + "visualGenInstancesCreated": { + "$ref": "#/definitions/types/PositiveInt" + }, + "activeVisualGenInstances": { + "$ref": "#/definitions/types/PositiveInt" + }, + "visualGenInitializationFailures": { + "$ref": "#/definitions/types/PositiveInt" + } + }, + "required": [ + "trtllmVersion", + "platform", + "pythonVersion", + "cpuArchitecture", + "cpuCount", + "gpuCount", + "gpuName", + "gpuMemoryMB", + "cudaVersion", + "runtimeKind", + "ingressPoint", + "modelId", + "pipelineClassName", + "resolvedPipelineClass", + "modality", + "launchMode", + "nodeCount", + "nWorkers", + "quantizationAlgo", + "dynamicWeightQuant", + "quantizedComponentsJson", + "cfgSize", + "ulyssesSize", + "asyncUlysses", + "ringSize", + "attn2dRowSize", + "attn2dColSize", + "tensorParallelSize", + "parallelVaeSize", + "parallelVaeSplitDim", + "parallelVae", + "stepCaching", + "cacheBackend", + "sparseAttention", + "sparseAttentionAlgorithm", + "sparseAttentionSparsity", + "quantizedAttention", + "attentionBackend", + "cudaGraphs", + "torchCompile", + "quantizedWeights", + "visualGenConfigJson", + "visualGenConfigMetaJson", + "visualGenInitializationAttempts", + "visualGenInstancesCreated", + "activeVisualGenInstances", + "visualGenInitializationFailures" + ] + }, + "trtllm_visual_gen_heartbeat": { + "eventMeta": { + "service": "telemetry", + "gdpr": { + "category": "functional", + "description": "TRT-LLM VisualGen periodic heartbeat with bounded topology and aggregate lifecycle counters." + } + }, + "additionalProperties": false, + "type": "object", + "properties": { + "seq": { + "$ref": "#/definitions/types/PositiveInt" + }, + "runtimeKind": { + "type": "string", + "enum": [ + "visual_gen" + ] + }, + "ingressPoint": { + "$ref": "#/definitions/types/ShortString" + }, + "nWorkers": { + "$ref": "#/definitions/types/PositiveInt" + }, + "gpuCount": { + "$ref": "#/definitions/types/PositiveInt" + }, + "visualGenInitializationAttempts": { + "$ref": "#/definitions/types/PositiveInt" + }, + "visualGenInstancesCreated": { + "$ref": "#/definitions/types/PositiveInt" + }, + "activeVisualGenInstances": { + "$ref": "#/definitions/types/PositiveInt" + }, + "visualGenInitializationFailures": { + "$ref": "#/definitions/types/PositiveInt" + } + }, + "required": [ + "seq", + "runtimeKind", + "ingressPoint", + "nWorkers", + "gpuCount", + "visualGenInitializationAttempts", + "visualGenInstancesCreated", + "activeVisualGenInstances", + "visualGenInitializationFailures" + ] + }, "trtllm_exit_report": { "eventMeta": { "service": "telemetry", @@ -275,6 +583,7 @@ "type": "string", "enum": [ "llm", + "visual_gen", "server", "engine_worker", "disagg_worker", @@ -312,6 +621,27 @@ }, "llmInitializationFailures": { "$ref": "#/definitions/types/PositiveInt" + }, + "visualGenInitializationAttempts": { + "$ref": "#/definitions/types/PositiveInt" + }, + "visualGenInstancesCreated": { + "$ref": "#/definitions/types/PositiveInt" + }, + "activeVisualGenInstances": { + "$ref": "#/definitions/types/PositiveInt" + }, + "visualGenInitializationFailures": { + "$ref": "#/definitions/types/PositiveInt" + }, + "runtimeKind": { + "type": "string", + "enum": [ + "llm", + "visual_gen", + "mixed", + "unknown" + ] } }, "required": [ @@ -329,7 +659,12 @@ "llmInstancesCreated", "activeLlmInstances", "maxConcurrentLlmInstances", - "llmInitializationFailures" + "llmInitializationFailures", + "visualGenInitializationAttempts", + "visualGenInstancesCreated", + "activeVisualGenInstances", + "visualGenInitializationFailures", + "runtimeKind" ] } } diff --git a/tensorrt_llm/usage/usage_lib.py b/tensorrt_llm/usage/usage_lib.py index cb342ca17d5c..958d4adee2c5 100644 --- a/tensorrt_llm/usage/usage_lib.py +++ b/tensorrt_llm/usage/usage_lib.py @@ -58,10 +58,16 @@ from tensorrt_llm.usage import schema from tensorrt_llm.usage.config import UsageContext -from tensorrt_llm.usage.llmapi_config import _failure_llm_api_config_payloads +from tensorrt_llm.usage.llmapi_config import ( + _failure_llm_api_config_payloads, + _failure_visual_gen_config_payloads, +) from tensorrt_llm.usage.llmapi_config import ( collect_llm_api_config_payloads as _collect_llm_api_config_payloads, ) +from tensorrt_llm.usage.llmapi_config import ( + collect_visual_gen_config_payloads as _collect_visual_gen_config_payloads, +) logger = logging.getLogger("tensorrt_llm") @@ -617,6 +623,105 @@ def _clamp_str(value: str, max_len: int) -> str: return value[:max_len] if len(value) > max_len else value +def _bounded_uint(value: Any, default: int = 0) -> int: + """Return a value inside NvTelemetry's unsigned 32-bit range.""" + if isinstance(value, int) and not isinstance(value, bool): + return min(max(value, 0), schema._UINT32_MAX) + return default + + +def _bounded_fraction(value: Any) -> float: + """Return a finite numeric value in the closed interval [0, 1].""" + if isinstance(value, (int, float)) and not isinstance(value, bool): + value = float(value) + if value == value and value not in (float("inf"), float("-inf")): + return min(max(value, 0.0), 1.0) + return 0.0 + + +def _visual_gen_initial_fields( + visual_gen_args: Any, + metadata: dict[str, Any], +) -> dict[str, Any]: + """Extract bounded static VisualGen fields from validated runtime state.""" + parallel = visual_gen_args.parallel_config + attention = visual_gen_args.attention_config + sparse = attention.sparse_attention_config + cache = visual_gen_args.cache_config + + modality = metadata.get("modality", "unknown") + if modality not in {"image", "video", "video_audio", "layered_image", "mixed"}: + modality = "unknown" + launch_mode = metadata.get("launch_mode", "unknown") + if launch_mode not in {"local_spawn", "torchrun", "slurm"}: + launch_mode = "unknown" + sparse_algorithm = getattr(sparse, "algorithm", "none") if sparse is not None else "none" + if sparse_algorithm not in {"skip_softmax", "vsa"}: + sparse_algorithm = "none" + if sparse_algorithm == "vsa": + sparse_sparsity = getattr(sparse, "vsa_sparsity", 0.0) + else: + sparse_sparsity = getattr(sparse, "target_sparsity", 0.0) or 0.0 + cache_backend = getattr(cache, "cache_backend", "none") if cache is not None else "none" + if cache_backend not in {"teacache", "cache_dit"}: + cache_backend = "none" + attention_backend = attention.backend + if attention_backend not in {"VANILLA", "TRTLLM", "FA4", "CUTEDSL"}: + attention_backend = "VANILLA" + + quantized_components = metadata.get("quantized_components", []) + if not isinstance(quantized_components, list): + quantized_components = [] + quantized_components = [ + component + for component in quantized_components + if component in ("transformer", "transformer_2") + ] + + attn2d_row, attn2d_col = parallel.attn2d_size + return { + "modelId": _clamp_str(str(metadata.get("model_id", "other")), schema._LONG_STR), + "pipelineClassName": _clamp_str( + str(metadata.get("pipeline_class_name", "unknown")), schema._LONG_STR + ), + "resolvedPipelineClass": _clamp_str( + str(metadata.get("resolved_pipeline_class", "unknown")), schema._LONG_STR + ), + "modality": modality, + "launchMode": launch_mode, + "nodeCount": _bounded_uint(metadata.get("node_count")), + "nWorkers": _bounded_uint(metadata.get("n_workers"), 1), + "quantizationAlgo": _clamp_str( + str(metadata.get("quantization_algo", "")), schema._SHORT_STR + ), + "dynamicWeightQuant": metadata.get("dynamic_weight_quant") is True, + "quantizedComponentsJson": json.dumps( + quantized_components, + separators=(",", ":"), + ), + "cfgSize": _bounded_uint(parallel.cfg_size, 1), + "ulyssesSize": _bounded_uint(parallel.ulysses_size, 1), + "asyncUlysses": parallel.async_ulysses is True, + "ringSize": _bounded_uint(parallel.ring_size, 1), + "attn2dRowSize": _bounded_uint(attn2d_row, 1), + "attn2dColSize": _bounded_uint(attn2d_col, 1), + "tensorParallelSize": _bounded_uint(parallel.tp_size, 1), + "parallelVaeSize": _bounded_uint(parallel.parallel_vae_size, 1), + "parallelVaeSplitDim": parallel.parallel_vae_split_dim, + "parallelVae": parallel.parallel_vae_size > 1, + "stepCaching": cache is not None, + "cacheBackend": cache_backend, + "sparseAttention": sparse is not None, + "sparseAttentionAlgorithm": sparse_algorithm, + "sparseAttentionSparsity": _bounded_fraction(sparse_sparsity), + "quantizedAttention": attention.quant_attention_config is not None, + "attentionBackend": attention_backend, + "cudaGraphs": visual_gen_args.cuda_graph_config.enable is True, + "torchCompile": visual_gen_args.torch_compile_config.enable is True, + "quantizedWeights": bool(quantized_components), + } + + def _background_reporter( llm_args: Any, pretrained_config: Any, @@ -737,6 +842,81 @@ def _background_reporter( _finish_background_reporter() +def _visual_gen_background_reporter( + visual_gen_args: Any, + metadata: dict[str, Any], + usage_context: str = "", +) -> None: + """Send the VisualGen initial report and periodic liveness snapshots.""" + try: + session = _get_session() + if session is None: + return + session_id = session.session_id + trtllm_version = session.trtllm_version + system_info = _collect_system_info() + gpu_info = _collect_gpu_info() + try: + config_json, config_meta_json = _collect_visual_gen_config_payloads(visual_gen_args) + except Exception: + config_json, config_meta_json = _failure_visual_gen_config_payloads( + args_class="VisualGenArgs" + ) + + event_snapshot = _event_snapshot(usage_context) + static_fields = _visual_gen_initial_fields(visual_gen_args, metadata) + initial_event = schema.TrtllmVisualGenInitialReport( + trtllmVersion=_clamp_str(trtllm_version or "", schema._SHORT_STR), + platform=_clamp_str(system_info.get("platform") or "", schema._LONG_STR), + pythonVersion=_clamp_str(system_info.get("python_version") or "", schema._SHORT_STR), + cpuArchitecture=_clamp_str( + system_info.get("cpu_architecture") or "", schema._SHORT_STR + ), + cpuCount=_bounded_uint(system_info.get("cpu_count")), + gpuCount=_bounded_uint(gpu_info.get("gpu_count")), + gpuName=_clamp_str(gpu_info.get("gpu_name") or "", schema._LONG_STR), + gpuMemoryMB=_bounded_uint(gpu_info.get("gpu_memory_mb")), + cudaVersion=_clamp_str(gpu_info.get("cuda_version") or "", schema._SHORT_STR), + visualGenConfigJson=config_json, + visualGenConfigMetaJson=config_meta_json, + **static_fields, + **_visual_gen_session_event_fields(event_snapshot), + ) + payload = schema.build_gxt_payload( + event=initial_event, + session_id=session_id, + trtllm_version=trtllm_version, + ) + if not session.claim_initial(): + return + _send_if_session_active(session, payload) + + heartbeat_interval = _get_heartbeat_interval() + for seq in range(_MAX_HEARTBEATS): + if _REPORTER_STOP.wait(timeout=heartbeat_interval): + return + try: + event_snapshot = _event_snapshot(usage_context) + heartbeat_event = schema.TrtllmVisualGenHeartbeat( + seq=seq, + nWorkers=static_fields["nWorkers"], + gpuCount=_bounded_uint(gpu_info.get("gpu_count")), + **_visual_gen_session_event_fields(event_snapshot), + ) + heartbeat_payload = schema.build_gxt_payload( + event=heartbeat_event, + session_id=session_id, + trtllm_version=trtllm_version, + ) + _send_if_session_active(session, heartbeat_payload) + except (urllib.error.URLError, OSError, ValueError, TypeError): + pass + except Exception: + pass + finally: + _finish_background_reporter() + + # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- @@ -778,6 +958,7 @@ def __init__( self.lifecycle_phase = lifecycle_phase self.observed_signal = 0 self.observed_outcome: Optional[TerminalOutcome] = None + self.runtime_kind: schema.RuntimeKind = "unknown" self.llm_initialization_attempts = 0 self.llm_instances_created = 0 @@ -785,6 +966,11 @@ def __init__( self.max_concurrent_llm_instances = 0 self.llm_initialization_failures = 0 + self.visual_gen_initialization_attempts = 0 + self.visual_gen_instances_created = 0 + self.active_visual_gen_instances = 0 + self.visual_gen_initialization_failures = 0 + self.disabled = False self.initial_reported = False self.terminal_reported = False @@ -809,6 +995,12 @@ def _refresh_metadata_unlocked(self) -> None: ): self.component = "disagg_worker" + def _observe_runtime_unlocked(self, runtime_kind: schema.RuntimeKind) -> None: + if self.runtime_kind == "unknown": + self.runtime_kind = runtime_kind + elif self.runtime_kind != runtime_kind: + self.runtime_kind = "mixed" + def refresh_metadata(self) -> None: """Promote correlation fields that become known after early startup.""" with self.lock: @@ -830,7 +1022,10 @@ def configure( "unknown", ): self.usage_context = usage_context - if self.component == "unknown" and component is not None: + if component is not None and ( + self.component == "unknown" + or (component == "visual_gen" and self.runtime_kind == "unknown") + ): self.component = component if lifecycle_phase is not None: self.lifecycle_phase = lifecycle_phase @@ -842,6 +1037,7 @@ def record_llm_initialization_attempt(self) -> bool: if self.disabled or self.terminal_reported: return False self.llm_initialization_attempts = self._increment(self.llm_initialization_attempts) + self._observe_runtime_unlocked("llm") self.lifecycle_phase = "model_initialization" if self.component == "unknown": self.component = "llm" @@ -874,6 +1070,45 @@ def record_llm_shutdown(self) -> None: if not self.disabled and self.active_llm_instances > 0: self.active_llm_instances -= 1 + def record_visual_gen_initialization_attempt(self) -> bool: + """Increment the attempt counter and enter VisualGen initialization.""" + with self.lock: + if self.disabled or self.terminal_reported: + return False + if self.runtime_kind == "unknown" and self.component in ("unknown", "llm"): + self.component = "visual_gen" + self._observe_runtime_unlocked("visual_gen") + self.visual_gen_initialization_attempts = self._increment( + self.visual_gen_initialization_attempts + ) + self.lifecycle_phase = "model_initialization" + self._refresh_metadata_unlocked() + return True + + def record_visual_gen_initialization_failure(self) -> None: + """Increment the VisualGen initialization failure counter.""" + with self.lock: + if not self.disabled and not self.terminal_reported: + self.visual_gen_initialization_failures = self._increment( + self.visual_gen_initialization_failures + ) + + def record_visual_gen_initialized(self) -> bool: + """Record a successfully constructed and active VisualGen object.""" + with self.lock: + if self.disabled or self.terminal_reported: + return False + self.visual_gen_instances_created = self._increment(self.visual_gen_instances_created) + self.active_visual_gen_instances = self._increment(self.active_visual_gen_instances) + self.lifecycle_phase = "serving" + return True + + def record_visual_gen_shutdown(self) -> None: + """Decrement the VisualGen active gauge without underflowing.""" + with self.lock: + if not self.disabled and self.active_visual_gen_instances > 0: + self.active_visual_gen_instances -= 1 + def set_lifecycle_phase(self, lifecycle_phase: schema.LifecyclePhase) -> None: """Update the best-known process lifecycle phase.""" with self.lock: @@ -920,6 +1155,11 @@ def _snapshot_unlocked(self) -> dict[str, Any]: "activeLlmInstances": self.active_llm_instances, "maxConcurrentLlmInstances": self.max_concurrent_llm_instances, "llmInitializationFailures": self.llm_initialization_failures, + "visualGenInitializationAttempts": self.visual_gen_initialization_attempts, + "visualGenInstancesCreated": self.visual_gen_instances_created, + "activeVisualGenInstances": self.active_visual_gen_instances, + "visualGenInitializationFailures": self.visual_gen_initialization_failures, + "runtimeKind": self.runtime_kind, "lifecyclePhase": self.lifecycle_phase, "component": self.component, "observedSignal": self.observed_signal, @@ -1049,6 +1289,11 @@ def _empty_event_snapshot(usage_context: str = "") -> dict[str, Any]: "activeLlmInstances": 0, "maxConcurrentLlmInstances": 0, "llmInitializationFailures": 0, + "visualGenInitializationAttempts": 0, + "visualGenInstancesCreated": 0, + "activeVisualGenInstances": 0, + "visualGenInitializationFailures": 0, + "runtimeKind": "unknown", "lifecyclePhase": "unknown", "component": "unknown", "observedSignal": 0, @@ -1062,7 +1307,7 @@ def _event_snapshot(usage_context: str = "") -> dict[str, Any]: def _session_event_fields(snapshot: dict[str, Any]) -> dict[str, Any]: - """Select correlation and lifecycle counters shared by all events.""" + """Select correlation and LLM counters for LLM events.""" return { "ingressPoint": snapshot["ingressPoint"], "disaggRole": snapshot["disaggRole"], @@ -1075,6 +1320,30 @@ def _session_event_fields(snapshot: dict[str, Any]) -> dict[str, Any]: } +def _visual_gen_session_event_fields(snapshot: dict[str, Any]) -> dict[str, Any]: + """Select ingress and VisualGen counters for VisualGen events.""" + return { + "ingressPoint": snapshot["ingressPoint"], + "visualGenInitializationAttempts": snapshot["visualGenInitializationAttempts"], + "visualGenInstancesCreated": snapshot["visualGenInstancesCreated"], + "activeVisualGenInstances": snapshot["activeVisualGenInstances"], + "visualGenInitializationFailures": snapshot["visualGenInitializationFailures"], + } + + +def _exit_session_event_fields(snapshot: dict[str, Any]) -> dict[str, Any]: + """Select correlation and all runtime counters for terminal events.""" + fields = _session_event_fields(snapshot) + fields.update( + runtimeKind=snapshot["runtimeKind"], + visualGenInitializationAttempts=snapshot["visualGenInitializationAttempts"], + visualGenInstancesCreated=snapshot["visualGenInstancesCreated"], + activeVisualGenInstances=snapshot["activeVisualGenInstances"], + visualGenInitializationFailures=snapshot["visualGenInitializationFailures"], + ) + return fields + + def _validated_usage_context(value: Any) -> str: """Return a bounded ingress category or the unset sentinel.""" if isinstance(value, UsageContext): @@ -1265,6 +1534,46 @@ def record_llm_shutdown() -> None: _session_call(lambda session: session.record_llm_shutdown(), None) +def record_visual_gen_initialization_attempt( + telemetry_config: Any = None, + *, + default_usage_context: str = "visual_gen_class", +) -> bool: + """Start local tracking and record entry into a VisualGen constructor.""" + if not apply_usage_session_config( + telemetry_config, + default_usage_context=default_usage_context, + component="visual_gen", + lifecycle_phase="model_initialization", + ): + return False + return _session_call( + lambda session: session.record_visual_gen_initialization_attempt(), + False, + ) + + +def record_visual_gen_initialization_failure() -> None: + """Record a handled Python exception from VisualGen construction.""" + _session_call( + lambda session: session.record_visual_gen_initialization_failure(), + None, + ) + + +def record_visual_gen_initialized() -> bool: + """Record one successfully constructed VisualGen object.""" + return _session_call( + lambda session: session.record_visual_gen_initialized(), + False, + ) + + +def record_visual_gen_shutdown() -> None: + """Mark one successfully tracked VisualGen object inactive.""" + _session_call(lambda session: session.record_visual_gen_shutdown(), None) + + def set_lifecycle_phase(lifecycle_phase: schema.LifecyclePhase) -> None: """Update the best-known phase for an outer process boundary.""" _session_call(lambda session: session.set_lifecycle_phase(lifecycle_phase), None) @@ -1382,6 +1691,7 @@ def report_exit( outcome.termination_kind == "clean" and snapshot["lifecyclePhase"] in ("cli_parsing", "config_validation") and snapshot["llmInitializationAttempts"] == 0 + and snapshot["visualGenInitializationAttempts"] == 0 ): return True @@ -1406,7 +1716,7 @@ def report_exit( lifecyclePhase=lifecycle_phase or snapshot["lifecyclePhase"], component=outcome.component or snapshot["component"], reportingSource=outcome.reporting_source, - **_session_event_fields(snapshot), + **_exit_session_event_fields(snapshot), ) payload = schema.build_gxt_payload( event=event, @@ -1447,28 +1757,17 @@ def report_exit( return claimed -def report_usage( - llm_args: Any = None, - pretrained_config: Any = None, - telemetry_config: Any = None, +def _start_background_usage_reporter( + target: Callable[..., None], + reporter_args: tuple[Any, ...], + telemetry_config: Any, + thread_name: str, ) -> None: - """Start background usage telemetry reporting. - - Call this once after model initialization. It spawns a daemon thread - that sends an initial report and periodic heartbeats. Subsequent calls - are no-ops (only one reporter thread per process). - - This function is fail-silent -- it will never raise an exception or - block the calling thread. - - Args: - llm_args: The LlmArgs object from BaseLLM (for config extraction). - pretrained_config: The pretrained model config (for architecture name). - telemetry_config: TelemetryConfig object (opt-out + usage context). - """ + """Start the process's single fail-silent usage reporter.""" global _PENDING_TERMINAL global _REPORTER_ACTIVE global _REPORTER_STARTED + claimed = False try: _, usage_context = _telemetry_settings(telemetry_config) if not apply_usage_session_config(telemetry_config): @@ -1483,18 +1782,19 @@ def report_usage( return _REPORTER_STARTED = True _REPORTER_ACTIVE = True + claimed = True _show_usage_notification() - thread = threading.Thread( - target=_background_reporter, - args=(llm_args, pretrained_config, usage_context), + target=target, + args=(*reporter_args, usage_context), daemon=True, - name="trtllm-usage-stats", + name=thread_name, ) thread.start() - except Exception: + if not claimed: + return pending = None with _REPORTER_LOCK: _REPORTER_STARTED = False @@ -1503,3 +1803,38 @@ def report_usage( _PENDING_TERMINAL = None if pending is not None: pending.completion.set() + + +def report_usage( + llm_args: Any = None, + pretrained_config: Any = None, + telemetry_config: Any = None, +) -> None: + """Start the LLM initial and heartbeat reporter after model startup.""" + _start_background_usage_reporter( + _background_reporter, + (llm_args, pretrained_config), + telemetry_config, + "trtllm-usage-stats", + ) + + +def report_visual_gen_usage( + visual_gen_args: Any, + metadata: Optional[dict[str, Any]] = None, + telemetry_config: Any = None, +) -> None: + """Start VisualGen initial and heartbeat reporting after pipeline startup. + + ``metadata`` must already contain only sanitized, bounded, allowlisted values. + """ + try: + reporter_metadata = dict(metadata or {}) + except Exception: + return + _start_background_usage_reporter( + _visual_gen_background_reporter, + (visual_gen_args, reporter_metadata), + telemetry_config, + "trtllm-visual-gen-usage-stats", + ) diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index 3559712867a8..a586d996b78d 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -30,6 +30,7 @@ from tensorrt_llm.llmapi.llm_args import Field from tensorrt_llm.llmapi.utils import StrictBaseModel, set_api_status from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.usage.config import TelemetryConfig from .sparse_attention import SkipSoftmaxAttentionConfig, VideoSparseAttentionConfig @@ -660,6 +661,13 @@ class VisualGenArgs(StrictBaseModel): ), ) + telemetry_config: TelemetryConfig = Field( + default_factory=TelemetryConfig, + telemetry=False, + status="prototype", + description="Configure anonymous usage telemetry for this VisualGen process.", + ) + revision: Optional[str] = Field( None, status="prototype", diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index a6274d353406..4066e324ba11 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -17,6 +17,7 @@ import itertools import secrets import sys +import threading import weakref from pathlib import Path from typing import Any, AsyncIterator, Dict, List, Literal, Optional, Union @@ -245,6 +246,8 @@ def __init__( ): self.model = str(model) self.args = (args or VisualGenArgs()).model_copy(update={"model": self.model}) + self._usage_lifecycle_active = False + self._usage_lifecycle_lock = threading.Lock() # In external-launch mode (torchrun/srun), ranks 1..N-1 run as pure # workers and never return to user code. @@ -252,13 +255,13 @@ def __init__( if ext is not None: rank, local_rank, world_size, master_addr, master_port = ext n_workers = self.args.parallel_config.n_workers - if world_size != n_workers: - raise ValueError( - f"Launcher world_size ({world_size}) does not match " - f"n_workers ({n_workers}). " - "Launch exactly n_workers tasks." - ) if rank != 0: + if world_size != n_workers: + raise ValueError( + f"Launcher world_size ({world_size}) does not match " + f"n_workers ({n_workers}). " + "Launch exactly n_workers tasks." + ) logger.info( f"VisualGen: rank {rank}/{world_size}, local_rank {local_rank} — " "starting as worker (external launch mode)" @@ -276,13 +279,61 @@ def __init__( local_rank=local_rank, ) sys.exit(0) - logger.info( - f"VisualGen: rank 0/{world_size} — coordinator + worker (external launch mode)" + + usage_attempt_tracked = False + try: + import tensorrt_llm.usage as _usage + from tensorrt_llm.usage import usage_lib as _usage_lib + + usage_attempt_tracked = _usage.record_visual_gen_initialization_attempt( + self.args.telemetry_config, + default_usage_context=_usage.UsageContext.VISUAL_GEN_CLASS.value, ) + except Exception as exc: + logger.debug(f"VisualGen telemetry initialization tracking failed: {exc}") + + try: + if ext is not None: + if world_size != n_workers: + raise ValueError( + f"Launcher world_size ({world_size}) does not match " + f"n_workers ({n_workers}). " + "Launch exactly n_workers tasks." + ) + logger.info( + f"VisualGen: rank 0/{world_size} — coordinator + worker (external launch mode)" + ) + self.executor = DiffusionRemoteClient(args=self.args) + except Exception: + try: + if usage_attempt_tracked: + _usage.record_visual_gen_initialization_failure() + except Exception as exc: + logger.debug(f"VisualGen telemetry failure tracking failed: {exc}") + raise + + try: + if usage_attempt_tracked: + self._usage_lifecycle_active = _usage.record_visual_gen_initialized() + telemetry_config = self.args.telemetry_config + if telemetry_config.usage_context == _usage.UsageContext.UNKNOWN: + telemetry_config = telemetry_config.model_copy( + update={"usage_context": _usage.UsageContext.VISUAL_GEN_CLASS} + ) + telemetry_metadata = dict(self.executor.telemetry_metadata) + telemetry_metadata.update( + launch_mode=self.executor.launch_mode, + node_count=self.executor.node_count, + n_workers=self.executor.n_workers, + ) + _usage_lib.report_visual_gen_usage( + self.args, + telemetry_metadata, + telemetry_config, + ) + except Exception as exc: + logger.debug(f"VisualGen telemetry setup failed: {exc}") - self.executor = DiffusionRemoteClient( - args=self.args, - ) self._req_counter = itertools.count() atexit.register(VisualGen._atexit_shutdown, weakref.ref(self)) @@ -484,6 +535,17 @@ def shutdown(self): logger.info("VisualGen: Shutting down") self.executor.shutdown() self.executor = None + usage_lifecycle_lock = getattr(self, "_usage_lifecycle_lock", None) + if usage_lifecycle_lock is not None: + with usage_lifecycle_lock: + if getattr(self, "_usage_lifecycle_active", False): + self._usage_lifecycle_active = False + try: + import tensorrt_llm.usage as _usage + + _usage.record_visual_gen_shutdown() + except Exception as exc: + logger.debug(f"VisualGen telemetry shutdown tracking failed: {exc}") @set_api_status("prototype") async def get_stats_async(self, timeout: Optional[float] = None) -> AsyncIterator[Dict]: diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py index 2de9a9ba02d0..710a6193e0af 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py @@ -216,6 +216,18 @@ def test_supported_models_no_duplicates(self): ids = VisualGen.supported_models() assert len(ids) == len(set(ids)), f"VisualGen.supported_models() returned duplicates: {ids}" + def test_registered_pipelines_declare_safe_telemetry_metadata(self): + """Built-in pipeline metadata is explicitly allowlisted and bounded.""" + from tensorrt_llm._torch.visual_gen.pipeline_registry import PIPELINE_REGISTRY + + allowed = {"image", "video", "video_audio", "layered_image", "mixed"} + missing = { + name: entry.modality + for name, entry in PIPELINE_REGISTRY.items() + if entry.modality not in allowed or not entry.telemetry_safe + } + assert not missing, f"Pipeline telemetry metadata is not allowlisted: {missing}" + class TestVisualGenArgsCacheBackend: def test_cache_dit_nested_config(self): diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py index c48d0fa963d1..ccefdf888631 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -723,6 +723,7 @@ def _build_ready_response(self, pipeline_cls, mock_self=None): "status": "READY", "default_generation_params": pipeline_cls.default_generation_params.fget(mock_self), "extra_param_specs": pipeline_cls.extra_param_specs.fget(mock_self), + "telemetry_metadata": {"modality": "video"}, }, ) @@ -732,6 +733,58 @@ def _roundtrip(self, response): return pickle.loads(pickle.dumps(response)) + @pytest.mark.parametrize( + ("external_launch", "environment", "expected"), + [ + (None, {}, ("local_spawn", 1)), + ((0, 0, 8, "localhost", 1234), {"GROUP_WORLD_SIZE": "2"}, ("torchrun", 2)), + ( + (0, 0, 8, "localhost", 1234), + {"SLURM_PROCID": "0", "SLURM_NNODES": "4"}, + ("slurm", 4), + ), + ], + ) + def test_launch_metadata_is_bounded(self, monkeypatch, external_launch, environment, expected): + from tensorrt_llm._torch.visual_gen.executor import _visual_gen_launch_metadata + + for name in ("SLURM_PROCID", "SLURM_NNODES", "GROUP_WORLD_SIZE", "LOCAL_WORLD_SIZE"): + monkeypatch.delenv(name, raising=False) + for name, value in environment.items(): + monkeypatch.setenv(name, value) + + assert _visual_gen_launch_metadata(external_launch) == expected + + def test_pipeline_telemetry_metadata_is_fail_silent(self, monkeypatch): + from tensorrt_llm._torch.visual_gen.pipeline_registry import ( + PIPELINE_REGISTRY, + AutoPipeline, + _PipelineEntry, + ) + + class _SlottedPipeline: + __slots__ = () + + @classmethod + def resolve_variant(cls, _config): + return cls + + def __init__(self, _config): + pass + + monkeypatch.setitem( + PIPELINE_REGISTRY, + "SlottedPipeline", + _PipelineEntry(pipeline_cls=_SlottedPipeline, telemetry_safe=True), + ) + monkeypatch.setattr( + AutoPipeline, + "_detect_from_checkpoint", + lambda _checkpoint_dir: "SlottedPipeline", + ) + + assert isinstance(AutoPipeline.from_config(object(), "/model"), _SlottedPipeline) + def test_ready_payload_pickle_roundtrip(self): """The READY dict survives pickle (the ZMQ transport layer).""" from tensorrt_llm._torch.visual_gen.executor import DiffusionResponse @@ -805,7 +858,10 @@ def test_flux_empty_specs_roundtrip(self): assert restored.output["default_generation_params"]["height"] == 1024 def test_client_extracts_metadata_from_ready(self): - """DiffusionRemoteClient stores metadata when processing a READY response.""" + """The real READY wait path stores pipeline and telemetry metadata.""" + import asyncio + from types import SimpleNamespace + from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2 import LTX2Pipeline from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema from tensorrt_llm.visual_gen.visual_gen import DiffusionRemoteClient @@ -816,20 +872,94 @@ def test_client_extracts_metadata_from_ready(self): ltx2_defaults = LTX2Pipeline.default_generation_params.fget(None) ltx2_specs = LTX2Pipeline.extra_param_specs.fget(None) - # Simulate what _wait_ready_async does: extract from the response payload - client = MagicMock(spec=DiffusionRemoteClient) - client.default_generation_params = {} - client.extra_param_specs = {} + async def wait_ready(): + client = SimpleNamespace( + lock=asyncio.Lock(), + completed_responses={-1: restored}, + default_generation_params={}, + extra_param_specs={}, + supports_image_edit=False, + ref_slot_specs={}, + telemetry_metadata={}, + worker_processes=[], + _ext_worker_thread=None, + response_event=asyncio.Event(), + ) + await DiffusionRemoteClient._wait_ready_async(client) + return client - payload = restored.output - if isinstance(payload, dict): - client.default_generation_params = payload.get("default_generation_params", {}) - client.extra_param_specs = payload.get("extra_param_specs", {}) + client = asyncio.run(wait_ready()) assert client.default_generation_params == ltx2_defaults assert set(client.extra_param_specs.keys()) == set(ltx2_specs.keys()) for spec in client.extra_param_specs.values(): assert isinstance(spec, ExtraParamSchema) + assert client.telemetry_metadata == {"modality": "video"} + + def test_worker_telemetry_metadata_allowlists_identity_and_components(self, monkeypatch): + """READY metadata never forwards an unknown model or component name.""" + from types import SimpleNamespace + + from tensorrt_llm._torch.visual_gen import executor + + class _Pipeline: + transformer_components = ("transformer", "private_component") + + def __init__(self): + self._telemetry_pipeline_class_name = "PublicPipeline" + self.pipeline_config = SimpleNamespace( + quant_config=SimpleNamespace(quant_algo=SimpleNamespace(name="NVFP4")), + dynamic_weight_quant=True, + get_quant_config=lambda component: SimpleNamespace( + quant_algo=( + SimpleNamespace(name="NVFP4") if component == "transformer" else None + ) + ), + ) + + registry = { + "PublicPipeline": SimpleNamespace( + pipeline_cls=_Pipeline, + hf_ids=["nvidia/public-model"], + modality="image", + telemetry_safe=True, + ) + } + monkeypatch.setattr(executor, "PIPELINE_REGISTRY", registry) + + public = executor._visual_gen_telemetry_metadata( + _Pipeline(), SimpleNamespace(model="nvidia/public-model") + ) + private = executor._visual_gen_telemetry_metadata( + _Pipeline(), SimpleNamespace(model="/private/customer/model") + ) + + assert public["model_id"] == "nvidia/public-model" + assert public["pipeline_class_name"] == "PublicPipeline" + assert public["resolved_pipeline_class"] == "_Pipeline" + assert public["modality"] == "image" + assert public["quantization_algo"] == "NVFP4" + assert public["quantized_components"] == ["transformer"] + assert private["model_id"] == "other" + + class _PrivatePipeline(_Pipeline): + def __init__(self): + super().__init__() + self._telemetry_pipeline_class_name = "CustomerPipeline" + + registry["CustomerPipeline"] = SimpleNamespace( + pipeline_cls=_PrivatePipeline, + hf_ids=["customer/private-model"], + modality="image", + telemetry_safe=False, + ) + private_pipeline = executor._visual_gen_telemetry_metadata( + _PrivatePipeline(), SimpleNamespace(model="customer/private-model") + ) + assert private_pipeline["model_id"] == "other" + assert private_pipeline["pipeline_class_name"] == "unknown" + assert private_pipeline["resolved_pipeline_class"] == "unknown" + assert private_pipeline["modality"] == "unknown" # ============================================================================= diff --git a/tests/unittest/usage/test_cli_telemetry.py b/tests/unittest/usage/test_cli_telemetry.py index 9ac1774ec1e2..1bd13d84dfe0 100644 --- a/tests/unittest/usage/test_cli_telemetry.py +++ b/tests/unittest/usage/test_cli_telemetry.py @@ -109,7 +109,7 @@ def _captured_terminal_parameters(payloads): """Return the parameters from one fully serialized terminal event.""" assert len(payloads) == 1 payload = payloads[0] - assert payload["eventSchemaVer"] == "0.7" + assert payload["eventSchemaVer"] == "0.8" assert len(payload["events"]) == 1 event = payload["events"][0] assert event["name"] == "trtllm_exit_report" @@ -171,6 +171,59 @@ def test_explicit_cli_telemetry_overrides_yaml_opt_out(self): apply_config.assert_not_called() + def test_visual_gen_bench_yaml_opt_out_precedes_prompt_validation( + self, + tmp_path, + captured_exit_payloads, + ): + """A valid YAML opt-out is honored before benchmark argument errors.""" + from click.testing import CliRunner + + from tensorrt_llm.commands.bench import main as bench_main + + config_path = tmp_path / "visual_gen.yaml" + config_path.write_text( + "telemetry_config:\n disabled: true\n", + encoding="utf-8", + ) + + result = CliRunner().invoke( + bench_main, + [ + "--model", + "test-model", + "visual-gen", + "--visual_gen_args", + str(config_path), + ], + ) + + assert result.exit_code == 2 + assert "Either --prompt or --prompt_file must be specified." in result.output + assert captured_exit_payloads == [] + + def test_visual_gen_bench_attributes_invalid_yaml_to_visual_gen( + self, + tmp_path, + captured_exit_payloads, + ): + from click.testing import CliRunner + + from tensorrt_llm.commands.bench import main as bench_main + + config_path = tmp_path / "visual_gen.yaml" + config_path.write_text("invalid: [", encoding="utf-8") + + result = CliRunner().invoke( + bench_main, + ["--model", "test-model", "visual-gen", "--visual_gen_args", str(config_path)], + ) + + assert result.exit_code != 0 + parameters = _captured_terminal_parameters(captured_exit_payloads) + assert parameters["component"] == "visual_gen" + assert parameters["ingressPoint"] == "cli_bench" + @pytest.mark.parametrize( ("yaml_config", "telemetry", "expected_disabled"), [ diff --git a/tests/unittest/usage/test_e2e_capture.py b/tests/unittest/usage/test_e2e_capture.py index 6a2bb928278d..365fb33fe3ab 100644 --- a/tests/unittest/usage/test_e2e_capture.py +++ b/tests/unittest/usage/test_e2e_capture.py @@ -330,7 +330,7 @@ def test_initial_report_captured(self, capture_server, monkeypatch): assert set(features.keys()) == expected_keys # Schema version - assert payload["eventSchemaVer"] == "0.7" + assert payload["eventSchemaVer"] == "0.8" # Disagg fields present (may be empty strings) assert "disaggRole" in params @@ -343,7 +343,7 @@ def test_initial_report_captured(self, capture_server, monkeypatch): # heartbeat must contain the current counter snapshot for the same # process-scoped telemetry session. assert heartbeat_payload["sessionId"] == payload["sessionId"] - assert heartbeat_payload["eventSchemaVer"] == "0.7" + assert heartbeat_payload["eventSchemaVer"] == "0.8" heartbeat = heartbeat_payload["events"][0] assert heartbeat["name"] == "trtllm_heartbeat" _assert_lifecycle_snapshot( @@ -382,7 +382,7 @@ def test_cli_serve_context_e2e(self, capture_server, monkeypatch): payload = CaptureHandler.captured_payloads[0] event = payload["events"][0] params = event["parameters"] - assert payload["eventSchemaVer"] == "0.7" + assert payload["eventSchemaVer"] == "0.8" _assert_lifecycle_snapshot( params, active_instances=1, @@ -442,7 +442,7 @@ def test_direct_llm_process_exit_e2e(self, capture_server, monkeypatch): initial_payload = _wait_for_event("trtllm_initial_report", timeout=5) exit_payload = _wait_for_event("trtllm_exit_report", timeout=5) assert exit_payload["sessionId"] == initial_payload["sessionId"] - assert exit_payload["eventSchemaVer"] == "0.7" + assert exit_payload["eventSchemaVer"] == "0.8" terminal_events = [ payload diff --git a/tests/unittest/usage/test_llmapi_config_capture.py b/tests/unittest/usage/test_llmapi_config_capture.py index 6b5a29fb7bec..550b02cb2a48 100644 --- a/tests/unittest/usage/test_llmapi_config_capture.py +++ b/tests/unittest/usage/test_llmapi_config_capture.py @@ -29,7 +29,10 @@ ) from tensorrt_llm.llmapi.utils import StrictBaseModel from tensorrt_llm.usage import usage_lib -from tensorrt_llm.usage.llmapi_config import collect_llm_api_config_payloads +from tensorrt_llm.usage.llmapi_config import ( + collect_llm_api_config_payloads, + collect_visual_gen_config_payloads, +) pytestmark = pytest.mark.cpu_only @@ -886,3 +889,41 @@ class _BigConfig(StrictBaseModel): config, meta = _loads_payloads(_BigConfig()) assert meta["payload_truncated"] is True assert len(rc._canonical_json(config).encode("utf-8")) <= 20 + + +def test_collect_visual_gen_config_uses_separate_source_and_excludes_model(): + """VisualGen capture reuses the sanitizer without collecting model paths.""" + from tensorrt_llm.visual_gen.args import VisualGenArgs + + config_json, meta_json = collect_visual_gen_config_payloads( + VisualGenArgs(model="/private/customer/model") + ) + config = json.loads(config_json) + meta = json.loads(meta_json) + + assert "model" not in config + assert "telemetry_config.disabled" not in config + assert config["parallel_config.cfg_size"] == 1 + assert config["attention_config.backend"] == "VANILLA" + assert meta["source"] == "effective_validated_visual_gen_args" + assert meta["args_class"] == "VisualGenArgs" + + +def test_collect_visual_gen_config_uses_canonical_manifest_for_subclasses(): + """Custom subclasses cannot add fields or expose their class name.""" + from tensorrt_llm.visual_gen.args import VisualGenArgs + + class _PrivateVisualGenArgs(VisualGenArgs): + customer_secret_enabled: bool = True + customer_private_limit: int = 7 + + config_json, meta_json = collect_visual_gen_config_payloads( + _PrivateVisualGenArgs(model="/private/customer/model") + ) + config = json.loads(config_json) + meta = json.loads(meta_json) + + assert "customer_secret_enabled" not in config + assert "customer_private_limit" not in config + assert config["parallel_config.cfg_size"] == 1 + assert meta["args_class"] == "VisualGenArgs" diff --git a/tests/unittest/usage/test_llmapi_config_telemetry_docs.py b/tests/unittest/usage/test_llmapi_config_telemetry_docs.py index e6e8f68b7c83..88cd7d185230 100644 --- a/tests/unittest/usage/test_llmapi_config_telemetry_docs.py +++ b/tests/unittest/usage/test_llmapi_config_telemetry_docs.py @@ -473,7 +473,8 @@ def test_renderer_emits_table_from_committed_golden(tmp_path): out = tmp_path / "telemetry.md" generator.generate_telemetry_reference(_repo_root(), out) text = out.read_text() - assert "## LLM API Configuration Fields" in text + assert "## Runtime Configuration Fields" in text + assert "### `VisualGenArgs`" in text assert "python3 scripts/generate_llm_args_golden_manifest.py" in text assert "explicitly marked" not in text # opt-in prose must be gone assert "`backend`" in text # a known captured key renders diff --git a/tests/unittest/usage/test_reporter.py b/tests/unittest/usage/test_reporter.py index d67135dd1af3..1e8f3e85d302 100644 --- a/tests/unittest/usage/test_reporter.py +++ b/tests/unittest/usage/test_reporter.py @@ -59,6 +59,32 @@ def reporter_session(enable_telemetry): assert usage_lib.apply_usage_session_config() +def _visual_gen_args(): + """Return the small validated-config surface consumed by the reporter.""" + parallel = SimpleNamespace( + cfg_size=1, + ulysses_size=1, + async_ulysses=False, + ring_size=1, + attn2d_size=(1, 1), + tp_size=1, + parallel_vae_size=1, + parallel_vae_split_dim="width", + ) + attention = SimpleNamespace( + backend="VANILLA", + sparse_attention_config=None, + quant_attention_config=None, + ) + return SimpleNamespace( + parallel_config=parallel, + attention_config=attention, + cache_config=None, + cuda_graph_config=SimpleNamespace(enable=False), + torch_compile_config=SimpleNamespace(enable=False), + ) + + # --------------------------------------------------------------------------- # Console notification tests # --------------------------------------------------------------------------- @@ -162,6 +188,101 @@ def test_get_trtllm_version_returns_string(self): result = usage_lib._get_trtllm_version() assert isinstance(result, str) + def test_report_visual_gen_usage_spawns_visual_gen_reporter( + self, monkeypatch, enable_telemetry + ): + """VisualGen starts its dedicated initial/heartbeat reporter.""" + usage_lib._NOTIFICATION_SHOWN.set() + mock_thread = MagicMock() + + with patch.object(usage_lib.threading, "Thread", return_value=mock_thread) as thread_cls: + usage_lib.report_visual_gen_usage(_visual_gen_args()) + + assert thread_cls.call_args.kwargs["target"] is usage_lib._visual_gen_background_reporter + assert thread_cls.call_args.kwargs["name"] == "trtllm-visual-gen-usage-stats" + assert thread_cls.call_args.kwargs["daemon"] is True + mock_thread.start.assert_called_once() + + def test_visual_gen_reporter_sends_initial_and_heartbeat(self, monkeypatch): + """VisualGen reports bounded runtime metadata and current lifecycle counters.""" + + class _OneHeartbeat: + def __init__(self): + self.wait_count = 0 + + def wait(self, timeout): + del timeout + self.wait_count += 1 + return self.wait_count > 1 + + monkeypatch.setenv("TRTLLM_USAGE_FORCE_ENABLED", "1") + assert usage_lib.record_visual_gen_initialization_attempt() + assert usage_lib.record_visual_gen_initialized() + sent = [] + visual_gen_args = _visual_gen_args() + visual_gen_args.parallel_config.cfg_size = 2 + visual_gen_args.parallel_config.tp_size = 2 + visual_gen_args.cache_config = SimpleNamespace(cache_backend="teacache") + metadata = { + "model_id": "nvidia/test-model", + "pipeline_class_name": "TestPipeline", + "resolved_pipeline_class": "ResolvedPipeline", + "modality": "image", + "launch_mode": "local_spawn", + "node_count": 2, + "n_workers": 3, + "quantization_algo": "NVFP4", + "dynamic_weight_quant": True, + "quantized_components": ["transformer"], + } + + with ( + patch.object(usage_lib, "_collect_system_info", return_value={}), + patch.object(usage_lib, "_collect_gpu_info", return_value={"gpu_count": 4}), + patch.object( + usage_lib, + "_collect_visual_gen_config_payloads", + return_value=("{}", "{}"), + ), + patch.object(usage_lib, "_send_to_gxt", side_effect=sent.append), + patch.object(usage_lib, "_REPORTER_STOP", _OneHeartbeat()), + ): + usage_lib._visual_gen_background_reporter(visual_gen_args, metadata, "visual_gen_class") + + assert [payload["events"][0]["name"] for payload in sent] == [ + "trtllm_visual_gen_initial_report", + "trtllm_visual_gen_heartbeat", + ] + initial = sent[0]["events"][0]["parameters"] + expected_initial = { + "modelId": "nvidia/test-model", + "pipelineClassName": "TestPipeline", + "resolvedPipelineClass": "ResolvedPipeline", + "modality": "image", + "launchMode": "local_spawn", + "nodeCount": 2, + "nWorkers": 3, + "quantizationAlgo": "NVFP4", + "dynamicWeightQuant": True, + "quantizedComponentsJson": '["transformer"]', + "cfgSize": 2, + "tensorParallelSize": 2, + "stepCaching": True, + "cacheBackend": "teacache", + } + assert {key: initial[key] for key in expected_initial} == expected_initial + + heartbeat = sent[1]["events"][0]["parameters"] + assert heartbeat["seq"] == 0 + assert heartbeat["runtimeKind"] == "visual_gen" + assert heartbeat["ingressPoint"] == "visual_gen_class" + assert heartbeat["nWorkers"] == 3 + assert heartbeat["gpuCount"] == 4 + assert heartbeat["visualGenInitializationAttempts"] == 1 + assert heartbeat["visualGenInstancesCreated"] == 1 + assert heartbeat["activeVisualGenInstances"] == 1 + assert heartbeat["visualGenInitializationFailures"] == 0 + # --------------------------------------------------------------------------- # Duplicate reporter guard tests @@ -877,6 +998,44 @@ def test_initialization_failure_updates_only_failure_counter(self, enable_teleme assert snapshot["llmInstancesCreated"] == 0 assert snapshot["activeLlmInstances"] == 0 + def test_visual_gen_lifecycle_updates_separate_counters(self, enable_telemetry): + """VisualGen lifecycle state is independent from LLM lifecycle state.""" + assert usage_lib.record_visual_gen_initialization_attempt() + assert usage_lib.record_visual_gen_initialized() + + snapshot = usage_lib._SESSION.snapshot() + assert snapshot["runtimeKind"] == "visual_gen" + assert snapshot["visualGenInitializationAttempts"] == 1 + assert snapshot["visualGenInstancesCreated"] == 1 + assert snapshot["activeVisualGenInstances"] == 1 + assert snapshot["llmInitializationAttempts"] == 0 + + usage_lib.record_visual_gen_shutdown() + assert usage_lib._SESSION.snapshot()["activeVisualGenInstances"] == 0 + + def test_mixed_session_exit_contains_both_runtime_counters(self, enable_telemetry): + """A shared process exit snapshot identifies mixed LLM/VisualGen use.""" + assert usage_lib.record_llm_initialization_attempt() + assert usage_lib.record_llm_initialized() + assert usage_lib.record_visual_gen_initialization_attempt() + assert usage_lib.record_visual_gen_initialized() + sent = [] + + with patch.object(usage_lib, "_send_to_gxt", side_effect=sent.append): + assert usage_lib.report_exit( + usage_lib.TerminalOutcome( + termination_kind="clean", + component="server", + exit_code_known=True, + exit_code=0, + ) + ) + + parameters = sent[0]["events"][0]["parameters"] + assert parameters["runtimeKind"] == "mixed" + assert parameters["llmInstancesCreated"] == 1 + assert parameters["visualGenInstancesCreated"] == 1 + def test_monotonic_counters_saturate_at_uint32(self, enable_telemetry): """Cumulative counters never exceed the SMS PositiveInt bound.""" assert usage_lib.apply_usage_session_config() diff --git a/tests/unittest/usage/test_schema.py b/tests/unittest/usage/test_schema.py index c911ad02b75c..b571775a89c2 100644 --- a/tests/unittest/usage/test_schema.py +++ b/tests/unittest/usage/test_schema.py @@ -251,7 +251,7 @@ class TestSchemaConstants: [ ("CLIENT_ID", "616561816355034"), ("EVENT_PROTOCOL", "1.6"), - ("EVENT_SCHEMA_VER", "0.7"), + ("EVENT_SCHEMA_VER", "0.8"), ("EVENT_SYS_VER", "trtllm-telemetry/1.0"), ], ) @@ -364,6 +364,7 @@ def test_expected_contexts_exist(self): "cli_bench", "cli_eval", "disaggregated", + "visual_gen_class", } actual = {ctx.value for ctx in _llm_args_mod.UsageContext} assert actual == expected @@ -422,13 +423,15 @@ def test_schema_has_draft07(self): sms_schema = self._load_sms_schema() assert sms_schema["$schema"] == "http://json-schema.org/draft-07/schema#" - def test_schema_has_three_events(self): + def test_schema_has_five_events(self): """SMS schema defines exactly the supported telemetry events.""" sms_schema = self._load_sms_schema() events = sms_schema["definitions"]["events"] assert set(events.keys()) == { "trtllm_initial_report", "trtllm_heartbeat", + "trtllm_visual_gen_initial_report", + "trtllm_visual_gen_heartbeat", "trtllm_exit_report", } @@ -437,6 +440,8 @@ def test_schema_has_three_events(self): [ ("trtllm_initial_report", schema.TrtllmInitialReport), ("trtllm_heartbeat", schema.TrtllmHeartbeat), + ("trtllm_visual_gen_initial_report", schema.TrtllmVisualGenInitialReport), + ("trtllm_visual_gen_heartbeat", schema.TrtllmVisualGenHeartbeat), ("trtllm_exit_report", schema.TrtllmExitReport), ], ) @@ -497,7 +502,13 @@ def test_initial_report_has_all_expected_fields(self): @pytest.mark.parametrize( "event_name", - ["trtllm_initial_report", "trtllm_heartbeat", "trtllm_exit_report"], + [ + "trtllm_initial_report", + "trtllm_heartbeat", + "trtllm_visual_gen_initial_report", + "trtllm_visual_gen_heartbeat", + "trtllm_exit_report", + ], ) def test_all_event_fields_are_required_and_closed(self, event_name): """Every SMS event requires its declared fields and rejects extras.""" @@ -551,12 +562,7 @@ def test_no_pii_fields_in_any_event(self): def test_events_have_gdpr_metadata(self): """Every event has GDPR metadata in eventMeta.""" sms_schema = self._load_sms_schema() - for event_name in ( - "trtllm_initial_report", - "trtllm_heartbeat", - "trtllm_exit_report", - ): - event = sms_schema["definitions"]["events"][event_name] + for event_name, event in sms_schema["definitions"]["events"].items(): assert "eventMeta" in event, f"{event_name} missing eventMeta" assert "gdpr" in event["eventMeta"], f"{event_name} missing gdpr in eventMeta" gdpr = event["eventMeta"]["gdpr"] @@ -623,73 +629,71 @@ def test_envelope_contains_all_gxt_v16_keys(self): # --- JSON schema validation (from collapsed TestSchemaDriftDetection) --- - def test_initial_report_validates_against_json_schema(self): - """A fully-populated TrtllmInitialReport must validate against the JSON schema.""" - import jsonschema - - sms_schema = json.loads(schemas.SMS_SCHEMA_PATH.read_text()) - report = schema.TrtllmInitialReport( - trtllmVersion="1.0", - platform="Linux", - pythonVersion="3.10", - cpuArchitecture="x86_64", - cpuCount=8, - gpuCount=1, - gpuName="H100", - gpuMemoryMB=81920, - cudaVersion="12.0", - architectureClassName="LlamaForCausalLM", - backend="pytorch", - tensorParallelSize=1, - pipelineParallelSize=1, - contextParallelSize=1, - moeExpertParallelSize=0, - moeTensorParallelSize=0, - dtype="float16", - quantizationAlgo="", - kvCacheDtype="", - ingressPoint="llm_class", - featuresJson='{"lora":false}', - llmApiConfigJson='{"tensor_parallel_size":1}', - llmApiConfigMetaJson='{"capture_succeeded":true}', - disaggRole="", - deploymentId="", - ) - payload = report.model_dump(by_alias=True) - initial_schema = sms_schema["definitions"]["events"]["trtllm_initial_report"].copy() - initial_schema["definitions"] = sms_schema["definitions"] - jsonschema.validate(instance=payload, schema=initial_schema) - - def test_heartbeat_validates_against_json_schema(self): - """A TrtllmHeartbeat must validate against the JSON schema.""" - import jsonschema - - sms_schema = json.loads(schemas.SMS_SCHEMA_PATH.read_text()) - heartbeat = schema.TrtllmHeartbeat(seq=0) - payload = heartbeat.model_dump(by_alias=True) - hb_schema = sms_schema["definitions"]["events"]["trtllm_heartbeat"].copy() - hb_schema["definitions"] = sms_schema["definitions"] - jsonschema.validate(instance=payload, schema=hb_schema) - - def test_exit_report_validates_against_json_schema(self): - """A TrtllmExitReport must validate against the JSON schema.""" + @pytest.mark.parametrize( + ("event_name", "event"), + [ + ( + "trtllm_initial_report", + schema.TrtllmInitialReport( + trtllmVersion="1.0", + platform="Linux", + pythonVersion="3.10", + cpuArchitecture="x86_64", + cpuCount=8, + gpuCount=1, + gpuName="H100", + gpuMemoryMB=81920, + cudaVersion="12.0", + architectureClassName="LlamaForCausalLM", + backend="pytorch", + tensorParallelSize=1, + pipelineParallelSize=1, + contextParallelSize=1, + moeExpertParallelSize=0, + moeTensorParallelSize=0, + dtype="float16", + quantizationAlgo="", + kvCacheDtype="", + ingressPoint="llm_class", + featuresJson='{"lora":false}', + llmApiConfigJson='{"tensor_parallel_size":1}', + llmApiConfigMetaJson='{"capture_succeeded":true}', + disaggRole="", + deploymentId="", + ), + ), + ("trtllm_heartbeat", schema.TrtllmHeartbeat(seq=0)), + ( + "trtllm_visual_gen_initial_report", + schema.TrtllmVisualGenInitialReport(), + ), + ( + "trtllm_visual_gen_heartbeat", + schema.TrtllmVisualGenHeartbeat(seq=0), + ), + ( + "trtllm_exit_report", + schema.TrtllmExitReport( + exitCodeKnown=True, + exitCode=130, + signalNumber=2, + terminationKind="signal", + lifecyclePhase="serving", + component="server", + reportingSource="self", + ingressPoint="cli_serve", + ), + ), + ], + ) + def test_events_validate_against_json_schema(self, event_name, event): + """Pydantic event payloads satisfy their SMS event definitions.""" import jsonschema sms_schema = json.loads(schemas.SMS_SCHEMA_PATH.read_text()) - report = schema.TrtllmExitReport( - exitCodeKnown=True, - exitCode=130, - signalNumber=2, - terminationKind="signal", - lifecyclePhase="serving", - component="server", - reportingSource="self", - ingressPoint="cli_serve", - ) - payload = report.model_dump(by_alias=True) - exit_schema = sms_schema["definitions"]["events"]["trtllm_exit_report"].copy() - exit_schema["definitions"] = sms_schema["definitions"] - jsonschema.validate(instance=payload, schema=exit_schema) + event_schema = sms_schema["definitions"]["events"][event_name].copy() + event_schema["definitions"] = sms_schema["definitions"] + jsonschema.validate(instance=event.model_dump(by_alias=True), schema=event_schema) # --------------------------------------------------------------------------- diff --git a/tests/unittest/visual_gen/test_telemetry.py b/tests/unittest/visual_gen/test_telemetry.py new file mode 100644 index 000000000000..e85ef53fd479 --- /dev/null +++ b/tests/unittest/visual_gen/test_telemetry.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import threading +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +import tensorrt_llm.usage as usage +from tensorrt_llm.usage import usage_lib +from tensorrt_llm.visual_gen.args import VisualGenArgs +from tensorrt_llm.visual_gen.visual_gen import VisualGen + +pytestmark = pytest.mark.cpu_only + + +def _executor(): + return SimpleNamespace( + telemetry_metadata={"model_id": "other", "modality": "image"}, + launch_mode="local_spawn", + node_count=1, + n_workers=1, + shutdown=MagicMock(), + ) + + +def test_visual_gen_reports_initialized_runtime_and_shutdown_once(): + executor = _executor() + args = VisualGenArgs(model="/private/model") + + with ( + patch( + "tensorrt_llm.visual_gen.visual_gen.DiffusionRemoteClient", + return_value=executor, + ), + patch.object(usage, "record_visual_gen_initialization_attempt", return_value=True), + patch.object(usage, "record_visual_gen_initialized", return_value=True), + patch.object(usage, "record_visual_gen_shutdown") as record_shutdown, + patch.object(usage_lib, "report_visual_gen_usage") as report_usage, + patch("tensorrt_llm.visual_gen.visual_gen.atexit.register"), + ): + visual_gen = VisualGen(model=args.model, args=args) + visual_gen.shutdown() + visual_gen.shutdown() + + metadata = report_usage.call_args.args[1] + assert metadata["launch_mode"] == "local_spawn" + assert metadata["node_count"] == 1 + assert metadata["n_workers"] == 1 + assert report_usage.call_args.args[2].usage_context is usage.UsageContext.VISUAL_GEN_CLASS + executor.shutdown.assert_called_once() + record_shutdown.assert_called_once() + + +def test_visual_gen_records_initialization_failure(): + args = VisualGenArgs(model="/private/model") + + with ( + patch( + "tensorrt_llm.visual_gen.visual_gen.DiffusionRemoteClient", + side_effect=RuntimeError("worker failed"), + ), + patch.object(usage, "record_visual_gen_initialization_attempt", return_value=True), + patch.object(usage, "record_visual_gen_initialization_failure") as record_failure, + ): + with pytest.raises(RuntimeError, match="worker failed"): + VisualGen(model=args.model, args=args) + + record_failure.assert_called_once() + + +def test_visual_gen_records_external_world_size_failure(): + args = VisualGenArgs(model="/private/model") + + with ( + patch( + "tensorrt_llm.visual_gen.visual_gen._detect_external_launch", + return_value=(0, 0, 2, "localhost", 1234), + ), + patch.object(usage, "record_visual_gen_initialization_attempt", return_value=True), + patch.object(usage, "record_visual_gen_initialization_failure") as record_failure, + patch("tensorrt_llm.visual_gen.visual_gen.DiffusionRemoteClient") as executor, + ): + with pytest.raises(ValueError, match=r"world_size \(2\) does not match n_workers \(1\)"): + VisualGen(model=args.model, args=args) + + record_failure.assert_called_once() + executor.assert_not_called() + + +def test_visual_gen_shutdown_failure_can_be_retried(): + visual_gen = VisualGen.__new__(VisualGen) + visual_gen.executor = _executor() + visual_gen.executor.shutdown.side_effect = [RuntimeError("shutdown failed"), None] + visual_gen._usage_lifecycle_active = True + visual_gen._usage_lifecycle_lock = threading.Lock() + + with ( + patch.object(usage, "record_visual_gen_shutdown") as record_shutdown, + pytest.raises(RuntimeError, match="shutdown failed"), + ): + visual_gen.shutdown() + + assert visual_gen.executor is not None + assert visual_gen._usage_lifecycle_active is True + record_shutdown.assert_not_called() + + with patch.object(usage, "record_visual_gen_shutdown") as record_shutdown: + visual_gen.shutdown() + + assert visual_gen.executor is None + record_shutdown.assert_called_once()