Serving the Trained Model

Starting the Triton Server

Training writes everything you need to serve alongside the trained weights. Assume your training output is at /path/to/output (the value you passed as output_dir in your config). The structure you will find there is:

/path/to/output/
├── model_final.pt
├── xgboost_fraud.json
├── metrics.json
└── infer/
    ├── triton/            ← Triton model repository — pass this as MODEL_REPO
    ├── serve.sh           ← launches the Triton container
    └── client.py          ← full inference client (real data, metrics, explain mode)

Step 1 — Install the client dependencies (host, once)

The inference client runs on the host (outside Docker). Install its dependencies once:

pip install "tritonclient[grpc]" pandas numpy torch

Step 2 — Start the Triton Server

serve.sh mounts the Triton model repository into the official nvcr.io/nvidia/tritonserver:26.04-py3 container and installs the Python backend dependencies at startup:

MODEL_REPO=/path/to/output/infer/triton \
  bash /path/to/output/infer/serve.sh

Run this in a dedicated terminal (it stays in the foreground). You will see a stream of startup logs as Triton loads each model.

What serve.sh executes internally:

docker run --rm \
  --gpus all \
  --shm-size=3g \
  -e LLM_API_KEY \
  -e LLM_BASE_URL \
  -e LLM_MODEL \
  -e NVIDIA_API_KEY \
  -p 127.0.0.1:8000:8000 \
  -p 127.0.0.1:8001:8001 \
  -p 127.0.0.1:8002:8002 \
  -v "/path/to/output/infer/triton:/models" \
  nvcr.io/nvidia/tritonserver:26.04-py3 \
  bash -c "pip install --quiet --root-user-action=ignore \
             'torch_geometric>=2.6' 'xgboost>=2.1,<3' 'captum>=0.7' 'openai>=1.0' && \
           tritonserver \
             --model-repository=/models \
             --backend-config=python,stub-timeout-seconds=120"

Key design choices:

• Python dependencies (torch_geometric, xgboost, captum, openai) are installed inside the container at every startup (torch is already in the Triton base image) — the model repository itself stays clean and portable.
• stub-timeout-seconds=120 allows the Python backend models up to 2 minutes to initialise. This is necessary on first startup when PyTorch and the GNN architecture are loaded for the first time.
• LLM API key environment variables are forwarded from the host shell into the container automatically — no secrets need to be baked into images or files.

Exposed ports:

Step 3 — Wait for the Server to Become Ready

The server takes 30–90 seconds on the first startup (pip install + model loading). Poll the health endpoint in a separate terminal:

until curl -sf http://localhost:8000/v2/health/ready > /dev/null; do
  echo "Waiting for Triton..."
  sleep 3
done
echo "Triton is ready."

The server is ready when the container logs show:

I0101 00:00:00.000000 1 grpc_server.cc:xxxx] Started GRPCInferenceService at 0.0.0.0:8001
I0101 00:00:00.000000 1 http_server.cc:xxxx] Started HTTPService at 0.0.0.0:8000
I0101 00:00:00.000000 1 server.cc:xxxx] Triton Inference Server is ready.

Verify all six models are loaded:

curl -s http://localhost:8000/v2/models | python3 -m json.tool

Expected: all of gnn_embedder, xgb_fraud, xgb_explainer, llm_explainer, fraud_pipeline, fraud_pipeline_explained show "state": "READY".

You can also check a specific model:

curl -s http://localhost:8000/v2/models/fraud_pipeline/ready

LLM Explainability at Serve Time

The LLM explainer calls an external OpenAI-compatible API. The API key must be set as an environment variable in the host shell before running serve.sh — the key is forwarded into the container automatically; you never need to write it to disk.

Option A — NVIDIA NIM (default, no LLM_BASE_URL needed):

export NVIDIA_API_KEY=nvapi-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
# Optionally override the model (default: nvidia/meta/llama-3.2-1b-instruct):
export LLM_MODEL=nvidia/meta/llama-3.1-70b-instruct
MODEL_REPO=/path/to/output/infer/triton \
  bash /path/to/output/infer/serve.sh

Default endpoint: https://inference-api.nvidia.com Default model: nvidia/meta/llama-3.2-1b-instruct

Option B — OpenAI or any compatible API endpoint:

export LLM_API_KEY=sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
export LLM_BASE_URL=https://api.openai.com/v1
export LLM_MODEL=gpt-4o-mini
MODEL_REPO=/path/to/output/infer/triton \
  bash /path/to/output/infer/serve.sh

Option C — Any other OpenAI-compatible endpoint:

export LLM_API_KEY=<your-api-key>
export LLM_BASE_URL=<your-openai-compatible-endpoint>
export LLM_MODEL=<model-name>
MODEL_REPO=/path/to/output/infer/triton \
  bash /path/to/output/infer/serve.sh

Confirm the LLM backend connected successfully by checking the Triton startup logs for:

[llm_explainer] LLM connectivity check passed.

Running without an API key: The server starts normally. fraud_pipeline (scoring only) works fully. fraud_pipeline_explained also starts, but llm_explainer disables itself gracefully — it logs a warning and returns a placeholder string ([LLM explanation disabled: set LLM_API_KEY to enable]) in place of a real explanation. This is not an error state; scoring is unaffected.

Production Deployment Checklist

Before routing production traffic to the server:

Multi-GPU Inference — Horizontal Scaling

By default, all models load on GPU 0 with 1 instance. To spread load across multiple GPUs, use the NUM_GPUS and INSTANCES_PER_GPU environment variables when calling serve.sh. The script automatically patches the config.pbtxt files before starting Triton — no manual file editing required. Training and deployment typically happen on different machines, so GPU count is set at serve time, not at training time.

Usage:

# Single GPU, 1 instance (default):
MODEL_REPO=... bash serve.sh
# 2 GPUs, 1 instance each (2 total per model):
NUM_GPUS=2 INSTANCES_PER_GPU=1 MODEL_REPO=... bash serve.sh
# 2 GPUs, 3 instances each (6 total per model):
NUM_GPUS=2 INSTANCES_PER_GPU=3 MODEL_REPO=... bash serve.sh
# 8 GPUs, 1 instance each (8 total per model):
NUM_GPUS=8 INSTANCES_PER_GPU=1 MODEL_REPO=... bash serve.sh

Triton load-balances incoming requests across all instances automatically.

Which models are patched:

Node prediction pipeline (GNN_XGBoost_NP):

Edge prediction pipeline (GNN_XGBoost):

What serve.sh writes to each patched config:

instance_group [ { kind: KIND_GPU  gpus: [0, 1, ...]  count: <INSTANCES_PER_GPU> } ]

count is per GPU, not total. gpus: [0, 1] count: 3 creates 6 instances total (3 per GPU).

Startup time with many instances: torch.compile warmup runs once per GNN embedder instance (gnn_embedder for EP, np_gnn_embedder for NP). With 8 GPUs × 3 instances = 24 instances, startup can take several minutes. The default stub-timeout-seconds=120 may need increasing — set it in serve.sh if models fail to load:

--backend-config=python,stub-timeout-seconds=600