[python] Support session based sticky routing in async mode

Author: xyang16

.github/workflows/integration.yml

@@ -173,7 +173,9 @@ jobs:
           - test: TestCorrectnessTrtLlm
             instance: g6
             failure-prefix: trtllm
-
+          - test: TestStickyRouting
+            instance: g6
+            failure-prefix: lmi
     outputs:
       failure_cpu: ${{ steps.test-failure.outputs.failure_cpu }}
       failure_gpu: ${{ steps.test-failure.outputs.failure_gpu }}
@@ -271,4 +273,4 @@ jobs:
           ./stop_instance.sh $instance_id
 
           instance_id=${{ needs.create-runners.outputs.cpu_instance_id }}
-          ./stop_instance.sh $instance_id
+          ./stop_instance.sh $instance_id

…/djl_python/custom_formatter_handling.py …/python/setup/djl_python/base_handler.pyengines/python/setup/djl_python/custom_formatter_handling.py renamed to engines/python/setup/djl_python/base_handler.py engines/python/setup/djl_python/custom_formatter_handling.py renamed to engines/python/setup/djl_python/base_handler.py

@@ -12,10 +12,14 @@
 # the specific language governing permissions and limitations under the License.
 import logging
 
+from djl_python.inputs import Input
+from djl_python.outputs import Output
 from djl_python.service_loader import get_annotated_function
 
 logger = logging.getLogger(__name__)
 
+HEADER_SAGEMAKER_SESSION_ID = "X-Amzn-SageMaker-Session-Id"
+
 
 class CustomFormatterError(Exception):
     """Exception raised when custom formatter code fails"""
@@ -26,11 +30,12 @@ def __init__(self, message: str, original_exception: Exception):
         self.__cause__ = original_exception
 
 
-class CustomFormatterHandler:
+class BaseHandler:
 
     def __init__(self):
         self.output_formatter = None
         self.input_formatter = None
+        self.session_manager = None
 
     def load_formatters(self, model_dir: str):
         """Load custom formatters from model.py"""
@@ -79,3 +84,34 @@ async def apply_output_formatter_streaming_raw(self, stream_generator):
             logger.exception("Streaming formatter failed")
             raise CustomFormatterError(
                 "Custom streaming formatter execution failed", e)
+
+    async def create_session(self, inputs: Input):
+        outputs = Output()
+        try:
+            session = self.session_manager.create_session()
+            outputs.add_property(HEADER_SAGEMAKER_SESSION_ID,
+                                 session.session_id)
+            outputs.add_property("Content-Type", "application/json")
+            outputs.add(Output.binary_encode(
+                {"result": f"Session {session.session_id} created"}),
+                        key="result")
+        except Exception as e:
+            return Output().error("create_session_error", message=str(e))
+
+        logger.info(f"Session {session.session_id} created")
+        return outputs
+
+    async def close_session(self, inputs: Input):
+        outputs = Output()
+        session_id = inputs.get_property(HEADER_SAGEMAKER_SESSION_ID)
+        try:
+            self.session_manager.close_session(session_id)
+            outputs.add_property("Content-Type", "application/json")
+            outputs.add(Output.binary_encode(
+                {"result": f"Session {session_id} closed"}),
+                        key="result")
+        except Exception as e:
+            return Output().error("close_session_error", message=str(e))
+
+        logger.info(f"Session {session_id} closed")
+        return outputs

engines/python/setup/djl_python/lmi_vllm/vllm_async_service.py

@@ -31,7 +31,8 @@
 from djl_python.outputs import Output
 from djl_python.encode_decode import decode
 from djl_python.async_utils import handle_streaming_response, create_non_stream_output
-from djl_python.custom_formatter_handling import CustomFormatterHandler, CustomFormatterError
+from djl_python.base_handler import BaseHandler, CustomFormatterError
+from djl_python.session_manager import SessionManager
 
 from .request_response_utils import (
     ProcessedRequest,
@@ -46,8 +47,10 @@
 
 logger = logging.getLogger(__name__)
 
+HEADER_SAGEMAKER_SESSION_ID = "X-Amzn-SageMaker-Session-Id"
 
-class VLLMHandler(CustomFormatterHandler):
+
+class VLLMHandler(BaseHandler):
 
     def __init__(self):
         super().__init__()
@@ -119,12 +122,20 @@ async def initialize(self, properties: dict):
             tool_parser=self.vllm_properties.tool_call_parser,
             reasoning_parser=self.vllm_properties.reasoning_parser,
         )
+        self.session_manager: SessionManager = SessionManager(properties)
         self.initialized = True
 
     def preprocess_request(self, inputs: Input) -> ProcessedRequest:
         batch = inputs.get_batches()
         assert len(batch) == 1, "only one request per batch allowed"
         raw_request = batch[0]
+
+        # Get session id
+        session_id = raw_request.get_property(HEADER_SAGEMAKER_SESSION_ID)
+        session = self.session_manager.get_session(session_id)
+        if session is None:
+            raise RuntimeError(f"Requested session {session_id} not found")
+
         content_type = raw_request.get_property("Content-Type")
         decoded_payload = decode(raw_request, content_type)
 
@@ -226,6 +237,24 @@ async def inference(
 service = VLLMHandler()
 
 
+async def create_session(inputs: Input) -> Output:
+    if not service.initialized:
+        await service.initialize(inputs.get_properties())
+        logger.info("vllm service initialized")
+
+    outputs = await service.create_session(inputs)
+    return outputs
+
+
+async def close_session(inputs: Input) -> Output:
+    if not service.initialized:
+        await service.initialize(inputs.get_properties())
+        logger.info("vllm service initialized")
+
+    outputs = await service.close_session(inputs)
+    return outputs
+
+
 async def handle(
         inputs: Input
 ) -> Optional[Union[Output, AsyncGenerator[Output, None]]]:

engines/python/setup/djl_python/session_manager.py

@@ -58,7 +58,7 @@ def remove(self):
             shutil.rmtree(self.files_path)
             return True
         else:
-            logging.warning(f"session not found: {self.session_id}")
+            logging.warning(f"session directory not found: {self.session_id}")
             return False
 
     def _path(self, key: str):
@@ -80,6 +80,7 @@ def __init__(self, properties: dict):
         self.sessions_s3url = properties.get("sessions_s3url", None)
         if not os.path.exists(self.sessions_path):
             os.makedirs(self.sessions_path)
+        self.sessions: dict[str, Session] = {}
 
         atexit.register(self._save_sessions_to_s3)
 
@@ -95,25 +96,32 @@ def create_session(self) -> Session:
         session.put(".creation_time", time.time())
 
         self.cloud_watch.post("create_session")
+        self.sessions[session_id] = session
         return session
 
     def get_session(self, session_id: str) -> Optional[Session]:
         if not session_id or not UUID_PATTERN.match(session_id):
             raise ValueError(f"invalid session_id: {session_id}")
 
-        session = Session(session_id, self.sessions_path)
+        if session_id not in self.sessions:
+            raise ValueError(f"Session not found: {session_id}")
+        session = self.sessions[session_id]
+
         if not os.path.exists(session.files_path):
             return self._recover_from_s3(session)
 
         return session
 
     def close_session(self, session_id):
         if not session_id or not UUID_PATTERN.match(session_id):
-            raise ValueError(f"invalid session_id: {session_id}")
+            raise ValueError(f"Invalid session_id: {session_id}")
 
-        session = Session(session_id, self.sessions_path)
+        if session_id not in self.sessions:
+            raise ValueError(f"Session not found: {session_id}")
+        session = self.sessions[session_id]
         if session.remove():
             self.cloud_watch.post("close_session")
+        self.sessions.pop(session_id)
 
     def _clean_expired_session(self):
         sessions = os.listdir(self.sessions_path)

engines/python/setup/djl_python/session_utils.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python
+#
+# Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file
+# except in compliance with the License. A copy of the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "LICENSE.txt" file accompanying this file. This file is distributed on an "AS IS"
+# BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied. See the License for
+# the specific language governing permissions and limitations under the License.
+
+HEADER_SAGEMAKER_SESSION_ID = "X-Amzn-SageMaker-Session-Id"
+
+
+def get_session(session_manager, request):
+    session_id = request.get_property(HEADER_SAGEMAKER_SESSION_ID)
+    if session_id is None:
+        return None
+    session = session_manager.get_session(session_id)
+    return session

engines/python/src/main/java/ai/djl/python/engine/PyProcess.java

@@ -158,7 +158,8 @@ Output predict(Input inputs, int timeout, boolean initialLoad) throws TranslateE
         // In RollingBatch, we queue adapter loading jobs to occur after the initial load.
         // Executing those in RollingBatch context doesn't work, so we need to handle them in the
         // 'standard' way.
-        if (initialLoad || inputs.getProperty("handler", null) != null) {
+        if (initialLoad
+                || (inputs.getProperty("handler", null) != null && asyncRequestManager == null)) {
             return predictStandard(inputs, timeout, initialLoad);
         }
         if (rollingBatch != null) {