使用Vertex AI TextEmbeddingModel的延迟问题及预热方案咨询

问题背景 #

使用Vertex AI的TextEmbeddingModel计算嵌入向量时，首次调用延迟远高于后续调用（如测试中首次耗时1.35秒，后续仅0.12秒左右），确认该现象并非上下文缓存导致（sdk_encode为单次独立调用），需要通过系统预热降低初始请求延迟。

测试代码及输出 #

import time
import google.generativeai as genai
from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
base_model_name = "text-embedding-004"
EMBED_TASK_TYPE = "RETRIEVAL_QUERY"

text_embedding_model = TextEmbeddingModel.from_pretrained(base_model_name)
model = TextEmbeddingModel.from_pretrained(base_model_name)

def sdk_encode( text):
        inputs = [TextEmbeddingInput(text.lower(), EMBED_TASK_TYPE) ]
        kwargs = {}
        embeddings = model.get_embeddings(inputs, **kwargs)
        text_embeddings = [embedding.values for embedding in embeddings]
        return text_embeddings[0] if len(text_embeddings) == 1 else text_embeddings

queries = ["I want to take pto Monday", "I want to take pto Tuesday", "I want to take pto Friday"]
for i in range(3): 
    query = queries[i]
    start_time = time.time()
    sdk_encode(query)
    end_time = time.time()
    sdk_delay = end_time - start_time
    print(f"Vertext SDK Latency for {query}: {sdk_delay}")

输出：

Vertext SDK Latency for I want to take pto Monday: 1.3506088256835938
Vertext SDK Latency for I want to take pto Tuesday: 0.12767696380615234
Vertext SDK Latency for I want to take pto Friday: 0.12481999397277832

核心原因 #

首次调用时需要完成模型资源加载、API连接握手、认证流程初始化等一次性操作，这些步骤会产生额外延迟；后续调用则复用已建立的连接和加载的模型，因此耗时大幅降低。

预热解决方案 #

1. 启动阶段执行预热调用 #

在模型初始化完成后，立即发起一个简单的测试调用，提前消耗初始化延迟，确保实际业务请求时模型已处于就绪状态。

修改后的示例代码：

import time
import google.generativeai as genai
from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
base_model_name = "text-embedding-004"
EMBED_TASK_TYPE = "RETRIEVAL_QUERY"

# 仅初始化一次模型实例（移除重复初始化）
model = TextEmbeddingModel.from_pretrained(base_model_name)

# 预热函数：触发模型加载与连接建立
def warm_up_model():
    test_input = [TextEmbeddingInput("warmup", EMBED_TASK_TYPE)]
    model.get_embeddings(test_input)
    print("模型预热完成")

# 启动时执行预热
warm_up_model()

def sdk_encode(text):
        inputs = [TextEmbeddingInput(text.lower(), EMBED_TASK_TYPE) ]
        embeddings = model.get_embeddings(inputs)
        text_embeddings = [embedding.values for embedding in embeddings]
        return text_embeddings[0] if len(text_embeddings) == 1 else text_embeddings

queries = ["I want to take pto Monday", "I want to take pto Tuesday", "I want to take pto Friday"]
for i in range(3): 
    query = queries[i]
    start_time = time.time()
    sdk_encode(query)
    end_time = time.time()
    sdk_delay = end_time - start_time
    print(f"Vertex SDK Latency for {query}: {sdk_delay}")

2. 复用模型实例 #

原代码中重复初始化了两次TextEmbeddingModel.from_pretrained，这会额外增加资源加载开销。确保全局仅初始化一次模型实例并复用，避免重复执行加载流程。

3. 优化客户端连接配置 #

通过配置客户端选项，提前建立并维持与Vertex AI服务的连接，减少首次调用的握手延迟。示例如下：

from google.api_core.client_options import ClientOptions

# 替换为你的服务区域对应的endpoint
client_options = ClientOptions(api_endpoint="us-central1-aiplatform.googleapis.com")
model = TextEmbeddingModel.from_pretrained(base_model_name, client_options=client_options)

4. 生产环境：利用Vertex AI端点预热功能 #

如果是生产部署场景，可在创建模型端点时配置warmup_requests参数，提交模拟请求让服务提前加载模型并保持就绪状态，避免实际流量到达时的初始延迟。

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