go.py

#!/usr/bin/env python3

import argparse
import secrets
import time
import sys
from pathlib import Path, PurePosixPath
import json
import signal

import boto3
from botocore.exceptions import ClientError
import paramiko
import yaml
from glom import glom


def arch_for_instance_type(ec2, instance_type):
    arches = glom(
        ec2.describe_instance_types(InstanceTypes=[instance_type]),
        "InstanceTypes.0.ProcessorInfo.SupportedArchitectures",
    )
    if len(arches) == 1:
        return arches[0]
    else:
        assert arches == ["i386", "x86_64"]
        return "x86_64"


def ami_for_arch(ec2, arch, ubuntu_version):
    ami_info = ec2.describe_images(
        Owners=["099720109477"],  # Canonical Inc.
        Filters=[
            {
                "Name": "architecture",
                "Values": [arch],
            },
            {
                "Name": "name",
                "Values": [f"*-{ubuntu_version}-*-server-*"],
            },
        ],
    )
    ami = sorted(ami_info["Images"], key=lambda x: x["CreationDate"])[-1]
    return glom(ami, "ImageId")


# Instead of just deleting *our* security group, we delete all security groups with the
# given tags. This way if one invocation crashes for some reason and fails to clean up
# after itself, the next invocation will clean things up for it.
#
# Technically there's a race condition here:
#   1) Invocation A (starting up): creates a new security group
#   2) Invocation B (finishing): deletes A's security group, because it's not in use yet
#   3) Invocation A: attempts to start using security group, but fails b/c B deleted it
#
# This requires some pretty precise timing and just re-running the script will sort it
# out, so ... I guess let's not worry about it for now.
def clean_up_unused_sgs(ec2, tag):
    response = ec2.describe_security_groups(
        Filters=[
            {
                "Name": "tag-key",
                "Values": [tag],
            },
        ]
    )
    KNOWN_ERRORS = [
        # Someone else deleted it
        "InvalidGroup.NotFound",
        # This group is still in use (maybe by another parallel invocation)
        "InvalidGroup.InUse",
        "DependencyViolation",
    ]
    for group_id in glom(response, ("SecurityGroups", ["GroupId"])):
        try:
            ec2.delete_security_group(GroupId=group_id)
        except ClientError as exc:
            if glom(exc.response, "Error.Code") in KNOWN_ERRORS:
                pass
            else:
                print(f"Unexpected error: {exc.response}")
        else:
            print(f"  Removed {group_id}")


SG_PREFIX = "tmp-wheel-builder-"
SG_TAG = "tmp-wheel-builder"


def create_our_sg(ec2):
    # We don't actually need secure randomness here; this is just an easy way to get a
    # unique printable string
    name = SG_PREFIX + secrets.token_urlsafe(16)

    # Just use the default VPC, no real reason not to.
    response = ec2.describe_vpcs(Filters=[{"Name": "is-default", "Values": ["true"]}])
    default_vpc = glom(response, "Vpcs.0.VpcId")
    response = ec2.create_security_group(
        GroupName=name,
        Description="Temporary sg for building wheels",
        VpcId=default_vpc,
        TagSpecifications=[
            {
                "ResourceType": "security-group",
                "Tags": [
                    {
                        "Key": SG_TAG,
                        "Value": "true",
                    },
                ],
            },
        ],
    )
    sg_id = glom(response, "GroupId")
    ec2.authorize_security_group_ingress(
        GroupId=sg_id,
        IpPermissions=[
            {
                "FromPort": 22,
                "ToPort": 22,
                "IpProtocol": "tcp",
                "IpRanges": [
                    {
                        "CidrIp": "0.0.0.0/0",
                    }
                ],
            },
        ],
    )
    ec2.authorize_security_group_egress(
        GroupId=sg_id,
        IpPermissions=[
            {
                "FromPort": 22,
                "ToPort": 22,
                "IpProtocol": "tcp",
                "IpRanges": [
                    {
                        "CidrIp": "0.0.0.0/0",
                    }
                ],
            },
        ],
    )
    return sg_id


def raise_keyboard_interrupt(*_):
    raise KeyboardInterrupt


def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("repo")
    parser.add_argument("commit")
    parser.add_argument(
        "--package-name",
        default=None,
        help="""The import name for the Python package (only required if it doesn't
        match the repo name)""",
    )
    parser.add_argument(
        "--non-interactive",
        default=False,
        action="store_true",
        help="If the build errors, don't pause for interactive debugging",
    )
    parser.add_argument("--region", default="eu-west-1", help="AWS region")
    parser.add_argument(
        "--instance-type", default="m6g.2xlarge", help="EC2 instance type"
    )
    parser.add_argument(
        "--disk-size", metavar="N", default=20, type=int, help="in gigabytes"
    )
    parser.add_argument(
        "-w",
        "--wheelhouse",
        metavar="PATH",
        default="./wheelhouse",
        help="where to put built wheels",
    )
    parser.add_argument(
        "--add-ssh-pubkey",
        metavar="PATH",
        action="append",
        default=[],
        help="""An additional ssh public key you want to add to the builder's
        authorized_keys. (Only needed if you don't have an ssh-agent running.)
        Can be used multiple times.""",
    )
    parser.add_argument(
        "--ubuntu-version",
        default="22.04",
        help="host OS version (doesn't affect actual build containers)",
    )
    args = parser.parse_args()

    if ":" not in args.repo:
        args.repo = f"https://github.com/{args.repo}"
    if args.package_name is None:
        args.package_name = args.repo.rsplit("/", 1)[-1]
        if args.package_name.endswith(".git"):
            args.package_name = args.package_name[:-4]

    print(f"Validating instance type {args.instance_type}...")
    ec2 = boto3.client("ec2", region_name=args.region)
    arch = arch_for_instance_type(ec2, args.instance_type)

    print(
        f"Looks like you're building {arch} wheels. Finding appropriate AMI in {args.region}..."
    )
    ami = ami_for_arch(ec2, arch, args.ubuntu_version)
    print(f"  Using {ami}")

    print("Opening firewall for SSH access...")
    sg_id = create_our_sg(ec2)
    print(f"  Created {sg_id}")

    print("Generating temporary SSH key...")
    # Have to use RSA, because Paramiko doesn't yet support generating Ed25519 keys
    # https://github.com/paramiko/paramiko/issues/1136
    tmp_key = paramiko.RSAKey.generate(bits=4096)

    print("Spawning virtual machine...")

    authorized_keys = []
    for key in [tmp_key] + list(paramiko.agent.Agent().get_keys()):
        authorized_keys.append(f"{key.get_name()} {key.get_base64()}")
    for key_path in args.add_ssh_pubkey:
        authorized_keys.append(Path(key_path).read_text())

    user_data = {
        "groups": ["docker"],
        "users": [
            {
                "name": "ubuntu",
                "ssh_authorized_keys": authorized_keys,
                "sudo": "ALL=(ALL) NOPASSWD:ALL",
                "groups": "docker",
            },
        ],
        # We install a time bomb in our new VM, to make sure that if something goes
        # wrong and it gets left running, it will clean itself up.
        "write_files": [
            {
                "path": "/etc/systemd/system/auto-terminate.timer",
                "content": """
[Timer]
OnBootSec=12h

[Install]
WantedBy=timers.target
                """,
            },
            {
                "path": "/etc/systemd/system/auto-terminate.service",
                "content": """
[Unit]
Description=Auto-terminate instance

[Service]
Type=oneshot
ExecStart=poweroff
                """,
            },
        ],
        "runcmd": [
            "systemctl daemon-reload",
            "systemctl enable --now auto-terminate.timer",
        ],
    }
    user_data_str = "#cloud-config\n" + yaml.dump(user_data)

    response = ec2.run_instances(
        ImageId=ami,
        InstanceType=args.instance_type,
        SecurityGroupIds=[sg_id],
        MinCount=1,
        MaxCount=1,
        BlockDeviceMappings=[
            {
                "DeviceName": "/dev/sda1",
                "Ebs": {
                    "DeleteOnTermination": True,
                    "VolumeSize": args.disk_size,
                    "VolumeType": "gp3",
                },
            }
        ],
        UserData=user_data_str,
        InstanceInitiatedShutdownBehavior="terminate",
    )
    instance_id = glom(response, "Instances.0.InstanceId")
    print(f"  Instance is {instance_id}")

    signal.signal(signal.SIGTERM, raise_keyboard_interrupt)
    try:
        print("Waiting for VM to be assigned a public IP...")
        while True:
            try:
                response = ec2.describe_instances(InstanceIds=[instance_id])
            except ClientError as exc:
                if glom(exc.response, "Error.Code") == "InvalidInstanceID.NotFound":
                    # Eventual consistency!
                    pass
                else:
                    raise
            else:
                instance_status = glom(response, "Reservations.0.Instances.0")
                public_ip = glom(instance_status, "PublicIpAddress", default=None)
                if public_ip is not None:
                    break
                state = glom(instance_status, "State.Name")
                if state not in ["pending", "running"]:
                    print(f"...Something went wrong; instance state is {state}")
                    print("Bailing out")
                    sys.exit(1)
            time.sleep(1)

        print(f"  Public IP: {public_ip}")

        print("Connecting to VM...")

        class IgnoreHostKey(paramiko.client.MissingHostKeyPolicy):
            def missing_host_key(self, client, hostname, key):
                pass

        ssh = paramiko.client.SSHClient()
        ssh.set_missing_host_key_policy(IgnoreHostKey())
        while True:
            try:
                ssh.connect(
                    public_ip,
                    username="ubuntu",
                    pkey=key,
                    timeout=1,
                    look_for_keys=False,
                )
            except Exception:
                time.sleep(1)
            else:
                break
        print("  Connection succeeded.")

        print("Setting up VM...")
        sftp = ssh.open_sftp()

        # Copy over files to run on remote
        sftp.put(Path(__file__).parent / "remote.sh", "remote.sh")
        build_info = {
            "repo_url": args.repo,
            "commit": args.commit,
            "package_name": args.package_name,
        }
        with sftp.open("build-info.json", mode="w") as f:
            json.dump(build_info, f)

        print("Running build...")
        print("\n-------- remote build output starts here --------")
        channel = ssh.get_transport().open_session()
        channel.set_combine_stderr(True)

        # TODO FIXME: pass through CIBW_* env vars?
        # channel.update_environment(...)

        channel.get_pty()  # try to reduce output buffering
        channel.exec_command("bash ~/remote.sh")

        # No stdin
        channel.shutdown_write()

        # Copy stdout to the terminal
        sys.stdout.flush()
        while True:
            data = channel.recv(4096)
            if not data:
                break
            sys.stdout.buffer.write(data)
            sys.stdout.buffer.flush()

        print("-------- remote build output ends here --------\n")

        status = channel.recv_exit_status()
        if status:
            print(f"Build failed! (exit status: {status})")
            if not args.non_interactive:
                if len(authorized_keys) > 1:
                    print(f"To investigate further: ssh ubuntu@{public_ip}")
                else:
                    print("Configure an ssh-agent to get access to failed builds")
                input("Press enter when ready to tear everything down...")
                print("Goodbye.")
            sys.exit(1)

        local_wheelhouse = Path(args.wheelhouse).absolute()
        remote_wheelhouse = PurePosixPath("./wheelhouse")
        print(
            f"Fetching wheels from remote:{remote_wheelhouse} into local:{local_wheelhouse}..."
        )
        local_wheelhouse.mkdir(parents=True, exist_ok=True)
        for name in sftp.listdir(str(remote_wheelhouse)):
            print(f"  {name}")
            sftp.get(str(remote_wheelhouse / name), str(local_wheelhouse / name))

        print(f"Success! Your wheels are in {local_wheelhouse}")

    finally:
        print("Terminating instance...")
        ec2.terminate_instances(InstanceIds=[instance_id])
        print("Cleaning up security groups...")
        clean_up_unused_sgs(ec2, SG_TAG)


if __name__ == "__main__":
    main()