Repository: getmoto/moto
Pull request #7936 (2024-08-05 15:33:48)

Feature: Add RGTAPI support for additional Sagemaker resource types
Adds ResourceGroupsTaggingAPI support for Sagemaker resource types: compilation-job, domain, model-explainability-job-definition, model-quality-job-definition, and hyper-parameter-tuning-job


Patch:
diff --git a/moto/resourcegroupstaggingapi/models.py b/moto/resourcegroupstaggingapi/models.py
index eb1f69d34..087693651 100644
--- a/moto/resourcegroupstaggingapi/models.py
+++ b/moto/resourcegroupstaggingapi/models.py
@@ -681,10 +681,15 @@ class ResourceGroupsTaggingAPIBackend(BaseBackend):
                     "Tags": tags,
                 }
 
-        # sagemaker cluster and automljob
+        # sagemaker cluster, automljob, compilation-job, domain, model-explainability-job-definition, model-quality-job-definition, and hyper-parameter-tuning-job currently supported
         sagemaker_resource_map: Dict[str, Dict[str, Any]] = {
             "sagemaker:cluster": self.sagemaker_backend.clusters,
             "sagemaker:automl-job": self.sagemaker_backend.auto_ml_jobs,
+            "sagemaker:compilation-job": self.sagemaker_backend.compilation_jobs,
+            "sagemaker:domain": self.sagemaker_backend.domains,
+            "sagemaker:model-explainability-job-definition": self.sagemaker_backend.model_explainability_job_definitions,
+            "sagemaker:model-quality-job-definition": self.sagemaker_backend.model_quality_job_definitions,
+            "sagemaker:hyper-parameter-tuning-job": self.sagemaker_backend.hyper_parameter_tuning_jobs,
         }
         for resource_type, resource_source in sagemaker_resource_map.items():
             if (
diff --git a/moto/sagemaker/models.py b/moto/sagemaker/models.py
index 5b97304ed..e0bdbb939 100644
--- a/moto/sagemaker/models.py
+++ b/moto/sagemaker/models.py
@@ -148,6 +148,12 @@ PAGINATION_MODEL = {
         "limit_default": 100,
         "unique_attribute": "model_card_arn",
     },
+    "list_model_bias_job_definitions": {
+        "input_token": "next_token",
+        "limit_key": "max_results",
+        "limit_default": 100,
+        "unique_attribute": "arn",
+    },
 }
 
 METRIC_INFO_TYPE = Dict[str, Union[str, int, float, datetime]]
@@ -2864,6 +2870,7 @@ class SageMakerModelBackend(BaseBackend):
         self.model_package_name_mapping: Dict[str, str] = {}
         self.feature_groups: Dict[str, FeatureGroup] = {}
         self.clusters: Dict[str, Cluster] = {}
+        self.model_bias_job_definitions: Dict[str, FakeModelBiasJobDefinition] = {}
         self.auto_ml_jobs: Dict[str, AutoMLJob] = {}
         self.compilation_jobs: Dict[str, CompilationJob] = {}
         self.domains: Dict[str, Domain] = {}
@@ -3002,6 +3009,7 @@ class SageMakerModelBackend(BaseBackend):
             "pipeline": self.pipelines,
             "model-package-group": self.model_package_groups,
             "cluster": self.clusters,
+            "model-bias-job-definition": self.model_bias_job_definitions,
             "automl-job": self.auto_ml_jobs,
             "compilation-job": self.compilation_jobs,
             "domain": self.domains,
@@ -4848,6 +4856,57 @@ class SageMakerModelBackend(BaseBackend):
             )
         return nodes_list
 
+    def create_model_bias_job_definition(
+        self,
+        account_id: str,
+        job_definition_name: str,
+        tags: List[Dict[str, str]] = [],
+        role_arn: str = "",
+        job_resources: Optional[Dict[str, Any]] = None,
+        stopping_condition: Optional[Dict[str, Any]] = None,
+        environment: Optional[Dict[str, str]] = None,
+        network_config: Optional[Dict[str, Any]] = None,
+        model_bias_baseline_config: Optional[Dict[str, Any]] = None,
+        model_bias_app_specification: Optional[Dict[str, Any]] = None,
+        model_bias_job_input: Optional[Dict[str, Any]] = None,
+        model_bias_job_output_config: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, str]:
+        job_definition = FakeModelBiasJobDefinition(
+            account_id=account_id,
+            region_name=self.region_name,
+            job_definition_name=job_definition_name,
+            tags=tags,
+            role_arn=role_arn,
+            job_resources=job_resources,
+            stopping_condition=stopping_condition,
+            environment=environment,
+            network_config=network_config,
+            model_bias_baseline_config=model_bias_baseline_config,
+            model_bias_app_specification=model_bias_app_specification,
+            model_bias_job_input=model_bias_job_input,
+            model_bias_job_output_config=model_bias_job_output_config,
+        )
+        self.model_bias_job_definitions[job_definition_name] = job_definition
+        return job_definition.response_create
+
+    @paginate(pagination_model=PAGINATION_MODEL)
+    def list_model_bias_job_definitions(self) -> List[Dict[str, str]]:
+        return [job.summary_object for job in self.model_bias_job_definitions.values()]
+
+    def describe_model_bias_job_definition(
+        self, job_definition_name: str
+    ) -> Dict[str, Any]:
+        job_definition = self.model_bias_job_definitions.get(job_definition_name)
+        if job_definition is None:
+            raise ResourceNotFound(f"Job definition {job_definition_name} not found")
+        return job_definition.response_object
+
+    def delete_model_bias_job_definition(self, job_definition_name: str) -> None:
+        if job_definition_name in self.model_bias_job_definitions:
+            del self.model_bias_job_definitions[job_definition_name]
+        else:
+            raise ResourceNotFound(f"Job definition {job_definition_name} not found")
+
     def create_auto_ml_job_v2(
         self,
         auto_ml_job_name: str,
@@ -5820,4 +5879,68 @@ class FakeTrialComponent(BaseObject):
         )
 
 
+class FakeModelBiasJobDefinition(BaseObject):
+    def __init__(
+        self,
+        account_id: str,
+        region_name: str,
+        job_definition_name: str,
+        tags: List[Dict[str, str]] = [],
+        role_arn: str = "",
+        job_resources: Optional[Dict[str, Any]] = None,
+        stopping_condition: Optional[Dict[str, Any]] = None,
+        environment: Optional[Dict[str, str]] = None,
+        network_config: Optional[Dict[str, Any]] = None,
+        model_bias_baseline_config: Optional[Dict[str, Any]] = None,
+        model_bias_app_specification: Optional[Dict[str, Any]] = None,
+        model_bias_job_input: Optional[Dict[str, Any]] = None,
+        model_bias_job_output_config: Optional[Dict[str, Any]] = None,
+    ):
+        self.job_definition_name = job_definition_name
+        self.job_definition_arn = FakeModelBiasJobDefinition.arn_formatter(
+            job_definition_name, account_id, region_name
+        )
+        self.tags = tags
+        self.role_arn = role_arn
+        self.job_resources = job_resources or {}
+        self.stopping_condition = stopping_condition or {}
+        self.environment = environment or {}
+        self.network_config = network_config or {}
+        self.model_bias_baseline_config = model_bias_baseline_config or {}
+        self.model_bias_app_specification = model_bias_app_specification or {}
+        self.model_bias_job_input = model_bias_job_input or {}
+        self.model_bias_job_output_config = model_bias_job_output_config or {}
+        self.creation_time = self.last_modified_time = datetime.now().strftime(
+            "%Y-%m-%d %H:%M:%S"
+        )
+
+    @property
+    def response_object(self) -> Dict[str, str]:
+        response_object = self.gen_response_object()
+        return {
+            k: v for k, v in response_object.items() if v is not None and v != [None]
+        }
+
+    @property
+    def response_create(self) -> Dict[str, str]:
+        return {"JobDefinitionArn": self.job_definition_arn}
+
+    @staticmethod
+    def arn_formatter(name: str, account_id: str, region: str) -> str:
+        return (
+            f"arn:aws:sagemaker:{region}:{account_id}:model-bias-job-definition/{name}"
+        )
+
+    @property
+    def summary_object(self) -> Dict[str, str]:
+        return {
+            "MonitoringJobDefinitionName": self.job_definition_name,
+            "MonitoringJobDefinitionArn": self.job_definition_arn,
+            "CreationTime": self.creation_time,
+            "EndpointName": self.model_bias_job_input.get("EndpointInput", {}).get(
+                "EndpointName", "EndpointName"
+            ),
+        }
+
+
 sagemaker_backends = BackendDict(SageMakerModelBackend, "sagemaker")
diff --git a/moto/sagemaker/responses.py b/moto/sagemaker/responses.py
index df44a2ce6..0226ce4ce 100644
--- a/moto/sagemaker/responses.py
+++ b/moto/sagemaker/responses.py
@@ -1024,6 +1024,52 @@ class SageMakerResponse(BaseResponse):
             dict(NextToken=next_token, ClusterNodeSummaries=cluster_node_summaries)
         )
 
+    def create_model_bias_job_definition(self) -> str:
+        account_id = self.current_account
+        job_definition_name = self._get_param("JobDefinitionName")
+        tags = self._get_param("Tags", [])
+        role_arn = self._get_param("RoleArn")
+        job_resources = self._get_param("JobResources")
+        stopping_condition = self._get_param("StoppingCondition")
+        environment = self._get_param("Environment", {})
+        network_config = self._get_param("NetworkConfig", {})
+        model_bias_baseline_config = self._get_param("ModelBiasBaselineConfig", {})
+        model_bias_app_specification = self._get_param("ModelBiasAppSpecification")
+        model_bias_job_input = self._get_param("ModelBiasJobInput")
+        model_bias_job_output_config = self._get_param("ModelBiasJobOutputConfig")
+
+        response = self.sagemaker_backend.create_model_bias_job_definition(
+            account_id=account_id,
+            job_definition_name=job_definition_name,
+            tags=tags,
+            role_arn=role_arn,
+            job_resources=job_resources,
+            stopping_condition=stopping_condition,
+            environment=environment,
+            network_config=network_config,
+            model_bias_baseline_config=model_bias_baseline_config,
+            model_bias_app_specification=model_bias_app_specification,
+            model_bias_job_input=model_bias_job_input,
+            model_bias_job_output_config=model_bias_job_output_config,
+        )
+        return json.dumps(response)
+
+    def list_model_bias_job_definitions(self) -> str:
+        result, next_token = self.sagemaker_backend.list_model_bias_job_definitions()
+        return json.dumps({"JobDefinitionSummaries": result, "NextToken": next_token})
+
+    def describe_model_bias_job_definition(self) -> str:
+        job_definition_name = self._get_param("JobDefinitionName")
+        job_definition = self.sagemaker_backend.describe_model_bias_job_definition(
+            job_definition_name
+        )
+        return json.dumps(job_definition)
+
+    def delete_model_bias_job_definition(self) -> str:
+        job_definition_name = self._get_param("JobDefinitionName")
+        self.sagemaker_backend.delete_model_bias_job_definition(job_definition_name)
+        return json.dumps({})
+
     def create_auto_ml_job_v2(self) -> str:
         auto_ml_job_name = self._get_param("AutoMLJobName")
         auto_ml_job_input_data_config = self._get_param("AutoMLJobInputDataConfig")


Test patch:
diff --git a/tests/test_sagemaker/test_sagemaker_model_bias_job_definitions.py b/tests/test_sagemaker/test_sagemaker_model_bias_job_definitions.py
new file mode 100644
index 000000000..013c34adc
--- /dev/null
+++ b/tests/test_sagemaker/test_sagemaker_model_bias_job_definitions.py
@@ -0,0 +1,252 @@
+"""Unit tests for sagemaker-supported APIs."""
+
+import boto3
+
+from moto import mock_aws
+from moto.core import DEFAULT_ACCOUNT_ID as ACCOUNT_ID
+
+
+@mock_aws
+def test_create_model_bias_job():
+    client = boto3.client("sagemaker", region_name="us-west-2")
+
+    response = client.create_model_bias_job_definition(
+        JobDefinitionName="test-bias-job",
+        RoleArn=f"arn:aws:iam::{ACCOUNT_ID}:role/SageMakerRole",
+        ModelBiasAppSpecification={
+            "ImageUri": f"{ACCOUNT_ID}.dkr.ecr.us-west-2.amazonaws.com/my-image:latest",
+            "ConfigUri": "s3://my-bucket/bias-config.json",
+        },
+        ModelBiasJobInput={
+            "EndpointInput": {
+                "EndpointName": "test-endpoint",
+                "LocalPath": "/opt/ml/processing/input/endpoint",
+                "S3InputMode": "File",
+                "S3DataDistributionType": "FullyReplicated",
+            },
+            "GroundTruthS3Input": {
+                "S3Uri": "s3://my-bucket/ground-truth-data",
+            },
+        },
+        ModelBiasJobOutputConfig={
+            "MonitoringOutputs": [
+                {
+                    "S3Output": {
+                        "S3Uri": "s3://my-bucket/bias-output",
+                        "LocalPath": "/opt/ml/processing/output",
+                        "S3UploadMode": "EndOfJob",
+                    },
+                },
+            ],
+        },
+        JobResources={
+            "ClusterConfig": {
+                "InstanceCount": 1,
+                "InstanceType": "ml.m5.xlarge",
+                "VolumeSizeInGB": 10,
+            },
+        },
+        NetworkConfig={
+            "EnableInterContainerTrafficEncryption": False,
+            "EnableNetworkIsolation": False,
+        },
+    )
+
+    expected_job_definition_arn = f"arn:aws:sagemaker:us-west-2:{ACCOUNT_ID}:model-bias-job-definition/test-bias-job"
+    assert response["JobDefinitionArn"] == expected_job_definition_arn
+
+
+@mock_aws
+def test_describe_model_bias_job():
+    client = boto3.client("sagemaker", region_name="us-east-1")
+
+    client.create_model_bias_job_definition(
+        JobDefinitionName="test-bias-job",
+        RoleArn=f"arn:aws:iam::{ACCOUNT_ID}:role/SageMakerRole",
+        ModelBiasAppSpecification={
+            "ImageUri": f"{ACCOUNT_ID}.dkr.ecr.us-west-2.amazonaws.com/my-image:latest",
+            "ConfigUri": "s3://my-bucket/bias-config.json",
+        },
+        ModelBiasJobInput={
+            "EndpointInput": {
+                "EndpointName": "test-endpoint",
+                "LocalPath": "/opt/ml/processing/input/endpoint",
+                "S3InputMode": "File",
+                "S3DataDistributionType": "FullyReplicated",
+            },
+            "GroundTruthS3Input": {
+                "S3Uri": "s3://my-bucket/ground-truth-data",
+            },
+        },
+        ModelBiasJobOutputConfig={
+            "MonitoringOutputs": [
+                {
+                    "S3Output": {
+                        "S3Uri": "s3://my-bucket/bias-output",
+                        "LocalPath": "/opt/ml/processing/output",
+                        "S3UploadMode": "EndOfJob",
+                    },
+                },
+            ],
+        },
+        JobResources={
+            "ClusterConfig": {
+                "InstanceCount": 1,
+                "InstanceType": "ml.m5.xlarge",
+                "VolumeSizeInGB": 10,
+            },
+        },
+        NetworkConfig={
+            "EnableInterContainerTrafficEncryption": False,
+            "EnableNetworkIsolation": False,
+        },
+    )
+    response = client.describe_model_bias_job_definition(
+        JobDefinitionName="test-bias-job"
+    )
+    assert response["JobDefinitionName"] == "test-bias-job"
+
+
+@mock_aws
+def test_list_model_bias_jobs():
+    client = boto3.client("sagemaker", region_name="us-east-1")
+    client.create_model_bias_job_definition(
+        JobDefinitionName="test-bias-job-1",
+        RoleArn=f"arn:aws:iam::{ACCOUNT_ID}:role/SageMakerRole",
+        ModelBiasAppSpecification={
+            "ImageUri": f"{ACCOUNT_ID}.dkr.ecr.us-west-2.amazonaws.com/my-image:latest",
+            "ConfigUri": "s3://my-bucket/bias-config.json",
+        },
+        ModelBiasJobInput={
+            "EndpointInput": {
+                "EndpointName": "test-endpoint",
+                "LocalPath": "/opt/ml/processing/input/endpoint",
+                "S3InputMode": "File",
+                "S3DataDistributionType": "FullyReplicated",
+            },
+            "GroundTruthS3Input": {
+                "S3Uri": "s3://my-bucket/ground-truth-data",
+            },
+        },
+        ModelBiasJobOutputConfig={
+            "MonitoringOutputs": [
+                {
+                    "S3Output": {
+                        "S3Uri": "s3://my-bucket/bias-output",
+                        "LocalPath": "/opt/ml/processing/output",
+                        "S3UploadMode": "EndOfJob",
+                    },
+                },
+            ],
+        },
+        JobResources={
+            "ClusterConfig": {
+                "InstanceCount": 1,
+                "InstanceType": "ml.m5.xlarge",
+                "VolumeSizeInGB": 10,
+            },
+        },
+        NetworkConfig={
+            "EnableInterContainerTrafficEncryption": False,
+            "EnableNetworkIsolation": False,
+        },
+    )
+    client.create_model_bias_job_definition(
+        JobDefinitionName="test-bias-job-2",
+        RoleArn=f"arn:aws:iam::{ACCOUNT_ID}:role/SageMakerRole",
+        ModelBiasAppSpecification={
+            "ImageUri": f"{ACCOUNT_ID}.dkr.ecr.us-west-2.amazonaws.com/my-image:latest",
+            "ConfigUri": "s3://my-bucket/bias-config.json",
+        },
+        ModelBiasJobInput={
+            "EndpointInput": {
+                "EndpointName": "test-endpoint",
+                "LocalPath": "/opt/ml/processing/input/endpoint",
+                "S3InputMode": "File",
+                "S3DataDistributionType": "FullyReplicated",
+            },
+            "GroundTruthS3Input": {
+                "S3Uri": "s3://my-bucket/ground-truth-data",
+            },
+        },
+        ModelBiasJobOutputConfig={
+            "MonitoringOutputs": [
+                {
+                    "S3Output": {
+                        "S3Uri": "s3://my-bucket/bias-output",
+                        "LocalPath": "/opt/ml/processing/output",
+                        "S3UploadMode": "EndOfJob",
+                    },
+                },
+            ],
+        },
+        JobResources={
+            "ClusterConfig": {
+                "InstanceCount": 1,
+                "InstanceType": "ml.m5.xlarge",
+                "VolumeSizeInGB": 10,
+            },
+        },
+        NetworkConfig={
+            "EnableInterContainerTrafficEncryption": False,
+            "EnableNetworkIsolation": False,
+        },
+    )
+    response = client.list_model_bias_job_definitions()
+
+    job_names = [
+        job["MonitoringJobDefinitionName"] for job in response["JobDefinitionSummaries"]
+    ]
+    assert "test-bias-job-1" in job_names
+    assert "test-bias-job-2" in job_names
+
+
+@mock_aws
+def test_delete_model_bias_job():
+    client = boto3.client("sagemaker", region_name="us-east-1")
+
+    client.create_model_bias_job_definition(
+        JobDefinitionName="test-bias-job",
+        RoleArn=f"arn:aws:iam::{ACCOUNT_ID}:role/SageMakerRole",
+        ModelBiasAppSpecification={
+            "ImageUri": f"{ACCOUNT_ID}.dkr.ecr.us-west-2.amazonaws.com/my-image:latest",
+            "ConfigUri": "s3://my-bucket/bias-config.json",
+        },
+        ModelBiasJobInput={
+            "EndpointInput": {
+                "EndpointName": "test-endpoint",
+                "LocalPath": "/opt/ml/processing/input/endpoint",
+                "S3InputMode": "File",
+                "S3DataDistributionType": "FullyReplicated",
+            },
+            "GroundTruthS3Input": {
+                "S3Uri": "s3://my-bucket/ground-truth-data",
+            },
+        },
+        ModelBiasJobOutputConfig={
+            "MonitoringOutputs": [
+                {
+                    "S3Output": {
+                        "S3Uri": "s3://my-bucket/bias-output",
+                        "LocalPath": "/opt/ml/processing/output",
+                        "S3UploadMode": "EndOfJob",
+                    },
+                },
+            ],
+        },
+        JobResources={
+            "ClusterConfig": {
+                "InstanceCount": 1,
+                "InstanceType": "ml.m5.xlarge",
+                "VolumeSizeInGB": 10,
+            },
+        },
+        NetworkConfig={
+            "EnableInterContainerTrafficEncryption": False,
+            "EnableNetworkIsolation": False,
+        },
+    )
+    client.delete_model_bias_job_definition(JobDefinitionName="test-bias-job")
+    response = client.list_model_bias_job_definitions()
+    job_names = [job["JobDefinitionName"] for job in response["JobDefinitionSummaries"]]
+    assert "test-bias-job" not in job_names