Enable toggling sticky routing - default false (#18)

* feat(sagemaker/sessions): add config to enable stateful sessions using env variable

* feat(sagemaker/sessions): add utility functions for getting session_manager

* chore(sagemaker/sessions): update tests

* feat: update env config to use pydantic model SageMakerConfig and use SAGEMAKER_  as prefix.

* feat(sagemaker/sessions): add validation layer so if session_manager 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

* Update way of setting sessions_path.

* import logger to sessions/transform.py

---------

Signed-off-by: Zuyi Zhao <[email protected]>

Author: zhaozuy

python/model_hosting_container_standards/sagemaker/config.py

@@ -1,5 +1,99 @@
 """SageMaker-specific configuration constants."""
 
+import os
+from typing import Any, Dict, Optional
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+
+SAGEMAKER_ENV_VAR_PREFIX = "SAGEMAKER_"
+
+
+class SageMakerConfig(BaseModel):
+    """Pydantic model for SageMaker configuration.
+
+    Automatically loads configuration from environment variables with SAGEMAKER_ prefix.
+    Example: SAGEMAKER_ENABLE_STATEFUL_SESSIONS=true -> enable_stateful_sessions=True
+
+    Only fields defined in this model are loaded. Other SAGEMAKER_* env vars
+    (like SAGEMAKER_MODEL_PATH) are ignored.
+
+    Usage:
+        # Create from environment variables
+        config = SagemakerConfig.from_env()
+
+        # Or just instantiate (automatically loads from env)
+        config = SagemakerConfig()
+
+        # Override specific values
+        config = SagemakerConfig(enable_stateful_sessions=True)
+    """
+
+    model_config = ConfigDict(extra="ignore")
+
+    # Stateful sessions configuration
+    enable_stateful_sessions: bool = Field(
+        default=False, description="Enable stateful sessions for the application"
+    )
+    sessions_expiration: int = Field(
+        default=1200,  # 20 minutes
+        description="Session expiration time in seconds",
+        gt=0,
+    )
+    sessions_path: Optional[str] = Field(
+        default=None,
+        description="Custom path for session storage (defaults to /dev/shm or temp)",
+    )
+
+    @classmethod
+    def from_env(cls) -> "SageMakerConfig":
+        """Create SagemakerConfig from environment variables.
+
+        Returns:
+            SagemakerConfig instance with values loaded from SAGEMAKER_* env vars
+        """
+        return cls()
+
+    @model_validator(mode="before")
+    @classmethod
+    def load_from_env_vars(cls, data: Any) -> Dict[str, Any]:
+        """Load configuration from environment variables.
+
+        Extracts SAGEMAKER_* environment variables and merges with any provided data.
+        Provided data takes precedence over environment variables.
+        Unknown SAGEMAKER_* variables are ignored (only defined fields are loaded).
+        """
+        # Extract env vars with SAGEMAKER_ prefix
+        env_config = {
+            key[len(SAGEMAKER_ENV_VAR_PREFIX) :].lower(): val
+            for key, val in os.environ.items()
+            if key.startswith(SAGEMAKER_ENV_VAR_PREFIX)
+        }
+
+        # If data is provided, merge with env config (data takes precedence)
+        if isinstance(data, dict):
+            return {**env_config, **data}
+        return env_config
+
+    @field_validator("enable_stateful_sessions", mode="before")
+    @classmethod
+    def parse_bool_string(cls, v: Any) -> bool:
+        """Convert string values from env vars to boolean."""
+        if isinstance(v, bool):
+            return v
+        if isinstance(v, str):
+            return v.lower() in ("true", "1")
+        return bool(v)
+
+    @field_validator("sessions_expiration", mode="before")
+    @classmethod
+    def parse_int_string(cls, v: Any) -> int:
+        """Convert string values from env vars to integer."""
+        if isinstance(v, int):
+            return v
+        if isinstance(v, str):
+            return int(v)
+        return int(v)
+
 
 class SageMakerEnvVars:
     """SageMaker environment variable names."""

python/model_hosting_container_standards/sagemaker/sessions/handlers.py

@@ -5,8 +5,13 @@
 from fastapi.exceptions import HTTPException
 
 from ...logging_config import logger
-from .manager import SessionManager
-from .models import SageMakerSessionHeader, SessionRequestType
+from .manager import get_session_manager
+from .models import (
+    SESSION_DISABLED_ERROR_DETAIL,
+    SESSION_DISABLED_LOG_MESSAGE,
+    SageMakerSessionHeader,
+    SessionRequestType,
+)
 from .utils import get_session_id_from_request
 
 
@@ -27,11 +32,10 @@ def get_handler_for_request_type(request_type: SessionRequestType):
         return None
 
 
-async def close_session(session_manager: SessionManager, raw_request: Request):
+async def close_session(raw_request: Request):
     """Close an existing session and clean up its resources.
 
     Args:
-        session_manager: SessionManager instance to manage the session lifecycle
         raw_request: FastAPI Request object containing session ID in headers
 
     Returns:
@@ -41,6 +45,13 @@ async def close_session(session_manager: SessionManager, raw_request: Request):
         HTTPException: If session closure fails with 424 FAILED_DEPENDENCY status
     """
     session_id = get_session_id_from_request(raw_request)
+    session_manager = get_session_manager()
+    if session_manager is None:
+        logger.error(SESSION_DISABLED_LOG_MESSAGE)
+        raise HTTPException(
+            status_code=HTTPStatus.BAD_REQUEST.value,
+            detail=SESSION_DISABLED_ERROR_DETAIL,
+        )
     try:
         session_manager.close_session(session_id)
         logger.info(f"Session {session_id} closed")
@@ -57,11 +68,10 @@ async def close_session(session_manager: SessionManager, raw_request: Request):
         )
 
 
-async def create_session(session_manager: SessionManager, raw_request: Request):
+async def create_session(raw_request: Request):
     """Create a new stateful session with expiration tracking.
 
     Args:
-        session_manager: SessionManager instance to manage the session lifecycle
         raw_request: FastAPI Request object (unused but part of handler signature)
 
     Returns:
@@ -70,6 +80,13 @@ async def create_session(session_manager: SessionManager, raw_request: Request):
     Raises:
         HTTPException: If session creation fails with 424 FAILED_DEPENDENCY status
     """
+    session_manager = get_session_manager()
+    if session_manager is None:
+        logger.error(SESSION_DISABLED_LOG_MESSAGE)
+        raise HTTPException(
+            status_code=HTTPStatus.BAD_REQUEST.value,
+            detail=SESSION_DISABLED_ERROR_DETAIL,
+        )
     try:
         session = session_manager.create_session()
         # expiration_ts is guaranteed to be set for newly created sessions

python/model_hosting_container_standards/sagemaker/sessions/manager.py

@@ -11,6 +11,8 @@
 from threading import RLock
 from typing import Optional
 
+from ..config import SageMakerConfig
+
 
 class Session:
     """Represents a single stateful session with file-based storage.
@@ -137,7 +139,7 @@ def __init__(self, properties: dict):
         else:
             session_dir = os.path.join(tempfile.gettempdir(), "sagemaker_sessions")
 
-        self.sessions_path = properties.get("sessions_path", session_dir)
+        self.sessions_path = properties.get("sessions_path") or session_dir
         self.sessions: dict[str, Session] = {}
         self._lock = RLock()  # Thread safety for concurrent session access
 
@@ -248,5 +250,49 @@ def _clean_expired_session(self):
                     self.close_session(session_id)
 
 
-# Global SessionManager instance
-session_manager = SessionManager({})
+def _init_session_manager(config: SageMakerConfig) -> SessionManager | None:
+    """Initialize a SessionManager if stateful sessions are enabled.
+
+    Args:
+        config: SagemakerConfig instance with session settings
+
+    Returns:
+        SessionManager instance if enabled, None otherwise
+    """
+    if config.enable_stateful_sessions:
+        # Convert config to dict for SessionManager
+        config_dict = {
+            "sessions_expiration": str(config.sessions_expiration),
+            "sessions_path": config.sessions_path,
+        }
+        return SessionManager(config_dict)
+    return None
+
+
+def get_session_manager() -> SessionManager | None:
+    """Get the global session manager instance.
+
+    Returns:
+        The global SessionManager instance, or None if not initialized
+    """
+    return session_manager
+
+
+def init_session_manager_from_env() -> SessionManager | None:
+    """Initialize the global session manager from environment variables.
+
+    This can be called to reinitialize the session manager after environment
+    variables have been set.
+
+    Returns:
+        The initialized SessionManager instance, or None if disabled
+    """
+    global session_manager
+    config = SageMakerConfig.from_env()
+    session_manager = _init_session_manager(config)
+    return session_manager
+
+
+# Global SessionManager instance - initialized from environment variables
+_config = SageMakerConfig.from_env()
+session_manager = _init_session_manager(_config)

python/model_hosting_container_standards/sagemaker/sessions/models.py

@@ -37,3 +37,11 @@ class SageMakerSessionHeader:
     SESSION_ID = "X-Amzn-SageMaker-Session-Id"
     NEW_SESSION_ID = "X-Amzn-SageMaker-New-Session-Id"
     CLOSED_SESSION_ID = "X-Amzn-SageMaker-Closed-Session-Id"
+
+
+# Error messages for session management
+SESSION_DISABLED_ERROR_DETAIL = "Invalid payload. stateful sessions not enabled"
+SESSION_DISABLED_LOG_MESSAGE = (
+    f"Invalid payload. stateful sessions not enabled, "
+    f"{SageMakerSessionHeader.SESSION_ID} header not supported"
+)

python/model_hosting_container_standards/sagemaker/sessions/transform.py

@@ -7,9 +7,14 @@
 from pydantic import ValidationError
 
 from ...common import BaseApiTransform, BaseTransformRequestOutput
+from ...logging_config import logger
 from .handlers import get_handler_for_request_type
-from .manager import SessionManager, session_manager
-from .models import SessionRequest
+from .manager import SessionManager, get_session_manager
+from .models import (
+    SESSION_DISABLED_ERROR_DETAIL,
+    SESSION_DISABLED_LOG_MESSAGE,
+    SessionRequest,
+)
 from .utils import get_session, get_session_id_from_request
 
 
@@ -38,7 +43,9 @@ def _parse_session_request(request_data: dict) -> Optional[SessionRequest]:
         return None
 
 
-def _validate_session_if_present(raw_request: Request, session_manager: SessionManager):
+def _validate_session_if_present(
+    raw_request: Request, session_manager: Optional[SessionManager]
+):
     """Validate that the session ID in the request exists and is not expired.
 
     Args:
@@ -60,7 +67,7 @@ def _validate_session_if_present(raw_request: Request, session_manager: SessionM
 
 
 def process_session_request(
-    request_data: dict, raw_request: Request, session_manager: SessionManager
+    request_data: dict, raw_request: Request, session_manager: Optional[SessionManager]
 ):
     """Process a potential session management request.
 
@@ -89,16 +96,22 @@ def process_session_request(
     # Not a session request - pass through for normal processing
     if session_request is None:
         return BaseTransformRequestOutput(
-            request=None,
             raw_request=raw_request,
             intercept_func=None,
         )
 
+    if session_manager is None:
+        logger.error(SESSION_DISABLED_LOG_MESSAGE)
+        raise HTTPException(
+            status_code=HTTPStatus.BAD_REQUEST.value,
+            detail=SESSION_DISABLED_ERROR_DETAIL,
+        )
+
     # Route to appropriate session management handler
     intercept_func = get_handler_for_request_type(session_request.requestType)
 
     return BaseTransformRequestOutput(
-        request=session_manager, raw_request=raw_request, intercept_func=intercept_func
+        raw_request=raw_request, intercept_func=intercept_func
     )
 
 
@@ -121,7 +134,7 @@ def __init__(self, request_shape, response_shape={}):
             The request/response shapes are passed to the parent class but not used
             for validation in this transform, as session requests use their own validation.
         """
-        self._session_manager = session_manager
+        self._session_manager = get_session_manager()
         super().__init__(request_shape, response_shape)
 
     async def transform_request(self, raw_request):

python/model_hosting_container_standards/sagemaker/sessions/utils.py

@@ -1,4 +1,5 @@
 from http import HTTPStatus
+from typing import Optional
 
 from fastapi import Request
 from fastapi.exceptions import HTTPException
@@ -22,7 +23,7 @@ def get_session_id_from_request(raw_request: Request):
     return raw_request.headers.get(SageMakerSessionHeader.SESSION_ID)
 
 
-def get_session(session_manager: SessionManager, raw_request: Request):
+def get_session(session_manager: Optional[SessionManager], raw_request: Request):
     """Retrieve the session associated with the request.
 
     Args: