tools/cluster-activity/arvados_cluster_activity/main.py

   1 #!/usr/bin/env python3
   2 # Copyright (C) The Arvados Authors. All rights reserved.
   3 #
   4 # SPDX-License-Identifier: AGPL-3.0
   5
   6 import argparse
   7 import sys
   8
   9 import arvados
  10 import arvados.util
  11 import ciso8601
  12 import csv
  13 import os
  14 import logging
  15 import re
  16
  17 from arvados_cluster_activity.report import ClusterActivityReport, aws_monthly_cost, format_with_suffix_base2
  18 from arvados_cluster_activity.prometheus import get_metric_usage, get_data_usage
  19
  20 from arvados_cluster_activity._version import __version__
  21
  22 from datetime import timedelta, timezone, datetime
  23 import base64
  24
  25 def parse_arguments(arguments):
  26     arg_parser = argparse.ArgumentParser()
  27     arg_parser.add_argument('--start', help='Start date for the report in YYYY-MM-DD format (UTC) (or use --days)')
  28     arg_parser.add_argument('--end', help='End date for the report in YYYY-MM-DD format (UTC), default "now"')
  29     arg_parser.add_argument('--days', type=int, help='Number of days before "end" to start the report (or use --start)')
  30     arg_parser.add_argument('--cost-report-file', type=str, help='Export cost report to specified CSV file')
  31     arg_parser.add_argument('--include-workflow-steps', default=False,
  32                             action="store_true", help='Include individual workflow steps (optional)')
  33     arg_parser.add_argument('--columns', type=str, help="""Cost report columns (optional), must be comma separated with no spaces between column names.
  34     Available columns are: Project, ProjectUUID, Workflow, WorkflowUUID, Step, StepUUID, Sample, SampleUUID, User, UserUUID, Submitted, Started, Runtime, Cost""")
  35     arg_parser.add_argument('--exclude', type=str, help="Exclude workflows containing this substring (may be a regular expression)")
  36
  37     arg_parser.add_argument('--html-report-file', type=str, help='Export HTML report to specified file')
  38     arg_parser.add_argument(
  39         '--version', action='version', version="%s %s" % (sys.argv[0], __version__),
  40         help='Print version and exit.')
  41
  42     arg_parser.add_argument('--cluster', type=str, help='Cluster to query for prometheus stats')
  43     arg_parser.add_argument('--prometheus-auth', type=str, help='Authorization file with prometheus info')
  44
  45     args = arg_parser.parse_args(arguments)
  46
  47     if args.days and args.start:
  48         arg_parser.print_help()
  49         print("Error: either specify --days or both --start and --end")
  50         exit(1)
  51
  52     if not args.days and not args.start:
  53         arg_parser.print_help()
  54         print("\nError: either specify --days or both --start and --end")
  55         exit(1)
  56
  57     if (args.start and not args.end):
  58         arg_parser.print_help()
  59         print("\nError: no start or end date found, either specify --days or both --start and --end")
  60         exit(1)
  61
  62     if args.end:
  63         try:
  64             to = datetime.strptime(args.end,"%Y-%m-%d")
  65         except:
  66             arg_parser.print_help()
  67             print("\nError: end date must be in YYYY-MM-DD format")
  68             exit(1)
  69     else:
  70         to = datetime.now(timezone.utc)
  71
  72     if args.days:
  73         since = to - timedelta(days=args.days)
  74
  75     if args.start:
  76         try:
  77             since = datetime.strptime(args.start,"%Y-%m-%d")
  78         except:
  79             arg_parser.print_help()
  80             print("\nError: start date must be in YYYY-MM-DD format")
  81             exit(1)
  82
  83
  84     if args.prometheus_auth:
  85         with open(args.prometheus_auth, "rt") as f:
  86             for line in f:
  87                 if line.startswith("export "):
  88                    line = line[7:]
  89                 sp = line.strip().split("=")
  90                 if sp[0].startswith("PROMETHEUS_"):
  91                     os.environ[sp[0]] = sp[1]
  92
  93     return args, since, to
  94
  95 def print_data_usage(prom, timestamp, cluster, label):
  96     value, dedup_ratio = get_data_usage(prom, timestamp, cluster)
  97
  98     if value is None:
  99         return
 100
 101     monthly_cost = aws_monthly_cost(value)
 102     print(label,
 103           "%s apparent," % (format_with_suffix_base2(value*dedup_ratio)),
 104           "%s actually stored," % (format_with_suffix_base2(value)),
 105           "$%.2f monthly S3 storage cost" % monthly_cost)
 106
 107 def print_container_usage(prom, start_time, end_time, metric, label, fn=None):
 108     cumulative = 0
 109
 110     for rs in get_metric_usage(prom, start_time, end_time, metric):
 111         # Calculate the sum of values
 112         #print(rs.sum()["y"])
 113         cumulative += rs.sum()["y"]
 114
 115     if fn is not None:
 116         cumulative = fn(cumulative)
 117
 118     print(label % cumulative)
 119
 120
 121 def get_prometheus_client():
 122     from prometheus_api_client import PrometheusConnect
 123
 124     prom_host = os.environ.get("PROMETHEUS_HOST")
 125     prom_token = os.environ.get("PROMETHEUS_APIKEY")
 126     prom_user = os.environ.get("PROMETHEUS_USER")
 127     prom_pw = os.environ.get("PROMETHEUS_PASSWORD")
 128
 129     headers = {}
 130     if prom_token:
 131         headers["Authorization"] = "Bearer %s" % prom_token
 132
 133     if prom_user:
 134         headers["Authorization"] = "Basic %s" % str(base64.b64encode(bytes("%s:%s" % (prom_user, prom_pw), 'utf-8')), 'utf-8')
 135
 136     try:
 137         return PrometheusConnect(url=prom_host, headers=headers)
 138     except Exception as e:
 139         logging.warn("Connecting to Prometheus failed, will not collect activity from Prometheus.  Error was: %s" % e)
 140         return None
 141
 142 def report_from_prometheus(prom, cluster, since, to):
 143
 144     if not cluster:
 145         arv_client = arvados.api()
 146         cluster = arv_client.config()["ClusterID"]
 147
 148     print(cluster, "between", since, "and", to, "timespan", (to-since))
 149
 150     try:
 151         print_data_usage(prom, since, cluster, "at start:")
 152     except:
 153         logging.exception("Failed to get start value")
 154
 155     try:
 156         print_data_usage(prom, to - timedelta(minutes=240), cluster, "current :")
 157     except:
 158         logging.exception("Failed to get end value")
 159
 160     print_container_usage(prom, since, to, "arvados_dispatchcloud_containers_running{cluster='%s'}" % cluster, '%.1f container hours', lambda x: x/60)
 161     print_container_usage(prom, since, to, "sum(arvados_dispatchcloud_instances_price{cluster='%s'})" % cluster, '$%.2f spent on compute', lambda x: x/60)
 162     print()
 163
 164
 165 def main(arguments=None):
 166     if arguments is None:
 167         arguments = sys.argv[1:]
 168
 169     args, since, to = parse_arguments(arguments)
 170
 171     logging.getLogger().setLevel(logging.INFO)
 172
 173     prom = None
 174     if "PROMETHEUS_HOST" in os.environ:
 175         prom = get_prometheus_client()
 176     else:
 177         logging.warn("PROMETHEUS_HOST not found, not collecting activity from Prometheus")
 178
 179     reporter = ClusterActivityReport(prom)
 180
 181     if args.cost_report_file:
 182         with open(args.cost_report_file, "wt") as f:
 183             reporter.csv_report(since, to, f, args.include_workflow_steps, args.columns, args.exclude)
 184     else:
 185         logging.info("Use --cost-report-file to get a CSV file of workflow runs")
 186
 187     if args.html_report_file:
 188         with open(args.html_report_file, "wt") as f:
 189             f.write(reporter.html_report(since, to, args.exclude, args.include_workflow_steps))
 190     else:
 191         logging.info("Use --html-report-file to get HTML report of cluster usage")
 192
 193     if not args.cost_report_file and not args.html_report_file:
 194         report_from_prometheus(prom, args.cluster, since, to)
 195
 196 if __name__ == "__main__":
 197     main()