diff --git a/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py b/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py index 738b47e309..0741d97018 100644 --- a/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py +++ b/sagemaker-core/src/sagemaker/core/telemetry/telemetry_logging.py @@ -88,7 +88,7 @@ def wrapper(*args, **kwargs): sagemaker_session = None if len(args) > 0 and hasattr(args[0], "sagemaker_session"): # Get the sagemaker_session from the instance method args - sagemaker_session = args[0].sagemaker_session or _get_default_sagemaker_session() + sagemaker_session = args[0].sagemaker_session elif len(args) > 0 and hasattr(args[0], "_sagemaker_session"): # Get the sagemaker_session from the instance method args (private attribute) sagemaker_session = args[0]._sagemaker_session @@ -98,6 +98,13 @@ def wrapper(*args, **kwargs): "sagemaker_session", _get_default_sagemaker_session() ) + # Fallback: check kwargs for sagemaker_session (e.g., classmethods where + # args[0] is the class and the session is passed as a keyword argument) + if not sagemaker_session: + sagemaker_session = kwargs.get("sagemaker_session") or ( + _get_default_sagemaker_session() + ) + if sagemaker_session: logger.debug("sagemaker_session found, preparing to emit telemetry...") logger.info(TELEMETRY_OPT_OUT_MESSAGING) diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 6f60d325ad..4cf07750b4 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -463,7 +463,7 @@ def _initialize_compute_config(self) -> None: def _initialize_network_config(self) -> None: """Initialize network configuration from Networking object.""" if self.network: - if self.network.vpc_config: + if hasattr(self.network, "vpc_config") and self.network.vpc_config: self.vpc_config = self.network.vpc_config else: self.vpc_config = ( diff --git a/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py b/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py index fafd132547..58121ea4ab 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py +++ b/sagemaker-serve/tests/unit/test_model_builder_coverage_boost.py @@ -399,6 +399,7 @@ def test_from_jumpstart_config_applies_network_isolation(self, mock_deploy_kwarg mock_session = Mock() mock_session.boto_region_name = "us-west-2" + mock_session.sagemaker_config = None mb = ModelBuilder.from_jumpstart_config( jumpstart_config=js_config,