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')

Author: jamesfyu

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()))

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)