Inference & Triton Deployment Pipeline

Once a clinical model achieves satisfactory validation metrics, the ColabBio pipeline automatically transitions the artifact from MLflow Registry to NVIDIA Triton Inference Server for real-time clinical consumption.

graph TD
    WEIGHTS["Trained Model Weights (.pt)"] --> REG["MLFLOW_REGISTER (Registry v1)"]
    REG --> EXPORT["Export / Convert to ONNX & TensorRT"]
    EXPORT --> CONFIG["Generate config.pbtxt (Batching & Dynamic dims)"]
    CONFIG --> REPO["Triton Model Repository: /models/abmil_tumor_classifier/1"]
    
    REPO --> HOT_RELOAD["TRITON_DEPLOY: POST /v2/repository/models/load"]
    HOT_RELOAD --> TRITON_SRV[("NVIDIA Triton Inference Server (GPU)")]
    
    TRITON_SRV --> VIEWER["Pathology Viewer (Heatmap Tiles)"]
    TRITON_SRV --> AGENT["Clinical Companion (LangGraph4j Tools)"]

---

1. Model Registration (MLFLOW_REGISTER)

The registration process publishes the model artifact, associates metadata, and increments the model version in the registry.

process MLFLOW_REGISTER {
    tag "MLflow Register: ${model_name}"
    container 'ghcr.io/colabbio/slidelab:latest'

    input:
    val model_name
    path weights

    output:
    val "mlflow-models:/${model_name}/1", emit: model_uri

    script:
    """
    python3 /opt/colabbio/mlops/register_model.py \
        --model-name "${model_name}" \
        --weights-path "${weights}" \
        --tracking-uri "${params.mlflow_tracking_uri}" \
        --stage "Staging"
    """
}

---

2. Triton Repository Structure & Configuration

Triton serves models from a strictly versioned directory containing the compiled model and its protobuf configuration (config.pbtxt):

/models/abmil_tumor_classifier/
├── config.pbtxt
└── 1/
    └── model.onnx

Example config.pbtxt

name: "abmil_tumor_classifier"
platform: "onnxruntime_onnx"
max_batch_size: 64

input [
  {
    name: "tile_embeddings"
    data_type: TYPE_FP32
    dims: [ -1, 512 ]
  }
]

output [
  {
    name: "tumor_probability"
    data_type: TYPE_FP32
    dims: [ 1 ]
  },
  {
    name: "attention_weights"
    data_type: TYPE_FP32
    dims: [ -1 ]
  }
]

instance_group [
  {
    count: 2
    kind: KIND_GPU
    gpus: [ 0 ]
  }
]

dynamic_batching {
  max_queue_delay_microseconds: 5000
}

---

3. Zero-Downtime Deployment (TRITON_DEPLOY)

The pipeline downloads the registered artifact, synchronizes the model repository, and instructs Triton to load the new version dynamically over its management API.

process TRITON_DEPLOY {
    tag "Triton Deploy: ${model_uri}"
    container 'ghcr.io/colabbio/slidelab:latest'

    input:
    val model_uri
    val triton_repo_path
    val triton_admin_url

    script:
    """
    # 1. Pull model from MLflow Registry
    python3 -m mlflow models download -m "${model_uri}" -d ./extracted_model

    # 2. Sync to Triton Model Repository volume
    cp -r ./extracted_model/data/model.onnx "${triton_repo_path}/abmil_tumor_classifier/1/"

    # 3. Hot-reload signal to Triton API without restarting the container
    curl -X POST "${triton_admin_url}/v2/repository/models/abmil_tumor_classifier/load"
    """
}

---

4. Real-Time Consumption

Once deployed:

• Pathology Viewer: Requests tile classification and attention weights to overlay real-time heatmap layers at 60 FPS.
• Clinical Companion: Evaluates tumor probability via triton_infer_biomarkers tool to support diagnostic conclusions.