Support BatchGetMetrics from sagemaker-metrics

resource_plan.csv

@@ -35,6 +35,7 @@ InferenceExperiment,resource,"['create', 'get', 'get_all', 'start']","['delete',
 InferenceRecommendationsJob,resource,"['create', 'get', 'get_all']","['refresh', 'stop', 'wait', 'wait_for_delete']",[],[],"['CreateInferenceRecommendationsJob', 'DescribeInferenceRecommendationsJob', 'ListInferenceRecommendationsJobs', 'StopInferenceRecommendationsJob']","[{'name': 'Status', 'shape_name': 'RecommendationJobStatus'}]","['PENDING', 'IN_PROGRESS', 'COMPLETED', 'FAILED', 'STOPPING', 'STOPPED', 'DELETING', 'DELETED']"
 LabelingJob,resource,"['create', 'get', 'get_all']","['refresh', 'stop', 'wait']",[],[],"['CreateLabelingJob', 'DescribeLabelingJob', 'ListLabelingJobs', 'StopLabelingJob']","[{'name': 'LabelingJobStatus', 'shape_name': 'LabelingJobStatus'}]","['Initializing', 'InProgress', 'Completed', 'Failed', 'Stopping', 'Stopped']"
 LineageGroup,resource,"['get', 'get_all']",['refresh'],[],[],"['DescribeLineageGroup', 'ListLineageGroups']",[],[]
+Metrics,resource,[],[],[],[],[],[],[]
 MlflowTrackingServer,resource,"['create', 'get', 'get_all', 'start']","['delete', 'refresh', 'stop', 'update', 'wait_for_delete', 'wait_for_status']",[],[],"['CreateMlflowTrackingServer', 'DeleteMlflowTrackingServer', 'DescribeMlflowTrackingServer', 'ListMlflowTrackingServers', 'StartMlflowTrackingServer', 'StopMlflowTrackingServer', 'UpdateMlflowTrackingServer']","[{'name': 'TrackingServerStatus', 'shape_name': 'TrackingServerStatus'}]","['Creating', 'Created', 'CreateFailed', 'Updating', 'Updated', 'UpdateFailed', 'Deleting', 'DeleteFailed', 'Stopping', 'Stopped', 'StopFailed', 'Starting', 'Started', 'StartFailed', 'MaintenanceInProgress', 'MaintenanceComplete', 'MaintenanceFailed']"
 Model,resource,"['create', 'get', 'get_all']","['delete', 'refresh']",[],[],"['CreateModel', 'DeleteModel', 'DescribeModel', 'ListModels']",[],[]
 ModelBiasJobDefinition,resource,"['create', 'get', 'get_all']","['delete', 'refresh']",[],[],"['CreateModelBiasJobDefinition', 'DeleteModelBiasJobDefinition', 'DescribeModelBiasJobDefinition', 'ListModelBiasJobDefinitions']",[],[]

src/sagemaker_core/main/resources.py

@@ -15898,6 +15898,61 @@ def get_policy(
         return GetLineageGroupPolicyResponse(**transformed_response)
 
 
+class Metrics(Base):
+    """
+    Class representing resource Metrics
+
+    """
+
+    @Base.add_validate_call
+    def batch_get_metrics(
+        self,
+        metric_queries: List[MetricQuery],
+        session: Optional[Session] = None,
+        region: Optional[str] = None,
+    ) -> Optional[BatchGetMetricsResponse]:
+        """
+        Used to retrieve training metrics from SageMaker.
+
+        Parameters:
+            metric_queries: Queries made to retrieve training metrics from SageMaker.
+            session: Boto3 session.
+            region: Region name.
+
+        Returns:
+            BatchGetMetricsResponse
+
+        Raises:
+            botocore.exceptions.ClientError: This exception is raised for AWS service related errors.
+                The error message and error code can be parsed from the exception as follows:
+                ```
+                try:
+                    # AWS service call here
+                except botocore.exceptions.ClientError as e:
+                    error_message = e.response['Error']['Message']
+                    error_code = e.response['Error']['Code']
+                ```
+        """
+
+        operation_input_args = {
+            "MetricQueries": metric_queries,
+        }
+        # serialize the input request
+        operation_input_args = serialize(operation_input_args)
+        logger.debug(f"Serialized input request: {operation_input_args}")
+
+        client = Base.get_sagemaker_client(
+            session=session, region_name=region, service_name="sagemaker-metrics"
+        )
+
+        logger.debug(f"Calling batch_get_metrics API")
+        response = client.batch_get_metrics(**operation_input_args)
+        logger.debug(f"Response: {response}")
+
+        transformed_response = transform(response, "BatchGetMetricsResponse")
+        return BatchGetMetricsResponse(**transformed_response)
+
+
 class MlflowTrackingServer(Base):
     """
     Class representing resource MlflowTrackingServer
@@ -16518,6 +16573,7 @@ def wrapper(*args, **kwargs):
                         "s3_data_source": {
                             "s3_uri": {"type": "string"},
                             "s3_data_type": {"type": "string"},
+                            "manifest_s3_uri": {"type": "string"},
                         }
                     }
                 },

src/sagemaker_core/main/shapes.py

@@ -406,6 +406,18 @@ class MetricQueryResult(Base):
     message: Optional[str] = Unassigned()
 
 
+class BatchGetMetricsResponse(Base):
+    """
+    BatchGetMetricsResponse
+
+    Attributes
+    ----------------------
+    metric_query_results: The results of a query to retrieve training metrics from SageMaker.
+    """
+
+    metric_query_results: Optional[List[MetricQueryResult]] = Unassigned()
+
+
 class BatchPutMetricsError(Base):
     """
     BatchPutMetricsError

src/sagemaker_core/tools/additional_operations.json

@@ -171,6 +171,16 @@
             "service_name": "sagemaker-metrics"
         }
     },
+    "Metrics": {
+        "BatchGetMetrics": {
+            "operation_name": "BatchGetMetrics",
+            "resource_name": "Metrics",
+            "method_name": "batch_get_metrics",
+            "return_type": "BatchGetMetricsResponse",
+            "method_type": "object",
+            "service_name": "sagemaker-metrics"
+        }
+    },
     "HubContent": {
         "ListHubContentVersions": {
             "operation_name": "ListHubContentVersions",