sagemaker: additional InferenceComponent support (#1686)

Signed-off-by: Tristan Burgess <[email protected]>

Author: tristanburgess

pkg/config/services.go

@@ -950,9 +950,11 @@ var SupportedServices = serviceConfigs{
 		Alias:     "sagemaker",
 		ResourceFilters: []*string{
 			aws.String("sagemaker:endpoint"),
+			aws.String("sagemaker:inference-component"),
 		},
 		DimensionRegexps: []*regexp.Regexp{
 			regexp.MustCompile(":endpoint/(?P<EndpointName>[^/]+)$"),
+			regexp.MustCompile(":inference-component/(?P<InferenceComponentName>[^/]+)$"),
 		},
 	},
 	{
@@ -965,6 +967,16 @@ var SupportedServices = serviceConfigs{
 			regexp.MustCompile(":endpoint/(?P<EndpointName>[^/]+)$"),
 		},
 	},
+	{
+		Namespace: "/aws/sagemaker/InferenceComponents",
+		Alias:     "sagemaker-inference-components",
+		ResourceFilters: []*string{
+			aws.String("sagemaker:inference-component"),
+		},
+		DimensionRegexps: []*regexp.Regexp{
+			regexp.MustCompile(":inference-component/(?P<InferenceComponentName>[^/]+)$"),
+		},
+	},
 	{
 		Namespace: "/aws/sagemaker/TrainingJobs",
 		Alias:     "sagemaker-training",

pkg/job/maxdimassociator/associator.go

@@ -255,10 +255,10 @@ func fixDimension(namespace string, dim model.Dimension) (model.Dimension, bool)
 		}
 	}
 
-	// AWS Sagemaker endpoint name may have upper case characters
+	// AWS Sagemaker endpoint name and inference component name may have upper case characters
 	// Resource ARN is only in lower case, hence transforming
-	// endpoint name value to be able to match the resource ARN
-	if namespace == "AWS/SageMaker" && dim.Name == "EndpointName" {
+	// name value to be able to match the resource ARN
+	if namespace == "AWS/SageMaker" && (dim.Name == "EndpointName" || dim.Name == "InferenceComponentName") {
 		dim.Value = strings.ToLower(dim.Value)
 		return dim, true
 	}

pkg/job/maxdimassociator/associator_sagemaker_inf_component_test.go

@@ -0,0 +1,75 @@
+// Copyright 2024 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package maxdimassociator
+
+import (
+	"testing"
+
+	"github.com/prometheus/common/promslog"
+	"github.com/stretchr/testify/require"
+
+	"github.com/prometheus-community/yet-another-cloudwatch-exporter/pkg/config"
+	"github.com/prometheus-community/yet-another-cloudwatch-exporter/pkg/model"
+)
+
+var sagemakerInfComponentJobOne = &model.TaggedResource{
+	ARN:       "arn:aws:sagemaker:us-west-2:123456789012:inference-component/example-inference-component-one",
+	Namespace: "/aws/sagemaker/InferenceComponents",
+}
+
+var sagemakerInfComponentJobResources = []*model.TaggedResource{
+	sagemakerInfComponentJobOne,
+}
+
+func TestAssociatorSagemakerInfComponentJob(t *testing.T) {
+	type args struct {
+		dimensionRegexps []model.DimensionsRegexp
+		resources        []*model.TaggedResource
+		metric           *model.Metric
+	}
+
+	type testCase struct {
+		name             string
+		args             args
+		expectedSkip     bool
+		expectedResource *model.TaggedResource
+	}
+
+	testcases := []testCase{
+		{
+			name: "1 dimension should not match but not skip",
+			args: args{
+				dimensionRegexps: config.SupportedServices.GetService("/aws/sagemaker/InferenceComponents").ToModelDimensionsRegexp(),
+				resources:        sagemakerInfComponentJobResources,
+				metric: &model.Metric{
+					MetricName: "CPUUtilizationNormalized",
+					Namespace:  "/aws/sagemaker/InferenceComponents",
+					Dimensions: []model.Dimension{
+						{Name: "InferenceComponentName", Value: "example-inference-component-one"},
+					},
+				},
+			},
+			expectedSkip:     false,
+			expectedResource: sagemakerInfComponentJobOne,
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			associator := NewAssociator(promslog.NewNopLogger(), tc.args.dimensionRegexps, tc.args.resources)
+			res, skip := associator.AssociateMetricToResource(tc.args.metric)
+			require.Equal(t, tc.expectedSkip, skip)
+			require.Equal(t, tc.expectedResource, res)
+		})
+	}
+}

pkg/job/maxdimassociator/associator_sagemaker_test.go

@@ -37,10 +37,22 @@ var sagemakerEndpointInvocationUpper = &model.TaggedResource{
 	Namespace: "AWS/SageMaker",
 }
 
+var sagemakerInferenceComponentInvocationOne = &model.TaggedResource{
+	ARN:       "arn:aws:sagemaker:us-west-2:123456789012:inference-component/example-inference-component-one",
+	Namespace: "AWS/SageMaker",
+}
+
+var sagemakerInferenceComponentInvocationUpper = &model.TaggedResource{
+	ARN:       "arn:aws:sagemaker:us-west-2:123456789012:inference-component/example-inference-component-upper",
+	Namespace: "AWS/SageMaker",
+}
+
 var sagemakerInvocationResources = []*model.TaggedResource{
 	sagemakerEndpointInvocationOne,
 	sagemakerEndpointInvocationTwo,
 	sagemakerEndpointInvocationUpper,
+	sagemakerInferenceComponentInvocationOne,
+	sagemakerInferenceComponentInvocationUpper,
 }
 
 func TestAssociatorSagemaker(t *testing.T) {
@@ -127,6 +139,38 @@ func TestAssociatorSagemaker(t *testing.T) {
 			expectedSkip:     false,
 			expectedResource: sagemakerEndpointInvocationUpper,
 		},
+		{
+			name: "inference component match",
+			args: args{
+				dimensionRegexps: config.SupportedServices.GetService("AWS/SageMaker").ToModelDimensionsRegexp(),
+				resources:        sagemakerInvocationResources,
+				metric: &model.Metric{
+					MetricName: "ModelLatency",
+					Namespace:  "AWS/SageMaker",
+					Dimensions: []model.Dimension{
+						{Name: "InferenceComponentName", Value: "example-inference-component-one"},
+					},
+				},
+			},
+			expectedSkip:     false,
+			expectedResource: sagemakerInferenceComponentInvocationOne,
+		},
+		{
+			name: "inference component match in Upper case",
+			args: args{
+				dimensionRegexps: config.SupportedServices.GetService("AWS/SageMaker").ToModelDimensionsRegexp(),
+				resources:        sagemakerInvocationResources,
+				metric: &model.Metric{
+					MetricName: "ModelLatency",
+					Namespace:  "AWS/SageMaker",
+					Dimensions: []model.Dimension{
+						{Name: "InferenceComponentName", Value: "Example-Inference-Component-Upper"},
+					},
+				},
+			},
+			expectedSkip:     false,
+			expectedResource: sagemakerInferenceComponentInvocationUpper,
+		},
 	}
 
 	for _, tc := range testcases {