fastembed/tests/test_multi_gpu.py

import pytest
from fastembed import (
    TextEmbedding,
    SparseTextEmbedding,
    LateInteractionTextEmbedding,
    ImageEmbedding,
)
from fastembed.rerank.cross_encoder import TextCrossEncoder
from tests.config import TEST_MISC_DIR

CACHE_DIR = "../model_cache"


@pytest.mark.skip(reason="Requires a multi-gpu server")
@pytest.mark.parametrize("device_id", [None, 0, 1])
def test_gpu_via_providers(device_id) -> None:
    docs = ["hello world", "flag embedding"]

    device_id = device_id if device_id is not None else 0
    providers = (
        ["CUDAExecutionProvider"]
        if device_id is None
        else [("CUDAExecutionProvider", {"device_id": device_id})]
    )
    embedding_model = TextEmbedding(
        "sentence-transformers/all-MiniLM-L6-v2",
        providers=providers,
        cache_dir=CACHE_DIR,
    )
    list(embedding_model.embed(docs))
    options = embedding_model.model.model.get_provider_options()
    assert options["CUDAExecutionProvider"]["device_id"] == str(device_id)

    embedding_model = SparseTextEmbedding(
        "prithvida/Splade_PP_en_v1",
        providers=providers,
        cache_dir=CACHE_DIR,
    )
    list(embedding_model.embed(docs))
    options = embedding_model.model.model.get_provider_options()
    assert options["CUDAExecutionProvider"]["device_id"] == str(device_id)

    embedding_model = SparseTextEmbedding(
        "Qdrant/bm42-all-minilm-l6-v2-attentions",
        providers=providers,
        cache_dir=CACHE_DIR,
    )
    list(embedding_model.embed(docs))
    options = embedding_model.model.model.get_provider_options()
    assert options["CUDAExecutionProvider"]["device_id"] == str(device_id)

    embedding_model = LateInteractionTextEmbedding(
        "colbert-ir/colbertv2.0",
        providers=providers,
        cache_dir=CACHE_DIR,
    )
    list(embedding_model.embed(docs))
    options = embedding_model.model.model.get_provider_options()
    assert options["CUDAExecutionProvider"]["device_id"] == str(device_id)

    embedding_model = ImageEmbedding(
        model_name="Qdrant/clip-ViT-B-32-vision",
        providers=providers,
        cache_dir=CACHE_DIR,
    )
    images = [
        TEST_MISC_DIR / "image.jpeg",
        str(TEST_MISC_DIR / "small_image.jpeg"),
    ]
    list(embedding_model.embed(images))
    options = embedding_model.model.model.get_provider_options()
    assert options["CUDAExecutionProvider"]["device_id"] == str(device_id)

    model = TextCrossEncoder(
        model_name="Xenova/ms-marco-MiniLM-L-6-v2",
        providers=providers,
        cache_dir=CACHE_DIR,
    )
    query = "What is the capital of France?"
    documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."]
    list(model.rerank(query, documents))
    options = embedding_model.model.model.get_provider_options()
    assert options["CUDAExecutionProvider"]["device_id"] == str(device_id)


@pytest.mark.skip(reason="Requires a multi-gpu server")
@pytest.mark.parametrize("device_ids", [None, [0], [1], [0, 1]])
def test_gpu_cuda_device_ids(device_ids) -> None:
    docs = ["hello world", "flag embedding"]
    device_id = device_ids[0] if device_ids else 0
    embedding_model = TextEmbedding(
        "sentence-transformers/all-MiniLM-L6-v2",
        cuda=True,
        device_ids=device_ids,
        cache_dir=CACHE_DIR,
    )
    list(embedding_model.embed(docs))
    options = embedding_model.model.model.get_provider_options()
    assert options["CUDAExecutionProvider"]["device_id"] == str(
        device_id
    ), f"Text embedding: {options}"

    embedding_model = SparseTextEmbedding(
        "prithvida/Splade_PP_en_v1",
        cuda=True,
        device_ids=device_ids,
        cache_dir=CACHE_DIR,
    )
    list(embedding_model.embed(docs))
    options = embedding_model.model.model.get_provider_options()
    assert options["CUDAExecutionProvider"]["device_id"] == str(
        device_id
    ), f"Sparse text embedding: {options}"

    embedding_model = SparseTextEmbedding(
        "Qdrant/bm42-all-minilm-l6-v2-attentions",
        cuda=True,
        device_ids=device_ids,
        cache_dir=CACHE_DIR,
    )
    list(embedding_model.embed(docs))
    options = embedding_model.model.model.get_provider_options()
    assert options["CUDAExecutionProvider"]["device_id"] == str(device_id), f"Bm42: {options}"

    embedding_model = LateInteractionTextEmbedding(
        "colbert-ir/colbertv2.0",
        cuda=True,
        device_ids=device_ids,
        cache_dir=CACHE_DIR,
    )
    list(embedding_model.embed(docs))
    options = embedding_model.model.model.get_provider_options()
    assert options["CUDAExecutionProvider"]["device_id"] == str(
        device_id
    ), f"Late interaction text embedding: {options}"

    embedding_model = ImageEmbedding(
        model_name="Qdrant/clip-ViT-B-32-vision",
        cuda=True,
        device_ids=device_ids,
        cache_dir=CACHE_DIR,
    )
    images = [
        TEST_MISC_DIR / "image.jpeg",
        str(TEST_MISC_DIR / "small_image.jpeg"),
    ]
    list(embedding_model.embed(images))
    options = embedding_model.model.model.get_provider_options()
    assert options["CUDAExecutionProvider"]["device_id"] == str(
        device_id
    ), f"Image embedding: {options}"

    if device_ids is None or len(device_ids) == 1:
        model = TextCrossEncoder(
            model_name="Xenova/ms-marco-MiniLM-L-6-v2",
            cuda=True,
            device_ids=device_ids,
            cache_dir=CACHE_DIR,
        )
        query = "What is the capital of France?"
        documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."]
        list(model.rerank(query, documents))
        options = embedding_model.model.model.get_provider_options()
        assert options["CUDAExecutionProvider"]["device_id"] == str(
            device_id
        ), f"Text cross encoder: {options}"


@pytest.mark.skip(reason="Requires a multi-gpu server")
@pytest.mark.parametrize(
    "device_ids,parallel", [(None, None), (None, 2), ([1], None), ([1], 1), ([1], 2), ([0, 1], 2)]
)
def test_multi_gpu_parallel_inference(device_ids, parallel) -> None:
    docs = ["hello world", "flag embedding"] * 100
    batch_size = 5

    embedding_model = TextEmbedding(
        "sentence-transformers/all-MiniLM-L6-v2",
        cuda=True,
        device_ids=device_ids,
        cache_dir=CACHE_DIR,
        lazy_load=True,
    )
    list(embedding_model.embed(docs, batch_size=batch_size, parallel=parallel))

    embedding_model = SparseTextEmbedding(
        "prithvida/Splade_PP_en_v1",
        cuda=True,
        device_ids=device_ids,
        cache_dir=CACHE_DIR,
    )
    list(embedding_model.embed(docs, batch_size=batch_size, parallel=parallel))

    embedding_model = SparseTextEmbedding(
        "Qdrant/bm42-all-minilm-l6-v2-attentions",
        cuda=True,
        device_ids=device_ids,
        cache_dir=CACHE_DIR,
    )
    list(embedding_model.embed(docs, batch_size=batch_size, parallel=parallel))

    embedding_model = LateInteractionTextEmbedding(
        "colbert-ir/colbertv2.0",
        cuda=True,
        device_ids=device_ids,
        cache_dir=CACHE_DIR,
    )
    list(embedding_model.embed(docs, batch_size=batch_size, parallel=parallel))

    embedding_model = ImageEmbedding(
        model_name="Qdrant/clip-ViT-B-32-vision",
        cuda=True,
        device_ids=device_ids,
        cache_dir=CACHE_DIR,
    )
    images = [
        TEST_MISC_DIR / "image.jpeg",
        str(TEST_MISC_DIR / "small_image.jpeg"),
    ] * 100
    list(embedding_model.embed(images, batch_size=batch_size, parallel=parallel))