integrate zmq

[aniketmaurya]

What does this PR do?

Before submitting

• Was this discussed/agreed via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure to update the docs?
• Did you write any new necessary tests?

Faster process communication with zmq.

TODO: A follow up PR to add a proxy and support multiple inference worker processes and multiple uvicorn processes.

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in GitHub issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

[bhimrazy]

Seems like a great addition, @aniketmaurya 🎉!
I remember llamastack also has zmq under the hood.
I'm looking forward to learning more about it and its uses over multiprocessing queues.
😊

[aniketmaurya]

thanks for taking a look @bhimrazy! yeah, zmq seem to improve the process communication time for sending the prediction results to the main process. This is gonna be very useful especially while serving streaming tokens from LLMs.

[codecov]

Codecov Report

Attention: Patch coverage is 61.06195% with 44 lines in your changes missing coverage. Please review.

Project coverage is 88%. Comparing base (747308a) to head (bf6969c).
Report is 2 commits behind head on main.

Additional details and impacted files

@@         Coverage Diff         @@
##           main   #403   +/-   ##
===================================
- Coverage    89%    88%   -1%     
===================================
  Files        30     30           
  Lines      1893   1976   +83     
===================================
+ Hits       1683   1734   +51     
- Misses      210    242   +32     

[lantiga]

The PR looks good in general, but not abstracting away the interprocess communication mechanism makes the code more complex (you have a socket rather than a queue so that has more opaque semantics), also a lot of details are guarded by if .. else conditions

I would take this opportunity to create a base class that takes care of the communication, and two concrete classes. You can then instantiate them in LitServer based on the one you want and you don't have to deal with conditionals and keep the code easily consumable.

Clear semantics is particularly important for people implementing loops, we need to keep it simple.

[aniketmaurya]

agree with your points @lantiga, I would put these inside the put_response method which can be used everywhere else. I will also create an encapsulation to hide these socket details for the next PR which would enable zmq for multiple workers.

[aniketmaurya]

creating followup as per above suggestion.