launch_ec2_cluster_hosts.py

#!/usr/bin/env python3
#
help_string = '''
Tool to launch a set of clean EC2 hosts which can be used as a MongoDB sharded cluster.

Since it interacts with AWS, a default region_name should be set in ~/.aws/config and the AWS
parameters should be specified either in the same config file or as environment variables.

To monitor the progress of the host initialization scripts, follow /var/log/cloud-init-output.log.

To allow traffic from your current IP to the security group used by the instances, run:
  aws ec2 authorize-security-group-ingress --group-name <your_user> --protocol all --cidr "$(curl -s https://ipv4.wtfismyip.com/text)/32"

Use --help for more information on the supported commands.
'''

import argparse
import boto3
import json
import logging
import os
import sys

from common.common import normalize_clustertag, yes_no
from common.ec2_instances import (CLIENT_HOST_TEMPLATE, create_and_attach_volumes,
                                  describe_all_instances, extract_data_volumes_from_template,
                                  filter_instances_by_role, launch_instances, load_template,
                                  make_client_driver_host_configuration,
                                  make_cluster_host_configuration, make_instance_tag_specifications,
                                  resolve_security_group_id, terminate_cluster_resources,
                                  wait_for_instances)
from common.version import CTOOLS_VERSION

# Ensure that the caller is using python 3
if (sys.version_info[0] < 3):
    raise Exception("Must be using Python 3")


def describe_cluster(ec2, clustertag, user):
    all_instances = describe_all_instances(ec2, clustertag, user)

    cluster_json = {
        "Name":
            clustertag,
        "Hosts":
            list(
                map(
                    lambda x: x["PublicDnsName"],
                    filter_instances_by_role(all_instances, 'config') +
                    filter_instances_by_role(all_instances, 'shard0') +
                    filter_instances_by_role(all_instances, 'shard1'))),
        "DriverHosts":
            list(
                map(lambda x: x["PublicDnsName"], filter_instances_by_role(all_instances,
                                                                           'driver'))),
        "MongoBinPath":
            "/home/ubuntu/workspace/mongo/bazel-bin/install-devcore/bin",
        "RemoteMongoDPath":
            "/mnt/data/mongod",
        "RemoteMongoSPath":
            "/mnt/data/mongos",
        "FeatureFlags": [],
        "MongoDParameters": ["--wiredTigerCacheSizeGB 18", ],
        "MongoSParameters": [],
    }

    return json.dumps(cluster_json, indent=2, separators=(', ', ': '))


def main_launch(args, ec2):
    """
    Implementation of the launch command
    """
    template = load_template(args.template)
    client_template = load_template(CLIENT_HOST_TEMPLATE)
    template['KeyName'] = args.user
    client_template['KeyName'] = args.user
    sg_id = resolve_security_group_id(ec2, args.user)
    template['SecurityGroupIds'] = [sg_id]
    client_template['SecurityGroupIds'] = [sg_id]

    use_volume_copy = getattr(args, 'use_volume_copy', None)
    template, data_volumes = extract_data_volumes_from_template(template)

    skip_format = bool(use_volume_copy)

    ##############################################################################################
    # DRIVER INSTANCES

    client_driver_instances = launch_instances(
        ec2,
        client_template,
        tag_specs=make_instance_tag_specifications(args.clustertag, 'driver', args.user),
        user_data=make_client_driver_host_configuration(args.clustertag),
        count=1,
    )

    ##############################################################################################
    # SHARD INSTANCES

    # Config instances
    config_instances = launch_instances(
        ec2,
        template,
        tag_specs=make_instance_tag_specifications(args.clustertag, 'config', args.user),
        user_data=make_cluster_host_configuration(args.clustertag, args.filesystem, 'config',
                                                  skip_format=skip_format),
        count=3,
    )

    # Shard(s) instances
    shard_instances = []
    for shard_id in ['shard0', 'shard1']:
        shard_instances += launch_instances(
            ec2,
            template,
            tag_specs=make_instance_tag_specifications(args.clustertag, shard_id, args.user),
            user_data=make_cluster_host_configuration(args.clustertag, args.filesystem, shard_id,
                                                      skip_format=skip_format),
            count=args.shard_repl_set_nodes,
        )

    all_cluster_instances = config_instances + shard_instances
    wait_for_instances(ec2, client_driver_instances + all_cluster_instances)

    create_and_attach_volumes(ec2, data_volumes, all_cluster_instances, args.clustertag, args.user,
                              source_volume_id=use_volume_copy)

    cluster_desc = describe_cluster(ec2, args.clustertag, args.user)
    output_dir = args.clustertag
    output_file = os.path.join(output_dir, 'deployment_description.json')
    os.makedirs(output_dir, exist_ok=True)
    with open(output_file, 'w') as f:
        f.write(cluster_desc)
    print(cluster_desc)

    logging.info(f'Cluster configuration written to {output_file}')
    logging.info(
        f'To deploy binaries to cluster, now run ./remote_control_cluster.py {args.clustertag} create  ...'
    )


def main_terminate(args, ec2):
    """
    Implementation of the terminate command
    """
    yes_no(f'About to terminate all resources for cluster tag {args.clustertag}')

    total = terminate_cluster_resources(ec2, args.clustertag, args.user)
    if total == 0:
        logging.warning(f'No resources found with cluster tag {args.clustertag}')


def main_describe(args, ec2):
    """
    Implementation of the describe command
    """
    cluster_desc = describe_cluster(ec2, args.clustertag, args.user)
    print(cluster_desc)


if __name__ == "__main__":
    argsParser = argparse.ArgumentParser(description=help_string)
    logging.basicConfig(format='%(asctime)s [%(levelname)s] %(message)s', level=logging.INFO)

    argsParser.add_argument(
        'clustertag', help=
        ('String with which to tag all the instances which will be spawned for this cluster so they '
         'can easily be identified. For example 5.0, 6.0 etc. There must not be any existing '
         'instances with that tag.'), type=normalize_clustertag)

    argsParser.add_argument(
        '--user', required=True, help=
        ('Your AWS key pair name and owner tag (e.g. firstname.lastname). The key pair must '
         'exist in the target AWS account and the private key must be at ~/.ssh/mongodb-aws-kernel-test.'
         ), type=str)

    subparsers = argsParser.add_subparsers(title='subcommands')

    ###############################################################################################
    # Arguments for the 'launch' command
    parser_launch = subparsers.add_parser(
        'launch', help='Launches the EC2 hosts which will comprise the cluster.')
    parser_launch.add_argument(
        '--template', required=True,
        help='Path to a JSON file with EC2 instance parameters (e.g. Atlas-M60.json).')
    parser_launch.add_argument('--shard-repl-set-nodes',
                               help='Number of nodes to use for the shard replica sets.', type=int,
                               default=1)
    parser_launch.add_argument('--filesystem', choices=['xfs', 'ext4'],
                               help='Filesystem to use for the data volume.', default='xfs')
    parser_launch.add_argument(
        '--use-volume-copy',
        help='EBS volume ID to snapshot and attach as data volume to each node.', type=str,
        default=None, metavar='vol-XXXX')
    parser_launch.set_defaults(func=main_launch)

    ###############################################################################################
    # Arguments for the 'terminate' command
    parser_terminate = subparsers.add_parser('terminate',
                                             help='Terminates the EC2 hosts for a cluster.')
    parser_terminate.set_defaults(func=main_terminate)

    ###############################################################################################
    # Arguments for the 'describe' command
    parser_describe = subparsers.add_parser(
        'describe', help='Describes all the hosts which comprise the cluster')
    parser_describe.set_defaults(func=main_describe)

    args = argsParser.parse_args()
    logging.info(f"CTools version {CTOOLS_VERSION} starting with arguments: '{args}'")

    ec2_instance = boto3.client('ec2', region_name='us-east-1')

    args.func(args, ec2_instance)