部署TorchScript模型至Google Cloud Function后出现500错误的排查求助

问题场景 #

将TorchScript模型部署到Google Cloud Function后，部署流程显示成功，但通过本地终端用cURL测试时触发500 Internal Server Error。查看云函数日志，报错信息为：TypeError: my_cloud_function() takes 0 positional arguments but 1 was given。

原模型代码 #

import torch
from transformers import MobileBertTokenizer

from google.cloud import storage
from flask import Flask, request

app = Flask(__name__)

def load_model_from_cloud_storage(bucket_name, model_filename):
    storage_client = storage.Client()
    bucket = storage_client.bucket(bucket_name)
    blob = bucket.blob(model_filename)

    # Download the model file to a temporary location
    temp_model_filename = "/tmp/" + model_filename
    blob.download_to_filename(temp_model_filename)

    # Load the TorchScript model
    model = torch.jit.load(temp_model_filename)

    return model

@app.route('/', methods=['POST'])
def my_cloud_function():
    # Handle your Cloud Function logic
    bucket_name = "pytorchmodel1"
    model_filename = "model.pt"

    # Load the TorchScript model
    model = load_model_from_cloud_storage(bucket_name, model_filename)

    # Load the BERT tokenizer
    tokenizer = MobileBertTokenizer.from_pretrained('google/mobilebert-uncased')

    # Retrieve the messages from the request
    request_data = request.get_json()
    messages = request_data.get('messages', [])

    results = []
    for message in messages:
        # Tokenize the message
        inputs = tokenizer.encode_plus(
            message,
            add_special_tokens=True,
            max_length=32,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )

        # Perform inference using the loaded model
        input_ids = inputs['input_ids']
        attention_mask = inputs['attention_mask']
        outputs = model(input_ids)
        probabilities = torch.softmax(outputs['logits'], dim=1).squeeze()

        # Get the probability of being a scam
        prob_scam = probabilities[1].item()

        results.append({
            'message': message,
            'probability_scam': prob_scam
        })

    return {
        'results': results
    }

原requirements.txt #

torch==2.0.1
transformers==4.30.2
google-cloud-storage==1.44.0
flask==2.2.4

原测试cURL命令 #

curl -X POST -H "Content-Type: application/json" -d "{\"messages\": \"Our systems indicate your Social Security number has been involved in illegal activities. You must urgently visit https://example.gov/reportfraud/ to report any unauthorized usage. This is an URGENT situation.\"}" <my-function-url>

报错原因分析 #

1. 函数参数不符合GCP规范：Google Cloud Function的HTTP触发函数必须接收一个请求对象参数（通常命名为request），原函数my_cloud_function()未定义该参数，导致调用时传参失败。
2. 冗余Flask配置：GCP Cloud Function的HTTP触发机制已内置请求处理逻辑，无需手动初始化Flask应用和添加@app.route装饰器，这部分代码属于冗余配置，会干扰函数的正常触发。
3. 请求数据格式错误：原cURL命令中messages字段值为字符串，但代码逻辑期望接收数组，会导致后续遍历处理出错。
4. 性能问题：原代码在每次请求时重新加载模型和tokenizer，会大幅增加响应时间，甚至可能触发超时。

修复方案 #

1. 修正函数定义与结构 #

移除Flask相关代码，让函数接收request参数，适配GCP Cloud Function的HTTP触发规范。同时将模型和tokenizer的加载逻辑移到全局作用域，实现冷启动后复用，提升性能。

2. 优化后的代码 #

import torch
from transformers import MobileBertTokenizer
from google.cloud import storage

# 全局加载模型和tokenizer，冷启动后复用
bucket_name = "pytorchmodel1"
model_filename = "model.pt"

def load_model_from_cloud_storage(bucket_name, model_filename):
    storage_client = storage.Client()
    bucket = storage_client.bucket(bucket_name)
    blob = bucket.blob(model_filename)

    temp_model_filename = "/tmp/" + model_filename
    blob.download_to_filename(temp_model_filename)
    model = torch.jit.load(temp_model_filename)
    return model

# 全局初始化模型和tokenizer
model = load_model_from_cloud_storage(bucket_name, model_filename)
tokenizer = MobileBertTokenizer.from_pretrained('google/mobilebert-uncased')

def my_cloud_function(request):
    # 解析请求数据
    request_data = request.get_json()
    messages = request_data.get('messages', [])
    
    # 兼容单条消息的情况
    if isinstance(messages, str):
        messages = [messages]

    results = []
    for message in messages:
        inputs = tokenizer.encode_plus(
            message,
            add_special_tokens=True,
            max_length=32,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )

        input_ids = inputs['input_ids']
        attention_mask = inputs['attention_mask']
        outputs = model(input_ids)
        probabilities = torch.softmax(outputs['logits'], dim=1).squeeze()
        prob_scam = probabilities[1].item()

        results.append({
            'message': message,
            'probability_scam': prob_scam
        })

    return {'results': results}

3. 更新requirements.txt #

移除不需要的flask依赖：

torch==2.0.1
transformers==4.30.2
google-cloud-storage==1.44.0

4. 修正测试cURL命令 #

将messages改为数组格式：

curl -X POST -H "Content-Type: application/json" -d "{\"messages\": [\"Our systems indicate your Social Security number has been involved in illegal activities. You must urgently visit https://example.gov/reportfraud/ to report any unauthorized usage. This is an URGENT situation.\"]}" <my-function-url>

验证步骤 #

1. 重新部署修改后的云函数
2. 使用修正后的cURL命令发起请求，确认返回正常的JSON结果
3. 查看云函数日志，确认无报错信息

内容的提问来源于stack exchange，提问作者Taskeen Jafri