fix(infra): fix some dependency hells and add some lazy loading to reduce celery worker RAM usage #2
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
https://linear.app/danswer/issue/DAN-2573/move-more-imports-onto-lazy-loading
so i have a script that i run to check memory usage for celery workers,
before this pr its ~600MB per worker
after its ~250 MB for some workers
docker container from 4.3GB -> 2.819GiB
to diagnose why, its not easy - its not that all pip dependencies are loaded into memory at worker start so i can just lazy load any, its specifically ones that get imported at runtime due to an actual import statement
this makes it very tricky to track down exactly what causes the 600 MB. literally had to trial and error suspiscious imports tracing starting from the celery worker main file
TBH the existing repo dependency graph is a little scuffed, one example branch that caused the worker to import llm stuff (there are like a dozen of these had to swift through to get all the memory offenders down):
app base
->
/Users/edwinluo/onyx/backend/onyx/background/celery/tasks/docprocessing/utils.py
->
redis connector
->
/Users/edwinluo/onyx/backend/onyx/redis/redis_connector_delete.py
->
/Users/edwinluo/onyx/backend/onyx/db/document.py
->
/Users/edwinluo/onyx/backend/onyx/db/feedback.py
->
/Users/edwinluo/onyx/backend/onyx/db/chat.py
->
/Users/edwinluo/onyx/backend/onyx/context/search/utils.py
->
/Users/edwinluo/onyx/backend/onyx/db/search_settings.py
->
/Users/edwinluo/onyx/backend/onyx/db/llm.py OR /Users/edwinluo/onyx/backend/onyx/natural_language_processing/search_nlp_models.py
->
/Users/edwinluo/onyx/backend/onyx/llm/utils.py (langchian, litellm, etc.)
How Has This Been Tested?
[Describe the tests you ran to verify your changes]
Backporting (check the box to trigger backport action)
Note: You have to check that the action passes, otherwise resolve the conflicts manually and tag the patches.
Summary by cubic
Moves heavy imports to lazy-loading across indexing, LLM, NLP, and file-processing code to reduce worker memory and speed up startup. Also consolidates search doc conversion into SearchDoc and extracts PromptSnapshot to a shared schema. Addresses Linear DAN-2573 (Reduce Memory usage in Onyx).
Refactors
Migration