IBM UFCG

Pesquisa e desenvolvimento em IA, hardware de alto desempenho e infraestrutura em parceria entre UFCG e IBM.

Construindo API para inferências de LLMs em um servidor IBM Power9

Contexto

Este é o terceiro post de uma série de tutoriais cujo objetivo é mostrar passo a passo como construir uma API de Modelos de Linguagem em um servidor Power9, desde a configuração do sistema operacional até a execução remota de inferências. Já configuramos o sistema operacional, os drivers NVIDIA, CUDA e cuDNN no primeiro post, e no segundo post instalamos Conda e PyTorch. Nesta etapa, vamos construir a API usando FastAPI e a biblioteca Transformers, baixando modelos do Hugging Face e executando o servidor web com uvicorn.

A API implementada terá as funcionalidades de gerar API Key, carregar modelos, realizar inferências, obter status e desccaregar modelos.

FastAPI: Framework web moderno para construção de APIs com Python 3.8+, baseado em tipagem estática e assíncrona. Foi projetado para ser rápido, fácil de usar e robusto, tornando o desenvolvimento de APIs mais eficiente.

Transformers: Biblioteca de código aberto desenvolvida pela Hugging Face. Fornece acesso prático e eficiente a uma ampla coleção de modelos pré-treinados de última geração para Processamento de Linguagem Natural (PLN), visão computacional e áudio.

Hugging Face: Hugging Face é uma plataforma focada em inteligência artificial, conhecida por hospedar modelos de NLP e outras tarefas. O Hugging Face Hub é um repositório colaborativo onde desenvolvedores e pesquisadores podem compartilhar, versionar e baixar modelos prontos para uso, facilitando o acesso e integração de modelos.

Uvicorn: Servidor web ASGI (Asynchronous Server Gateway Interface). O Uvicorn é um servidor de alta performance para aplicações Python assíncronas.

TL;DR

  • Este post apresenta o passo a passo para implementar uma API que realiza inferências de Grandes Modelos de Linguagem.
  • Usaremos FastAPI e Transformers para desenvolver essa API e Hugging Face para baixar os modelos.

Configuração do Ambiente

Estrutura de Diretórios

Primeiro, vamos criar a estrutura básica do projeto:

model_api/
├── requirements.txt
├── app/
│   ├── __init__.py
│   ├── main.py
│   ├── schemas.py
│   ├── auth.py
│   ├── model_manager.py
│   ├── utils.py
│   └── apikey_store.json
└── README.md (opcional)

Arquivo requirements.txt

Vamos usar FastAPI e Transformers para implementar a API. Além disso, usaremos uvicorn para executar o servidor, pydantic para validação de dados de entrada e torch, que já instalamos no tutorial anterior.

Primeiro, vamos instalar as bibliotecas necessárias e depois preencher o arquivo requirements.txt. Lembre-se de ativar o ambiente conda se você o criou, para garantir o uso correto do pytorch.

conda activate llm_api
pip install fastapi uvicorn transformers

O arquivo requirements.txt ficará assim:

requirements.txt

1fastapi>=0.104.0
2uvicorn>=0.24.0
3torch>=2.0.0
4transformers>=4.35.0
5pydantic>=2.0.0

Arquivo de Armazenamento de API Keys

O arquivo apikey_store.json será usado para armazenar as chaves de API geradas. Vamos iniciá-lo vazio, contendo apenas {}.

apikey_store.json

1{}

Schemas e validação de dados

Os schemas são essenciais para validar os dados de entrada e saída da API. Eles garantem que os dados estejam no formato correto e permitem a geração automática de documentação.

Vamos criar o arquivo app/schemas.py com todos os modelos de dados. Teremos quatro modelos: GenerateRequest, LoadModelRequest, ApiKeyResponse e LDAPUserRequest.

schemas.py

 1from pydantic import BaseModel, Field
 2from typing import Optional
 3
 4class GenerateRequest(BaseModel):
 5    model_name: str = Field(..., description="The name of the model to use for generation.")
 6    prompt: str = Field(..., description="The input text to generate a response for.")
 7    max_tokens: Optional[int] = Field(300, description="The maximum length of the generated response.")
 8    temperature: Optional[float] = Field(1.0, description="The sampling temperature for generation.")
 9    top_p: Optional[float] = Field(1.0, description="The cumulative probability for nucleus sampling.")
10    hf_token: Optional[str] = Field(None, description="The Hugging Face tokenizer to use, if applicable.")
11
12
13class LoadModelRequest(BaseModel):
14    model_name: str = Field(..., description="The name of the model to load.")
15    device: Optional[str] = Field("cuda", description="The device to load the model on (e.g., 'cpu', 'cuda').")
16    hf_token: Optional[str] = Field(None, description="The Hugging Face tokenizer to use, if applicable.")
17
18class ApiKeyResponse(BaseModel):
19    api_key: str = Field(..., description="The API key for accessing the model API.")
20
21class LDAPUserRequest(BaseModel):
22    username: str = Field(..., description="The username for LDAP authentication.")
  • Todas as classes herdam da classe BaseModel da biblioteca pydantic, obtendo funcionalidades de validação, serialização e documentação automática.
  • O campo Field(...) define um campo obrigatório sem valor padrão.
  • O campo Field(value) define um campo obrigatório com value como valor padrão.
  • O tipo Optional[type] indica que o campo é opcional, mas deve ser do tipo type se fornecido.

Com os schemas definidos, vamos criar o arquivo responsável pela autenticação via API Key.

Autenticação e API Keys

O sistema de autenticação protege a API, garantindo que apenas usuários autorizados possam acessar os endpoints. Vamos implementar um mecanismo baseado em API Keys.

Vamos criar o arquivo app/auth.py com todas as funcionalidades de autenticação.

auth.py

 1import secrets 
 2import json
 3from fastapi import HTTPException, Request
 4
 5APIKEY_STORE_FILE = "app/apikey_store.json"
 6
 7def load_apikeys():
 8    try:
 9        with open(APIKEY_STORE_FILE, "r") as f:
10            return json.load(f)
11    except FileNotFoundError:
12        raise HTTPException(
13            status_code=404,
14            detail=f"Arquivo de API keys não encontrado: {APIKEY_STORE_FILE}")
15    
16def save_apikeys(keys: dict):
17    with open(APIKEY_STORE_FILE, "w") as f:
18        json.dump(keys, f, indent=4)
19
20def generate_apikey(user:str) -> str:
21    key = secrets.token_hex(32)
22    keys = load_apikeys()
23    keys[user] = key
24    save_apikeys(keys)
25    return key
26
27async def verify_apikey(request: Request) -> bool:
28    apikey = request.headers.get("x-API-Key")
29    if not apikey:
30        raise HTTPException(
31            status_code=401,
32            detail="API key não fornecida.")
33    try:
34        keys = load_apikeys()
35        if apikey in keys.values():
36            return True
37    
38    except json.JSONDecodeError:
39        raise HTTPException(
40        status_code=403,
41        detail="API key inválida.")
  • A função load_apikeys carrega as informações armazenadas no arquivo app/apikey_store.json.
  • save_apikeys é responsável por salvar o conteúdo no formato JSON.
  • A função generate_apikey cria uma chave para um usuário e a adiciona ao dicionário, usando o username como chave.
  • verify_apikey será chamada sempre que uma requisição chegar, para realizar a validação.

Gerenciador de Modelos e GPU

O app/model_manager.py é o coração da API, responsável por carregar, gerenciar e executar os modelos de linguagem. Ele otimiza o uso de GPU/CPU e garante eficiência na geração do texto.

model_manager.py

 1import torch 
 2from transformers import AutoTokenizer, AutoModelForCausalLM
 3from fastapi import HTTPException
 4import gc
 5from .utils import is_model_on_gpu
 6
 7DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 8
 9class ModelManager:
10    def __init__(self):
11        self.model = None
12        self.tokenizer = None
13        self.model_name = None
14
15    def load_model(self, model_name: str, hf_token:str = None, device: str = DEVICE):
16        if self.model_name != None and self.model_name != model_name:
17            print("Removendo modelo carregado anteriormente...")
18
19        self.unload_model()        
20        print(f"Carregando modelo {model_name} no dispositivo {device}...")
21       
22        if self.model_name != model_name:
23            try:            
24                if hf_token:           
25                    self.tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
26                    self.model = AutoModelForCausalLM.from_pretrained(model_name, device_map="balanced", token=hf_token)
27                else:
28                    self.tokenizer = AutoTokenizer.from_pretrained(model_name)
29                    self.model = AutoModelForCausalLM.from_pretrained(model_name, device_map="balanced")
30                self.model.eval()
31                self.model_name = model_name
32                print(is_model_on_gpu(self.model.hf_device_map, self.model_name))
33                
34            except Exception as e:
35                raise HTTPException(status_code=500, detail=f"Erro ao carregar modelo: {str(e)}")
36        else:
37            print(f"O modelo {model_name} já está carregado.")
38
39    def generate(self, model_name:str, hf_token: str, prompt:str, max_tokens:int = 300, temperature:float = 1.0, top_p:float = 1.0) -> str:
40        
41        if self.model_name != model_name:
42            self.load_model(model_name, hf_token, device=DEVICE)
43
44        if self.model is None or self.tokenizer is None:
45            raise HTTPException(status_code=400, detail="Nenhum modelo carregado.")
46
47        try:
48            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
49            with torch.no_grad(): 
50                outputs = self.model.generate(**inputs, max_new_tokens=max_tokens,temperature=temperature, top_p=top_p, eos_token_id=self.tokenizer.eos_token_id)
51            return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
52        except Exception as e:
53            raise HTTPException(status_code=500, detail=f"Erro ao gerar texto: {str(e)}")
54    
55    def get_status(self) -> str:        
56        if self.model is None:
57            self.unload_model()
58            return "Nenhum modelo carregado."       
59        return is_model_on_gpu(self.model.hf_device_map, self.model_name)
60
61    def unload_model(self):
62        self.model = None
63        self.tokenizer = None
64        old_model = self.model_name if self.model_name else False
65        self.model_name = None
66
67        gc.collect()
68        torch.cuda.empty_cache()
69        return f"Modelo {old_model} descarregado com sucesso." if old_model else "Nenhum modelo carregado para descarregar."
70
71manager = ModelManager()
  • A função load_model carrega o novo modelo na memória, removendo algum modelo que foi carregado anteriormente.
  • generate é a principal função da API, ela é responsável por realizar a inferência do modelo. Permite alterar os parâmetros: temperature, top_p e max_tokens.
  • get_status é responsável por informar se existe modelo carregado e se está em GPU ou CPU.
  • A função unload_model remove o modelo da memória, limpando o cache do CUDA e utilizando o garbage collector do python para não restar resquícios que possam atrapalhar futuros carregamentos.

Endpoints da API FastAPI

O arquivo app/main.py é onde todos os componentes se conectam. Nele definimos todos os endpoints e a lógica de roteamento da API.

main.py

 1from fastapi import FastAPI, Request, HTTPException, Depends
 2from fastapi.responses import JSONResponse
 3from app import schemas, model_manager, auth
 4
 5app = FastAPI()
 6
 7async def require_api_key(request: Request) -> schemas.LDAPUserRequest:
 8    user = await auth.verify_apikey(request)
 9    if not user:
10        raise HTTPException(status_code=401, detail="API key invalida.")
11    return user
12
13@app.post("/generate_apikey")
14async def generate_apikey(payload: schemas.LDAPUserRequest) -> JSONResponse:
15    key = auth.generate_apikey(payload.username)
16    return JSONResponse(status_code=200, content={"api_key": key})
17
18@app.post("/load_model", dependencies=[Depends(require_api_key)])
19async def load_model(payload: schemas.LoadModelRequest) -> JSONResponse:
20    try:
21        model_manager.manager.load_model(payload.model_name, payload.hf_token, payload.device)
22        return JSONResponse(content={"message": f"Modelo {payload.model_name} carregado com sucesso."})
23    except Exception as e:
24        raise HTTPException(status_code=500, content={"error": str(e)})
25    
26@app.post("/generate", dependencies=[Depends(require_api_key)])
27async def generate(payload: schemas.GenerateRequest)-> JSONResponse:
28    try:
29        result = model_manager.manager.generate(payload.model_name, payload.hf_token,payload.prompt, payload.max_tokens, payload.temperature, payload.top_p)
30        return {"result": result}
31    except Exception as e:
32        return JSONResponse(status_code=500, content={"error": str(e)})
33    
34@app.get("/status", dependencies=[Depends(require_api_key)])
35async def status()-> JSONResponse:
36    str_status = model_manager.manager.get_status()
37    return JSONResponse(content={"status": str_status})
38
39@app.post("/unload_model", dependencies=[Depends(require_api_key)])
40async def unload_model() -> JSONResponse:
41    try:
42        str_unload = model_manager.manager.unload_model()
43        return JSONResponse(content={"message":str_unload})
44    except Exception as e:
45        raise HTTPException(status_code=500, content={"error": str(e)})
  • A função require_api_key verifica a API Key sempre que chega uma requisição e retorna o usuário autenticado ou gera erro 401.
  • generate_apikey gera e retorna uma nova chave de API para o usuário informado.
  • load_model carrega o modelo especificado. Caso o modelo necessite de um token Hugging Face, a função também recebe esse parâmetro.
  • A função generate é responsável por fazer o modelo realizar a inferência a partir do prompt e os parâmetros passados.
  • Ao chamar o endpoint status o usuário recebe o status atual do gerenciador de modelos.
  • unload_model descarrega o modelo atualmente carregado e retorna uma mensagem de sucesso caso tenha concluído corretamente.

Arquivo utils.py

O arquivo app/utils.py contém a função que verifica se o modelo carregado está totalmente/parcialmente em GPU ou foi carregado em CPU.

utils.py

1def is_model_on_gpu(hf_device_map: dict, model_name: str) -> str:
2    if '' in hf_device_map.keys() and hf_device_map[''] == 'cpu':
3        return f"Modelo {model_name} carregado totalmente na CPU."
4    elif 'cpu' in hf_device_map.values():
5        return f"Algumas camadas do modelo {model_name} estão carregadas na CPU."
6    else:
7        return f"Modelo {model_name} carregado totalmente na GPU."

Executando a API

Para executar a API com o uvicorn é muito simples, basta executar um comando com as informações de host e porta para o serviço rodar.

uvicorn app.main:app --host 0.0.0.0 --port 8000 --reload

  • app:main se refere ao arquivo app/main.py responsável por conectar todos os componentes e receber as requisições realizadas pelo usuário.

  • --host 0.0.0.0 define o endereço IP no qual o servidor Uvicorn irá escutar as requisições. O valor 0.0.0.0 define que este servidor estará acessível de qualquer interface de rede disponível na máquina Power9.

  • --port 8000 especifica a porta na qual o servidor irá escutar as requisições.

  • --reload flag para ser utilizada em desenvolvimento. Recarrega a aplicação sempre que uma mudança é realizada.

Seguindo estas implementações, você terá uma API capaz de realizar inferências com Modelos de Linguagem baixados do Hugging Face. No próximo tutorial será demonstrado como enviar requisições para a API via curl e python.