From 4ab06416a1532a925b3e5e396075f76f76db221c Mon Sep 17 00:00:00 2001 From: Claire Peters Date: Fri, 10 Nov 2023 17:46:38 -0800 Subject: [PATCH] more refactoring --- .../xdmod/management/commands/xdmod_usage.py | 82 +++++++++++-------- coldfront/plugins/xdmod/utils.py | 50 ++++++----- 2 files changed, 78 insertions(+), 54 deletions(-) diff --git a/coldfront/plugins/xdmod/management/commands/xdmod_usage.py b/coldfront/plugins/xdmod/management/commands/xdmod_usage.py index fd1948554..d4835b17a 100644 --- a/coldfront/plugins/xdmod/management/commands/xdmod_usage.py +++ b/coldfront/plugins/xdmod/management/commands/xdmod_usage.py @@ -6,6 +6,7 @@ from django.db.models import Q from coldfront.core.allocation.models import Allocation +from coldfront.core.resource.models import Resource from coldfront.plugins.xdmod.utils import (XDMOD_ACCOUNT_ATTRIBUTE_NAME, XDMOD_CLOUD_CORE_TIME_ATTRIBUTE_NAME, XDMOD_CLOUD_PROJECT_ATTRIBUTE_NAME, @@ -57,23 +58,20 @@ def write(self, data): os.dup2(devnull, sys.stdout.fileno()) sys.exit(1) - def id_resources(self, s): + def run_resource_checks(self, resource): + rname = resource.get_attribute(XDMOD_RESOURCE_ATTRIBUTE_NAME) + if not rname and resource.parent_resource: + rname = resource.parent_resource.get_attribute( + XDMOD_RESOURCE_ATTRIBUTE_NAME) + if self.filter_resource and self.filter_resource != rname: + return None + return rname + + def id_allocation_resources(self, s): resources = [] for r in s.resources.all(): - rname = r.get_attribute(XDMOD_RESOURCE_ATTRIBUTE_NAME) - if not rname and r.parent_resource: - rname = r.parent_resource.get_attribute( - XDMOD_RESOURCE_ATTRIBUTE_NAME) - - if not rname: - continue - if self.filter_resource and self.filter_resource != rname: - continue - + rname = self.run_resource_checks(r) resources.append(rname) - if len(resources) == 0: - logger.warning("%s attribute not found on any resouces for allocation: %s", - XDMOD_RESOURCE_ATTRIBUTE_NAME, s) return resources def attribute_check(self, s, attr_name, num=False): @@ -88,9 +86,18 @@ def attribute_check(self, s, attr_name, num=False): return attr def filter_allocations(self, allocations, account_attr_name=XDMOD_ACCOUNT_ATTRIBUTE_NAME): - allocations = allocations.select_related('project').prefetch_related( - 'resources', 'allocationattribute_set', 'allocationuser_set' - ) + cleared_resources = Resource.objects.filter( + Q(resourceattribute__resource_attribute_type__name=XDMOD_RESOURCE_ATTRIBUTE_NAME) | + Q(parent_resource__resourceattribute__resource_attribute_type__name=XDMOD_RESOURCE_ATTRIBUTE_NAME) + ) + + allocations = ( + allocations.select_related('project') + .prefetch_related( + 'resources', 'allocationattribute_set', 'allocationuser_set' + ) + .filter(resources__in=cleared_resources) + ) if self.fetch_expired: allocations = allocations.filter(~Q(status__name='Active')) else: @@ -133,7 +140,8 @@ def process_total_storage(self): ] ) - allocations = self.filter_allocations(allocations, account_attr_name=XDMOD_STORAGE_GROUP_ATTRIBUTE_NAME) + allocations = self.filter_allocations( + allocations, account_attr_name=XDMOD_STORAGE_GROUP_ATTRIBUTE_NAME) for s in allocations.distinct(): account_name = self.attribute_check(s, XDMOD_STORAGE_GROUP_ATTRIBUTE_NAME) @@ -141,13 +149,11 @@ def process_total_storage(self): if None in [account_name, cpu_hours]: continue - resources = self.id_resources(s) - if len(resources) == 0: - continue + resources = self.id_allocation_resources(s) fetcher = XDModFetcher(s.start_date, s.end_date, resources=resources) try: - usage = fetcher.xdmod_fetch_total_storage( + usage = fetcher.xdmod_fetch_storage( account_name, statistics='avg_physical_usage' ) except XdmodNotFoundError: @@ -199,13 +205,11 @@ def process_total_gpu_hours(self): if None in [account_name, cpu_hours]: continue - resources = self.id_resources(s) - if len(resources) == 0: - continue + resources = self.id_allocation_resources(s) fetcher = XDModFetcher(s.start_date, s.end_date, resources=resources) try: - usage = fetcher.xdmod_fetch_total_cpu_hours( + usage = fetcher.xdmod_fetch_cpu_hours( account_name, statistics='total_gpu_hours' ) except XdmodNotFoundError: @@ -250,19 +254,33 @@ def process_total_cpu_hours(self): ) allocations = self.filter_allocations(allocations) + resource_filter = ( + Q(resourceattribute__resource_attribute_type__name=XDMOD_RESOURCE_ATTRIBUTE_NAME) | + Q(parent_resource__resourceattribute__resource_attribute_type__name=XDMOD_RESOURCE_ATTRIBUTE_NAME) + ) + resources_all = Resource.objects.filter( + id__in=[r.id for a in allocations for r in a.resources.all()]).filter( + resource_filter + ) + + # fetcher = XDModFetcher(resources=resources_all) + # try: + # usage = fetcher.xdmod_fetch_all_project_usages('total_cpu_hours') + # except XdmodNotFoundError: + # raise XdmodNotFoundError( + # "No data in XDMoD found for resources %s", resources_all) + for s in allocations.distinct(): account_name = self.attribute_check(s, XDMOD_ACCOUNT_ATTRIBUTE_NAME) cpu_hours = self.attribute_check(s, XDMOD_CPU_HOURS_ATTRIBUTE_NAME, num=True) if None in [account_name, cpu_hours]: continue - resources = self.id_resources(s) - if len(resources) == 0: - continue + resources = self.id_allocation_resources(s) fetcher = XDModFetcher(s.start_date, s.end_date, resources=resources) try: - usage = fetcher.xdmod_fetch_total_cpu_hours(account_name) + usage = fetcher.xdmod_fetch_cpu_hours(account_name) except XdmodNotFoundError: logger.warning( "No data in XDMoD found for allocation %s account %s resources %s", @@ -316,9 +334,7 @@ def process_cloud_core_time(self): s, XDMOD_CLOUD_CORE_TIME_ATTRIBUTE_NAME, num=True) if None in [project_name, core_time]: continue - resources = self.id_resources(s) - if len(resources) == 0: - continue + resources = self.id_allocation_resources(s) fetcher = XDModFetcher(s.start_date, s.end_date, resources=resources) try: diff --git a/coldfront/plugins/xdmod/utils.py b/coldfront/plugins/xdmod/utils.py index c027ce126..cbfb23cee 100644 --- a/coldfront/plugins/xdmod/utils.py +++ b/coldfront/plugins/xdmod/utils.py @@ -48,6 +48,9 @@ logger = logging.getLogger(__name__) +QUARTER_START, QUARTER_END = get_quarter_start_end() + + class XdmodError(Exception): pass @@ -55,22 +58,20 @@ class XdmodNotFoundError(XdmodError): pass class XDModFetcher: - def __init__(self, start, end, resources=None,): + def __init__(self, start=QUARTER_START, end=QUARTER_END, resources=None): self.url = f'{XDMOD_API_URL}{_ENDPOINT_CORE_HOURS}' if resources is None: resources = [] payload = _DEFAULT_PARAMS - # payload['start_date'] = start - # payload['end_date'] = end - start_date, end_date = get_quarter_start_end() - payload['start_date'] = start_date - payload['end_date'] = end_date + payload['start_date'] = start + payload['end_date'] = end payload['resource_filter'] = f'"{",".join(resources)}"' payload['operation'] = 'get_data' self.payload = payload + self.group_by = {'total':'pi', 'per-user':'user'} - def fetch_data(self, search_item, payload): + def fetch_data(self, payload, search_item=None): r = requests.get( self.url, params=payload, auth=HTTPBasicAuth(XDMOD_USER, XDMOD_PASS) ) @@ -93,7 +94,7 @@ def fetch_data(self, search_item, payload): rows = root.find('rows') if len(rows) != 1: raise XdmodNotFoundError( - f'Rows not found for {search_item} - {self.payload["resources"]}' + f'Rows not found for {search_item} - {self.payload["resource_filter"]}' ) cells = rows.find('row').findall('cell') if len(cells) != 2: @@ -102,26 +103,33 @@ def fetch_data(self, search_item, payload): stats = cells[1].find('value').text return stats - def xdmod_fetch_total_cpu_hours(self, account, statistics='total_cpu_hours'): - """fetch total cpu hours.""" + def xdmod_fetch(self, account, statistics, realm, group_by='total'): + """fetch either total or per-user usage stats""" payload = dict(self.payload) payload['pi_filter'] = f'"{account}"' - payload['group_by'] = 'pi' - payload['realm'] = 'Jobs' + payload['group_by'] = self.group_by[group_by] payload['statistic'] = statistics - - core_hours = self.fetch_data(account, payload) + payload['realm'] = realm + core_hours = self.fetch_data(payload, search_item=account) return core_hours - def xdmod_fetch_total_storage(self, account, statistics='physical_usage'): - """fetch total storage.""" + def xdmod_fetch_all_project_usages(self, statistic): + """return usage statistics for all projects""" payload = dict(self.payload) - payload['pi_filter'] = f'"{account}"' payload['group_by'] = 'pi' - payload['realm'] = 'Storage' - payload['statistic'] = statistics + payload['realm'] = 'Jobs' + payload['statistic'] = statistic + stats = self.fetch_data(payload) + return stats + + def xdmod_fetch_cpu_hours(self, account, group_by='total', statistics='total_cpu_hours'): + """fetch either total or per-user cpu hours""" + core_hours = self.xdmod_fetch(account, statistics, 'Jobs', group_by=group_by) + return core_hours - stats = self.fetch_data(account, payload) + def xdmod_fetch_storage(self, account, group_by='total', statistics='physical_usage'): + """fetch total or per-user storage stats.""" + stats = self.xdmod_fetch(account, statistics, 'Storage', group_by=group_by) physical_usage = float(stats) / 1E9 return physical_usage @@ -133,5 +141,5 @@ def xdmod_fetch_cloud_core_time(self, project): payload['realm'] = 'Cloud' payload['statistic'] = 'cloud_core_time' - core_hours = self.fetch_data(project, payload) + core_hours = self.fetch_data(payload, search_item=project) return core_hours