MCP Server Triton Inference Server — tensor inputs, config.pbtxt, HTTP vs gRPC, backends

Triton HTTP REST infer endpoint and tensor format

Triton's HTTP REST API follows the KServe v2 inference protocol. Every request to POST /v2/models/{model_name}/infer wraps inputs in a list of tensor objects. The tensor format is verbose compared to OpenAI-style APIs — it exists to support arbitrary model I/O shapes across PyTorch, TensorFlow, ONNX, and Python backends. Fetch the model's config once at startup to learn the exact tensor names and datatypes, then cache them — Triton configs do not change without a model reload.

import { z } from 'zod';

const TRITON_HTTP = process.env.TRITON_HTTP_URL ?? 'http://localhost:8000';

// Tensor type mapping from Triton config.pbtxt datatype to JS
type TritonDatatype = 'BYTES' | 'FP32' | 'FP64' | 'INT32' | 'INT64' | 'BOOL';

interface TritonInput {
  name:     string;
  shape:    number[];
  datatype: TritonDatatype;
  data:     unknown[];
}

interface TritonOutput {
  name: string;
}

interface TritonInferRequest {
  inputs:  TritonInput[];
  outputs: TritonOutput[];
}

// Fetch model config to discover tensor names — cache the result
const modelConfigCache = new Map<string, { inputs: Array<{name:string;data_type:string}>; outputs: Array<{name:string;data_type:string}> }>();

async function getModelConfig(modelName: string) {
  if (modelConfigCache.has(modelName)) return modelConfigCache.get(modelName)!;

  const res = await fetch(`${TRITON_HTTP}/v2/models/${modelName}/config`, {
    signal: AbortSignal.timeout(10_000),
  });
  if (!res.ok) throw new Error(`Triton config fetch failed for model '${modelName}': HTTP ${res.status}`);

  const config = await res.json() as {
    input:  Array<{ name: string; data_type: string; dims: number[] }>;
    output: Array<{ name: string; data_type: string; dims: number[] }>;
  };

  const result = {
    inputs:  config.input,
    outputs: config.output,
  };
  modelConfigCache.set(modelName, result);
  return result;
}

// MCP tool — text inference via Triton Python backend
server.tool(
  'triton_text_infer',
  {
    model:  z.string().min(1),
    text:   z.string().min(1).max(8000),
  },
  async ({ model, text }) => {
    // Fetch config to get actual tensor names
    const config = await getModelConfig(model);

    const inputName  = config.inputs[0]?.name;
    const outputName = config.outputs[0]?.name;

    if (!inputName || !outputName) {
      throw new Error(`Model '${model}' has no inputs or outputs in config`);
    }

    // Triton BYTES datatype: send string data as an array of strings (not base64 on input)
    const requestBody: TritonInferRequest = {
      inputs: [{
        name:     inputName,
        shape:    [1],          // batch size 1
        datatype: 'BYTES',
        data:     [text],       // plain string — Triton handles encoding on ingress
      }],
      outputs: [{ name: outputName }],
    };

    const res = await fetch(`${TRITON_HTTP}/v2/models/${model}/infer`, {
      method:  'POST',
      headers: { 'Content-Type': 'application/json' },
      body:    JSON.stringify(requestBody),
      signal:  AbortSignal.timeout(30_000),
    });

    if (!res.ok) {
      const body = await res.text();
      throw new Error(`Triton infer failed for '${model}': HTTP ${res.status} — ${body}`);
    }

    const responseBody = await res.json() as {
      outputs: Array<{ name: string; datatype: string; data: unknown[] }>;
    };

    const outputTensor = responseBody.outputs.find(o => o.name === outputName);
    if (!outputTensor) throw new Error(`Output tensor '${outputName}' not found in Triton response`);

    // BYTES output is base64-encoded — must decode
    const rawValue = outputTensor.data[0];
    let outputText: string;

    if (outputTensor.datatype === 'BYTES' && typeof rawValue === 'string') {
      outputText = Buffer.from(rawValue, 'base64').toString('utf-8');
    } else {
      outputText = String(rawValue);
    }

    return { content: [{ type: 'text', text: outputText }] };
  }
);

Triton Python backends (the most common for NLP models) use the triton_python_backend_utils.pb_utils.Tensor API internally, which maps to BYTES datatype for string I/O. Always assume BYTES output is base64-encoded — even if the model author's Python code writes a plain string, Triton's HTTP serialization encodes it at the boundary.

Batch inference and multi-input models

Triton is designed for batch throughput — sending multiple inputs in one request is more efficient than N sequential calls. The shape field in each tensor specifies [batch_size, ...]. For NLP models with tokenized inputs, you may need multiple input tensors (e.g., input_ids, attention_mask, token_type_ids) that must all have matching batch dimensions.

// MCP tool — batch text inference via Triton (multiple texts, one request)
server.tool(
  'triton_batch_infer',
  {
    model: z.string().min(1),
    texts: z.array(z.string().min(1).max(2000)).min(1).max(32),
  },
  async ({ model, texts }) => {
    const config     = await getModelConfig(model);
    const inputName  = config.inputs[0]?.name  ?? 'text_input';
    const outputName = config.outputs[0]?.name ?? 'text_output';

    const batchSize = texts.length;

    const requestBody: TritonInferRequest = {
      inputs: [{
        name:     inputName,
        shape:    [batchSize],   // batch dimension = number of texts
        datatype: 'BYTES',
        data:     texts,         // array of strings — one per batch item
      }],
      outputs: [{ name: outputName }],
    };

    const res = await fetch(`${TRITON_HTTP}/v2/models/${model}/infer`, {
      method:  'POST',
      headers: { 'Content-Type': 'application/json' },
      body:    JSON.stringify(requestBody),
      signal:  AbortSignal.timeout(60_000),
    });

    if (!res.ok) {
      const body = await res.text();
      throw new Error(`Triton batch infer failed: HTTP ${res.status} — ${body}`);
    }

    const responseBody = await res.json() as {
      outputs: Array<{ name: string; datatype: string; data: unknown[] }>;
    };

    const outputTensor = responseBody.outputs.find(o => o.name === outputName);
    if (!outputTensor) throw new Error(`Output tensor '${outputName}' not in Triton response`);

    // Decode each BYTES output value from base64
    const results = outputTensor.data.map(raw => {
      if (outputTensor.datatype === 'BYTES' && typeof raw === 'string') {
        return Buffer.from(raw, 'base64').toString('utf-8');
      }
      return String(raw);
    });

    return {
      content: [{
        type: 'text',
        text: JSON.stringify({ results, count: results.length }),
      }],
    };
  }
);

// MCP tool — embedding inference with FP32 tensor output (e.g. sentence-bert on Triton)
server.tool(
  'triton_embed',
  {
    model: z.string().min(1),
    text:  z.string().min(1).max(4000),
  },
  async ({ model, text }) => {
    const config     = await getModelConfig(model);
    const inputName  = config.inputs[0]?.name  ?? 'text';
    const outputName = config.outputs[0]?.name ?? 'embedding';

    const res = await fetch(`${TRITON_HTTP}/v2/models/${model}/infer`, {
      method:  'POST',
      headers: { 'Content-Type': 'application/json' },
      body:    JSON.stringify({
        inputs:  [{ name: inputName, shape: [1], datatype: 'BYTES', data: [text] }],
        outputs: [{ name: outputName }],
      }),
      signal:  AbortSignal.timeout(15_000),
    });

    if (!res.ok) throw new Error(`Triton embed failed: HTTP ${res.status}`);

    const body = await res.json() as {
      outputs: Array<{ name: string; datatype: string; data: number[]; shape: number[] }>;
    };

    const embedTensor = body.outputs.find(o => o.name === outputName);
    if (!embedTensor) throw new Error(`Embedding tensor '${outputName}' not found`);

    // FP32 outputs are NOT base64-encoded — they come as JSON number arrays
    const embedding = embedTensor.data;  // number[]

    return {
      content: [{
        type: 'text',
        text: JSON.stringify({
          embedding,
          dimensions: embedding.length,
          shape:      embedTensor.shape,
        }),
      }],
    };
  }
);

Model repository management and health probes

Triton's /v2/health/ready endpoint returns 200 only when the server is ready to accept inference requests — all models in the model repository with version_policy: ALL must be loaded. If a model fails to load (missing weights, wrong backend), Triton still returns 200 on /health/ready but the failed model will return an error on individual infer calls. Use GET /v2/models/{model}/ready for per-model readiness — it returns 200 if that model is loaded and 400 if it failed to load.

// MCP resource — Triton server and per-model health
server.resource('triton_health', 'triton://health', async () => {
  try {
    // Server-level readiness — all configured models should be loaded
    const serverRes = await fetch(`${TRITON_HTTP}/v2/health/ready`, {
      signal: AbortSignal.timeout(5_000),
    });

    if (serverRes.status !== 200) {
      return {
        contents: [{
          uri: 'triton://health',
          text: JSON.stringify({
            status:      serverRes.status === 503 ? 'loading' : 'error',
            http_status: serverRes.status,
            checked_at:  new Date().toISOString(),
          }),
        }],
      };
    }

    // Fetch index of all loaded models
    const indexRes = await fetch(`${TRITON_HTTP}/v2/models`, {
      signal: AbortSignal.timeout(5_000),
    });

    const models = indexRes.ok
      ? (await indexRes.json() as { models: Array<{name: string; version: string; state: string}> }).models
      : [];

    const ready    = models.filter(m => m.state === 'READY');
    const notReady = models.filter(m => m.state !== 'READY');

    return {
      contents: [{
        uri: 'triton://health',
        text: JSON.stringify({
          status:          notReady.length === 0 ? 'ok' : 'degraded',
          models_ready:    ready.length,
          models_not_ready: notReady.map(m => ({ name: m.name, state: m.state })),
          checked_at:      new Date().toISOString(),
        }),
      }],
    };
  } catch (err: any) {
    return {
      contents: [{
        uri: 'triton://health',
        text: JSON.stringify({ status: 'unreachable', message: err.message }),
      }],
    };
  }
});

// MCP tool — load a model on demand (requires --model-control-mode explicit at startup)
server.tool(
  'triton_load_model',
  { model: z.string().min(1) },
  async ({ model }) => {
    const res = await fetch(`${TRITON_HTTP}/v2/repository/models/${model}/load`, {
      method: 'POST',
      signal: AbortSignal.timeout(120_000),  // loading can take up to 2min for large models
    });

    if (!res.ok) {
      const body = await res.text();
      throw new Error(`Triton model load failed for '${model}': HTTP ${res.status} — ${body}`);
    }

    return {
      content: [{
        type: 'text',
        text: JSON.stringify({ loaded: true, model }),
      }],
    };
  }
);

Triton's model control mode must be set at startup: --model-control-mode=explicit disables automatic loading and requires explicit POST /v2/repository/models/{model}/load calls; --model-control-mode=poll (default) auto-loads models when files appear in the model repository. For MCP servers that need to load models on agent request, use explicit mode — otherwise Triton auto-loads every model in the repository at startup regardless of whether they will be used.