[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: TestStatefulModel
+            instance: g6
+            failure-prefix: lmi
     outputs:
       failure_cpu: ${{ steps.test-failure.outputs.failure_cpu }}
       failure_gpu: ${{ steps.test-failure.outputs.failure_gpu }}

engines/python/setup/djl_python/async_utils.py

@@ -18,12 +18,14 @@
 
 
 def create_non_stream_output(data: Union[str, dict],
+                             properties: Optional[dict] = None,
                              error: Optional[str] = None,
                              code: Optional[int] = None) -> Output:
     return _create_output(
         data,
         True,
         "application/json",
+        properties=properties,
         error=error,
         code=code,
     )
@@ -46,6 +48,7 @@ def _create_output(
     data: Union[str, dict],
     last_chunk: bool,
     content_type: str,
+    properties: Optional[dict] = None,
     error: Optional[str] = None,
     code: Optional[int] = None,
 ) -> Output:
@@ -65,6 +68,9 @@ def _create_output(
         response_dict["code"] = code
     output = Output()
     output.add_property("Content-Type", content_type)
+    if properties:
+        for k, v in properties.items():
+            output.add_property(k, v)
     output.add(Output.binary_encode(response_dict))
     return output
 

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

@@ -31,9 +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 .request_response_utils import (
+from djl_python.custom_formatter_handling import CustomFormatterError, CustomFormatterHandler
+from djl_python.lmi_vllm.request_response_utils import (
     ProcessedRequest,
     vllm_stream_output_formatter,
     vllm_non_stream_output_formatter,
@@ -43,9 +42,15 @@
     lmi_with_details_non_stream_output_formatter,
     lmi_non_stream_output_formatter,
 )
+from djl_python.session_manager import SessionManager
+from djl_python.session_utils import (create_session, close_session,
+                                      get_session,
+                                      session_non_stream_output_formatter)
 
 logger = logging.getLogger(__name__)
 
+SESSION_REQUESTS = {"NEW_SESSION": create_session, "CLOSE": close_session}
+
 
 class VLLMHandler(CustomFormatterHandler):
 
@@ -119,12 +124,15 @@ async def initialize(self, properties: dict):
             tool_parser=self.vllm_properties.tool_call_parser,
             reasoning_parser=self.vllm_properties.reasoning_parser,
         )
+        if properties.get("enable_stateful_sessions", "true") == "true":
+            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]
+        session = get_session(self.session_manager, raw_request)
         content_type = raw_request.get_property("Content-Type")
         decoded_payload = decode(raw_request, content_type)
 
@@ -160,6 +168,20 @@ def preprocess_request(self, inputs: Input) -> ProcessedRequest:
             vllm_invoke_function = self.chat_completion_service.create_chat_completion
             non_stream_output_formatter = vllm_non_stream_output_formatter
             stream_output_formatter = vllm_stream_output_formatter
+        elif "requestType" in decoded_payload:
+            request_type = decoded_payload["requestType"]
+            if request_type not in SESSION_REQUESTS.keys():
+                raise RuntimeError(
+                    f"invalid payload. request type must be one of {SESSION_REQUESTS.keys()}"
+                )
+            if self.session_manager is None:
+                raise RuntimeError(
+                    f"invalid payload. stateful sessions not enabled, {request_type} not supported"
+                )
+            vllm_request = self.session_manager, inputs
+            vllm_invoke_function = SESSION_REQUESTS[request_type]
+            non_stream_output_formatter = session_non_stream_output_formatter
+            stream_output_formatter = vllm_stream_output_formatter
         else:
             raise RuntimeError(
                 "invalid payload. must contain prompt, inputs, or messages")

engines/python/setup/djl_python/session_manager.py

@@ -53,13 +53,11 @@ def get(self, key: str, d=None):
             return pickle.load(f)
 
     def remove(self):
-        if os.path.exists(self.files_path):
-            logging.info(f"closing session: {self.session_id}")
-            shutil.rmtree(self.files_path)
-            return True
-        else:
-            logging.warning(f"session not found: {self.session_id}")
-            return False
+        if not os.path.exists(self.files_path):
+            raise ValueError(f"session not found: {self.session_id}")
+        logging.info(f"closing session: {self.session_id}")
+        shutil.rmtree(self.files_path)
+        return True
 
     def _path(self, key: str):
         return os.path.join(self.files_path, key.replace("/", "-"))
@@ -92,18 +90,32 @@ def create_session(self) -> Session:
         session_id = str(uuid.uuid4())
         session = Session(session_id, self.sessions_path)
         os.makedirs(session.files_path)
-        session.put(".creation_time", time.time())
+        expiration_time = time.time() + self.expiration
+        session.put(".expiration_time", expiration_time)
 
         self.cloud_watch.post("create_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}")
+        if session_id == "NEW_SESSION" or not session_id:
+            return None
 
+        if not UUID_PATTERN.match(session_id):
+            logging.warning(f"invalid session_id: {session_id}")
+            return None
+
+        # Session expired
         session = Session(session_id, self.sessions_path)
+        if session.get(".expiration_time") is not None \
+                and time.time() > session.get(".expiration_time"):
+            return None
+
+        # Session not found, try to recover from s3 bucket
         if not os.path.exists(session.files_path):
-            return self._recover_from_s3(session)
+            session = self._recover_from_s3(session)
+            if session is None:
+                raise ValueError(f"session not found: {session_id}")
+            return session
 
         return session
 
@@ -119,7 +131,8 @@ def _clean_expired_session(self):
         sessions = os.listdir(self.sessions_path)
         for session_id in sessions:
             session = Session(session_id, self.sessions_path)
-            if time.time() - session.get(".creation_time") > self.expiration:
+            if session.get(".expiration_time") is None \
+                    or time.time() > session.get(".expiration_time"):
                 self.close_session(session_id)
 
     def _recover_from_s3(self, session: Session) -> Optional[Session]:

engines/python/setup/djl_python/session_utils.py

@@ -0,0 +1,85 @@
+#!/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.
+import datetime
+import logging
+
+from djl_python.async_utils import create_non_stream_output
+from djl_python.outputs import Output
+
+logger = logging.getLogger(__name__)
+
+HEADER_SAGEMAKER_SESSION_ID = "X-Amzn-SageMaker-Session-Id"
+HEADER_SAGEMAKER_CLOSED_SESSION_ID = "X-Amzn-SageMaker-Closed-Session-Id"
+
+
+async def create_session(request):
+    session_manager, inputs = request
+    try:
+        session = session_manager.create_session()
+        expiration_ts = datetime.datetime.fromtimestamp(
+            session.get(".expiration_time")).strftime("%Y-%m-%dT%H:%M:%SZ")
+        logger.info(f"Session {session.session_id} created")
+        return {
+            "data": {
+                "result": f"Session {session.session_id} created"
+            },
+            "properties": {
+                HEADER_SAGEMAKER_SESSION_ID:
+                f"{session.session_id}; Expires={expiration_ts}"
+            }
+        }
+    except Exception as e:
+        return {"error": f"Failed to create session: {str(e)}", "code": 424}
+
+
+async def close_session(request):
+    session_manager, inputs = request
+    session_id = inputs.get_property(HEADER_SAGEMAKER_SESSION_ID)
+    try:
+        session_manager.close_session(session_id)
+        logger.info(f"Session {session_id} closed")
+        return {
+            "data": {
+                "result": f"Session {session_id} closed"
+            },
+            "properties": {
+                HEADER_SAGEMAKER_CLOSED_SESSION_ID: f"{session_id}"
+            }
+        }
+    except Exception as e:
+        return {"error": f"Failed to close session: {str(e)}", "code": 424}
+
+
+def get_session(session_manager, request):
+    session_id = request.get_property(HEADER_SAGEMAKER_SESSION_ID)
+    if session_manager is None:
+        if session_id is not None:
+            raise RuntimeError(
+                f"invalid payload. stateful sessions not enabled, {HEADER_SAGEMAKER_SESSION_ID} header not supported"
+            )
+        return None
+    session = session_manager.get_session(session_id)
+    return session
+
+
+def session_non_stream_output_formatter(
+    response: dict,
+    **_,
+) -> Output:
+    if "error" in response:
+        return create_non_stream_output("",
+                                        error=response["error"],
+                                        code=response["code"])
+
+    return create_non_stream_output(response["data"],
+                                    properties=response.get("properties"))

engines/python/src/test/java/ai/djl/python/engine/PySessionsTests.java

@@ -84,10 +84,23 @@ public void testLocalLoadSave()
 
             // test session timeout
             Thread.sleep(1000);
+            regular = new Input();
+            regular.addProperty("Content-Type", "application/json");
+            regular.addProperty("X-Amzn-SageMaker-Session-Id", sessionId);
+            regular.add(BytesSupplier.wrapAsJson(Map.of("action", "regular")));
+            ret = predictor.predict(regular);
+            Assert.assertEquals(ret.getProperty("Content-Type", null), "application/json");
+            Assert.assertTrue(ret.getAsString(0).contains("session not found"));
+            long count;
+            try (Stream<Path> files = Files.list(path)) {
+                count = files.count();
+            }
+            Assert.assertEquals(count, 1);
+
+            // create a new session
             ret = predictor.predict(createSession);
             sessionId = ret.getProperty("X-Amzn-SageMaker-Session-Id", null);
             Assert.assertNotNull(sessionId);
-            long count;
             try (Stream<Path> files = Files.list(path)) {
                 count = files.count();
             }