Optimizing Multi-LoRA Serving in Production: Dynamic Adapter Swapping with Triton Inference Server — code snippets

snippet-1.txt

name: "ensemble_base_model"
backend: "vllm"
max_batch_size: 0

input [
  {
    name: "prompt"
    data_type: TYPE_STRING
    dims: [ 1 ]
  },
  {
    name: "sampling_parameters"
    data_type: TYPE_STRING
    dims: [ 1 ]
  },
  {
    name: "lora_name"
    data_type: TYPE_STRING
    dims: [ 1 ]
  },
  {
    name: "lora_path"
    data_type: TYPE_STRING
    dims: [ 1 ]
  }
]

output [
  {
    name: "text_output"
    data_type: TYPE_STRING
    dims: [ 1 ]
  }
]

parameters: {
  key: "model"
  value: { string_value: "/opt/triton/models/base_llama3_8b" }
}
parameters: {
  key: "enable_lora"
  value: { string_value: "true" }
}
parameters: {
  key: "max_loras"
  value: { string_value: "32" }
}
parameters: {
  key: "max_lora_rank"
  value: { string_value: "16" }
}
parameters: {
  key: "lora_extra_vocab_size"
  value: { string_value: "256" }
}
parameters: {
  key: "gpu_memory_utilization"
  value: { string_value: "0.80" }
}
parameters: {
  key: "max_cpu_loras"
  value: { string_value: "128" }
}

snippet-2.txt

name: "multilora_router"
backend: "python"
max_batch_size: 0

input [
  {
    name: "prompt"
    data_type: TYPE_STRING
    dims: [ 1 ]
  },
  {
    name: "adapter_id"
    data_type: TYPE_STRING
    dims: [ 1 ]
  }
]

output [
  {
    name: "generated_text"
    data_type: TYPE_STRING
    dims: [ 1 ]
  }
]

instance_group [
  {
    count: 8
    kind: KIND_CPU
  }
]

snippet-3.py

import json
import numpy as np
import triton_python_backend_utils as pb_utils
from cache_manager import AdapterCacheManager

class TritonPythonModel:
    def initialize(self, args):
        self.model_config = json.loads(args['model_config'])
        self.base_model_name = "ensemble_base_model"
        
        # Configure tiered storage location
        self.s3_bucket = "enterprise-lora-registry"
        self.local_cache_dir = "/tmp/lora_cache"
        
        # Local NVMe pool capacity capped at 128 adapters
        self.cache_manager = AdapterCacheManager(
            s3_bucket=self.s3_bucket,
            local_dir=self.local_cache_dir,
            max_disk_adapters=128
        )

    def execute(self, requests):
        responses = []
        for request in requests:
            prompt_tensor = pb_utils.get_input_tensor_by_name(request, "prompt")
            adapter_id_tensor = pb_utils.get_input_tensor_by_name(request, "adapter_id")
            
            if prompt_tensor is None or adapter_id_tensor is None:
                responses.append(pb_utils.InferenceResponse(
                    output_tensors=[],
                    error=pb_utils.TritonModelError("Missing required inputs: prompt and adapter_id")
                ))
                continue
                
            prompt = prompt_tensor.as_numpy()[0].decode('utf-8')
            adapter_id = adapter_id_tensor.as_numpy()[0].decode('utf-8')
            
            try:
                # Stage adapter to local NVMe if not present
                adapter_path = self.cache_manager.get_adapter_path(adapter_id)
                
                # Define vLLM sampling parameters
                sampling_params = json.dumps({
                    "temperature": 0.0,
                    "max_tokens": 256,
                    "stop": ["<|eot_id|>"]
                })
                
                # Construct NumPy inputs for BLS call
                sampling_params_np = np.array([sampling_params], dtype=np.object_)
                lora_name_np = np.array([adapter_id], dtype=np.object_)
                lora_path_np = np.array([adapter_path], dtype=np.object_)
                
                vllm_inputs = [
                    pb_utils.Tensor("prompt", prompt_tensor.as_numpy()),
                    pb_utils.Tensor("sampling_parameters", sampling_params_np),
                    pb_utils.Tensor("lora_name", lora_name_np),
                    pb_utils.Tensor("lora_path", lora_path_np)
                ]
                
                # Non-blocking execution within Triton's C API context
                inference_request = pb_utils.InferenceRequest(
                    model_name=self.base_model_name,
                    requested_output_names=["text_output"],
                    inputs=vllm_inputs
                )
                
                inference_response = inference_request.exec()
                
                if inference_response.has_error():
                    responses.append(pb_utils.InferenceResponse(
                        output_tensors=[],
                        error=inference_response.get_error()
                    ))
                else:
                    text_output_tensor = pb_utils.get_output_tensor_by_name(
                        inference_response, "text_output"
                    )
                    output_tensor = pb_utils.Tensor(
                        "generated_text", text_output_tensor.as_numpy()
                    )
                    responses.append(pb_utils.InferenceResponse(
                        output_tensors=[output_tensor]
                    ))
                    
            except Exception as e:
                responses.append(pb_utils.InferenceResponse(
                    output_tensors=[],
                    error=pb_utils.TritonModelError(f"Multi-LoRA router failed: {str(e)}")
                ))
                
        return responses

    def finalize(self):
        pass

snippet-4.py

import os
import shutil
import threading
from collections import OrderedDict
import boto3

class AdapterCacheManager:
    def __init__(self, s3_bucket, local_dir, max_disk_adapters=128):
        self.s3_bucket = s3_bucket
        self.local_dir = local_dir
        self.max_disk_adapters = max_disk_adapters
        self.lock = threading.Lock()
        
        # Thread-safe LRU tracking
        self.cache = OrderedDict()
        self.s3_client = boto3.client('s3')
        
        # Guard active downloads to prevent duplicate network requests
        self.active_downloads = {}
        
        os.makedirs(self.local_dir, exist_ok=True)
        self._scan_existing_adapters()

    def _scan_existing_adapters(self):
        for item in os.listdir(self.local_dir):
            item_path = os.path.join(self.local_dir, item)
            if os.path.isdir(item_path):
                self.cache[item] = item_path

    def get_adapter_path(self, adapter_id):
        download_event = None
        
        with self.lock:
            if adapter_id in self.cache:
                self.cache.move_to_end(adapter_id)
                return self.cache[adapter_id]
            
            # Check if another thread is already downloading this adapter
            if adapter_id in self.active_downloads:
                download_event = self.active_downloads[adapter_id]
            else:
                download_event = threading.Event()
                self.active_downloads[adapter_id] = download_event
                
        # If another thread is downloading, block and wait for it
        if download_event is not None and adapter_id not in self.active_downloads:
            # We are the downloader thread
            adapter_path = os.path.join(self.local_dir, adapter_id)
            try:
                self._download_from_s3(adapter_id, adapter_path)
                with self.lock:
                    self.cache[adapter_id] = adapter_path
                    if len(self.cache) > self.max_disk_adapters:
                        self._evict_lru()
            finally:
                with self.lock:
                    download_event.set()
                    self.active_downloads.pop(adapter_id, None)
            return adapter_path
        else:
            # We are a waiter thread
            download_event.wait()
            with self.lock:
                if adapter_id in self.cache:
                    return self.cache[adapter_id]
                raise RuntimeError(f"Parallel download of adapter '{adapter_id}' failed.")

    def _download_from_s3(self, adapter_id, destination_path):
        os.makedirs(destination_path, exist_ok=True)
        prefix = f"adapters/{adapter_id}/"
        
        try:
            paginator = self.s3_client.get_paginator('list_objects_v2')
            pages = paginator.paginate(Bucket=self.s3_bucket, Prefix=prefix)
            
            file_downloaded = False
            for page in pages:
                if 'Contents' not in page:
                    continue
                for obj in page['Contents']:
                    key = obj['Key']
                    relative_path = os.path.relpath(key, prefix)
                    dest_file_path = os.path.join(destination_path, relative_path)
                    
                    os.makedirs(os.path.dirname(dest_file_path), exist_ok=True)
                    self.s3_client.download_file(self.s3_bucket, key, dest_file_path)
                    file_downloaded = True
                    
            if not file_downloaded:
                raise FileNotFoundError(f"Adapter '{adapter_id}' not found in S3.")
                
        except Exception as e:
            if os.path.exists(destination_path):
                shutil.rmtree(destination_path)
            raise RuntimeError(f"S3 download failed for {adapter_id}: {str(e)}")

    def _evict_lru(self):
        lru_key, lru_path = self.cache.popitem(last=False)
        try:
            if os.path.exists(lru_path):
                shutil.rmtree(lru_path)
        except Exception as e:
            # Re-insert to retry eviction later
            self.cache[lru_key] = lru_path
            self.cache.move_to_end(lru_key, last=False)
            raise RuntimeError(f"Failed to evict {lru_key}: {str(e)}")

snippet-5.py

import numpy as np
import tritonclient.grpc as grpcclient

def run_multilora_inference(server_url, prompt, adapter_id):
    try:
        # Establish a single TCP connection with HTTP/2 keep-alive via gRPC
        client = grpcclient.InferenceServerClient(url=server_url)
    except Exception as e:
        raise ConnectionError(f"Failed to connect to Triton server: {e}")

    # Pack values into NumPy object arrays
    prompt_data = np.array([prompt], dtype=np.object_)
    adapter_data = np.array([adapter_id], dtype=np.object_)

    # Prepare inputs matching Triton's expected signature
    inputs = [
        grpcclient.InferInput("prompt", [1], "BYTES"),
        grpcclient.InferInput("adapter_id", [1], "BYTES")
    ]
    
    inputs[0].set_data_from_numpy(prompt_data)
    inputs[1].set_data_from_numpy(adapter_data)

    try:
        # Execute gRPC call
        response = client.infer(
            model_name="multilora_router",
            inputs=inputs,
            timeout=15.0 # Account for cold-start staging timeouts
        )
        generated_text = response.as_numpy("generated_text")[0].decode('utf-8')
        return generated_text
    except Exception as e:
        raise RuntimeError(f"Triton inference execution failed: {e}")

if __name__ == "__main__":
    SERVER_URL = "10.0.12.45:8001" # Internal Load Balancer IP
    res = run_multilora_inference(
        server_url=SERVER_URL,
        prompt="Synthesize the transaction logs for anomalous behavior:",
        adapter_id="lora_secops_detection_v4"
    )
    print(f"Generated Output: {res}")

snippet-6.sh

#!/usr/bin/env bash
set -euo pipefail

# Production pre-warming script to populate the local NVMe cache
S3_BUCKET="enterprise-lora-registry"
LOCAL_CACHE_DIR="/tmp/lora_cache"
HOT_ADAPTERS=("lora_finance_prod" "lora_legal_prod" "lora_security_core")

echo "[+] Initializing NVMe cache directory: ${LOCAL_CACHE_DIR}"
mkdir -p "${LOCAL_CACHE_DIR}"

for adapter in "${HOT_ADAPTERS[@]}"; do
    TARGET_DIR="${LOCAL_CACHE_DIR}/${adapter}"
    
    if [ -d "${TARGET_DIR}" ]; then
        echo "[*] Adapter '${adapter}' is already staged. Skipping."
        continue
    fi
    
    echo "[*] Staging '${adapter}' from S3 (s3://${S3_BUCKET}/adapters/${adapter}/)..."
    aws s3 cp "s3://${S3_BUCKET}/adapters/${adapter}/" "${TARGET_DIR}/" --recursive --quiet
    
    if [ ! -f "${TARGET_DIR}/adapter_config.json" ] || [ ! -f "${TARGET_DIR}/adapter_model.safetensors" ]; then
        echo "[!] Error: Staged adapter '${adapter}' is corrupted or incomplete."
        rm -rf "${TARGET_DIR}"
        exit 1
    fi
    echo "[+] Staged '${adapter}' successfully."
done

# Wait for Triton to report readiness
TRITON_READY_URL="http://localhost:8000/v1/health/ready"
echo "[+] Waiting for Triton Server health validation..."

for i in {1..30}; do
    STATUS_CODE=$(curl -s -o /dev/null -w "%{http_code}" "${TRITON_READY_URL}" || true)
    if [ "$STATUS_CODE" -eq 200 ]; then
        echo "[+] Triton Server is ready. Pre-warming sequence completed."
        exit 0
    fi
    sleep 2
done

echo "[!] Timeout: Triton failed to reach ready state."
exit 1