From b17b493b7574854a028048dc31cff6025180497b Mon Sep 17 00:00:00 2001 From: Ramit Kataria Date: Wed, 5 Aug 2026 14:50:31 -0700 Subject: [PATCH] Eliminate trust policy race by migrating emr_eks test to Pod Identity The test shares one EMR job execution role across runs while each run creates its own EKS cluster. With IRSA, every run had to rewrite the shared role's trust policy to federate its own cluster's OIDC provider, and pruning old statements to stay under the trust policy size limit meant concurrent runs deleted each other's trust entries, failing jobs with AssumeRoleWithWebIdentity AccessDenied. IAM offers no compare-and-swap on trust policies, so this race cannot be fixed while runs mutate shared IAM state. EKS Pod Identity (supported by EMR on EKS since emr-7.3.0) removes the shared mutable state: the role keeps one static trust policy and each run creates Pod Identity associations scoped to its own cluster, also avoiding the 4096-char trust policy cap that limits IRSA to roughly a dozen clusters per role. Pod Identity is also AWS's recommended mechanism for pod credentials: > Unless you have specific usecases for IRSA, we recommend you use > EKS Pod Identities when using EKS. https://docs.aws.amazon.com/eks/latest/best-practices/identity-and-access-management.html#_eks_pod_identities_compared_to_irsa --- .../system/amazon/aws/example_emr_eks.py | 143 ++++++------------ 1 file changed, 46 insertions(+), 97 deletions(-) diff --git a/providers/amazon/tests/system/amazon/aws/example_emr_eks.py b/providers/amazon/tests/system/amazon/aws/example_emr_eks.py index 56af4e1f931fa..12a6ad1ca6778 100644 --- a/providers/amazon/tests/system/amazon/aws/example_emr_eks.py +++ b/providers/amazon/tests/system/amazon/aws/example_emr_eks.py @@ -16,14 +16,11 @@ # under the License. from __future__ import annotations -import json import logging import subprocess -import time from datetime import datetime, timedelta import boto3 -from tenacity import retry, retry_if_exception_type, stop_after_delay, wait_exponential from airflow.providers.amazon.aws.hooks.eks import ClusterStates, NodegroupStates from airflow.providers.amazon.aws.operators.eks import EksCreateClusterOperator, EksDeleteClusterOperator @@ -111,8 +108,7 @@ def run_eksctl_commands(cluster_name, ns): file = "https://github.com/weaveworks/eksctl/releases/latest/download/eksctl_$(uname -s)_amd64.tar.gz" commands = f""" curl --silent --location "{file}" | tar xz -C /tmp && - /tmp/eksctl create iamidentitymapping --cluster {cluster_name} --namespace {ns} --service-name "emr-containers" && - /tmp/eksctl utils associate-iam-oidc-provider --cluster {cluster_name} --approve + /tmp/eksctl create iamidentitymapping --cluster {cluster_name} --namespace {ns} --service-name "emr-containers" """ build = subprocess.Popen( @@ -127,43 +123,37 @@ def run_eksctl_commands(cluster_name, ns): raise RuntimeError(err) -@task(trigger_rule=TriggerRule.ALL_DONE) -def delete_iam_oidc_identity_provider(cluster_name): - oidc_provider_issuer_url = boto3.client("eks").describe_cluster( - name=cluster_name, - )["cluster"]["identity"]["oidc"]["issuer"] - oidc_provider_issuer_endpoint = oidc_provider_issuer_url.replace("https://", "") - - account_id = boto3.client("sts").get_caller_identity()["Account"] - boto3.client("iam").delete_open_id_connect_provider( - OpenIDConnectProviderArn=f"arn:aws:iam::{account_id}:oidc-provider/{oidc_provider_issuer_endpoint}" - ) +@task +def install_pod_identity_agent(cluster_name): + """Install and wait for the EKS Pod Identity Agent add-on. + Each run creates Pod Identity associations scoped to its own cluster. The + role's trust policy stays static. + """ + log = logging.getLogger(__name__) + eks_client = boto3.client("eks") + addon_name = "eks-pod-identity-agent" -@task -def update_trust_policy_execution_role(cluster_name, cluster_namespace, role_name): - # Remove any already existing trusted entities added with "update-role-trust-policy" - # Prevent getting an error "Cannot exceed quota for ACLSizePerRole" - client = boto3.client("iam") - role_trust_policy = client.get_role(RoleName=role_name)["Role"]["AssumeRolePolicyDocument"] - # We assume if the action is sts:AssumeRoleWithWebIdentity, the statement had been added with - # "update-role-trust-policy". Removing it to not exceed the quota - role_trust_policy["Statement"] = [ - statement - for statement in role_trust_policy["Statement"] - if statement["Action"] != "sts:AssumeRoleWithWebIdentity" - ] - - client.update_assume_role_policy( - RoleName=role_name, - PolicyDocument=json.dumps(role_trust_policy), + eks_client.create_addon(clusterName=cluster_name, addonName=addon_name) + + waiter = eks_client.get_waiter("addon_active") + waiter.wait( + clusterName=cluster_name, + addonName=addon_name, + WaiterConfig={"Delay": 10, "MaxAttempts": 60}, ) + log.info("Pod Identity Agent addon is ACTIVE on cluster %s", cluster_name) + - # See https://docs.aws.amazon.com/emr/latest/EMR-on-EKS-DevelopmentGuide/setting-up-trust-policy.html - # The action "update-role-trust-policy" is not available in boto3, thus we need to do it using AWS CLI +@task +def create_pod_identity_role_associations(cluster_name, namespace, role_name): + # `emr-containers create-role-associations` is a CLI-only helper — it creates + # the three Pod Identity associations (client, driver, executor) that the EMR + # job pods use. There is no boto3 equivalent. Requires AWS CLI > 2.24.0. + # See https://docs.aws.amazon.com/emr/latest/EMR-on-EKS-DevelopmentGuide/setting-up-enable-IAM.html commands = ( - f"aws emr-containers update-role-trust-policy --cluster-name {cluster_name} " - f"--namespace {cluster_namespace} --role-name {role_name}" + f"aws emr-containers create-role-associations --cluster-name {cluster_name} " + f"--namespace {namespace} --role-name {role_name}" ) build = subprocess.Popen( @@ -178,64 +168,23 @@ def update_trust_policy_execution_role(cluster_name, cluster_namespace, role_nam raise RuntimeError(err) -@task -def wait_for_trust_policy_propagation(cluster_name, role_name): - """Validate that the IAM trust policy has propagated by checking the role's - trust policy contains the expected OIDC provider. - - Uses exponential backoff retries (up to 5 minutes) instead of a fixed sleep, - which avoids both wasting time when propagation is fast and failing when it's slow. - """ - log = logging.getLogger(__name__) +@task(trigger_rule=TriggerRule.ALL_DONE) +def delete_pod_identity_role_associations(cluster_name, namespace, role_name): + commands = ( + f"aws emr-containers delete-role-associations --cluster-name {cluster_name} " + f"--namespace {namespace} --role-name {role_name}" + ) - # Determine the expected OIDC provider ARN from the EKS cluster - eks_client = boto3.client("eks") - oidc_issuer_url = eks_client.describe_cluster(name=cluster_name)["cluster"]["identity"]["oidc"]["issuer"] - oidc_issuer_endpoint = oidc_issuer_url.replace("https://", "") - account_id = boto3.client("sts").get_caller_identity()["Account"] - expected_oidc_provider_arn = f"arn:aws:iam::{account_id}:oidc-provider/{oidc_issuer_endpoint}" - - @retry( - retry=retry_if_exception_type(RuntimeError), - wait=wait_exponential(multiplier=1, min=5, max=30), - stop=stop_after_delay(300), - reraise=True, + build = subprocess.Popen( + commands, + shell=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, ) - def _validate_trust_policy(): - iam_client = boto3.client("iam") - - # Verify the trust policy document contains the expected OIDC provider - role = iam_client.get_role(RoleName=role_name)["Role"] - trust_policy = role["AssumeRolePolicyDocument"] - - has_oidc_statement = False - for statement in trust_policy.get("Statement", []): - if statement.get("Action") != "sts:AssumeRoleWithWebIdentity": - continue - principal = statement.get("Principal", {}) - federated = principal.get("Federated", "") - if oidc_issuer_endpoint in federated: - has_oidc_statement = True - break - - if not has_oidc_statement: - log.info( - "Trust policy does not yet contain OIDC provider %s, retrying...", - expected_oidc_provider_arn, - ) - raise RuntimeError( - f"Trust policy for role {role_name} does not yet contain " - f"the expected OIDC provider: {expected_oidc_provider_arn}" - ) - - log.info("Trust policy document confirmed for role %s", role_name) - - _validate_trust_policy() - - # Brief buffer after IAM confirms the trust policy document — cross-service - # caches (EKS/EMR) may still serve the old policy for a few seconds. - time.sleep(15) - log.info("Trust policy validation complete, proceeding.") + _, err = build.communicate() + + if build.returncode != 0: + raise RuntimeError(err) @task(trigger_rule=TriggerRule.ALL_DONE) @@ -332,7 +281,7 @@ def delete_virtual_cluster(virtual_cluster_id): task_id="start_job", virtual_cluster_id=str(create_emr_eks_cluster.output), execution_role_arn=job_role_arn, - release_label="emr-7.0.0-latest", + release_label="emr-7.13.0-latest", job_driver=job_driver_arg, configuration_overrides=configuration_overrides_arg, name="pi.py", @@ -385,15 +334,15 @@ def delete_virtual_cluster(virtual_cluster_id): create_cluster_and_nodegroup, await_create_nodegroup, run_eksctl_commands(eks_cluster_name, eks_namespace), - update_trust_policy_execution_role(eks_cluster_name, eks_namespace, job_role_name), - wait_for_trust_policy_propagation(eks_cluster_name, job_role_name), + install_pod_identity_agent(eks_cluster_name), + create_pod_identity_role_associations(eks_cluster_name, eks_namespace, job_role_name), # TEST BODY create_emr_eks_cluster, job_starter, job_waiter, describe_pod, # TEST TEARDOWN - delete_iam_oidc_identity_provider(eks_cluster_name), + delete_pod_identity_role_associations(eks_cluster_name, eks_namespace, job_role_name), delete_virtual_cluster(str(create_emr_eks_cluster.output)), delete_eks_cluster, await_delete_eks_cluster,