def test_aws_to_dict_and_back():
    from dask_cloudprovider.providers.aws.helper import aws_to_dict, dict_to_aws

    aws_dict = [{"key": "hello", "value": "world"}]
    aws_upper_dict = [{"Key": "hello", "Value": "world"}]
    py_dict = {"hello": "world"}

    assert dict_to_aws(py_dict) == aws_dict
    assert dict_to_aws(py_dict, upper=True) == aws_upper_dict
    assert aws_to_dict(aws_dict) == py_dict

    assert aws_to_dict(dict_to_aws(py_dict, upper=True)) == py_dict
    assert aws_to_dict(dict_to_aws(py_dict)) == py_dict
    assert dict_to_aws(aws_to_dict(aws_dict)) == aws_dict
    assert dict_to_aws(aws_to_dict(aws_upper_dict),
                       upper=True) == aws_upper_dict
Exemple #2
0
async def _cleanup_stale_resources():
    """ Clean up any stale resources which are tagged with 'createdBy': 'dask-cloudprovider'.

    This function will scan through AWS looking for resources that were created
    by the ``ECSCluster`` class. Any ECS clusters which do not have any running
    tasks will be deleted and then any supporting resources such as task definitions
    security groups and IAM roles that are not associated with an active cluster
    will also be deleted.

    The ``ECSCluster`` should clean up after itself when it is garbage collected
    however if the Python process is terminated without notice this may not happen.
    Therefore this is useful to remove shrapnel from past failures.

    """
    # Clean up clusters (clusters with no running tasks)
    session = aiobotocore.get_session()
    async with session.create_client("ecs") as ecs:
        active_clusters = []
        clusters_to_delete = []
        async for page in ecs.get_paginator("list_clusters").paginate():
            clusters = (
                await ecs.describe_clusters(
                    clusters=page["clusterArns"], include=["TAGS"]
                )
            )["clusters"]
            for cluster in clusters:
                if DEFAULT_TAGS.items() <= aws_to_dict(cluster["tags"]).items():
                    if cluster["runningTasksCount"] == 0:
                        clusters_to_delete.append(cluster["clusterArn"])
                    else:
                        active_clusters.append(cluster["clusterName"])
        for cluster_arn in clusters_to_delete:
            await ecs.delete_cluster(cluster=cluster_arn)

        # Clean up task definitions (with no active clusters)
        async for page in ecs.get_paginator("list_task_definitions").paginate():
            for task_definition_arn in page["taskDefinitionArns"]:
                response = await ecs.describe_task_definition(
                    taskDefinition=task_definition_arn, include=["TAGS"]
                )
                task_definition = response["taskDefinition"]
                task_definition["tags"] = response["tags"]
                task_definition_cluster = aws_to_dict(task_definition["tags"]).get(
                    "cluster"
                )
                if (
                    task_definition_cluster is None
                    or task_definition_cluster not in active_clusters
                ):
                    await ecs.deregister_task_definition(
                        taskDefinition=task_definition_arn
                    )

    # Clean up security groups (with no active clusters)
    async with session.create_client("ec2") as ec2:
        async for page in ec2.get_paginator("describe_security_groups").paginate(
            Filters=[{"Name": "tag:createdBy", "Values": ["dask-cloudprovider"]}]
        ):
            for group in page["SecurityGroups"]:
                sg_cluster = aws_to_dict(group["Tags"]).get("cluster")
                if sg_cluster is None or sg_cluster not in active_clusters:
                    await ec2.delete_security_group(
                        GroupName=group["GroupName"], DryRun=False
                    )

    # Clean up roles (with no active clusters)
    async with session.create_client("iam") as iam:
        async for page in iam.get_paginator("list_roles").paginate():
            for role in page["Roles"]:
                role["Tags"] = (
                    await iam.list_role_tags(RoleName=role["RoleName"])
                ).get("Tags")
                if DEFAULT_TAGS.items() <= aws_to_dict(role["Tags"]).items():
                    role_cluster = aws_to_dict(role["Tags"]).get("cluster")
                    if role_cluster is None or role_cluster not in active_clusters:
                        attached_policies = (
                            await iam.list_attached_role_policies(
                                RoleName=role["RoleName"]
                            )
                        )["AttachedPolicies"]
                        for policy in attached_policies:
                            await iam.detach_role_policy(
                                RoleName=role["RoleName"], PolicyArn=policy["PolicyArn"]
                            )
                        await iam.delete_role(RoleName=role["RoleName"])