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