Usage monitoring

usage-monitoring/config.json

@@ -0,0 +1,10 @@
+{
+  "token_file": "~/.config/usage-monitoring/os-token.json",
+  "allocation_resources": [
+    "jetstream2.indiana.xsede.org",
+    "jetstream2-gpu.indiana.xsede.org",
+    "jetstream2-lm.indiana.xsede.org"
+  ],
+  "data_file": "~/usage_monitoring.csv",
+  "test_csv_file": "/tmp/usage_monitoring_test.csv"
+}

usage-monitoring/install.sh

@@ -0,0 +1,11 @@
+# /usr/bin/env /usr/bin/bash
+
+set -x
+
+INSTALL_PATH=$HOME/.local/bin/
+CONFIG_PATH=$HOME/.config/usage-monitoring/
+
+mkdir -p $INSTALL_PATH $CONFIG_PATH
+
+ln -s $PWD/usage_monitoring.py $INSTALL_PATH/usage_monitoring.py
+ln -s $PWD/config.json $CONFIG_PATH/config.json

usage-monitoring/usage_monitoring.py

@@ -1,6 +1,7 @@
-from os import system
+#! /usr/bin/env conda run -n usage-monitoring python
+
 from subprocess import run
-from os.path import isfile
+from os.path import isfile, expanduser
 
 import argparse
 
@@ -11,7 +12,17 @@
 import requests
 from matplotlib import pyplot as plt
 
-from usage_monitoring_config import *
+def load_config(config_path):
+    try:
+        with open(expanduser(config_path)) as config:
+            c = json.load(config)
+            c['token_file'] = expanduser(c['token_file'])
+            c['data_file'] = expanduser(c['data_file'])
+            c['test_csv_file'] = expanduser(c['test_csv_file'])
+            return c
+    except FileNotFoundError:
+        print(f'File {config_path} not found! Exiting ...')
+        exit(1)
 
 def create_os_token(token_file):
     token = run(
@@ -30,21 +41,25 @@ def token_expired(token_file):
         now = datetime.now(UTC).timestamp()
         expire < now    
 
-def get_os_token(token_file='/tmp/os-token.json', force_new_token=False):
+def get_os_token(token_file, force_new_token=False):
     if not isfile(token_file) or force_new_token or token_expired(token_file):
         create_os_token(token_file)
     with open(token_file, 'r') as f:
-        json.load(f)['id']
+        return json.load(f)['id']
 
 def query_accounting_api(token):
     url = 'https://js2.jetstream-cloud.org:9001'
     headers = { 'X-Auth-Token': f'{token}' }
     response = requests.get(url, headers=headers)
-    response.raise_for_status()
+    try:
+        response.raise_for_status()
+    except Exception as ex:
+        print(ex)
+        exit(1)
     query = json.loads(response.text)
     return query
 
-def get_js2_resources(query):
+def get_js2_resources(query,allocation_resources):
     now = datetime.now()
     date_format = '%Y-%m-%d'
     all_resources = [ 
@@ -206,7 +221,10 @@ def usage_analysis(data,days_prior):
 
         # tot_sus - s1 = remaining_sus = r*(t2 - t1) --> t2 = remaining_sus/r + t1
         exhausted_ts = remaining_sus/r + cur_ts
-        exhausted_date = datetime.fromtimestamp(exhausted_ts)
+        try:
+            exhausted_date = datetime.fromtimestamp(exhausted_ts)
+        except OverflowError:
+            exhausted_date = None
 
         # s2 - s1 = r*(t2 - t1) --> s2 = r*(t2 - t1) + s1
         date_format = '%Y-%m-%d'
@@ -242,10 +260,13 @@ def usage_analysis(data,days_prior):
 
     return analysis
 
-def generate_usage_plot(resources, analyses):
+def generate_usage_plot(resources, analyses, allocation_resources):
     fig, ax = plt.subplots()
     for resource_type in allocation_resources:
         data = get_data_by_resource(resources, resource_type)
+        if data.empty:
+            print(f'No available data for {resource_type}')
+            continue
 
         timestamps = pd.array(data['timestamp'])
         dates = [ datetime.fromtimestamp(ts) for ts in timestamps ]
@@ -258,44 +279,55 @@ def generate_usage_plot(resources, analyses):
     plt.show()
     return 0
 
-def main(data_file):
+def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('-n', '--force-new-token', help='Force the creation of a new openstack token before query', action='store_true')
-    parser.add_argument('-w', '--write', help=f'Query Jetstream2 for new allocation data and write to data file: {data_file}', action='store_true')
-    parser.add_argument('-c', '--dump-csv', help=f'Dump the data from {data_file} in csv format', action='store_true')
-    parser.add_argument('-j', '--dump-json', help=f'Dump the data from {data_file} in json format', action='store_true')
+    parser.add_argument('-w', '--write', help='Query Jetstream2 for new allocation data and write to data file', action='store_true')
+    parser.add_argument('-c', '--dump-csv', help='Dump the data from data_file in csv format', action='store_true')
+    parser.add_argument('-j', '--dump-json', help='Dump the data from data_file in json format', action='store_true')
     parser.add_argument('-p', '--plot', help='Generate an interactive plot of SU usage data', action='store_true')
     parser.add_argument('-a', '--analysis-days', help='Days prior for which to perform an analysis', action='extend', nargs='+', type=int)
     parser.add_argument('-d', '--devel', help='Use test_csv_file for development work', action='store_true')
+    parser.add_argument('--config', help='Configuration file path', type=str, default="~/.config/usage-monitoring/config.json")
     args = vars(parser.parse_args())
 
+    c = load_config(args['config'])
+
     if not any([ args[key] for key in args.keys() ]):
         parser.parse_args(['--help'])
 
     if args['devel']:
-        data_file = test_csv_file
+        c['data_file'] = c['test_csv_file']
 
     if args['write']:
-        token = get_os_token(token_file,force_new_token=args['force_new_token'])
+        token = get_os_token(c['token_file'],force_new_token=args['force_new_token'])
         query = query_accounting_api(token)
-        resources = get_js2_resources(query)
-        write_resource_csv(resources, data_file)
+        resources = get_js2_resources(query,c['allocation_resources'])
+        write_resource_csv(resources, c['data_file'])
 
     if args['dump_csv']:
-        system(f'cat {data_file}')
+        dump = run(
+            ['cat', f'{c['data_file']}'],
+            check=True
+        )
 
     if args['dump_json']:
-        resources = read_resource_csv(data_file)
-        print(json.dumps(resources, indent=2))
+        resources = read_resource_csv(c['data_file'])
+        print(resources.to_json(orient='records', indent=2))
 
     if args['analysis_days']:
         # Get resources
-        resources = read_resource_csv(data_file)
+        resources = read_resource_csv(c['data_file'])
 
         analyses = []
         # Loop over resources to get each type of data found in allocation_resources
-        for resource_type in allocation_resources:
+        for resource_type in c['allocation_resources']:
             data = get_data_by_resource(resources, resource_type)
+            if data.empty:
+                print(f'No available data for {resource_type}. Skipping ...')
+                continue
+            if len(data) < 2:
+                print(f'Not enough data for {resource_type}: len(data) = {len(data)}. Skipping ...')
             # Perform analysis (usage rates, "forecast", )
             analyses.append(usage_analysis(data,args['analysis_days']))
 
@@ -304,8 +336,8 @@ def main(data_file):
     if args['plot']:
         if 'analyses' not in locals():
             analyses = None
-        resources = read_resource_csv(data_file)
-        generate_usage_plot(resources, analyses)
+        resources = read_resource_csv(c['data_file'])
+        generate_usage_plot(resources, analyses, c['allocation_resources'])
 
 if __name__ == "__main__":
-    main(data_file)
+    main()