#!/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"] # /// """Delete managed EC2 CI runners, or sweep expired instances by ownership tags.""" from __future__ import annotations import argparse import importlib.util import os import sys import time from pathlib import Path from datetime import datetime, timedelta, timezone from runner_utils import deregister_runner _SPEC = importlib.util.spec_from_file_location("aws_ec2_runner_provision", Path(__file__).with_name("aws-ec2-runner-provision.py")) assert _SPEC is not None and _SPEC.loader is not None provision = importlib.util.module_from_spec(_SPEC) sys.modules[_SPEC.name] = provision _SPEC.loader.exec_module(provision) make_client = provision.make_client dump_console = provision.dump_console MANAGED_BY = provision.MANAGED_BY REPO_TAG = provision.REPO_TAG RUN_TAG = provision.RUN_TAG RUNNER_TAG = provision.RUNNER_TAG TTL_TAG_KEY = provision.TTL_TAG_KEY SWEEP_TTL_HOURS = 4 def managed_instances(client, repo: str, run_id: str | None = None, instance_id: str | None = None) -> list[dict]: filters = [{"Name": "tag:managed-by", "Values": [MANAGED_BY]}, {"Name": f"tag:{REPO_TAG}", "Values": [repo]}] if instance_id: # A retried teardown may have a newer workflow attempt than creation. filters.append({"Name": "instance-id", "Values": [instance_id]}) elif run_id: filters.append({"Name": f"tag:{RUN_TAG}", "Values": [run_id]}) return [instance for page in client.get_paginator("describe_instances").paginate(Filters=filters) for reservation in page["Reservations"] for instance in reservation["Instances"]] def delete_instance(client, instance_id: str) -> bool: dump_console(client, instance_id) for attempt in range(12): try: client.terminate_instances(InstanceIds=[instance_id]) client.get_waiter("instance_terminated").wait( InstanceIds=[instance_id], WaiterConfig={"Delay": 5, "MaxAttempts": 60}) return True except Exception as error: code = getattr(error, "response", {}).get("Error", {}).get("Code") if code == "InvalidInstanceID.NotFound": if attempt < 11: time.sleep(5) continue # The ownership lookup observed this instance. Treat sustained # NotFound after bounded retries as an already-deleted VM. return True print(f"Failed to terminate {instance_id}: {error}") return False return False def teardown(client, repo: str, token: str, run_id: str | None = None, instance_id: str | None = None, now: datetime | None = None, runner_name: str | None = None) -> int: # Lookup by ownership/run tags also recovers a create whose response or job # outputs were lost. Sweep never touches legacy build/JSONBench instances. instances = managed_instances(client, repo, run_id, instance_id) # Tags can lag RunInstances. Retry before concluding a targeted VM is gone. if run_id and not instances: for _ in range(11): time.sleep(5) instances = managed_instances(client, repo, run_id, instance_id) if instances: break ok = True found = False for instance in instances: identity = instance["InstanceId"] if instance_id and identity != instance_id: continue found = True tags = {t["Key"]: t["Value"] for t in instance.get("Tags", [])} name = tags.get(RUNNER_TAG, "") if not name.startswith("ci-ec2-"): print(f"Skipping {identity}: missing managed runner name") ok = False continue if not run_id: ttl = tags.get(TTL_TAG_KEY) try: hours = SWEEP_TTL_HOURS if ttl is None else int(ttl) if not 1 <= hours <= 168: raise ValueError("TTL outside 1..168 hours") except (TypeError, ValueError): print(f"Skipping {identity}: invalid TTL {ttl!r}") continue if (now or datetime.now(timezone.utc)) - instance["LaunchTime"] < timedelta(hours=hours): continue # A GitHub outage must not prevent EC2 deletion, or vice versa. ok &= delete_instance(client, identity) try: ok &= deregister_runner(token, repo, name) except (Exception, SystemExit) as error: print(f"Failed to unregister {name}: {error}") ok = False if run_id and not found: # EC2 may already be gone while its GitHub registration remains. name = runner_name or f"ci-ec2-{run_id}" if not name.startswith("ci-ec2-"): raise ValueError("Invalid CI runner name") ok &= deregister_runner(token, repo, name) return 0 if ok else 1 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"), ("instance-id", "AWS_EC2_INSTANCE_ID"), ("runner-name", "AWS_EC2_RUNNER_NAME")): parser.add_argument(f"--{name}", default=os.environ.get(env)) parser.add_argument("--sweep", action="store_true") args = parser.parse_args() for name in ("region", "repo", "github_token"): if not getattr(args, name): parser.error(f"Missing {name.replace('_', '-')}") if not args.sweep and not args.run_id: parser.error("Missing run-id for targeted teardown") return teardown(make_client(args.region), args.repo, args.github_token, None if args.sweep else args.run_id, None if args.sweep else args.instance_id, runner_name=None if args.sweep else args.runner_name) if __name__ == "__main__": raise SystemExit(main())