[wip] Custom session handlers #25
Open
+3,286
−220
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
… SAGEMAKER_ as prefix.
…ds into toggle-sticky-routing
…is None but request is a session request, raise 400 error - change session api transform's transform_request to never return request field in output (previously used to pass session_manager - change handler code to only take raw_request as a parameter and use new utility function get_session_manager - add new integration tests for expected errors when sessions is disabled
…ds into toggle-sticky-routing
…eate/close session apis
|
Logic looks good to me in general. I want to understand the CX before we move further. Can you add docs/examples to help understand the CX ? |
…ds into toggle-sticky-routing Signed-off-by: Zuyi Zhao <[email protected]>
…container-standards into custom-session-handlers Signed-off-by: Zuyi Zhao <[email protected]>
…ds into custom-session-handlers Signed-off-by: Zuyi Zhao <[email protected]>
- Fix session ID injection logic to work with default session manager - Changed conditional from elif to separate if statements - Session ID now properly injected into request body when session_id_path is specified - Validation and injection can now happen independently - Update unit tests to match refactored transform API - Remove import of non-existent _validate_session_if_present function - Update test signatures to match _process_session_request parameters - Add missing SessionRequest import - Replace outdated test logic with current implementation - Add integration tests for session_id_path injection feature - Test session ID injection into flat and nested request body paths - Test injection with multiple requests and different sessions - Verify existing body fields are preserved during injection
- Rename session_id_path to request_session_id_path in stateful_session_manager - Rename session_id_path to response_session_id_path in register handlers - Add docstring clarifying request_session_id_path usage - Provide default content messages for create/close session handlers - Remove specific engine references (vLLM, TGI) from documentation - Delete obsolete test_registration.py - Improve integration tests
…ds into custom-session-handlers
Lokiiiiii
requested changes
Dec 8, 2025
Lokiiiiii
left a comment
Lokiiiiii
left a comment
There was a problem hiding this comment.
Summarizing -
- The parameter names for the handlers are un-intuitive. I recommend renaming it to be obvious in terms of the functionality or the input the developer needs to provide -
path_to_session_id_in_incoming_requestor something similar. - The documentation has inconsistent ordering of decorators. This will cause issues when overriding handlers.
- We need to improve the defaults and documentation needs to indicate defaults for handler arguments
| # Register custom close session handler | ||
| @register_close_session_handler( | ||
| request_shape={ | ||
| "session_id": 'headers."X-Amzn-SageMaker-Session-Id"' # Extract from header |
There was a problem hiding this comment.
Within sagemaker.sessions X-Amzn-SageMaker-Session-Id should be the default. We should remove this from the documentation for ease of understanding.
| capacity: int | ||
| session_id: Optional[str] = None | ||
|
|
||
| @register_create_session_handler( |
There was a problem hiding this comment.
I see inconsistent ordering of decorators.
I believe it is supposed to be reversed ie.., handlers should be registered before FastAPI route configurations ?
| # Register custom create session handler | ||
| @register_create_session_handler( | ||
| request_shape={ | ||
| "capacity": "`1024`" # JMESPath literal value |
There was a problem hiding this comment.
Is this a required new field ? Is there going to be a default ?
Lokiiiiii
reviewed
Dec 8, 2025
|
|
||
| @register_create_session_handler( | ||
| request_shape={ | ||
| "capacity": "`1024`", |
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.
Issue #, if available:
Description of changes:
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.