mirror of
https://github.com/ClickHouse/ClickHouse.git
synced 2024-12-11 17:02:25 +00:00
218 lines
7.8 KiB
Python
218 lines
7.8 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import requests
|
|
import argparse
|
|
import jwt
|
|
import sys
|
|
import json
|
|
import time
|
|
from collections import namedtuple
|
|
import boto3
|
|
|
|
def get_dead_runners_in_ec2(runners):
|
|
ids = {runner.name: runner for runner in runners if runner.offline == True and runner.busy == False}
|
|
if not ids:
|
|
return []
|
|
|
|
client = boto3.client('ec2')
|
|
|
|
print("Checking ids", list(ids.keys()))
|
|
instances_statuses = client.describe_instance_status(InstanceIds=list(ids.keys()))
|
|
found_instances = set([])
|
|
print("Response", instances_statuses)
|
|
for instance_status in instances_statuses['InstanceStatuses']:
|
|
if instance_status['InstanceState']['Name'] in ('pending', 'running'):
|
|
found_instances.add(instance_status['InstanceId'])
|
|
|
|
print("Found instances", found_instances)
|
|
result_to_delete = []
|
|
for instance_id, runner in ids.items():
|
|
if instance_id not in found_instances:
|
|
print("Instance", instance_id, "is not alive, going to remove it")
|
|
result_to_delete.append(runner)
|
|
return result_to_delete
|
|
|
|
def get_key_and_app_from_aws():
|
|
import boto3
|
|
secret_name = "clickhouse_github_secret_key"
|
|
session = boto3.session.Session()
|
|
client = session.client(
|
|
service_name='secretsmanager',
|
|
)
|
|
get_secret_value_response = client.get_secret_value(
|
|
SecretId=secret_name
|
|
)
|
|
data = json.loads(get_secret_value_response['SecretString'])
|
|
return data['clickhouse-app-key'], int(data['clickhouse-app-id'])
|
|
|
|
def handler(event, context):
|
|
private_key, app_id = get_key_and_app_from_aws()
|
|
main(private_key, app_id, True, True)
|
|
|
|
def get_installation_id(jwt_token):
|
|
headers = {
|
|
"Authorization": f"Bearer {jwt_token}",
|
|
"Accept": "application/vnd.github.v3+json",
|
|
}
|
|
response = requests.get("https://api.github.com/app/installations", headers=headers)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
return data[0]['id']
|
|
|
|
def get_access_token(jwt_token, installation_id):
|
|
headers = {
|
|
"Authorization": f"Bearer {jwt_token}",
|
|
"Accept": "application/vnd.github.v3+json",
|
|
}
|
|
response = requests.post(f"https://api.github.com/app/installations/{installation_id}/access_tokens", headers=headers)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
return data['token']
|
|
|
|
|
|
RunnerDescription = namedtuple('RunnerDescription', ['id', 'name', 'tags', 'offline', 'busy'])
|
|
|
|
def list_runners(access_token):
|
|
headers = {
|
|
"Authorization": f"token {access_token}",
|
|
"Accept": "application/vnd.github.v3+json",
|
|
}
|
|
response = requests.get("https://api.github.com/orgs/ClickHouse/actions/runners?per_page=100", headers=headers)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
total_runners = data['total_count']
|
|
runners = data['runners']
|
|
|
|
total_pages = int(total_runners / 100 + 1)
|
|
print("Total pages", total_pages)
|
|
for i in range(2, total_pages + 1):
|
|
response = requests.get(f"https://api.github.com/orgs/ClickHouse/actions/runners?page={i}&per_page=100", headers=headers)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
runners += data['runners']
|
|
|
|
print("Total runners", len(runners))
|
|
result = []
|
|
for runner in runners:
|
|
tags = [tag['name'] for tag in runner['labels']]
|
|
desc = RunnerDescription(id=runner['id'], name=runner['name'], tags=tags,
|
|
offline=runner['status']=='offline', busy=runner['busy'])
|
|
result.append(desc)
|
|
|
|
return result
|
|
|
|
def group_runners_by_tag(listed_runners):
|
|
result = {}
|
|
|
|
RUNNER_TYPE_LABELS = ['style-checker', 'builder', 'func-tester', 'stress-tester', 'fuzzer-unit-tester']
|
|
for runner in listed_runners:
|
|
for tag in runner.tags:
|
|
if tag in RUNNER_TYPE_LABELS:
|
|
if tag not in result:
|
|
result[tag] = []
|
|
result[tag].append(runner)
|
|
break
|
|
else:
|
|
if 'unlabeled' not in result:
|
|
result['unlabeled'] = []
|
|
result['unlabeled'].append(runner)
|
|
return result
|
|
|
|
|
|
def push_metrics_to_cloudwatch(listed_runners, namespace):
|
|
client = boto3.client('cloudwatch')
|
|
metrics_data = []
|
|
busy_runners = sum(1 for runner in listed_runners if runner.busy and not runner.offline)
|
|
metrics_data.append({
|
|
'MetricName': 'BusyRunners',
|
|
'Value': busy_runners,
|
|
'Unit': 'Count',
|
|
})
|
|
total_active_runners = sum(1 for runner in listed_runners if not runner.offline)
|
|
metrics_data.append({
|
|
'MetricName': 'ActiveRunners',
|
|
'Value': total_active_runners,
|
|
'Unit': 'Count',
|
|
})
|
|
total_runners = len(listed_runners)
|
|
metrics_data.append({
|
|
'MetricName': 'TotalRunners',
|
|
'Value': total_runners,
|
|
'Unit': 'Count',
|
|
})
|
|
if total_active_runners == 0:
|
|
busy_ratio = 100
|
|
else:
|
|
busy_ratio = busy_runners / total_active_runners * 100
|
|
|
|
metrics_data.append({
|
|
'MetricName': 'BusyRunnersRatio',
|
|
'Value': busy_ratio,
|
|
'Unit': 'Percent',
|
|
})
|
|
|
|
client.put_metric_data(Namespace=namespace, MetricData=metrics_data)
|
|
|
|
def delete_runner(access_token, runner):
|
|
headers = {
|
|
"Authorization": f"token {access_token}",
|
|
"Accept": "application/vnd.github.v3+json",
|
|
}
|
|
|
|
response = requests.delete(f"https://api.github.com/orgs/ClickHouse/actions/runners/{runner.id}", headers=headers)
|
|
response.raise_for_status()
|
|
print(f"Response code deleting {runner.name} is {response.status_code}")
|
|
return response.status_code == 204
|
|
|
|
def main(github_secret_key, github_app_id, push_to_cloudwatch, delete_offline_runners):
|
|
payload = {
|
|
"iat": int(time.time()) - 60,
|
|
"exp": int(time.time()) + (10 * 60),
|
|
"iss": github_app_id,
|
|
}
|
|
|
|
encoded_jwt = jwt.encode(payload, github_secret_key, algorithm="RS256")
|
|
installation_id = get_installation_id(encoded_jwt)
|
|
access_token = get_access_token(encoded_jwt, installation_id)
|
|
runners = list_runners(access_token)
|
|
grouped_runners = group_runners_by_tag(runners)
|
|
for group, group_runners in grouped_runners.items():
|
|
if push_to_cloudwatch:
|
|
print(group)
|
|
push_metrics_to_cloudwatch(group_runners, 'RunnersMetrics/' + group)
|
|
else:
|
|
print(group, f"({len(group_runners)})")
|
|
for runner in group_runners:
|
|
print('\t', runner)
|
|
|
|
if delete_offline_runners:
|
|
print("Going to delete offline runners")
|
|
dead_runners = get_dead_runners_in_ec2(runners)
|
|
for runner in dead_runners:
|
|
print("Deleting runner", runner)
|
|
delete_runner(access_token, runner)
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description='Get list of runners and their states')
|
|
parser.add_argument('-p', '--private-key-path', help='Path to file with private key')
|
|
parser.add_argument('-k', '--private-key', help='Private key')
|
|
parser.add_argument('-a', '--app-id', type=int, help='GitHub application ID', required=True)
|
|
parser.add_argument('--push-to-cloudwatch', action='store_true', help='Store received token in parameter store')
|
|
parser.add_argument('--delete-offline', action='store_true', help='Remove offline runners')
|
|
|
|
args = parser.parse_args()
|
|
|
|
if not args.private_key_path and not args.private_key:
|
|
print("Either --private-key-path or --private-key must be specified", file=sys.stderr)
|
|
|
|
if args.private_key_path and args.private_key:
|
|
print("Either --private-key-path or --private-key must be specified", file=sys.stderr)
|
|
|
|
if args.private_key:
|
|
private_key = args.private_key
|
|
else:
|
|
with open(args.private_key_path, 'r') as key_file:
|
|
private_key = key_file.read()
|
|
|
|
main(private_key, args.app_id, args.push_to_cloudwatch, args.delete_offline)
|