feat: Add SageMaker token generator to sagemaker-core (#5868)

shrijeetj · jamesfyu · web-flow · commit 37b43b7e5aac · 2026-05-19T13:11:34.000-07:00
* feat: Add SageMaker token generator to sagemaker-core Embed the aws-sagemaker-token-generator library into sagemaker.core so users can generate SageMaker bearer tokens without installing a separate wheel. Usage: from sagemaker.core.aws_sagemaker_token_generator import provide_token token = provide_token(region='us-east-1') * test: Add unit tests for SageMaker token generator Add unit tests for the token_generator module covering token format, base64 encoding, version info, multiple regions, session tokens, credential validation, expiry handling, and API consistency. Tests are modeled after the Bedrock token generator test suite: https://github.com/aws/aws-bedrock-token-generator-python/blob/main/tests/test_token_generator.py --------- Co-authored-by: James Yu <jamesfyu@amazon.com>
diff --git a/sagemaker-core/src/sagemaker/core/token_generator/__init__.py b/sagemaker-core/src/sagemaker/core/token_generator/__init__.py
@@ -0,0 +1,70 @@
+"""AWS SageMaker Token Generator.
+
+A lightweight module for generating short-term bearer tokens for AWS SageMaker
+API authentication. Provides the ``generate_token`` helper and the lower-level
+``SageMakerTokenGenerator`` class.
+
+Example::
+
+    >>> from sagemaker.core.token_generator import generate_token
+    >>> token = generate_token(region="us-east-1")
+"""
+
+from __future__ import annotations
+
+import os
+from datetime import timedelta
+
+from botocore.credentials import CredentialProvider
+from botocore.session import Session
+
+from sagemaker.core.token_generator.token_generator import (
+    TOKEN_DURATION,
+    SageMakerTokenGenerator,
+    _generate_token,
+)
+
+__all__ = ["SageMakerTokenGenerator", "generate_token"]
+
+
+def generate_token(
+    region: str | None = None,
+    aws_credentials_provider: CredentialProvider | None = None,
+    expiry: timedelta = timedelta(hours=12),
+) -> str:
+    """Generate a short-lived AWS SageMaker bearer token.
+
+    Args:
+        region (str): AWS region. Falls back to the ``AWS_REGION``
+            environment variable when not provided.
+        aws_credentials_provider (CredentialProvider): Optional credential
+            provider. Uses the default AWS credential chain when omitted.
+        expiry (timedelta): Token lifetime. Must be between 1 second and
+            12 hours inclusive. Defaults to 12 hours.
+
+    Returns:
+        str: A bearer token string.
+
+    Raises:
+        ValueError: If *region* is missing or *expiry* is out of range.
+        RuntimeError: If no valid AWS credentials are found.
+    """
+    region = region or os.environ.get("AWS_REGION")
+    if not region:
+        raise ValueError("Region must be provided or set via the AWS_REGION environment variable.")
+
+    if expiry.total_seconds() <= 0 or expiry.total_seconds() > TOKEN_DURATION:
+        raise ValueError(
+            "Token expiry must be greater than zero and less than or equal to 12 hours"
+        )
+
+    credentials = (
+        aws_credentials_provider.load() if aws_credentials_provider else Session().get_credentials()
+    )
+
+    if credentials is None:
+        raise RuntimeError(
+            "No AWS credentials found. Check your environment or credential provider."
+        )
+
+    return _generate_token(credentials, region, int(expiry.total_seconds()))
diff --git a/sagemaker-core/src/sagemaker/core/token_generator/token_generator.py b/sagemaker-core/src/sagemaker/core/token_generator/token_generator.py
@@ -0,0 +1,71 @@
+"""SageMaker Token Generator core signing logic.
+
+Generates short-term bearer tokens for AWS SageMaker API authentication
+using SigV4 signed pre-signed URLs.
+"""
+
+from __future__ import annotations
+
+import base64
+
+from botocore.auth import SigV4QueryAuth
+from botocore.awsrequest import AWSRequest
+from botocore.credentials import Credentials
+
+DEFAULT_HOST: str = "sagemaker.amazonaws.com"
+DEFAULT_URL: str = f"https://{DEFAULT_HOST}/"
+SERVICE_NAME: str = "sagemaker"
+AUTH_PREFIX: str = "sagemaker-api-key-"
+TOKEN_VERSION: str = "&Version=1"
+TOKEN_DURATION: int = 43200  # 12 hours in seconds
+
+
+def _generate_token(credentials: Credentials, region: str, expires: int) -> str:
+    """Build a presigned bearer token.
+
+    Args:
+        credentials (Credentials): AWS credentials.
+        region (str): AWS region.
+        expires (int): Expiry time in seconds.
+
+    Returns:
+        str: A base64-encoded bearer token string.
+    """
+    request = AWSRequest(
+        method="POST",
+        url=DEFAULT_URL,
+        headers={"host": DEFAULT_HOST},
+        params={"Action": "CallWithBearerToken"},
+    )
+
+    auth = SigV4QueryAuth(credentials, SERVICE_NAME, region, expires=expires)
+    auth.add_auth(request)
+
+    presigned_url = request.url.replace("https://", "") + TOKEN_VERSION
+    encoded_token = base64.b64encode(presigned_url.encode("utf-8")).decode("utf-8")
+
+    return f"{AUTH_PREFIX}{encoded_token}"
+
+
+class SageMakerTokenGenerator:
+    """Generate short-lived AWS SageMaker bearer tokens."""
+
+    def get_token(self, credentials: Credentials, region: str) -> str:
+        """Generate a token using provided credentials and region.
+
+        Args:
+            credentials (Credentials): AWS credentials to sign the request.
+            region (str): AWS region.
+
+        Returns:
+            str: A bearer token string.
+
+        Raises:
+            ValueError: If inputs are invalid.
+        """
+        if not credentials:
+            raise ValueError("Credentials cannot be None")
+        if not region or not isinstance(region, str):
+            raise ValueError("Region must be a non-empty string")
+
+        return _generate_token(credentials, region, TOKEN_DURATION)
diff --git a/sagemaker-core/tests/unit/test_token_generator.py b/sagemaker-core/tests/unit/test_token_generator.py
@@ -0,0 +1,170 @@
+# Copyright 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" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Unit tests for sagemaker.core.token_generator module."""
+
+from __future__ import absolute_import
+
+import base64
+from datetime import timedelta
+
+import pytest
+from unittest.mock import Mock
+
+from botocore.credentials import Credentials
+
+from sagemaker.core.token_generator import generate_token, SageMakerTokenGenerator
+from sagemaker.core.token_generator.token_generator import AUTH_PREFIX
+
+
+class TestSageMakerTokenGenerator:
+    """Tests for the SageMakerTokenGenerator class."""
+
+    @pytest.fixture(autouse=True)
+    def setup(self):
+        """Setup test credentials and token generator instance."""
+        self.token_generator = SageMakerTokenGenerator()
+        self.credentials = Credentials(
+            access_key="AKIAIOSFODNN7EXAMPLE",
+            secret_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
+        )
+        self.region = "us-west-2"
+
+    def test_get_token_returns_non_null_token(self):
+        """Test that get_token returns a non-null token."""
+        token = self.token_generator.get_token(self.credentials, self.region)
+
+        assert token is not None
+        assert len(token) > 0
+
+    def test_get_token_starts_with_correct_prefix(self):
+        """Test that the token starts with the correct prefix."""
+        token = self.token_generator.get_token(self.credentials, self.region)
+
+        assert token.startswith(AUTH_PREFIX)
+
+    def test_get_token_with_different_regions(self):
+        """Test token generation with different regions."""
+        regions = ["us-east-1", "us-west-2", "eu-west-1", "ap-northeast-1"]
+
+        for region in regions:
+            token = self.token_generator.get_token(self.credentials, region)
+
+            assert token is not None, f"Token should not be null for region: {region}"
+            assert token.startswith(
+                AUTH_PREFIX
+            ), f"Token should start with the correct prefix for region: {region}"
+
+    def test_get_token_is_base64_encoded(self):
+        """Test that the token is properly Base64 encoded."""
+        token = self.token_generator.get_token(self.credentials, self.region)
+
+        token_without_prefix = token[len(AUTH_PREFIX) :]
+        decoded = base64.b64decode(token_without_prefix)
+        assert decoded is not None
+
+    def test_get_token_contains_version_info(self):
+        """Test that the decoded token contains version information."""
+        token = self.token_generator.get_token(self.credentials, self.region)
+
+        token_without_prefix = token[len(AUTH_PREFIX) :]
+        decoded_string = base64.b64decode(token_without_prefix).decode("utf-8")
+        assert "&Version=1" in decoded_string
+
+    def test_get_token_different_credentials_produce_different_tokens(self):
+        """Test that different credentials produce different tokens."""
+        credentials1 = Credentials(
+            access_key="AKIAIOSFODNN7EXAMPLE",
+            secret_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
+        )
+        credentials2 = Credentials(
+            access_key="AKIAI44QH8DHBEXAMPLE",
+            secret_key="je7MtGbClwBF/2Zp9Utk/h3yCo8nvbEXAMPLEKEY",
+        )
+
+        token1 = self.token_generator.get_token(credentials1, self.region)
+        token2 = self.token_generator.get_token(credentials2, self.region)
+
+        assert token1 != token2
+
+    def test_get_token_with_session_token(self):
+        """Test token generation with session token (temporary credentials)."""
+        credentials_with_token = Credentials(
+            access_key="AKIAIOSFODNN7EXAMPLE",
+            secret_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
+            token="AQoDYXdzEJr...<remainder of security token>",
+        )
+
+        token = self.token_generator.get_token(credentials_with_token, self.region)
+
+        assert token is not None
+        assert token.startswith(AUTH_PREFIX)
+
+    def test_get_token_no_credentials_raises_error(self):
+        """Test that get_token raises ValueError when credentials are None."""
+        with pytest.raises(ValueError, match="Credentials cannot be None"):
+            self.token_generator.get_token(None, self.region)
+
+    def test_get_token_no_region_raises_error(self):
+        """Test that get_token raises ValueError when region is None or empty."""
+        with pytest.raises(ValueError, match="Region must be a non-empty string"):
+            self.token_generator.get_token(self.credentials, None)
+
+        with pytest.raises(ValueError, match="Region must be a non-empty string"):
+            self.token_generator.get_token(self.credentials, "")
+
+    def test_get_token_contains_correct_expiry(self):
+        """Test that the decoded token has the correct expiry duration (12 hours)."""
+        token = self.token_generator.get_token(self.credentials, self.region)
+
+        token_without_prefix = token[len(AUTH_PREFIX) :]
+        decoded_string = base64.b64decode(token_without_prefix).decode("utf-8")
+        assert "X-Amz-Expires=43200" in decoded_string
+
+    def test_get_token_vs_generate_token_consistency(self):
+        """Test that get_token and generate_token produce identical tokens for same inputs."""
+        mock_provider = Mock()
+        mock_provider.load.return_value = self.credentials
+
+        token1 = self.token_generator.get_token(self.credentials, self.region)
+
+        token2 = generate_token(
+            region=self.region,
+            aws_credentials_provider=mock_provider,
+            expiry=timedelta(hours=12),
+        )
+
+        assert token1 == token2
+        assert token1.startswith(AUTH_PREFIX)
+        assert token2.startswith(AUTH_PREFIX)
+        assert len(token1) == len(token2)
+
+    def test_generate_token_with_custom_expiry_produces_different_token(self):
+        """Test that different expiry durations produce different tokens."""
+        mock_provider = Mock()
+        mock_provider.load.return_value = self.credentials
+
+        token_default = generate_token(
+            region=self.region,
+            aws_credentials_provider=mock_provider,
+            expiry=timedelta(hours=12),
+        )
+
+        token_custom = generate_token(
+            region=self.region,
+            aws_credentials_provider=mock_provider,
+            expiry=timedelta(hours=6),
+        )
+
+        assert token_default != token_custom
+        assert token_default.startswith(AUTH_PREFIX)
+        assert token_custom.startswith(AUTH_PREFIX)
