텍스트 임베딩 서버를 뛰워서 임베딩하기

RAG를 구현해 보려고, 수집기를 작성해서 24시간 동작시켜 놨는데
수집기에서 수집된 데이터를 엘라스틱서치에 색인 하는 과정에서 embedding을 하기 위해 emdebbing 모델을 한번 사용하고 나면 해당 모듈이 끝나고 난후에도 Cuda메모리 사용량이 줄어 들지가 않는다.

vectorstoredb = ~ 
del vectorstoredb 

이런식으로 해당 객체를 지워도 보고.

impoer gc
vectorstoredb = ~ 
del vectorstoredb 
gc.collect()

이런식으로 가비지 컬렉터를 명시적으로 호출을 해봐도 도무지 내려가지가 않아서.

수집기 내부에서 직접적으로 cuda 자원을 사용하는 방식을 포기하고, 외부 api 호출 방식으로 바꾸기로 함.

임베딩이 가능한 외부 api를 찾아보니
tei(Text Embedding Inference)라는 도커 서비서를 이용해서 사용이 가능하다고 하여 기존에 쓰고 있는 BAAI/bge모델을 tei에 올려서 임베딩에 활용.

 bge-embedding:
    container_name: bge-embedding
    image: ghcr.io/huggingface/text-embeddings-inference:latest
    platform: linux/amd64
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
    command: --model-id BAAI/bge-m3
    environment:
      - CORS_ALLOW_ORIGIN=http://0.0.0.0:3000
      - PORT=3000
      - MAX_CLIENT_BATCH_SIZE=32
      - USE_FLASH_ATTENTION=True
    volumes:
      - ./embedding:/data
    ports:
      - "3000:3000"
    networks:
      - ps_network

기존 수집기가 celery, django. 조합으로 된 docker-compose 구성이어서 tei 도커 구성을 추가.

해당 임베딩 api를 사용하는 루틴으로 임베딩 모듈 수정.

text_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=250)
        embeddings = LocalAIEmbeddings(openai_api_base='http://localhost:3000', model="BAAI/bge-m3")

        vectorstore = ElasticsearchStore(
            index_name=settings.NEWS_EMBEDDING_INDEX,
            embedding=embeddings,
            es_connection=es
        )

정상적으로 동작 확인.

위와 같이 했을때 임베딩 이후 cuda 사용량에 변화는 없지만, 해당 도커가 항상 떠 있어서 메모리 점유율이 항상 100%가 됨..

스컈쥴러나 크론탭을 통해 해당 도커를 올렸다가 내리는 모듈도 만들어야 될듯.