diff --git a/sagemaker-core/src/sagemaker/core/resources.py b/sagemaker-core/src/sagemaker/core/resources.py index 19a314fc3f..a780c829eb 100644 --- a/sagemaker-core/src/sagemaker/core/resources.py +++ b/sagemaker-core/src/sagemaker/core/resources.py @@ -22984,7 +22984,11 @@ def refresh( """ operation_input_args = { - "ModelPackageName": self.model_package_name, + "ModelPackageName": ( + self.model_package_name + if not isinstance(self.model_package_name, Unassigned) + else self.model_package_arn + ), } # serialize the input request operation_input_args = serialize(operation_input_args) diff --git a/sagemaker-core/src/sagemaker/core/tools/resources_codegen.py b/sagemaker-core/src/sagemaker/core/tools/resources_codegen.py index d6cc9ba325..27b5ccea8c 100644 --- a/sagemaker-core/src/sagemaker/core/tools/resources_codegen.py +++ b/sagemaker-core/src/sagemaker/core/tools/resources_codegen.py @@ -98,6 +98,14 @@ REQUIRED = "required" GLOBAL_DEFAULTS = "GlobalDefaults" +# Some resources do not return their *Name identifier from the List API for certain +# variants (e.g. a versioned ModelPackage returns only ModelPackageArn). Because the +# Describe API accepts either the name or the ARN for the *Name parameter, refresh() +# falls back to the ARN attribute when the name identifier is Unassigned. See issue #5606. +REFRESH_IDENTIFIER_FALLBACKS = { + "ModelPackage": {"ModelPackageName": "model_package_arn"}, +} + class ResourcesCodeGen: """A class for generating resources based on a service JSON file. @@ -761,7 +769,10 @@ def _generate_operation_input_args( return operation_input_args def _generate_operation_input_necessary_args( - self, resource_operation: dict, resource_attributes: list + self, + resource_operation: dict, + resource_attributes: list, + identifier_fallbacks: dict = None, ) -> str: """Generate the operation input arguments string. @@ -770,7 +781,10 @@ def _generate_operation_input_necessary_args( Args: resource_operation (dict): The resource operation dictionary. - is_class_method (bool): Indicates method is class method, else object method. + resource_attributes (list): The resource object attributes available on self. + identifier_fallbacks (dict, optional): Maps an input member name to a resource + attribute to fall back to when the member's own attribute is Unassigned + (e.g. {"ModelPackageName": "model_package_arn"}). See issue #5606. Returns: str: The formatted operation input arguments string. @@ -778,12 +792,22 @@ def _generate_operation_input_necessary_args( input_shape_name = resource_operation["input"]["shape"] input_shape_members = list(self.shapes[input_shape_name]["members"].keys()) + identifier_fallbacks = identifier_fallbacks or {} args = list() for member in input_shape_members: - if convert_to_snake_case(member) in resource_attributes: - args.append(f"'{member}': self.{convert_to_snake_case(member)}") + snake_member = convert_to_snake_case(member) + if snake_member in resource_attributes: + if member in identifier_fallbacks: + fallback_attr = identifier_fallbacks[member] + args.append( + f"'{member}': self.{snake_member} " + f"if not isinstance(self.{snake_member}, Unassigned) " + f"else self.{fallback_attr}" + ) + else: + args.append(f"'{member}': self.{snake_member}") else: - args.append(f"'{member}': {convert_to_snake_case(member)}") + args.append(f"'{member}': {snake_member}") operation_input_args = ",\n".join(args) operation_input_args += "," @@ -1300,7 +1324,9 @@ def generate_refresh_method(self, resource_name: str, **kwargs) -> str: ) operation_input_args = self._generate_operation_input_necessary_args( - operation_metadata, kwargs["resource_attributes"] + operation_metadata, + kwargs["resource_attributes"], + identifier_fallbacks=REFRESH_IDENTIFIER_FALLBACKS.get(resource_name), ) operation = convert_to_snake_case(operation_name) diff --git a/sagemaker-core/tests/unit/generated/test_resources.py b/sagemaker-core/tests/unit/generated/test_resources.py index 260408df19..d192443dec 100644 --- a/sagemaker-core/tests/unit/generated/test_resources.py +++ b/sagemaker-core/tests/unit/generated/test_resources.py @@ -935,3 +935,40 @@ def test_no_synthesis_without_output_data_config(self): ) assert isinstance(job.model_artifacts, Unassigned) + + +class TestModelPackageRefreshIdentifierFallback(unittest.TestCase): + """ModelPackage.refresh() must work for versioned packages (issue #5606). + + ``ListModelPackages`` does not return ``ModelPackageName`` for versioned packages + inside a model package group, only ``ModelPackageArn``. Since ``DescribeModelPackage`` + accepts either for the ``ModelPackageName`` parameter, refresh() falls back to the ARN + when the name is ``Unassigned``. + """ + + _ARN = "arn:aws:sagemaker:us-west-2:123456789012:model-package/my-group/1" + + def _refresh_with(self, **attrs): + from sagemaker.core.resources import ModelPackage + + pkg = ModelPackage.model_construct(**attrs) + with ( + patch("sagemaker.core.resources.transform", return_value={}), + patch.object(Base, "get_sagemaker_client") as mock_get_client, + ): + client = mock_get_client.return_value + client.describe_model_package.return_value = {} + pkg.refresh() + return client + + def test_refresh_falls_back_to_arn_when_name_unassigned(self): + from sagemaker.core.utils.utils import Unassigned + + client = self._refresh_with(model_package_name=Unassigned(), model_package_arn=self._ARN) + client.describe_model_package.assert_called_once_with(ModelPackageName=self._ARN) + + def test_refresh_uses_name_when_present(self): + client = self._refresh_with( + model_package_name="my-model-package", model_package_arn=self._ARN + ) + client.describe_model_package.assert_called_once_with(ModelPackageName="my-model-package")