[API] Improve async functions #9039
Merged
Conversation
This file contains hidden or 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
…d improve log collection handling
There was a problem hiding this comment.
Pull request overview
This PR improves async function handling and database session lifecycle management by introducing async context managers for database sessions. The main goal is to prevent potential resource leaks and improve code maintainability through proper session management patterns.
Key Changes:
- Added
get_db_session_async()andget_db_session()context managers for automatic session commit/rollback/cleanup - Refactored
run_async_function_with_new_db_session()to support both coroutine functions and sync functions in async contexts using the new context managers - Replaced all
fastapi.concurrency.run_in_threadpoolcalls withmlrun.utils.run_in_threadpoolfor better standardization - Added
run_with_time_window_tracker_sync()to handle synchronous periodic functions with time window tracking
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
server/py/framework/db/session.py |
Introduces async and sync context managers for database sessions; refactors session management functions to use context managers |
server/py/framework/utils/time_window_tracker.py |
Refactors async time window tracker to use new session management; adds sync version of time window tracker function |
server/py/services/api/main.py |
Refactors log collection and monitoring functions to use async/sync context managers; replaces fastapi.concurrency with mlrun.utils; converts _monitor_runs from async to sync with proper session handling |
server/py/services/alerts/main.py |
Converts _generate_events from async to sync to work with new sync time window tracker |
server/py/framework/service/__init__.py |
Updates imports to use mlrun.utils.run_in_threadpool instead of fastapi.concurrency |
server/py/services/api/tests/unit/test_collect_runs_logs.py |
Adds comprehensive test coverage for async session usage in log collection and verification functions |
server/py/framework/tests/unit/db/test_session.py |
New test file with comprehensive tests for async session management functions |
pipeline-adapters/mlrun-pipelines-kfp-v1-8/src/mlrun_pipelines/client.py |
Removes verbose debug logging statements |
pipeline-adapters/mlrun-pipelines-kfp-v1-8/pyproject.toml |
Version bump to 0.6.2 |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
TomerShor
reviewed
Dec 11, 2025
Comment on lines
123
to
130
| await framework.db.session.run_async_function_with_new_db_session( | ||
| callback, last_update_time, *args, **kwargs | ||
| ) | ||
| await run_in_threadpool( | ||
| framework.db.session.run_function_with_new_db_session, | ||
| cycle_tracker.update_window, | ||
| now, | ||
| ) |
Member
There was a problem hiding this comment.
Can't both of those calls run with the same db session? why do we need a new one for each?
server/py/services/api/main.py
Outdated
| async def _base_handler( | ||
| self, | ||
| request: fastapi.Request, | ||
| request, |
Member
Author
There was a problem hiding this comment.
removed the import, I will add for type checking
…indow_update parameter and improve async handling
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 9 out of 9 changed files in this pull request and generated 5 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
pipeline-adapters/mlrun-pipelines-kfp-v1-8/src/mlrun_pipelines/client.py
Show resolved
Hide resolved
pipeline-adapters/mlrun-pipelines-kfp-v1-8/src/mlrun_pipelines/client.py
Show resolved
Hide resolved
TomerShor
approved these changes
Dec 14, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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 addresses ML-11432 by ensuring proper session lifecycle management in async contexts, preventing potential resource leaks and improving code maintainability.
🛠️ Changes Made
get_db_session_async()async context manager for proper database session management in async contexts_initiate_logs_collection()to use async session context manager instead of manual session handlingrun_async_function_with_new_db_session()to support both coroutine functions and sync functions in async contextsTimeWindowTrackerto work with async session context managers✅ Checklist
🧪 Testing
test_collect_runs_logs.pycovering:🔗 References
🚨 Breaking Changes?