diff --git a/sagemaker-core/src/sagemaker/core/__init__.py b/sagemaker-core/src/sagemaker/core/__init__.py index f25f18009d..4b4f0e397c 100644 --- a/sagemaker-core/src/sagemaker/core/__init__.py +++ b/sagemaker-core/src/sagemaker/core/__init__.py @@ -14,6 +14,7 @@ # Partner App from sagemaker.core.partner_app.auth_provider import PartnerAppAuthProvider # noqa: F401 +from sagemaker.core.partner_app.auth_provider import RequestsAuth # noqa: F401 # Attribution from sagemaker.core.telemetry.attribution import Attribution, set_attribution # noqa: F401 diff --git a/sagemaker-core/src/sagemaker/core/partner_app/__init__.py b/sagemaker-core/src/sagemaker/core/partner_app/__init__.py index 87ab21acec..8bea7d7dcc 100644 --- a/sagemaker-core/src/sagemaker/core/partner_app/__init__.py +++ b/sagemaker-core/src/sagemaker/core/partner_app/__init__.py @@ -11,6 +11,8 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. """__init__ file for sagemaker.core.partner_app""" -from __future__ import absolute_import +from __future__ import annotations from sagemaker.core.partner_app.auth_provider import PartnerAppAuthProvider # noqa: F401 +from sagemaker.core.partner_app.auth_provider import RequestsAuth # noqa: F401 +from sagemaker.core.partner_app.auth_utils import PartnerAppAuthUtils # noqa: F401 diff --git a/sagemaker-core/src/sagemaker/core/partner_app/auth_provider.py b/sagemaker-core/src/sagemaker/core/partner_app/auth_provider.py index 7abdb71e0b..3b485e503a 100644 --- a/sagemaker-core/src/sagemaker/core/partner_app/auth_provider.py +++ b/sagemaker-core/src/sagemaker/core/partner_app/auth_provider.py @@ -12,8 +12,9 @@ # language governing permissions and limitations under the License. """The SageMaker partner application SDK auth module""" -from __future__ import absolute_import +from __future__ import annotations +import logging import os import re from typing import Dict, Tuple @@ -25,6 +26,8 @@ from requests.models import PreparedRequest from sagemaker.core.partner_app.auth_utils import PartnerAppAuthUtils +logger = logging.getLogger(__name__) + SERVICE_NAME = "sagemaker" AWS_PARTNER_APP_ARN_REGEX = r"arn:aws[a-z\-]*:sagemaker:[a-z0-9\-]*:[0-9]{12}:partner-app\/.*" @@ -94,6 +97,7 @@ def __init__(self, credentials: Credentials = None): credentials if credentials is not None else boto3.Session().get_credentials() ) self.sigv4 = SigV4Auth(self.credentials, SERVICE_NAME, self.region) + logger.info("PartnerAppAuthProvider initialized for region: %s", self.region) def get_signed_request( self, url: str, method: str, headers: dict, body: object @@ -109,6 +113,7 @@ def get_signed_request( Returns: tuple: (url, headers) """ + logger.debug("Signing request: %s %s", method, url) return PartnerAppAuthUtils.get_signed_request( sigv4=self.sigv4, app_arn=self.app_arn, diff --git a/sagemaker-core/src/sagemaker/core/partner_app/auth_utils.py b/sagemaker-core/src/sagemaker/core/partner_app/auth_utils.py index eb1dcacaa9..5ba17a6cca 100644 --- a/sagemaker-core/src/sagemaker/core/partner_app/auth_utils.py +++ b/sagemaker-core/src/sagemaker/core/partner_app/auth_utils.py @@ -13,8 +13,9 @@ """Partner App Auth Utils Module""" -from __future__ import absolute_import +from __future__ import annotations +import logging from hashlib import sha256 import functools from typing import Tuple, Dict @@ -22,6 +23,8 @@ from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest +logger = logging.getLogger(__name__) + HEADER_CONNECTION = "Connection" HEADER_X_AMZ_TARGET = "X-Amz-Target" HEADER_AUTHORIZATION = "Authorization"