Merge branch 'master' into MultiICNameFix

Author: sage-maker

src/sagemaker/image_uris.py

@@ -101,6 +101,8 @@ def retrieve(
             https://github.com/aws/deep-learning-containers/blob/master/available_images.md
             (default: None).
         distribution (dict): A dictionary with information on how to run distributed training
+        base_framework_version (str): The base version number of PyTorch or Tensorflow.
+            (default: None).
         training_compiler_config (:class:`~sagemaker.training_compiler.TrainingCompilerConfig`):
             A configuration class for the SageMaker Training Compiler
             (default: None).

src/sagemaker/model_monitor/model_monitoring.py

@@ -2413,7 +2413,12 @@ def _update_data_quality_monitoring_schedule(
         )
         self.sagemaker_session.sagemaker_client.create_data_quality_job_definition(**request_dict)
         try:
-            self._update_monitoring_schedule(new_job_definition_name, schedule_cron_expression)
+            self._update_monitoring_schedule(
+                job_definition_name=new_job_definition_name,
+                schedule_cron_expression=schedule_cron_expression,
+                data_analysis_start_time=data_analysis_start_time,
+                data_analysis_end_time=data_analysis_end_time,
+            )
             self.job_definition_name = new_job_definition_name
             if role is not None:
                 self.role = role

src/sagemaker/session.py

@@ -5286,7 +5286,7 @@ def get_tagging_resources(self, tag_filters, resource_type_filters):
                 resource_tag_response = self.resource_group_tagging_client.get_resources(
                     TagFilters=tag_filters,
                     ResourceTypeFilters=resource_type_filters,
-                    NextToken=next_token,
+                    PaginationToken=next_token,
                 )
                 resource_list = resource_list + resource_tag_response["ResourceTagMappingList"]
                 next_token = resource_tag_response.get("PaginationToken")

tests/integ/sagemaker/serve/test_base_model_builder_deploy.py

@@ -12,38 +12,72 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
-import pytest
-
-from sagemaker import get_execution_role
-from sklearn.datasets import load_iris
-from sklearn.model_selection import train_test_split
-
 import os
+import uuid
+from typing import Generator
 
+import numpy as np
+import pandas as pd
+import pytest
+from sagemaker_core.main.resources import TrainingJob
 from sagemaker_core.main.shapes import (
     AlgorithmSpecification,
     Channel,
     DataSource,
-    S3DataSource,
     OutputDataConfig,
     ResourceConfig,
+    S3DataSource,
     StoppingCondition,
 )
-import uuid
-from sagemaker.serve.builder.model_builder import ModelBuilder
-import pandas as pd
-import numpy as np
-from sagemaker.serve import InferenceSpec, SchemaBuilder
-from sagemaker_core.main.resources import TrainingJob
+from sklearn.datasets import load_iris
+from sklearn.model_selection import train_test_split
 from xgboost import XGBClassifier
 
-from sagemaker.serverless.serverless_inference_config import ServerlessInferenceConfig
-
-from sagemaker.s3_utils import s3_path_join
+from sagemaker import get_execution_role
 from sagemaker.async_inference import AsyncInferenceConfig
+from sagemaker.s3_utils import s3_path_join
+from sagemaker.serve import InferenceSpec, SchemaBuilder
+from sagemaker.serve.builder.model_builder import ModelBuilder
+from sagemaker.serverless.serverless_inference_config import ServerlessInferenceConfig
 from tests.integ.utils import cleanup_model_resources
 
 
+@pytest.fixture(autouse=True)
+def cleanup_endpoints(mb_sagemaker_session) -> Generator[None, None, None]:
+    """Clean up any existing endpoints before and after tests."""
+    sagemaker_client = mb_sagemaker_session.sagemaker_client
+
+    # Pre-test cleanup
+    try:
+        endpoints = sagemaker_client.list_endpoints()
+        for endpoint in endpoints["Endpoints"]:
+            try:
+                sagemaker_client.delete_endpoint(EndpointName=endpoint["EndpointName"])
+                sagemaker_client.delete_endpoint_config(
+                    EndpointConfigName=endpoint["EndpointConfigName"]
+                )
+            except Exception as e:
+                print(f"Error cleaning up endpoint {endpoint['EndpointName']}: {e}")
+    except Exception as e:
+        print(f"Error listing endpoints: {e}")
+
+    yield
+
+    # Post-test cleanup
+    try:
+        endpoints = sagemaker_client.list_endpoints()
+        for endpoint in endpoints["Endpoints"]:
+            try:
+                sagemaker_client.delete_endpoint(EndpointName=endpoint["EndpointName"])
+                sagemaker_client.delete_endpoint_config(
+                    EndpointConfigName=endpoint["EndpointConfigName"]
+                )
+            except Exception as e:
+                print(f"Error cleaning up endpoint {endpoint['EndpointName']}: {e}")
+    except Exception as e:
+        print(f"Error listing endpoints: {e}")
+
+
 @pytest.fixture(scope="module")
 def xgboost_model_builder(mb_sagemaker_session):
     sagemaker_session = mb_sagemaker_session