Optimize Memory Handling in Embedding Computations and Refactor EmbeddingService #103
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Enhanced the embedding computation function to dynamically manage RAM by estimating the maximum embeddings that can fit and automatically saving them to optimize memory usage
- Extract core embedding service logic into embedding_service method - Add special case handling for ultra-long reads: - Immediately save ultra-long read embeddings to disk - Avoid loading additional sequences when ultra-long read detected - Dynamically calculate max embeddings that fit in available memory - Use this to determine when to flush embeddings to disk - Code cleanup: - Use type hints for improved readability - Docstrings for key methods
SebieF
added
enhancement
SebieF
requested changes
Aug 19, 2024
There was a problem hiding this comment.
Very nice work, thank you for looking into this quite complex problem! :)
Most of my comments are minor suggestions for improvements, the only major change that I would suggest is to unify the handling of the ultra-long reads and normal reads if at all possible, and removing the while loops altogether.
Other than that, the approach seems to work fine, I ran the tests locally and executed a run with ProtT5 for 2377 sequences successfully.
…ogging - Renamed embedding_service to _do_embeddings_computation for clarity. - Replaced while loop with for loop to prevent infinite loops and simplify processing. - Combined handling of ultra-long and normal reads into a single loop to reduce complexity. - Updated logging levels for better clarity and reduced unnecessary logs. - Optimized garbage collection by consolidating gc.collect() calls.
…gement - Replacing tearDown and manual cleanup with tempfile.TemporaryDirectory() for automatic resource management. - Adjusting long_length calculation to use a fixed value for local testing and a memory-based value for CI environments. - Enhancing _run_embedding_test to support both sequence_to_class and residue_to_class protocols. - Updating _verify_result to validate output based on the protocol used. - Adding new tests for comprehensive coverage of both embedding protocols.
|
Thank you for the suggestions on improving the changes I made. I just updated the code based on your feedback! :) |
SebieF
requested changes
Aug 25, 2024
- Moving progress bar initialization to ensure visibility during the first embedding calculation and updated its description. - Consolidating memory cleanup: moving del self._embedder and gc.collect() to avoid repeated calls. - Updating docstring in _max_embedding_fit to explain constants (0.75 and 18).
|
Sebastian, I just updated the code as you suggested 😉 |
Merged
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Description:
This PR enhances the
EmbeddingServiceby optimizing memory usage during embedding computations, refactoring core logic, and adding unit tests for robustness.Key Changes:
Dynamic Memory Management:
Refactor
EmbeddingService:SAVE_AFTER_N_EMBEDDINGS.Unit Testing:
Issue Reference:
Closes [#98]