greptimedb/.github/scripts/aws-ec2-runner-provision.py

#!/usr/bin/env python3
# Copyright 2023 Greptime Team
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# /// script
# requires-python = ">=3.10"
# dependencies = ["boto3>=1.35,<2"]
# ///

"""Provision an ephemeral EC2 CI runner from an Ubuntu x86_64 AMI."""

from __future__ import annotations

import argparse
import os
import re
import shlex
import time
from dataclasses import dataclass

from runner_utils import append_github_output, create_registration_token, find_runner_by_name, github_api

MANAGED_BY = "query-regression-ci"
REPO_TAG = "github-repository"
RUN_TAG = "query-regression-run-id"
RUNNER_TAG = "github-runner-name"
TTL_TAG_KEY = "runner-ttl-hours"


def make_client(region: str):
    import boto3
    from botocore.config import Config

    return boto3.client("ec2", region_name=region, config=Config(
        connect_timeout=10, read_timeout=30,
        retries={"mode": "standard", "max_attempts": 3},
    ))


@dataclass(frozen=True)
class ProvisionConfig:
    region: str
    image_id: str
    instance_type: str
    subnet_id: str
    security_group_id: str
    repo: str
    run_id: str
    github_token: str
    system_disk_gib: int = 500
    ttl_hours: int = 8

    def __post_init__(self):
        if not re.fullmatch(r"[0-9]+-[0-9]+", self.run_id):
            raise ValueError("run-id must be GitHub run_id-run_attempt")
        if not 20 <= self.system_disk_gib <= 2048:
            raise ValueError("system-disk-gib must be between 20 and 2048")
        if not 1 <= self.ttl_hours <= 168:
            raise ValueError("ttl-hours must be between 1 and 168")

    @property
    def runner_name(self) -> str:
        return f"ci-ec2-{self.run_id}"


def resolve_image(client, image_id: str) -> str:
    if image_id != "auto":
        return image_id
    images = client.describe_images(
        Owners=["099720109477"],
        Filters=[
            {"Name": "name", "Values": ["ubuntu/images/hvm-ssd-gp3/ubuntu-noble-24.04-amd64-server-*"]},
            {"Name": "state", "Values": ["available"]},
            {"Name": "architecture", "Values": ["x86_64"]},
        ],
    )["Images"]
    if not images:
        raise ValueError("No Canonical Ubuntu 24.04 amd64 AMI found in this region")
    image = max(images, key=lambda image: image["CreationDate"])
    print(f"Using AMI {image['ImageId']}: {image['Name']}")
    return image["ImageId"]


def root_disk(client, image_id: str, size: int) -> dict:
    images = client.describe_images(ImageIds=[image_id])["Images"]
    if len(images) != 1:
        raise ValueError(f"AMI not found: {image_id}")
    image = images[0]
    if image["State"] != "available" or image["Architecture"] != "x86_64":
        raise ValueError("Runner AMI must be available and x86_64")
    root = next((d for d in image["BlockDeviceMappings"]
                 if d["DeviceName"] == image["RootDeviceName"] and "Ebs" in d), None)
    if root is None:
        raise ValueError("Runner AMI must have an EBS root disk")
    minimum = root["Ebs"]["VolumeSize"]
    if size < minimum:
        raise ValueError(f"Requested root disk {size} GiB is smaller than AMI minimum {minimum} GiB")
    return {"DeviceName": image["RootDeviceName"], "Ebs": {
        "VolumeSize": size, "VolumeType": "gp3", "Iops": 3000,
        "Throughput": 125, "DeleteOnTermination": True,
    }}


def runner_download(token: str, repo: str) -> str:
    runners = github_api(token, "GET", f"/repos/{repo}/actions/runners/downloads")
    return next(r["download_url"] for r in runners if r["os"] == "linux" and r["architecture"] == "x64")


def render_user_data(name: str, token: str, repo: str, download_url: str) -> str:
    # Quote every external value before embedding it in the root shell script.
    name, token, repo, download_url = map(shlex.quote, (name, token, repo, download_url))
    return f"""#!/bin/bash
set -euo pipefail
for tool in git curl tar sudo; do command -v "$tool"; done
# Mask systemd-oomd before enabling swap: Ubuntu 24.04 kills the whole
# service cgroup on PSI pressure, and swap thrashing looks like pressure.
systemctl disable --now systemd-oomd.socket systemd-oomd.service || true
systemctl mask systemd-oomd.socket systemd-oomd.service || true
# apt-daily-upgrade / unattended-upgrades can `systemctl restart` services
# whose libraries were updated. That SIGTERM-s the runner mid-job; GitHub
# records UserCancelled (job still valid, child exit 143). These VMs live
# ~1h; security updates belong in the image, not at job time.
systemctl disable --now unattended-upgrades.service apt-daily.timer apt-daily-upgrade.timer apt-daily.service apt-daily-upgrade.service || true
systemctl mask unattended-upgrades.service apt-daily.timer apt-daily-upgrade.timer apt-daily.service apt-daily-upgrade.service || true
systemctl stop unattended-upgrades.service || true
killall -9 unattended-upgr || true
cat > /etc/apt/apt.conf.d/99disable-auto-updates <<'APTEOF'
APT::Periodic::Update-Package-Lists "0";
APT::Periodic::Unattended-Upgrade "0";
APT::Periodic::Download-Upgradeable-Packages "0";
APTEOF
if [[ ! -f "/swapfile" ]]; then
  fallocate --length 16G "/swapfile"
  chmod 600 "/swapfile"
  mkswap "/swapfile"
fi
swapon "/swapfile"
sysctl --write vm.swappiness=10
swapon --show
free --human

id runner >/dev/null 2>&1 || useradd --create-home --shell /bin/bash runner
# Establish membership before the runner service starts; setup installs Docker later.
getent group docker >/dev/null || groupadd --system docker
usermod -aG docker runner
printf 'runner ALL=(ALL) NOPASSWD: ALL\\n' > /etc/sudoers.d/ci-runner
chmod 0440 /etc/sudoers.d/ci-runner
visudo -cf /etc/sudoers.d/ci-runner
install -d -o runner -g "$(id -gn runner)" /home/runner/benchmark-data /opt/ci-github-runner
cd /opt/ci-github-runner
curl --fail --location --retry 3 {download_url} | tar xz
./bin/installdependencies.sh
chown -R runner:"$(id -gn runner)" /opt/ci-github-runner
runuser -u runner -- ./config.sh --unattended --ephemeral --disableupdate --url https://github.com/{repo} --token {token} --name {name} --labels {name}
./svc.sh install runner
# Match the ECS runner's behavior when a child process is OOM-killed.
service_name=$(cat .service)
mkdir -p "/etc/systemd/system/$service_name.d"
cat > "/etc/systemd/system/$service_name.d/oom.conf" <<'EOF'
[Service]
OOMPolicy=continue
EOF
systemctl daemon-reload
# Package installation in a job must not restart its runner service.
install -d /etc/needrestart/conf.d
cat > /etc/needrestart/conf.d/actions_runner_services.conf <<'EOF'
$nrconf{{override_rc}}{{qr(^actions\\.runner\\..+\\.service$)}} = 0;
EOF
./svc.sh start
"""


def run_instance(client, config: ProvisionConfig) -> str:
    image_id = resolve_image(client, config.image_id)
    disk = root_disk(client, image_id, config.system_disk_gib)
    instance = client.describe_instance_types(InstanceTypes=[config.instance_type])["InstanceTypes"][0]
    if "x86_64" not in instance["ProcessorInfo"]["SupportedArchitectures"]:
        raise ValueError("Runner instance type must support x86_64")
    token = create_registration_token(config.github_token, config.repo)
    print(f"::add-mask::{token}")
    script = render_user_data(config.runner_name, token, config.repo, runner_download(config.github_token, config.repo))
    tags = [
        {"Key": "Name", "Value": config.runner_name},
        {"Key": "managed-by", "Value": MANAGED_BY},
        {"Key": REPO_TAG, "Value": config.repo},
        {"Key": RUN_TAG, "Value": config.run_id},
        {"Key": RUNNER_TAG, "Value": config.runner_name},
        {"Key": TTL_TAG_KEY, "Value": str(config.ttl_hours)},
    ]
    response = client.run_instances(
        ImageId=image_id, InstanceType=config.instance_type,
        MinCount=1, MaxCount=1, SubnetId=config.subnet_id,
        SecurityGroupIds=[config.security_group_id],
        BlockDeviceMappings=[disk], UserData=script,
        ClientToken=config.runner_name,
        MetadataOptions={"HttpTokens": "required", "HttpEndpoint": "enabled"},
        TagSpecifications=[{"ResourceType": kind, "Tags": tags} for kind in ("instance", "volume")],
    )
    return response["Instances"][0]["InstanceId"]


def dump_console(client, instance_id: str) -> None:
    try:
        # Boto3 decodes the API's base64 console output.
        output = client.get_console_output(InstanceId=instance_id, Latest=True).get("Output", "")
        print(f"::group::EC2 console {instance_id}")
        print("\n".join(output.splitlines()[-80:]))
        print("::endgroup::")
    except Exception as error:
        print(f"Unable to read console for {instance_id}: {error}")


def provision(client, config: ProvisionConfig) -> int:
    instance_id = run_instance(client, config)
    # Publish before waiting: failed registration must still reach teardown.
    for key, value in (("instance_id", instance_id),
                       ("label", config.runner_name), ("runner_name", config.runner_name)):
        append_github_output(key, value)
    try:
        client.get_waiter("instance_running").wait(
            InstanceIds=[instance_id], WaiterConfig={"Delay": 5, "MaxAttempts": 60})
        deadline = time.monotonic() + 10 * 60
        while time.monotonic() < deadline:
            runner = find_runner_by_name(config.github_token, config.repo, config.runner_name)
            if runner and runner.get("status") == "online":
                print(f"Runner {config.runner_name} online: {instance_id}")
                return 0
            time.sleep(5)
        raise TimeoutError(f"Runner {config.runner_name} did not come online")
    finally:
        dump_console(client, instance_id)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    for name, env in (("region", "AWS_REGION"), ("repo", "GITHUB_REPOSITORY"),
                      ("github-token", "GH_PERSONAL_ACCESS_TOKEN"),
                      ("run-id", "AWS_EC2_RUN_ID"), ("image-id", "AWS_EC2_IMAGE_ID"),
                      ("instance-type", "AWS_EC2_INSTANCE_TYPE"),
                      ("subnet-id", "AWS_EC2_SUBNET_ID"),
                      ("security-group-id", "AWS_EC2_SECURITY_GROUP_ID")):
        parser.add_argument(f"--{name}", default=os.environ.get(env) or ("auto" if name == "image-id" else None))
    parser.add_argument("--system-disk-gib", type=int, default=os.environ.get("AWS_EC2_SYSTEM_DISK_GIB", "500"))
    parser.add_argument("--ttl-hours", type=int, default=os.environ.get("AWS_EC2_TTL_HOURS", "8"))
    args = parser.parse_args()
    for name, value in vars(args).items():
        if value is None or value == "":
            parser.error(f"Missing {name.replace('_', '-')}")
    config = ProvisionConfig(**vars(args))
    return provision(make_client(args.region), config)


if __name__ == "__main__":
    raise SystemExit(main())