feat: add grpc support for the huggingface text embedding inference (TEI) components

[maxdswain]

Related Issues

• fixes Support gRPC interface of TEI haystack#7522

Proposed Changes:

Add grpc support through optional grpc extras for the hugging face text embedding inference (TEI). To support this implementation, I generated a grpc client from the latest TEI protobuf. Instructions on how to regenerate the files from the protobuf have been added to the README.

This needed to be done as the existing huggingface hub inference client does not support grpc, and does not intend to.

How did you test it?

Added unit and integration tests. You need to run docker compose up -d before running the integration tests.

Notes for the reviewer

The TEI grpc API doesn't have an API where it takes a list of inputs to embed, so instead for the document embedder components, I used the EmbedStream interface. In my testing, this is faster then using multiple streams one after the other and batching requests.

I added grpc client initialisation in the warm_up/warm_up_async methods and client closing in the close/close_async methods to align with the recent changes to lifecycle management for HTTP clients.

Checklist

• I have read the contributors guidelines and the code of conduct
• I have updated the related issue with new insights and changes
• I added unit tests and updated the docstrings
• I've used one of the conventional commit types for my PR title: fix:, feat:, build:, chore:, ci:, docs:, style:, refactor:, perf:, test:.

[github-actions]

Heads-up for maintainers

This PR is from a fork and touches integrations whose integration tests require API keys.
Those tests are skipped in CI because fork PRs don't have access to repo secrets for security reasons.

Affected integrations:

• huggingface_api

Please run the integration tests locally (hatch run test:integration inside each folder) before approving.

[socket-security]

Review the following changes in direct dependencies. Learn more about Socket for GitHub.

Diff	Package	Supply Chain Security	Vulnerability	Quality	Maintenance	License
	pypi/​grpcio@​1.84.0

View full report

[github-actions]

Coverage report (huggingface_api)

Click to see where and how coverage changed

File	Statements	Missing	Coverage	Coverage (new stmts)	Lines missing
integrations/huggingface_api/src/haystack_integrations/components/embedders/huggingface_api
document_embedder.py					190-191, 341, 358-359, 392-421, 485, 494, 523, 532
sparse_document_embedder.py					200, 224, 248-249
sparse_text_embedder.py
text_embedder.py					153-154, 301-305, 342-346
integrations/huggingface_api/src/haystack_integrations/components/embedders/huggingface_api/_grpc
tei_pb2.py					33-97
tei_pb2_grpc.py					15-16, 19, 37, 49-51, 55-65, 83, 146-148, 152-154, 158-160, 164-166, 170-172, 176-178, 182-217, 235, 262, 289, 316, 343, 370, 396-411, 423-425, 429-431, 435-437, 441-443, 447-472, 490, 517, 544, 571, 597-602, 614-616, 620-622, 626-641, 659, 686, 712-727, 739-741, 745-747, 751-753, 757-759, 763-788, 806, 833, 860, 887
Project Total

This report was generated by python-coverage-comment-action

[maxdswain]

Sorry to bother you @bogdankostic, this should be ready to review and I've resolved the previously existing merge conflicts.

[anakin87]

Hey @maxdswain, thanks for this contribution!

Since this requires much more code to maintain, I'd like to see reproducible proof that it improves performance significantly compared to HTTP.

I quickly tried benchmarking this myself but I could not find conclusive results. Could you try and let me know?

[maxdswain]

Hey @maxdswain, thanks for this contribution!

Since this requires much more code to maintain, I'd like to see reproducible proof that it improves performance significantly compared to HTTP.

I quickly tried benchmarking this myself but I could not find conclusive results. Could you try and let me know?

Hi @anakin87, thanks for the reply. Most of the code in this PR is in the _grpc directory which is generated from the grpc protobuf so shouldn't require much of a maintenance burden. I've documented how to re-generate this on .proto changes in the README.md.

curl -Lo src/haystack_integrations/components/embedders/huggingface_api/_grpc/tei.proto https://raw.githubusercontent.com/huggingface/text-embeddings-inference/main/proto/tei.proto
pip install grpcio-tools mypy-protobuf
python -m grpc_tools.protoc --proto_path=src --python_out=src --pyi_out=src --grpc_python_out=src --mypy_grpc_out=src haystack_integrations/components/embedders/huggingface_api/_grpc/tei.proto

I've found there to be large performance differences at high concurrency (i.e., when ingesting documents or when many users are asking questions) when running on GPUs, as this is when inference time takes up a much smaller portion of the request time. I contribute to OS on a old personal laptop without access to GPUs so unfortunately I'm unable to give you more solid evidence on the performance differences.

[anakin87]

I see, thanks for clarifying.

Could you try exploring server reflection as an alternative? We might be able to build the client dynamically and avoid shipping and regenerating the protobuf bindings when TEI changes.

I tried a prototype using grpc-requests.

Could you try doing something similar and see if the performance benefits still hold? This solution would be simpler and more maintainable.