/scripts/launch.py

import datetime
import json
import os

import click

AMI_MAP = {
    "us-west-1": "FILL IN YOUR AMI HERE",
}


def highlight(x):
    if not isinstance(x, str):
        x = json.dumps(x, sort_keys=True, indent=2)
    click.secho(x, fg='green')


def upload_archive(exp_name, archive_excludes, s3_bucket):
    import hashlib, os.path as osp, subprocess, tempfile, uuid, sys

    # Archive this package
    thisfile_dir = osp.dirname(osp.abspath(__file__))
    pkg_parent_dir = osp.abspath(osp.join(thisfile_dir, '..', '..'))
    pkg_subdir = osp.basename(osp.abspath(osp.join(thisfile_dir, '..')))
    assert osp.abspath(__file__) == osp.join(pkg_parent_dir, pkg_subdir, 'scripts', 'launch.py'), 'You moved me!'

    # Run tar
    tmpdir = tempfile.TemporaryDirectory()
    local_archive_path = osp.join(tmpdir.name, '{}.tar.gz'.format(uuid.uuid4()))
    tar_cmd = ["tar", "-zcvf", local_archive_path, "-C", pkg_parent_dir]
    for pattern in archive_excludes:
        tar_cmd += ["--exclude", pattern]
    tar_cmd += ["-h", pkg_subdir]
    highlight(" ".join(tar_cmd))

    if sys.platform == 'darwin':
        # Prevent Mac tar from adding ._* files
        env = os.environ.copy()
        env['COPYFILE_DISABLE'] = '1'
        subprocess.check_call(tar_cmd, env=env)
    else:
        subprocess.check_call(tar_cmd)

    # Construct remote path to place the archive on S3
    with open(local_archive_path, 'rb') as f:
        archive_hash = hashlib.sha224(f.read()).hexdigest()
    remote_archive_path = '{}/{}_{}.tar.gz'.format(s3_bucket, exp_name, archive_hash)

    # Upload
    upload_cmd = ["aws", "s3", "cp", local_archive_path, remote_archive_path]
    highlight(" ".join(upload_cmd))
    subprocess.check_call(upload_cmd)

    presign_cmd = ["aws", "s3", "presign", remote_archive_path, "--expires-in", str(60 * 60 * 24 * 30)]
    highlight(" ".join(presign_cmd))
    remote_url = subprocess.check_output(presign_cmd).decode("utf-8").strip()
    return remote_url


def make_disable_hyperthreading_script():
    return """
# disable hyperthreading
# https://forums.aws.amazon.com/message.jspa?messageID=189757
for cpunum in $(
    cat /sys/devices/system/cpu/cpu*/topology/thread_siblings_list |
    sed 's/-/,/g' | cut -s -d, -f2- | tr ',' '\n' | sort -un); do
        echo 0 > /sys/devices/system/cpu/cpu$cpunum/online
done
"""


def make_download_and_run_script(code_url, cmd):
    return """su -l ubuntu <<'EOF'
set -x
cd ~
wget --quiet "{code_url}" -O code.tar.gz
tar xvaf code.tar.gz
rm code.tar.gz
cd es-distributed
{cmd}
EOF
""".format(code_url=code_url, cmd=cmd)


def make_master_script(code_url, exp_str):
    cmd = """
cat > ~/experiment.json <<< '{exp_str}'
python -m es_distributed.main master \
    --master_socket_path /var/run/redis/redis.sock \
    --log_dir ~ \
    --exp_file ~/experiment.json
    """.format(exp_str=exp_str)
    return """#!/bin/bash
{
set -x

%s

# Disable redis snapshots
echo 'save ""' >> /etc/redis/redis.conf

# Make the unix domain socket available for the master client
# (TCP is still enabled for workers/relays)
echo "unixsocket /var/run/redis/redis.sock" >> /etc/redis/redis.conf
echo "unixsocketperm 777" >> /etc/redis/redis.conf
mkdir -p /var/run/redis
chown ubuntu:ubuntu /var/run/redis

systemctl restart redis

%s
} >> /home/ubuntu/user_data.log 2>&1
""" % (make_disable_hyperthreading_script(), make_download_and_run_script(code_url, cmd))


def make_worker_script(code_url, master_private_ip):
    cmd = ("MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=1 "
           "python -m es_distributed.main workers "
           "--master_host {} "
           "--relay_socket_path /var/run/redis/redis.sock").format(master_private_ip)
    return """#!/bin/bash
{
set -x

%s

# Disable redis snapshots
echo 'save ""' >> /etc/redis/redis.conf

# Make redis use a unix domain socket and disable TCP sockets
sed -ie "s/port 6379/port 0/" /etc/redis/redis.conf
echo "unixsocket /var/run/redis/redis.sock" >> /etc/redis/redis.conf
echo "unixsocketperm 777" >> /etc/redis/redis.conf
mkdir -p /var/run/redis
chown ubuntu:ubuntu /var/run/redis

systemctl restart redis

%s
} >> /home/ubuntu/user_data.log 2>&1
""" % (make_disable_hyperthreading_script(), make_download_and_run_script(code_url, cmd))


@click.command()
@click.argument('exp_files', nargs=-1, type=click.Path(), required=True)
@click.option('--key_name', default=lambda: os.environ["KEY_NAME"])
@click.option('--aws_access_key_id', default=os.environ.get("AWS_ACCESS_KEY", None))
@click.option('--aws_secret_access_key', default=os.environ.get("AWS_ACCESS_SECRET", None))
@click.option('--archive_excludes', default=(".git", "__pycache__", ".idea", "scratch"))
@click.option('--s3_bucket')
@click.option('--spot_price')
@click.option('--region_name')
@click.option('--zone')
@click.option('--cluster_size', type=int, default=1)
@click.option('--spot_master', is_flag=True, help='Use a spot instance as the master')
@click.option('--master_instance_type')
@click.option('--worker_instance_type')
@click.option('--security_group')
@click.option('--yes', is_flag=True, help='Skip confirmation prompt')
def main(exp_files,
         key_name,
         aws_access_key_id,
         aws_secret_access_key,
         archive_excludes,
         s3_bucket,
         spot_price,
         region_name,
         zone,
         cluster_size,
         spot_master,
         master_instance_type,
         worker_instance_type,
         security_group,
         yes
         ):

    highlight('Launching:')
    highlight(locals())

    import boto3
    ec2 = boto3.resource(
        "ec2",
        region_name=region_name,
        aws_access_key_id=aws_access_key_id,
        aws_secret_access_key=aws_secret_access_key
    )
    as_client = boto3.client(
        'autoscaling',
        region_name=region_name,
        aws_access_key_id=aws_access_key_id,
        aws_secret_access_key=aws_secret_access_key
    )

    for i_exp_file, exp_file in enumerate(exp_files):
        with open(exp_file, 'r') as f:
            exp = json.loads(f.read())
        highlight('Experiment [{}/{}]:'.format(i_exp_file + 1, len(exp_files)))
        highlight(exp)
        if not yes:
            click.confirm('Continue?', abort=True)

        exp_prefix = exp['exp_prefix']
        exp_str = json.dumps(exp)

        exp_name = '{}_{}'.format(exp_prefix, datetime.datetime.now().strftime('%Y%m%d-%H%M%S'))

        code_url = upload_archive(exp_name, archive_excludes, s3_bucket)
        highlight("code_url: " + code_url)

        image_id = AMI_MAP[region_name]
        highlight('Using AMI: {}'.format(image_id))

        if spot_master:
            import base64
            requests = ec2.meta.client.request_spot_instances(
                SpotPrice=spot_price,
                InstanceCount=1,
                LaunchSpecification=dict(
                    ImageId=image_id,
                    KeyName=key_name,
                    InstanceType=master_instance_type,
                    EbsOptimized=True,
                    SecurityGroups=[security_group],
                    Placement=dict(
                        AvailabilityZone=zone,
                    ),
                    UserData=base64.b64encode(make_master_script(code_url, exp_str).encode()).decode()
                )
            )['SpotInstanceRequests']
            assert len(requests) == 1
            request_id = requests[0]['SpotInstanceRequestId']
            # Wait for fulfillment
            highlight('Waiting for spot request {} to be fulfilled'.format(request_id))
            ec2.meta.client.get_waiter('spot_instance_request_fulfilled').wait(SpotInstanceRequestIds=[request_id])
            req = ec2.meta.client.describe_spot_instance_requests(SpotInstanceRequestIds=[request_id])
            master_instance_id = req['SpotInstanceRequests'][0]['InstanceId']
            master_instance = ec2.Instance(master_instance_id)
        else:
            master_instance = ec2.create_instances(
                ImageId=image_id,
                KeyName=key_name,
                InstanceType=master_instance_type,
                EbsOptimized=True,
                SecurityGroups=[security_group],
                MinCount=1,
                MaxCount=1,
                Placement=dict(
                    AvailabilityZone=zone,
                ),
                UserData=make_master_script(code_url, exp_str)
            )[0]
        master_instance.create_tags(
            Tags=[
                dict(Key="Name", Value=exp_name + "-master"),
                dict(Key="es_dist_role", Value="master"),
                dict(Key="exp_prefix", Value=exp_prefix),
                dict(Key="exp_name", Value=exp_name),
            ]
        )
        highlight("Master created. IP: %s" % master_instance.public_ip_address)

        config_resp = as_client.create_launch_configuration(
            ImageId=image_id,
            KeyName=key_name,
            InstanceType=worker_instance_type,
            EbsOptimized=True,
            SecurityGroups=[security_group],
            LaunchConfigurationName=exp_name,
            UserData=make_worker_script(code_url, master_instance.private_ip_address),
            SpotPrice=spot_price,
        )
        assert config_resp["ResponseMetadata"]["HTTPStatusCode"] == 200

        asg_resp = as_client.create_auto_scaling_group(
            AutoScalingGroupName=exp_name,
            LaunchConfigurationName=exp_name,
            MinSize=cluster_size,
            MaxSize=cluster_size,
            DesiredCapacity=cluster_size,
            AvailabilityZones=[zone],
            Tags=[
                dict(Key="Name", Value=exp_name + "-worker"),
                dict(Key="es_dist_role", Value="worker"),
                dict(Key="exp_prefix", Value=exp_prefix),
                dict(Key="exp_name", Value=exp_name),
            ]
            # todo: also try placement group to see if there is increased networking performance
        )
        assert asg_resp["ResponseMetadata"]["HTTPStatusCode"] == 200
        highlight("Scaling group created")

        highlight("%s launched successfully." % exp_name)
        highlight("Manage at %s" % (
            "https://%s.console.aws.amazon.com/ec2/v2/home?region=%s#Instances:sort=tag:Name" % (
            region_name, region_name)
        ))


if __name__ == '__main__':
    main()