bug: Worker Exits with code 132! when using custom AzureOpenAI Embeddings

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

1. Using these configs deploy to Azure webapps.
2. Make a request to the Endpoint,
3. 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:

Start with config.yml and the custom AzureEmbeddingModel implementation, then reproduce the endpoint request in the Python 3.11 Docker/Azure Web App environment while collecting the worker's exit logs. Done means the same configuration returns HTTP 201 reliably instead of a 502 with worker exit code 132.