Triton PBText

Triton PBText is a VS Code extension for Triton Inference Server config.pbtxt files. It focuses on ensemble pipeline authoring while still helping with ordinary model configs.

Features

• Syntax highlighting for protobuf text format used by Triton .pbtxt files.
• Snippets for model configs, ensemble configs, ensemble steps, tensors, dynamic batching, instance groups, and Python backend models.
• Context-aware completions for top-level fields, tensor fields, ensemble scheduling, map entries, enums, platforms, and backends.
• Hover documentation for common Triton config fields.
• Lightweight diagnostics for bracket balance, invalid enum values, malformed dims, missing ensemble_scheduling, incomplete ensemble steps, and incomplete input_map / output_map blocks.
• Model repository scanning for model_name completions in ensemble steps.

Supported Scope

This project currently implements the planned v0.1, v0.2, and v0.3 scope:

• v0.1: language registration, TextMate highlighting, snippets, common completions, and basic diagnostics.
• v0.2: Triton-aware field and enum tables, hover documentation, improved diagnostics, and local model repository scanning for ensemble model_name completions.
• v0.3: model repository metadata indexing, pipeline-aware input_map / output_map completions, and ensemble data-flow diagnostics across referenced model configs.

It does not yet implement a full protobuf parser or formatter. The current parser is intentionally lightweight and focused on Triton model, tensor, and ensemble pipeline structure.

Pipeline-Aware Behavior

When a file is inside a Triton model repository, the extension reads sibling config.pbtxt files and uses their tensor declarations:

• model_name: completes model directory names.
• input_map.key completes input tensor names from the referenced model.
• output_map.key completes output tensor names from the referenced model.
• input_map.value completes tensors available before the current step, including ensemble inputs and outputs from earlier steps.
• output_map.value completes ensemble output names and tensors already known in the pipeline.

Diagnostics also flag common ensemble issues, including unknown referenced models, map keys that do not exist on the referenced model, values consumed before they are produced, missing required input maps, and declared ensemble outputs that no step produces.

Usage

Open a config.pbtxt or *.pbtxt file. VS Code should select the Triton PBText language automatically.

Useful snippet prefixes include:

• triton-model
• triton-ensemble
• triton-step
• triton-input
• triton-output
• triton-dynamic
• triton-instance
• triton-python

Settings

{
  "tritonPbtxt.diagnostics.enabled": true,
  "tritonPbtxt.completions.scanModelRepository": true,
  "tritonPbtxt.completions.maxScannedModels": 200
}

Development

Run the static checks from this directory:

npm run check

To try the extension locally, open this folder in VS Code and launch the extension host with F5.

To package a .vsix, install the VS Code Extension Manager and run:

npm install -g @vscode/vsce
vsce package

Design Notes

The extension intentionally starts with lightweight text analysis instead of a full language server. Triton config.pbtxt files are small, and the most common authoring mistakes are local enough to catch with fast scanning. The code separates pure analysis modules from VS Code provider registration so the behavior can be verified without launching VS Code.