getmoto / moto

A library that allows you to easily mock out tests based on AWS infrastructure.
http://docs.getmoto.org/en/latest/
Apache License 2.0
7.6k stars 2.03k forks source link

Inconsistent behavior of SageMaker Client Search Method for Model Packages #8155

Open malmeida1 opened 4 days ago

malmeida1 commented 4 days ago

Environment Information

Python Version: 3.10.14 boto3 == 1.24.4 pytest == 8.3.2 moto == 5.0.13

Problem Description

I believe this issue is very similar to #7898. I'm also using moto to test the retrieval of models, this time model packages instead of model package groups. I've created some model packages using create_model_package and when using the Resource ModelPackage in sagemaker_client.search, no model packages were returned. On the other hand, I had results returned when using the Resource ModelPackageGroup.

I also tried to list_model_packages with the same set of created model packages and had two outcomes:

  1. When using ModelPackageGroupName: the models that were under the provided model package group were returned.
  2. When not using any parameters: an empty list was returned.

Minimal Example

search_model_package.py

import boto3

class ModelPackageFinder:
    def __init__(self):
        self.sm_client = boto3.client('sagemaker')

    def get_model_groups_with_search_method(self):
        search_results = self.sm_client.search(
            Resource='ModelPackage',
            SearchExpression={
                'Filters': [{
                    'Name': 'ModelPackageName', 
                    'Operator': 'Contains', 
                    'Value': 'model-test'}]}
        )
        return search_results

    def get_model_packages_with_list_model_packages_method(self):
        list_model_results = self.sm_client.list_model_packages()
        return list_model_results

    def get_model_packages_with_list_model_packages_method_and_group(self):
        list_model_results = self.sm_client.list_model_packages(ModelPackageGroupName="model-test_DEV")
        return list_model_results

test_model_groups.py

import boto3
from moto import mock_aws
from search_model_package import ModelPackageFinder
from pprint import pprint

MOCK_SOURCE_ROLE_ARN = 'arn:aws:iam::123456789012:role/DSSourceRole'

def mock_model_groups(sm_client):
    sm_client.create_model_package_group(
        ModelPackageGroupName='model-test_DEV',
        ModelPackageGroupDescription='Model group in DEV environment.',
        Tags=[
            {
                'Key': 'RegisteredBySourceRoleArn',
                'Value': MOCK_SOURCE_ROLE_ARN
            }
        ]
    )
    sm_client.create_model_package_group(
        ModelPackageGroupName='model-test_PROD',
        ModelPackageGroupDescription='Model group in PROD environment.',
        Tags=[
            {
                'Key': 'RegisteredBySourceRoleArn',
                'Value': MOCK_SOURCE_ROLE_ARN
            }
        ]
    )

def mock_model_packages(sm_client):
    sm_client.create_model_package(
        ModelPackageGroupName='model-test_DEV',
        ModelPackageDescription="Model package version in DEV environment.",
        Tags=[
            {"Key": "RegisteredUsingSofiaSDK", "Value": "true"},
            {"Key": "RegisteredBySourceRoleArn", "Value": MOCK_SOURCE_ROLE_ARN},
        ],
    )
    sm_client.create_model_package(
        ModelPackageGroupName='model-test_PROD',
        ModelPackageDescription="Model package version in PROD environment.",
        Tags=[
            {"Key": "RegisteredUsingSofiaSDK", "Value": "true"},
            {"Key": "RegisteredBySourceRoleArn", "Value": MOCK_SOURCE_ROLE_ARN},
        ],
    )
    sm_client.create_model_package(
        ModelPackageGroupName='model-test_DEV',
        ModelPackageDescription="Model package version in DEV environment with change in hyperparameter.",
        Tags=[
            {"Key": "RegisteredUsingSofiaSDK", "Value": "true"},
            {"Key": "RegisteredBySourceRoleArn", "Value": MOCK_SOURCE_ROLE_ARN},
        ]
    )

@mock_aws
def test_search_function():
    sm_client = boto3.client('sagemaker')
    mock_model_groups(sm_client)
    mock_model_packages(sm_client)
    model_group_finder = ModelPackageFinder()
    model_groups = model_group_finder.get_model_groups_with_search_method()
    print('Sagemaker Search Function: ')
    pprint(model_groups)

@mock_aws
def test_list_function():
    sm_client = boto3.client('sagemaker')
    mock_model_groups(sm_client)
    mock_model_packages(sm_client)
    model_group_finder = ModelPackageFinder()
    model_groups = model_group_finder.get_model_packages_with_list_model_packages_method()
    print('List Model Package Function')
    pprint(model_groups)

@mock_aws
def test_list_function_with_parameter():
    sm_client = boto3.client('sagemaker')
    mock_model_groups(sm_client)
    mock_model_packages(sm_client)
    model_group_finder = ModelPackageFinder()
    model_groups = model_group_finder.get_model_packages_with_list_model_packages_method_and_group()
    print('List Model Package Function With Package Group as Parameter')
    pprint(model_groups)

When running pytest -s test_model_packages.py, the following print statements were shown in the console:

tests/test_model_packages.py Sagemaker Search Function:

{'ResponseMetadata': {'HTTPHeaders': {'date': 'Tue, 24 Sep 2024 16:37:25 GMT',
                                      'server': 'amazon.com',
                                      'x-amzn-requestid': 'LHRdZh1KJTBwIxP4aRRTLsFWN06Git5j9i3tx0ZoroLkJ8AC8aXZ'},
                      'HTTPStatusCode': 200,
                      'RequestId': 'LHRdZh1KJTBwIxP4aRRTLsFWN06Git5j9i3tx0ZoroLkJ8AC8aXZ',
                      'RetryAttempts': 0},
 'Results': []}

.List Model Package Function


{'ModelPackageSummaryList': [],
 'ResponseMetadata': {'HTTPHeaders': {'date': 'Tue, 24 Sep 2024 16:37:25 GMT',
                                      'server': 'amazon.com',
                                      'x-amzn-requestid': 'A7ZD2v5G8YWUvjCBRYa66aqmdQG2qWiSdCA72g0CD2nDBMIvvpQp'},
                      'HTTPStatusCode': 200,
                      'RequestId': 'A7ZD2v5G8YWUvjCBRYa66aqmdQG2qWiSdCA72g0CD2nDBMIvvpQp',
                      'RetryAttempts': 0}}

.List Model Package Function With Package Group as Parameter

{'ModelPackageSummaryList': [{'CreationTime': datetime.datetime(2024, 9, 24, 19, 37, 25, 956457, tzinfo=tzutc()),
                              'ModelPackageArn': 'arn:aws:sagemaker:us-east-1:123456789012:model-package/model-test_dev/1',
                              'ModelPackageDescription': 'Model package '
                                                         'version in DEV '
                                                         'environment.',
                              'ModelPackageGroupName': 'model-test_DEV',
                              'ModelPackageStatus': 'Completed',
                              'ModelPackageVersion': 1},
                             {'CreationTime': datetime.datetime(2024, 9, 24, 19, 37, 25, 957802, tzinfo=tzutc()),
                              'ModelPackageArn': 'arn:aws:sagemaker:us-east-1:123456789012:model-package/model-test_dev/2',
                              'ModelPackageDescription': 'Model package '
                                                         'version in DEV '
                                                         'environment with '
                                                         'change in '
                                                         'hyperparameter.',
                              'ModelPackageGroupName': 'model-test_DEV',
                              'ModelPackageStatus': 'Completed',
                              'ModelPackageVersion': 2}],
 'ResponseMetadata': {'HTTPHeaders': {'date': 'Tue, 24 Sep 2024 16:37:25 GMT',
                                      'server': 'amazon.com',
                                      'x-amzn-requestid': 'OTR8HEhDYNVT6N5Z7AaC5fn8ZCJAfBfbmjiJeFvhzopSJFWNcI8e'},
                      'HTTPStatusCode': 200,
                      'RequestId': 'OTR8HEhDYNVT6N5Z7AaC5fn8ZCJAfBfbmjiJeFvhzopSJFWNcI8e',
                      'RetryAttempts': 0}}

As seen above, the sm_client.search did not return any results when using the ModelPackage Resource. The sm_client.list_model_packages method also did not return anything when passing no parameters, but it did when passing the ModelPackageGroupName parameter.

bblommers commented 3 days ago

Hi @malmeida1, thank you for raising this, and welcome to Moto! I'll mark it as a bug.

Would you like to open a PR to fix this yourself?