bug: Worker Exits with code 132! when using custom AzureOpenAI Embeddings #1492
Description
Did you check docs and existing issues?
- I have read all the NeMo-Guardrails docs
- I have updated the package to the latest version before submitting this issue
- (optional) I have used the develop branch
- I have searched the existing issues of NeMo-Guardrails
Python version (python --version)
3.11
Operating system/version
Docker container with python 3.11-slim image and Azure Webapp p2mv3
NeMo-Guardrails version (if you must use a specific version and not the latest
0.15.0
Describe the bug
Most of the times when I deploy my webapp and make an API request to nemoguardrails configuration, I get an error that worker exited with code 132! I have to rerun pipelines multiple times to make the same code work. I was able to narrow down the problem to embeddings model. I have overwrtitten the methods for embeddings provider to use AzureOpenAI embeddings. These are the logs that I get when I run rails.
config.yml
#dummy values to override with db values
models:
- type: main
engine: openai
model: gpt-4o
parameters:
deployment_name:
openai_api_key:
api_version:
azure_endpoint:
temperature:
- type: content_safety
engine: openai
model: gpt-4o
parameters:
deployment_name:
api_key:
api_version:
azure_endpoint:
temperature:
- type: topic_control
engine: openai
model: gpt-4o
parameters:
deployment_name:
api_key:
api_version:
azure_endpoint:
temperature:
- type: embeddings
engine: openai
model: gpt-4o
parameters:
deployment_name:
api_key:
api_version:
azure_endpoint:
rails:
input:
parallel: True
flows:
- ascii smuggle check
- self check input
- content safety check input $model=content_safety
- topic safety check input $model=topic_control
- facts check
- pii redact
AzureOpenAI Embedding Support Added using
import asyncio
from contextvars import ContextVar
from typing import List
from langchain_openai import AzureOpenAIEmbeddings
from nemoguardrails.embeddings.providers.base import EmbeddingModel
from app.core.config import settings
import time
import os
import logging
async_client_var: ContextVar = ContextVar("async_client", default=None)
class AzureEmbeddingModel(EmbeddingModel):
engine_name = "azure"
def __init__(
self,
embedding_model: str,
**kwargs,
):
self.model = embedding_model
self.client = AzureOpenAIEmbeddings(
azure_deployment=settings.AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME,
openai_api_version=settings.AZURE_OPENAI_API_EMBEDDING_VERSION,
azure_endpoint=settings.AZURE_OPENAI_ENDPOINT,
api_key=settings.AZURE_OPENAI_API_KEY,
)
self.embedding_size_dict = {
"text-embedding-ada-002": 1536,
"text-embedding-3-small": 1536,
"text-embedding-3-large": 3072,
}
if self.model in self.embedding_size_dict:
self.embedding_size = self.embedding_size_dict[self.model]
else:
# Perform a first encoding to get the embedding size
self.embedding_size = len(self.encode(["test"])[0])
async def encode_async(self, documents: List[str]) -> List[List[float]]:
start_time = time.time()
loop = asyncio.get_running_loop()
embeddings = await loop.run_in_executor(None, self.encode, documents)
end_time = time.time()
print(f"TIME - initialize_rails_config --> encode_async latency: {end_time - start_time} sec")
return embeddings
def encode(self, documents: List[str]) -> List[List[float]]:
embeddings = self.client.embed_documents(documents)
return embeddings
Steps To Reproduce
- Using these configs deploy to Azure webapps.
- Make a request to the Endpoint,
- Returns 502 server error.
Expected Behavior
Returns 201 response when API call to the endpoint with this config is made.
Expected Logs:
Actual Behavior
API call returns 502 error and logs show:
