# 用 PyTorch Profiler 看耗时分布
torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA
    ],
    on_trace_ready=torch.profiler.tensorboard_trace_handler('./log')
) as prof:
    outputs = model(input_ids, attention_mask)
print(prof.key_averages().table(sort_by="cuda_time_total"))

# 环境信息确认命令
python -c "import torch; print(torch.__version__)"
python -c "import onnxruntime; print(onnxruntime.__version__)"
/usr/src/tensorrt/bin/trtexec --version
nvidia-smi --query-gpu=name,driver_version,memory.total --format=csv

# export_onnx.py
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

# 模型路径，这里用的是微调后的模型
model_path = "./models/bert-base-intent/checkpoint-5000"
model = AutoModelForSequenceClassification.from_pretrained(model_path, torchscript=True)
model.eval()

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

# 构造示例输入，这里的维度必须和实际推理时一致
dummy_input = tokenizer(
    "how to pay my bill",
    return_tensors="pt",
    padding="max_length",
    truncation=True,
    max_length=128
)

# 导出ONNX，关键是dynamic_axes
torch.onnx.export(
    model,
    (dummy_input["input_ids"], dummy_input["attention_mask"]),
    "./models/bert-intent.onnx",
    input_names=["input_ids", "attention_mask"],
    output_names=["logits"],
    dynamic_axes={
        "input_ids": {0: "batch_size", 1: "seq_len"},
        "attention_mask": {0: "batch_size", 1: "seq_len"},
        "logits": {0: "batch_size"}
    },
    opset_version=14,
    do_constant_folding=True,
    verbose=False
)
print("ONNX exported to ./models/bert-intent.onnx")

# onnx_inference.py
import numpy as np
import onnxruntime as ort
import time

class ONNXPredictor:
    def __init__(self, onnx_path, use_gpu=False):
        self.session = ort.InferenceSession(
            onnx_path,
            providers=[
                "CUDAExecutionProvider" if use_gpu else "CPUExecutionProvider",
                "CPUExecutionProvider"
            ]
        )
        # 打印实际使用的provider，确认GPU是否生效
        print("Using providers:", self.session.get_providers())
        
        self.input_names = [inp.name for inp in self.session.get_inputs()]
        self.output_names = [out.name for out in self.session.get_outputs()]
        
    def predict(self, input_ids, attention_mask):
        inputs = {
            "input_ids": input_ids.astype(np.int64),
            "attention_mask": attention_mask.astype(np.int64)
        }
        start = time.perf_counter()
        outputs = self.session.run(self.output_names, inputs)
        infer_ms = (time.perf_counter() - start) * 1000
        return outputs[0], infer_ms

# 使用示例
if __name__ == "__main__":
    import onnxruntime as ort
    
    predictor = ONNXPredictor("./models/bert-intent.onnx", use_gpu=True)
    
    # 构造输入
    input_ids = np.ones((1, 128), dtype=np.int64)
    attention_mask = np.ones((1, 128), dtype=np.int64)
    
    logits, latency = predictor.predict(input_ids, attention_mask)
    print(f"Latency: {latency:.2f}ms")
    print(f"Logits shape: {logits.shape}")

# 用trtexec转ONNX为TensorRT引擎
# 关键参数：--fp16开启半精度，--maxBatchSize设batch上限
# 注意：这里生成的engine文件和GPU型号绑定，换机器需要重新生成
/usr/src/tensorrt/bin/trtexec \
    --onnx=./models/bert-intent.onnx \
    --saveEngine=./models/bert-intent-fp16.engine \
    --fp16 \
    --maxBatchSize=32 \
    --workspace=1024 \
    --verbose

# convert_trt.py
import tensorrt as trt

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def build_engine(onnx_path, engine_path, fp16=True):
    """ONNX转TensorRT engine"""
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, TRT_LOGGER)
    
    with open(onnx_path, "rb") as f:
        if not parser.parse(f.read()):
            print("Failed to parse ONNX file")
            for i in range(parser.num_errors):
                print(parser.get_error(i))
            return None
    
    config = builder.create_builder_config()
    config.max_workspace_size = 1 << 30  # 1GB
    
    if fp16:
        config.set_flag(trt.BuilderFlag.FP16)
    
    # 动态shape需要设置优化配置文件
    profile = builder.create_optimization_profile()
    profile.set_shape("input_ids", (1, 64), (1, 128), (32, 256))
    profile.set_shape("attention_mask", (1, 64), (1, 128), (32, 256))
    config.add_optimization_profile(profile)
    
    engine = builder.build_engine(network, config)
    if engine:
        with open(engine_path, "wb") as f:
            f.write(engine.serialize())
        print(f"Engine saved to {engine_path}")
    return engine

if __name__ == "__main__":
    build_engine(
        onnx_path="./models/bert-intent.onnx",
        engine_path="./models/bert-intent-fp16.engine",
        fp16=True
    )

# trt_inference.py
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import time

class TRTPredictor:
    def __init__(self, engine_path):
        self.logger = trt.Logger(trt.Logger.WARNING)
        self.runtime = trt.Runtime(self.logger)
        
        with open(engine_path, "rb") as f:
            self.engine = self.runtime.deserialize_cuda_engine(f.read())
        self.context = self.engine.create_execution_context()
        
        # 分配显存
        self.allocate_buffers()
        
    def allocate_buffers(self):
        self.inputs = []
        self.outputs = []
        self.bindings = []
        
        for binding in self.engine:
            shape = self.engine.get_binding_shape(binding)
            size = trt.volume(shape)
            dtype = trt.nptype(self.engine.get_binding_dtype(binding))
            
            host_mem = cuda.pagelocked_empty(size, dtype)
            device_mem = cuda.mem_alloc(host_mem.nbytes)
            self.bindings.append(int(device_mem))
            
            if self.engine.binding_is_input(binding):
                self.inputs.append((host_mem, device_mem))
            else:
                self.outputs.append((host_mem, device_mem))
    
    def infer(self, input_ids, attention_mask):
        # 拷贝输入到显存
        np.copyto(self.inputs[0][0], input_ids.ravel())
        np.copyto(self.inputs[1][0], attention_mask.ravel())
        cuda.memcpy_htod(self.inputs[0][1], self.inputs[0][0])
        cuda.memcpy_htod(self.inputs[1][1], self.inputs[1][0])
        
        start = time.perf_counter()
        self.context.execute_v2(bindings=self.bindings)
        cuda.Context.synchronize()
        infer_ms = (time.perf_counter() - start) * 1000
        
        # 拷贝输出回内存
        cuda.memcpy_dtoh(self.outputs[0][0], self.outputs[0][1])
        
        logits = self.outputs[0][0].copy().reshape(1, -1)
        return logits, infer_ms

if __name__ == "__main__":
    predictor = TRTPredictor("./models/bert-intent-fp16.engine")
    
    input_ids = np.ones((1, 128), dtype=np.int32)
    attention_mask = np.ones((1, 128), dtype=np.int32)
    
    logits, latency = predictor.infer(input_ids, attention_mask)
    print(f"TensorRT latency: {latency:.2f}ms")

# model_repository/bert_intent/1/model.onnx  (实际文件)
# 把ONNX文件放在版本号目录下

# model_repository/bert_intent/config.pbtxt
name: "bert_intent"
platform: "onnxruntime_onnx"
max_batch_size: 32
input [
  {
    name: "input_ids"
    data_type: TYPE_INT64
    dims: [-1, 128]
  },
  {
    name: "attention_mask"
    data_type: TYPE_INT64
    dims: [-1, 128]
  }
]
output [
  {
    name: "logits"
    data_type: TYPE_FP32
    dims: [-1, 2]
  }
]

dynamic_batching {
  preferred_batch_size: [4, 8, 16]
  max_queue_delay_microseconds: 100
}

# 启动Triton服务
# 用GPU跑需要加 --gpus 1，拉取镜像的时候注意版本
docker pull nvcr.io/nvidia/tritonserver:23.05-py3

# 挂载模型仓库并启动
docker run --gpus=1 --rm -p 8000:8000 -p 8001:8001 -p 8002:8002 \
  -v /data/model_repository:/models \
  nvcr.io/nvidia/tritonserver:23.05-py3 \
  tritonserver --model-repository=/models \
  --model-control-mode=poll \
  --metrics-interval-ms=5000

# triton_client.py
import numpy as np
import tritonclient.http as httpclient
import time

client = httpclient.InferenceServerClient(url="localhost:8000")

# 确认模型状态
print(client.get_model_repository_index())

input_ids = np.ones((1, 128), dtype=np.int64)
attention_mask = np.ones((1, 128), dtype=np.int64)

inputs = [
    httpclient.InferInput("input_ids", input_ids.shape, "INT64"),
    httpclient.InferInput("attention_mask", attention_mask.shape, "INT64")
]
inputs[0].set_data_from_numpy(input_ids)
inputs[1].set_data_from_numpy(attention_mask)

outputs = [httpclient.InferRequestedOutput("logits")]

# 压测100次取平均
latencies = []
for _ in range(100):
    start = time.perf_counter()
    result = client.infer("bert_intent", inputs=inputs, outputs=outputs)
    latencies.append((time.perf_counter() - start) * 1000)

logits = result.as_numpy("logits")
print(f"Average latency: {np.mean(latencies):.2f}ms")
print(f"P99 latency: {np.percentile(latencies, 99):.2f}ms")
print(f"Logits: {logits}")