fix: deduplicate principals in sagemaker-templates ModelPackageGroupPolicy (#397)

* fix: deduplicate principals in sagemaker-templates ModelPackageGroupPolicy

Fixes "Invalid policy provided: Duplicate principal" error when dev,
pre-prod, and prod account IDs resolve to the same AWS account in
single-account deployments.

* style: fix ruff formatting

* style: fix ruff formatting for older version

Author: weklund

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### **Changed**
 
+- fixed duplicate principal error in `sagemaker-templates` module when dev, pre-prod, and prod account IDs resolve to the same AWS account
 - fixed `sagemaker-templates` Model Deploy seed code incorrectly granting S3 permissions to a ManagedPolicy instead of a Role, which caused deployment failures with CDK 2.174.0+
 - update qs to 6.14.1 via npm override to address security vulnerability
 - pin @cdklabs/generative-ai-cdk-constructs to 0.1.311 to fix build compatibility

modules/sagemaker/sagemaker-templates/templates/finetune_llm_evaluation/product_stack.py

@@ -48,6 +48,9 @@ def __init__(
         pre_prod_account_id = Aws.ACCOUNT_ID if not pre_prod_account_id else pre_prod_account_id
         prod_account_id = Aws.ACCOUNT_ID if not prod_account_id else prod_account_id
 
+        # Deduplicate account IDs to avoid "Duplicate principal" errors in single-account deployments
+        unique_account_ids = list(dict.fromkeys([dev_account_id, pre_prod_account_id, prod_account_id]))
+
         Tags.of(self).add("sagemaker:project-id", sagemaker_project_id)
         Tags.of(self).add("sagemaker:project-name", sagemaker_project_name)
         if sagemaker_domain_id:
@@ -147,9 +150,8 @@ def __init__(
                         )
                     ],
                     principals=[
-                        iam.ArnPrincipal(f"arn:{Aws.PARTITION}:iam::{dev_account_id}:root"),
-                        iam.ArnPrincipal(f"arn:{Aws.PARTITION}:iam::{pre_prod_account_id}:root"),
-                        iam.ArnPrincipal(f"arn:{Aws.PARTITION}:iam::{prod_account_id}:root"),
+                        iam.ArnPrincipal(f"arn:{Aws.PARTITION}:iam::{account_id}:root")
+                        for account_id in unique_account_ids
                     ],
                 ),
                 iam.PolicyStatement(
@@ -167,9 +169,8 @@ def __init__(
                         )
                     ],
                     principals=[
-                        iam.ArnPrincipal(f"arn:{Aws.PARTITION}:iam::{dev_account_id}:root"),
-                        iam.ArnPrincipal(f"arn:{Aws.PARTITION}:iam::{pre_prod_account_id}:root"),
-                        iam.ArnPrincipal(f"arn:{Aws.PARTITION}:iam::{prod_account_id}:root"),
+                        iam.ArnPrincipal(f"arn:{Aws.PARTITION}:iam::{account_id}:root")
+                        for account_id in unique_account_ids
                     ],
                 ),
             ]

modules/sagemaker/sagemaker-templates/templates/hf_import_models/product_stack.py

@@ -45,6 +45,9 @@ def __init__(
         pre_prod_account_id = Aws.ACCOUNT_ID if not pre_prod_account_id else pre_prod_account_id
         prod_account_id = Aws.ACCOUNT_ID if not prod_account_id else prod_account_id
 
+        # Deduplicate account IDs to avoid "Duplicate principal" errors in single-account deployments
+        unique_account_ids = list(dict.fromkeys([pre_prod_account_id, prod_account_id]))
+
         Tags.of(self).add("sagemaker:project-id", sagemaker_project_id)
         Tags.of(self).add("sagemaker:project-name", sagemaker_project_name)
         if sagemaker_domain_id:
@@ -129,10 +132,7 @@ def __init__(
                     resources=[
                         f"arn:{Aws.PARTITION}:sagemaker:{Aws.REGION}:{Aws.ACCOUNT_ID}:model-package-group/{model_package_group_name}"
                     ],
-                    principals=[
-                        iam.AccountPrincipal(pre_prod_account_id),
-                        iam.AccountPrincipal(prod_account_id),
-                    ],
+                    principals=[iam.AccountPrincipal(account_id) for account_id in unique_account_ids],
                 ),
                 iam.PolicyStatement(
                     sid="ModelPackage",
@@ -145,10 +145,7 @@ def __init__(
                     resources=[
                         f"arn:{Aws.PARTITION}:sagemaker:{Aws.REGION}:{Aws.ACCOUNT_ID}:model-package/{model_package_group_name}/*"
                     ],
-                    principals=[
-                        iam.AccountPrincipal(pre_prod_account_id),
-                        iam.AccountPrincipal(prod_account_id),
-                    ],
+                    principals=[iam.AccountPrincipal(account_id) for account_id in unique_account_ids],
                 ),
             ]
         ).to_json()

modules/sagemaker/sagemaker-templates/templates/xgboost_abalone/product_stack.py

@@ -50,6 +50,9 @@ def __init__(
         pre_prod_account_id = Aws.ACCOUNT_ID if not pre_prod_account_id else pre_prod_account_id
         prod_account_id = Aws.ACCOUNT_ID if not prod_account_id else prod_account_id
 
+        # Deduplicate account IDs to avoid "Duplicate principal" errors in single-account deployments
+        unique_account_ids = list(dict.fromkeys([dev_account_id, pre_prod_account_id, prod_account_id]))
+
         dev_vpc = None
         if dev_vpc_id:
             dev_vpc = ec2.Vpc.from_lookup(self, "dev-vpc", vpc_id=dev_vpc_id)
@@ -153,9 +156,8 @@ def __init__(
                         )
                     ],
                     principals=[
-                        iam.ArnPrincipal(f"arn:{Aws.PARTITION}:iam::{dev_account_id}:root"),
-                        iam.ArnPrincipal(f"arn:{Aws.PARTITION}:iam::{pre_prod_account_id}:root"),
-                        iam.ArnPrincipal(f"arn:{Aws.PARTITION}:iam::{prod_account_id}:root"),
+                        iam.ArnPrincipal(f"arn:{Aws.PARTITION}:iam::{account_id}:root")
+                        for account_id in unique_account_ids
                     ],
                 ),
                 iam.PolicyStatement(
@@ -173,9 +175,8 @@ def __init__(
                         )
                     ],
                     principals=[
-                        iam.ArnPrincipal(f"arn:{Aws.PARTITION}:iam::{dev_account_id}:root"),
-                        iam.ArnPrincipal(f"arn:{Aws.PARTITION}:iam::{pre_prod_account_id}:root"),
-                        iam.ArnPrincipal(f"arn:{Aws.PARTITION}:iam::{prod_account_id}:root"),
+                        iam.ArnPrincipal(f"arn:{Aws.PARTITION}:iam::{account_id}:root")
+                        for account_id in unique_account_ids
                     ],
                 ),
             ]

modules/sagemaker/sagemaker-templates/tests/test_stack.py

@@ -157,3 +157,92 @@ def test_no_cdk_nag_errors(stack: cdk.Stack, project_template_type) -> None:
         Match.string_like_regexp(r"AwsSolutions-.*"),
     )
     assert not nag_errors, f"Found {len(nag_errors)} CDK nag errors"
+
+
+@pytest.fixture(scope="function")
+def stack_single_account(stack_defaults, project_template_type) -> cdk.Stack:
+    import stack
+
+    app = cdk.App()
+    same_account_id = os.environ["CDK_DEFAULT_ACCOUNT"]
+    sagemaker_domain_id = "domain_id"
+
+    xgboost_abalone_settings = None
+    hf_import_models_settings = None
+
+    if project_template_type == ProjectTemplateType.XGBOOST_ABALONE:
+        xgboost_abalone_settings = XGBoostAbaloneProjectSettings(
+            enable_network_isolation="False", encrypt_inter_container_traffic="False"
+        )
+    elif project_template_type == ProjectTemplateType.HF_IMPORT_MODELS:
+        hf_import_models_settings = HfImportModelsProjectSettings(
+            hf_access_token_secret="test-hf-token-secret", hf_model_id="test-model-id"
+        )
+
+    return stack.ProjectStack(
+        app,
+        "test-single-account-stack",
+        project_template_type=project_template_type,
+        sagemaker_project_name="test-project",
+        sagemaker_project_id="test-project-id",
+        dev_vpc_id="",
+        dev_subnet_ids=[],
+        dev_security_group_ids=[],
+        pre_prod_account_id=same_account_id,
+        pre_prod_region="us-east-1",
+        pre_prod_vpc_id="",
+        pre_prod_subnet_ids=[],
+        pre_prod_security_group_ids=[],
+        prod_account_id=same_account_id,
+        prod_region="us-east-1",
+        prod_vpc_id="",
+        prod_subnet_ids=[],
+        prod_security_group_ids=[],
+        env=cdk.Environment(
+            account=same_account_id,
+            region=os.environ["CDK_DEFAULT_REGION"],
+        ),
+        sagemaker_domain_id=sagemaker_domain_id,
+        sagemaker_domain_arn=f"arn:aws:sagemaker:::domain/{sagemaker_domain_id}",
+        repository_type=RepositoryType.CODECOMMIT,
+        access_token_secret_name=None,
+        aws_codeconnection_arn=None,
+        repository_owner=None,
+        xgboost_abalone_project_settings=xgboost_abalone_settings,
+        model_deploy_project_settings=None,
+        hf_import_models_project_settings=hf_import_models_settings,
+        batch_inference_project_settings=None,
+    )
+
+
+@pytest.mark.parametrize(
+    "project_template_type",
+    [
+        ProjectTemplateType.XGBOOST_ABALONE,
+        ProjectTemplateType.FINETUNE_LLM_EVALUATION,
+        ProjectTemplateType.HF_IMPORT_MODELS,
+    ],
+    indirect=True,
+)
+def test_no_duplicate_principals_in_model_package_group_policy(
+    stack_single_account: cdk.Stack, project_template_type
+) -> None:
+    import json
+
+    from aws_cdk.assertions import Template
+
+    template = Template.from_stack(stack_single_account)
+    model_package_groups = template.find_resources("AWS::SageMaker::ModelPackageGroup")
+
+    for logical_id, resource in model_package_groups.items():
+        policy = resource.get("Properties", {}).get("ModelPackageGroupPolicy")
+        if policy and isinstance(policy, str):
+            policy_doc = json.loads(policy)
+            for statement in policy_doc.get("Statement", []):
+                principal = statement.get("Principal", {})
+                if isinstance(principal, dict):
+                    aws_principals = principal.get("AWS", [])
+                    if isinstance(aws_principals, list):
+                        assert len(aws_principals) == len(
+                            set(aws_principals)
+                        ), f"Duplicate principals in {logical_id}: {aws_principals}"