Update crc-usage, crc-sus, and crc_proposal_end to gather information…

… from keystone (#233)

* Start on building a usage table with info from keystone instead of crc-bank

* utility functions for interacting with keystone

* pull function for getting allocations by primary key

* add requests

* verify provided slurm account against keystone as well as slurm

* properly handle dates and pretty table rows

* split output into summary and usage tables, keep track of per cluster usage

* PEP fixes and comments

* Update crc_sus.py

Output remaining SUs on each clsuter using keystone requests and sreport

* adjustments from meeting

* filter on active and expiration date in request for allocations

* add and move slurm functions for gathering usage data into utils.system_info

* move slurm functions out of utils.keystone

* add usage table setup

* small fix with field_names

* fix imports, use account_name in checking for whether slurm account exists

* fix adding total to out_data

* clean up usage output formatting

* small formatting changes to crc_sus

* formatting adjustment

* add request primary key as query param for requests

* use cluster totals instead of allocation totals, other review items

* use request pk when gathering allocations

* use check_slurm_account_exists

* working state for crc-sus and crc-usage

* fix crc_usage test that wa failing

* pull duplicate functionality for running subprocesses

* small typo in crc_usage test string

* crc-proposal-end (#232)

* Update crc_proposal_end.py

Added authentication procedure

* Update crc_proposal_end.py

* Update crc_proposal_end.py

* Update crc_proposal_end.py

Updated requests with group_id

* Update crc_proposal_end.py

Update requests

* working version of crc_proposal_end

* remove unused imports

---------

Co-authored-by: Nickolas Comeau <[email protected]>

* fix expected output string

* codacy / PEP8 items

* pull crc-bank from dependencies

* add prettytable to deps

* needs to be an f string

* fix conversion to percent that was broken

* more broken division

* Rework utility functions to be more specific, adjust wrappers to utilize them

* use reworked utility functions in crc_usage and build/print summary and usage tables in separate functions

* codacy fixes

* fix bug where all allocations appear under each request

* researchgroup -> researchgroup ID, per_request option for per cluster totals

* user per_cluster_totals

* fix bare except

* small codacy fixes

---------

Co-authored-by: chnixi <[email protected]>

apps/crc_proposal_end.py

@@ -7,10 +7,11 @@
 import grp
 import os
 from argparse import Namespace
-
-from bank.account_logic import AccountServices
+from getpass import getpass
 
 from .utils.cli import BaseParser
+from .utils.keystone import *
+from .utils.system_info import Slurm
 
 
 class CrcProposalEnd(BaseParser):
@@ -31,8 +32,16 @@ def app_logic(self, args: Namespace) -> None:
             args: Parsed command line arguments
         """
 
-        acct = AccountServices(args.account)
-        end_date = acct._get_active_proposal_end_date()
+        Slurm.check_slurm_account_exists(args.account)
+        auth_header = get_auth_header(KEYSTONE_URL,
+                                      {'username': os.environ["USER"],
+                                       'password': getpass("Please enter your CRC login password:\n")})
+        keystone_group_id = get_researchgroup_id(KEYSTONE_URL, args.account, auth_header)
+        alloc_requests = get_active_requests(KEYSTONE_URL, keystone_group_id, auth_header)
+
+        if not (keystone_group_id and requests):
+            print(f"No active allocation information found in accounting system for '{args.account}'")
+            exit()
 
-        # Format the account name and end date as an easy-to-read string
-        print(f"The active proposal for account {args.account} ends on {end_date}")
+        for request in alloc_requests:
+            print(f"Resource Allocation Request: '{request['title']}' ends on {request['expire']} ")

apps/crc_sus.py

@@ -1,16 +1,17 @@
-"""Print an account's service unit allocation.
+"""Print an account's service unit allocation, with a usage value relative to that amount.
 
-This application is designed to interface with the CRC banking application
-and will not work without a running bank installation.
+This application is designed to interface with the Keystone banking system
+and will not work without a running keystone installation.
 """
+
 import grp
 import os
 from argparse import Namespace
-from typing import Dict
-
-from bank.account_logic import AccountServices
+from getpass import getpass
 
 from .utils.cli import BaseParser
+from .utils.keystone import *
+from .utils.system_info import Slurm
 
 
 class CrcSus(BaseParser):
@@ -24,47 +25,28 @@ def __init__(self) -> None:
         help_text = f"SLURM account name [defaults to your primary group: {default_group}]"
         self.add_argument('account', nargs='?', default=default_group, help=help_text)
 
-    def get_allocation_info(self, account: str) -> Dict[str, int]:
-        """Return the service unit allocation for a given account name
-
-        Args:
-            account: The name of the account
-
-        Returns:
-            A dictionary mapping cluster names to the number of service units
-        """
-
-        acct = AccountServices(account)
-        allocs = acct._get_active_proposal_allocation_info()
-
-        allocations = dict()
-        for cluster in allocs:
-            allocations[cluster.cluster_name] = cluster.service_units_total - cluster.service_units_used
-
-        return allocations
-
     @staticmethod
-    def build_output_string(account: str, **allocation: int) -> str:
+    def build_output_string(account: str, used: int, total: int, cluster: str) -> str:
         """Build a string describing an account's service unit allocation
 
         Args:
             account: The name of the account
-            allocation: The number of service units allocated for each cluster
-
+            total: The number of service units allocated for each cluster
+            used: number of SUs used on the cluster
+            cluster: name of cluster
         Returns:
             A string summarizing the account allocation
         """
 
         output_lines = [f'Account {account}']
 
-        # Right justify cluster names to the same length
-        cluster_name_length = max(len(cluster) for cluster in allocation)
-        for cluster, sus in allocation.items():
-            if sus > 0:
-                out = f' cluster {cluster:>{cluster_name_length}} has {sus:,} SUs remaining'
-            else:
-                out = f" cluster {cluster:>{cluster_name_length}} is LOCKED due to exceeding usage limits"
-            output_lines.append(out)
+        remaining = total - used
+
+        if remaining > 0:
+            out = f' cluster {cluster} has {remaining} SUs remaining'
+        else:
+            out = f" cluster {cluster} is LOCKED due to reaching usage limits"
+        output_lines.append(out)
 
         return '\n'.join(output_lines)
 
@@ -75,6 +57,25 @@ def app_logic(self, args: Namespace) -> None:
             args: Parsed command line arguments
         """
 
-        account_info = self.get_allocation_info(args.account)
-        output_string = self.build_output_string(args.account, **account_info)
-        print(output_string)
+        Slurm.check_slurm_account_exists(account_name=args.account)
+        auth_header = get_auth_header(KEYSTONE_URL,
+                                      {'username': os.environ["USER"],
+                                       'password': getpass("Please enter your CRC login password:\n")})
+        # Determine if provided or default account is in Keystone
+        keystone_group_id = get_researchgroup_id(KEYSTONE_URL, args.account, auth_header)
+        alloc_requests = get_active_requests(KEYSTONE_URL, keystone_group_id, auth_header)
+        if not (keystone_group_id and alloc_requests):
+            print(f"No active allocation information found in accounting system for '{args.account}'")
+            exit()
+
+        per_cluster_totals = get_per_cluster_totals(alloc_requests, auth_header)
+        earliest_date = get_earliest_startdate(alloc_requests)
+
+        for cluster in per_cluster_totals:
+            used = Slurm.get_cluster_usage_by_user(args.account, earliest_date, cluster)
+            if not used:
+                used = 0
+            else:
+                used = int(used['total'])
+
+            print(self.build_output_string(args.account, used, per_cluster_totals[cluster], cluster))

apps/crc_usage.py

@@ -1,21 +1,23 @@
-"""User facing application for returning user system usage from the CRC banking application.
+"""User facing application for returning user system usage from the Keystone application.
 
-This application is designed to interface with the CRC banking application
-and will not work without a running bank installation.
+This application is designed to interface with keystone
+and will not work without the application running on keystone.crc.pitt.edu.
 """
 
 import grp
 import os
 from argparse import Namespace
+from getpass import getpass
 
-from bank.account_logic import AccountServices
+from prettytable import PrettyTable
 
 from .utils.cli import BaseParser
-from .utils.system_info import Shell
+from .utils.keystone import *
+from .utils.system_info import Slurm
 
 
 class CrcUsage(BaseParser):
-    """Display a Slurm account's cluster usage."""
+    """Display a Slurm account's allocation usage"""
 
     def __init__(self) -> None:
         """Define the application commandline interface"""
@@ -26,16 +28,89 @@ def __init__(self) -> None:
         help_text = "slurm account name (defaults to the current user's primary group name)"
         self.add_argument('account', nargs='?', default=default_group, help=help_text)
 
+    @staticmethod
+    def print_summary_table(alloc_requests: [dict], account_name: str, per_request_totals: dict) -> None:
+        """Build and print a human-readable summary table for the slurm account with info from Keystone"""
+
+        # Initialize table for summary of requests and allocations
+        summary_table = PrettyTable(header=True, padding_width=2, max_table_width=79, min_table_width=79)
+        summary_table.title = f"Resource Allocation Request Information for '{account_name}'"
+        summary_table.field_names = ["ID", "TITLE", "EXPIRATION DATE"]
+
+        # Print request and allocation information for active allocations from the provided group
+        for request in alloc_requests:
+
+            summary_table.add_row([f"{request['id']}", f"{request['title']}", f"{request['expire']}"], divider=True)
+            summary_table.add_row(["", "CLUSTER", "SERVICE UNITS"])
+            summary_table.add_row(["", "----", "----"])
+            awarded_totals = per_request_totals[request['id']]
+
+            for cluster, total in awarded_totals.items():
+                summary_table.add_row(["", f"{cluster}", f"{total}"])
+
+            summary_table.add_row(["", "", ""], divider=True)
+
+        print(summary_table)
+
+    @staticmethod
+    def print_usage_table(account_name: str, awarded_totals: dict, earliest_date: date) -> None:
+        """Build and print a human-readable usage table for the slurm account with info from Keystone and
+        sreport"""
+
+        # Initialize table for summary of usage
+        usage_table = PrettyTable(header=False, padding_width=2, max_table_width=79, min_table_width=79)
+        usage_table.title = f"Summary of Usage Across All Clusters"
+
+        for cluster, total_awarded in awarded_totals.items():
+            usage_by_user = Slurm.get_cluster_usage_by_user(account_name=account_name,
+                                                            start_date=earliest_date,
+                                                            cluster=cluster)
+            if not usage_by_user:
+                usage_table.add_row([f"{cluster}", f"TOTAL USED: 0", f"AWARDED: {total_awarded}", f"% USED: 0"],
+                                    divider=True)
+                usage_table.add_row(["", "", "", ""], divider=True)
+                continue
+
+            total_used = int(usage_by_user.pop('total'))
+            percent_used = int((total_used / total_awarded * 100) // 1)
+            usage_table.add_row(
+                [f"{cluster}", f"TOTAL USED: {total_used}", f"AWARDED: {total_awarded}", f"% USED: {percent_used}"],
+                divider=True)
+            usage_table.add_row(["", "USER", "USED", "% USED"])
+            usage_table.add_row(["", "----", "----", "----"])
+            for user, usage in sorted(usage_by_user.items(), key=lambda item: item[1], reverse=True):
+                percent = int((usage / total_awarded * 100) // 1)
+                if percent == 0:
+                    percent = '<1'
+                usage_table.add_row(["", user, int(usage), percent])
+
+            usage_table.add_row(["", "", "", ""], divider=True)
+
+        print(usage_table)
+
     def app_logic(self, args: Namespace) -> None:
         """Logic to evaluate when executing the application
 
         Args:
             args: Parsed command line arguments
         """
 
-        account_exists = Shell.run_command(f'sacctmgr -n list account account={args.account} format=account%30')
-        if not account_exists:
-            raise RuntimeError(f"No slurm account was found with the name '{args.account}'.")
+        Slurm.check_slurm_account_exists(account_name=args.account)
+        auth_header = get_auth_header(KEYSTONE_URL,
+                                      {'username': os.environ["USER"],
+                                       'password': getpass("Please enter your CRC login password:\n")})
+
+        # Gather AllocationRequests from Keystone
+        keystone_group_id = get_researchgroup_id(KEYSTONE_URL, args.account, auth_header)
+        alloc_requests = get_active_requests(KEYSTONE_URL, keystone_group_id, auth_header)
+        if not (keystone_group_id and alloc_requests):
+            print(f"No active allocation data found in the accounting system for '{args.account}'")
+            exit()
+
+        self.print_summary_table(alloc_requests,
+                                 args.account,
+                                 get_per_cluster_totals(alloc_requests, auth_header, per_request=True))
 
-        account = AccountServices(args.account)
-        print(account._build_usage_table())
+        self.print_usage_table(args.account,
+                               get_per_cluster_totals(alloc_requests, auth_header),
+                               get_earliest_startdate(alloc_requests))

apps/utils/keystone.py

@@ -0,0 +1,79 @@
+"""Utility functions used across various wrappers for interacting with keystone"""
+
+from datetime import date
+
+import requests
+
+KEYSTONE_URL = "https://keystone.crc.pitt.edu"
+CLUSTERS = {1: 'MPI', 2: 'SMP', 3: 'HTC', 4: 'GPU'}
+
+
+def get_auth_header(keystone_url: str, auth_header: dict) -> dict:
+    """ Generate an authorization header to be used for accessing information from keystone"""
+
+    response = requests.post(f"{keystone_url}/authentication/new/", json=auth_header)
+    response.raise_for_status()
+    tokens = response.json()
+    return {"Authorization": f"Bearer {tokens['access']}"}
+
+
+def get_request_allocations(keystone_url: str, request_pk: int, auth_header: dict) -> dict:
+    """Get All Allocation information from keystone for a given request"""
+
+    response = requests.get(f"{keystone_url}/allocations/allocations/?request={request_pk}", headers=auth_header)
+    response.raise_for_status()
+    return response.json()
+
+
+def get_active_requests(keystone_url: str, group_pk: int, auth_header: dict) -> [dict]:
+    """Get all active AllocationRequest information from keystone for a given group"""
+
+    response = requests.get(f"{keystone_url}/allocations/requests/?group={group_pk}&status=AP", headers=auth_header)
+    response.raise_for_status()
+    return [request for request in response.json()
+            if date.fromisoformat(request['active']) <= date.today() < date.fromisoformat(request['expire'])]
+
+
+def get_researchgroup_id(keystone_url: str, account_name: str, auth_header: dict) -> int:
+    """Get the Researchgroup ID from keystone for the specified Slurm account"""
+
+    response = requests.get(f"{keystone_url}/users/researchgroups/?name={account_name}", headers=auth_header)
+    response.raise_for_status()
+
+    try:
+        group_id = int(response.json()[0]['id'])
+    except IndexError:
+        group_id = None
+
+    return group_id
+
+
+def get_earliest_startdate(alloc_requests: [dict]) -> date:
+    """Given a number of requests, determine the earliest start date across them"""
+
+    earliest_date = date.today()
+    for request in alloc_requests:
+        start = date.fromisoformat(request['active'])
+        if start < earliest_date:
+            earliest_date = start
+
+    return earliest_date
+
+
+def get_per_cluster_totals(alloc_requests: [dict], auth_header: dict, per_request: bool = False) -> dict:
+    """Gather the awarded totals across the given requests on each cluster into a dictionary"""
+
+    per_cluster_totals = {}
+    for request in alloc_requests:
+        if per_request:
+            per_cluster_totals[request['id']] = {}
+        for allocation in get_request_allocations(KEYSTONE_URL, request['id'], auth_header):
+            cluster = CLUSTERS[allocation['cluster']]
+            if per_request:
+                per_cluster_totals[request['id']].setdefault(cluster, 0)
+                per_cluster_totals[request['id']][cluster] += allocation['awarded']
+            else:
+                per_cluster_totals.setdefault(cluster, 0)
+                per_cluster_totals[cluster] += allocation['awarded']
+
+    return per_cluster_totals