fix: duplicate princpals for dev nonprod and prod (#402)

Author: weklund

CHANGELOG.md

@@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### **Changed**
 
+- fixed duplicate principal error in `sagemaker-templates` KMS key and S3 bucket policies when dev, pre-prod, and prod account IDs resolve to the same AWS account
+
 ## v3.2.0
 
 ### **Added**

modules/sagemaker/sagemaker-templates/settings.py

@@ -53,6 +53,7 @@ class ModuleSettings(CdkBaseSettings):
     sagemaker_project_name: str
     sagemaker_project_id: str
 
+    dev_account_id: Optional[str] = Field(default=None)
     dev_vpc_id: Optional[str] = Field(default=None)
     dev_subnet_ids: List[str] = Field(default=[])
     dev_security_group_ids: List[str] = Field(default=[])

modules/sagemaker/sagemaker-templates/stack.py

@@ -27,6 +27,7 @@ def __init__(
         sagemaker_domain_arn: str,
         sagemaker_project_name: str,
         sagemaker_project_id: str,
+        dev_account_id: str,
         dev_vpc_id: str,
         dev_subnet_ids: List[str],
         dev_security_group_ids: List[str],
@@ -61,6 +62,7 @@ def __init__(
                 self,
                 "XGBoostAbaloneProject",
                 build_app_asset=cast(s3_assets.Asset, build_app_asset),
+                dev_account_id=dev_account_id,
                 pre_prod_account_id=pre_prod_account_id,
                 prod_account_id=prod_account_id,
                 sagemaker_domain_id=sagemaker_domain_id,
@@ -100,6 +102,7 @@ def __init__(
                 build_app_asset=cast(s3_assets.Asset, build_app_asset),
                 sagemaker_project_name=sagemaker_project_name,
                 sagemaker_project_id=sagemaker_project_id,
+                dev_account_id=dev_account_id,
                 pre_prod_account_id=pre_prod_account_id,
                 prod_account_id=prod_account_id,
                 sagemaker_domain_id=sagemaker_domain_id,
@@ -118,6 +121,7 @@ def __init__(
                 sagemaker_domain_arn=sagemaker_domain_arn,
                 sagemaker_project_name=sagemaker_project_name,
                 sagemaker_project_id=sagemaker_project_id,
+                dev_account_id=dev_account_id,
                 pre_prod_account_id=pre_prod_account_id,
                 prod_account_id=prod_account_id,
                 hf_access_token_secret=hf_import_models_project_settings.hf_access_token_secret,

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

@@ -32,6 +32,7 @@ def __init__(
         build_app_asset: s3_assets.Asset,
         sagemaker_project_name: str,
         sagemaker_project_id: str,
+        dev_account_id: str,
         pre_prod_account_id: str,
         prod_account_id: str,
         sagemaker_domain_id: str,
@@ -44,12 +45,14 @@ def __init__(
     ) -> None:
         super().__init__(scope, id)
 
-        dev_account_id = Aws.ACCOUNT_ID
+        dev_account_id = Aws.ACCOUNT_ID if not dev_account_id else dev_account_id
         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]))
+        # Deduplicate cross-account IDs (pre-prod and prod) for KMS and S3 policies
+        unique_cross_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)
@@ -83,10 +86,7 @@ def __init__(
                         resources=[
                             "*",
                         ],
-                        principals=[
-                            iam.AccountPrincipal(pre_prod_account_id),
-                            iam.AccountPrincipal(prod_account_id),
-                        ],
+                        principals=[iam.AccountPrincipal(account_id) for account_id in unique_cross_account_ids],
                     ),
                 ]
             ),
@@ -126,10 +126,7 @@ def __init__(
                     model_bucket.arn_for_objects(key_pattern="*"),
                     model_bucket.bucket_arn,
                 ],
-                principals=[
-                    iam.AccountPrincipal(pre_prod_account_id),
-                    iam.AccountPrincipal(prod_account_id),
-                ],
+                principals=[iam.AccountPrincipal(account_id) for account_id in unique_cross_account_ids],
             )
         )
 

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

@@ -30,6 +30,7 @@ def __init__(
         sagemaker_domain_arn: str,
         sagemaker_project_name: str,
         sagemaker_project_id: str,
+        dev_account_id: str,
         pre_prod_account_id: str,
         prod_account_id: str,
         hf_access_token_secret: str,
@@ -42,11 +43,14 @@ def __init__(
     ) -> None:
         super().__init__(scope, construct_id)
 
+        dev_account_id = Aws.ACCOUNT_ID if not dev_account_id else dev_account_id
         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]))
+        unique_account_ids = list(dict.fromkeys([dev_account_id, pre_prod_account_id, prod_account_id]))
+        # Deduplicate cross-account IDs (pre-prod and prod) for KMS and S3 policies
+        unique_cross_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)
@@ -80,10 +84,7 @@ def __init__(
                         resources=[
                             "*",
                         ],
-                        principals=[
-                            iam.AccountPrincipal(pre_prod_account_id),
-                            iam.AccountPrincipal(prod_account_id),
-                        ],
+                        principals=[iam.AccountPrincipal(account_id) for account_id in unique_cross_account_ids],
                     ),
                 ]
             ),
@@ -117,8 +118,8 @@ def __init__(
         s3_artifact.grant_read_write(iam.AccountRootPrincipal())
 
         # PROD account access to objects in the bucket
-        s3_artifact.grant_read_write(iam.AccountPrincipal(pre_prod_account_id))
-        s3_artifact.grant_read_write(iam.AccountPrincipal(prod_account_id))
+        for account_id in unique_cross_account_ids:
+            s3_artifact.grant_read_write(iam.AccountPrincipal(account_id))
 
         # cross account model registry resource policy
         model_package_group_name = f"{sagemaker_project_name}-{sagemaker_project_id}"

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

@@ -28,6 +28,7 @@ def __init__(
         scope: Construct,
         id: str,
         build_app_asset: s3_assets.Asset,
+        dev_account_id: str,
         pre_prod_account_id: str,
         prod_account_id: str,
         sagemaker_domain_id: str,
@@ -46,12 +47,14 @@ def __init__(
     ) -> None:
         super().__init__(scope, id)
 
-        dev_account_id = Aws.ACCOUNT_ID
+        dev_account_id = Aws.ACCOUNT_ID if not dev_account_id else dev_account_id
         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]))
+        # Deduplicate cross-account IDs (pre-prod and prod) for KMS and S3 policies
+        unique_cross_account_ids = list(dict.fromkeys([pre_prod_account_id, prod_account_id]))
 
         dev_vpc = None
         if dev_vpc_id:
@@ -89,10 +92,7 @@ def __init__(
                         resources=[
                             "*",
                         ],
-                        principals=[
-                            iam.AccountPrincipal(pre_prod_account_id),
-                            iam.AccountPrincipal(prod_account_id),
-                        ],
+                        principals=[iam.AccountPrincipal(account_id) for account_id in unique_cross_account_ids],
                     ),
                 ]
             ),
@@ -132,10 +132,7 @@ def __init__(
                     model_bucket.arn_for_objects(key_pattern="*"),
                     model_bucket.bucket_arn,
                 ],
-                principals=[
-                    iam.AccountPrincipal(pre_prod_account_id),
-                    iam.AccountPrincipal(prod_account_id),
-                ],
+                principals=[iam.AccountPrincipal(account_id) for account_id in unique_cross_account_ids],
             )
         )
 

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

@@ -50,6 +50,7 @@ def stack(stack_defaults, project_template_type) -> cdk.Stack:
     project_name = "test-project"
     dep_name = "test-deployment"
     mod_name = "test-module"
+    dev_account_id = "dev_account_id"
     dev_vpc_id = "vpc"
     dev_subnet_ids = ["sub"]
     dev_security_group_ids = ["sg"]
@@ -108,6 +109,7 @@ def stack(stack_defaults, project_template_type) -> cdk.Stack:
         project_template_type=project_template_type,
         sagemaker_project_name=sagemaker_project_name,
         sagemaker_project_id=sagemaker_project_id,
+        dev_account_id=dev_account_id,
         dev_vpc_id=dev_vpc_id,
         dev_subnet_ids=dev_subnet_ids,
         dev_security_group_ids=dev_security_group_ids,
@@ -185,6 +187,7 @@ def stack_single_account(stack_defaults, project_template_type) -> cdk.Stack:
         project_template_type=project_template_type,
         sagemaker_project_name="test-project",
         sagemaker_project_id="test-project-id",
+        dev_account_id=same_account_id,
         dev_vpc_id="",
         dev_subnet_ids=[],
         dev_security_group_ids=[],
@@ -227,22 +230,161 @@ def stack_single_account(stack_defaults, project_template_type) -> cdk.Stack:
 def test_no_duplicate_principals_in_model_package_group_policy(
     stack_single_account: cdk.Stack, project_template_type
 ) -> None:
+    """Test that single-account deployments don't create duplicate principals.
+
+    When all account IDs (dev, pre-prod, prod) are the same, the policy should
+    deduplicate them to avoid SageMaker's "Duplicate principal" validation error.
+
+    This test checks each policy statement individually, as SageMaker rejects
+    policies where the same principal appears multiple times within a single statement.
+    """
     import json
 
     from aws_cdk.assertions import Template
 
+    def stringify_principal(principal) -> str:
+        """Convert a principal to a comparable string representation."""
+        if isinstance(principal, str):
+            return principal
+        elif isinstance(principal, dict):
+            # Handle Fn::Join and other intrinsic functions
+            return json.dumps(principal, sort_keys=True)
+        return str(principal)
+
+    def check_statement_for_duplicates(statement: dict, logical_id: str, statement_sid: str) -> None:
+        """Check a single policy statement for duplicate principals."""
+        principal = statement.get("Principal", {})
+        if isinstance(principal, dict):
+            aws_principals = principal.get("AWS", [])
+            if isinstance(aws_principals, list):
+                principal_strs = [stringify_principal(p) for p in aws_principals]
+                assert len(principal_strs) == len(
+                    set(principal_strs)
+                ), f"Duplicate principals in {logical_id} statement '{statement_sid}': {aws_principals}"
+
     template = Template.from_stack(stack_single_account)
     model_package_groups = template.find_resources("AWS::SageMaker::ModelPackageGroup")
 
+    assert model_package_groups, "Expected at least one ModelPackageGroup resource"
+
     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}"
+        if policy:
+            if isinstance(policy, str):
+                policy = json.loads(policy)
+
+            if isinstance(policy, dict):
+                for statement in policy.get("Statement", []):
+                    statement_sid = statement.get("Sid", "unknown")
+                    check_statement_for_duplicates(statement, logical_id, statement_sid)
+
+
+@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_kms_key_policy(stack_single_account: cdk.Stack, project_template_type) -> None:
+    """Test that single-account deployments don't create duplicate principals in KMS key policies.
+
+    When pre-prod and prod account IDs are the same, the KMS key policy should
+    deduplicate them to avoid "Duplicate principal" validation errors.
+    """
+    import json
+
+    from aws_cdk.assertions import Template
+
+    def stringify_principal(principal) -> str:
+        """Convert a principal to a comparable string representation."""
+        if isinstance(principal, str):
+            return principal
+        elif isinstance(principal, dict):
+            return json.dumps(principal, sort_keys=True)
+        return str(principal)
+
+    def check_statement_for_duplicates(statement: dict, logical_id: str, statement_sid: str) -> None:
+        """Check a single policy statement for duplicate principals."""
+        principal = statement.get("Principal", {})
+        if isinstance(principal, dict):
+            aws_principals = principal.get("AWS", [])
+            if isinstance(aws_principals, list):
+                principal_strs = [stringify_principal(p) for p in aws_principals]
+                assert len(principal_strs) == len(
+                    set(principal_strs)
+                ), f"Duplicate principals in KMS key {logical_id} statement '{statement_sid}': {aws_principals}"
+
+    template = Template.from_stack(stack_single_account)
+    kms_keys = template.find_resources("AWS::KMS::Key")
+
+    assert kms_keys, "Expected at least one KMS Key resource"
+
+    for logical_id, resource in kms_keys.items():
+        policy = resource.get("Properties", {}).get("KeyPolicy")
+        if policy:
+            if isinstance(policy, str):
+                policy = json.loads(policy)
+
+            if isinstance(policy, dict):
+                for statement in policy.get("Statement", []):
+                    statement_sid = statement.get("Sid", "unknown")
+                    check_statement_for_duplicates(statement, logical_id, statement_sid)
+
+
+@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_s3_bucket_policy(stack_single_account: cdk.Stack, project_template_type) -> None:
+    """Test that single-account deployments don't create duplicate principals in S3 bucket policies.
+
+    When pre-prod and prod account IDs are the same, the S3 bucket policy should
+    deduplicate them to avoid "Duplicate principal" validation errors.
+    """
+    import json
+
+    from aws_cdk.assertions import Template
+
+    def stringify_principal(principal) -> str:
+        """Convert a principal to a comparable string representation."""
+        if isinstance(principal, str):
+            return principal
+        elif isinstance(principal, dict):
+            return json.dumps(principal, sort_keys=True)
+        return str(principal)
+
+    def check_statement_for_duplicates(statement: dict, logical_id: str, statement_sid: str) -> None:
+        """Check a single policy statement for duplicate principals."""
+        principal = statement.get("Principal", {})
+        if isinstance(principal, dict):
+            aws_principals = principal.get("AWS", [])
+            if isinstance(aws_principals, list):
+                principal_strs = [stringify_principal(p) for p in aws_principals]
+                assert len(principal_strs) == len(set(principal_strs)), (
+                    f"Duplicate principals in S3 bucket policy {logical_id} "
+                    f"statement '{statement_sid}': {aws_principals}"
+                )
+
+    template = Template.from_stack(stack_single_account)
+    bucket_policies = template.find_resources("AWS::S3::BucketPolicy")
+
+    assert bucket_policies, "Expected at least one S3 BucketPolicy resource"
+
+    for logical_id, resource in bucket_policies.items():
+        policy = resource.get("Properties", {}).get("PolicyDocument")
+        if policy:
+            if isinstance(policy, str):
+                policy = json.loads(policy)
+
+            if isinstance(policy, dict):
+                for statement in policy.get("Statement", []):
+                    statement_sid = statement.get("Sid", "unknown")
+                    check_statement_for_duplicates(statement, logical_id, statement_sid)