From fefffb4195cb321724c7c383397293a447f66265 Mon Sep 17 00:00:00 2001 From: Claire Peters Date: Tue, 7 Nov 2023 16:52:36 -0800 Subject: [PATCH] refactor xdmod code to improve readability --- .../commands/add_allocation_defaults.py | 2 +- .../allocation/allocation_detail.html | 24 +- coldfront/core/allocation/views.py | 3 +- .../slurm/management/commands/slurm_check.py | 6 +- .../xdmod/management/commands/xdmod_usage.py | 303 ++++++------------ coldfront/plugins/xdmod/utils.py | 223 +++++-------- 6 files changed, 197 insertions(+), 364 deletions(-) diff --git a/coldfront/core/allocation/management/commands/add_allocation_defaults.py b/coldfront/core/allocation/management/commands/add_allocation_defaults.py index c295b8280..116111ec6 100644 --- a/coldfront/core/allocation/management/commands/add_allocation_defaults.py +++ b/coldfront/core/allocation/management/commands/add_allocation_defaults.py @@ -48,7 +48,7 @@ def handle(self, *args, **options): # UBCCR defaults ('Cloud Account Name', 'Text', False, False), # ('CLOUD_USAGE_NOTIFICATION', 'Yes/No', False, True), - ('Core Usage (Hours)', 'Int', True, False), + ('Core Usage (Hours)', 'Float', True, False), # ('Accelerator Usage (Hours)', 'Int', True, False), # ('Cloud Storage Quota (TB)', 'Float', True, False), # ('EXPIRE NOTIFICATION', 'Yes/No', False, True), diff --git a/coldfront/core/allocation/templates/allocation/allocation_detail.html b/coldfront/core/allocation/templates/allocation/allocation_detail.html index 379eb68ec..a1709100a 100644 --- a/coldfront/core/allocation/templates/allocation/allocation_detail.html +++ b/coldfront/core/allocation/templates/allocation/allocation_detail.html @@ -143,17 +143,19 @@

Allocation Information

Last Synced {{user_sync_dt}} - - Quota (TB): - {{ allocation.size|floatformat:2 }} - - - Total Amount Due: - {% cost_tb allocation.size as cost %} - {% if cost %} - {{ cost }} - {% endif %} - + {% if "Storage" in allocation.get_parent_resource.resource_type.name %} + + Quota (TB): + {{ allocation.size|floatformat:2 }} + + + Total Amount Due: + {% cost_tb allocation.size as cost %} + {% if cost %} + {{ cost }} + {% endif %} + + {% endif %} Start Date: diff --git a/coldfront/core/allocation/views.py b/coldfront/core/allocation/views.py index bb3fdb0c2..5222d2701 100644 --- a/coldfront/core/allocation/views.py +++ b/coldfront/core/allocation/views.py @@ -429,7 +429,8 @@ def get_queryset(self): # Allocation Attribute Name if data.get('allocation_attribute_name') and data.get('allocation_attribute_value'): allocations = allocations.filter( - Q(allocationattribute__allocation_attribute_type=data.get('allocation_attribute_name')) & + Q(allocationattribute__allocation_attribute_type=data.get( + 'allocation_attribute_name')) & Q(allocationattribute__value=data.get('allocation_attribute_value')) ) diff --git a/coldfront/plugins/slurm/management/commands/slurm_check.py b/coldfront/plugins/slurm/management/commands/slurm_check.py index 9512b743a..caf89f581 100644 --- a/coldfront/plugins/slurm/management/commands/slurm_check.py +++ b/coldfront/plugins/slurm/management/commands/slurm_check.py @@ -38,9 +38,11 @@ def add_arguments(self, parser): parser.add_argument("-c", "--cluster", help="Run sacctmgr dump [cluster] as input") parser.add_argument("-s", "--sync", - help="Remove associations in Slurm that no longer exist in ColdFront", action="store_true") + help="Remove associations in Slurm that no longer exist in ColdFront", + action="store_true") parser.add_argument("-n", "--noop", - help="Print commands only. Do not run any commands.", action="store_true") + help="Print commands only. Do not run any commands.", + action="store_true") parser.add_argument("-u", "--username", help="Check specific username") parser.add_argument("-a", "--account", help="Check specific account") parser.add_argument("-x", "--header", diff --git a/coldfront/plugins/xdmod/management/commands/xdmod_usage.py b/coldfront/plugins/xdmod/management/commands/xdmod_usage.py index f29d401fd..fd1948554 100644 --- a/coldfront/plugins/xdmod/management/commands/xdmod_usage.py +++ b/coldfront/plugins/xdmod/management/commands/xdmod_usage.py @@ -15,8 +15,7 @@ XDMOD_STORAGE_ATTRIBUTE_NAME, XDMOD_STORAGE_GROUP_ATTRIBUTE_NAME, XdmodNotFoundError, - xdmod_fetch_cloud_core_time, - xdmod_fetch_total_cpu_hours, xdmod_fetch_total_storage) + XDModFetcher) logger = logging.getLogger(__name__) @@ -58,6 +57,62 @@ def write(self, data): os.dup2(devnull, sys.stdout.fileno()) sys.exit(1) + def id_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 + + 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): + attr = s.get_attribute(attr_name) + check_pass = attr + if num: + check_pass = attr is not None + if not check_pass: + logger.warning("%s attribute not found for allocation: %s", + attr_name, s) + return None + 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' + ) + if self.fetch_expired: + allocations = allocations.filter(~Q(status__name='Active')) + else: + allocations = allocations.filter(status__name='Active') + + if self.filter_user: + allocations = allocations.filter(project__pi__username=self.filter_user) + + if self.filter_project: + allocations = allocations.filter( + Q(allocationattribute__allocation_attribute_type__name=XDMOD_CLOUD_PROJECT_ATTRIBUTE_NAME) + & Q(allocationattribute__value=self.filter_project) + ) + + if self.filter_account: + allocations = allocations.filter( + Q(allocationattribute__allocation_attribute_type__name=account_attr_name) + & Q(allocationattribute__value=self.filter_account) + ) + + return allocations + def process_total_storage(self): header = [ 'allocation_id', @@ -71,81 +126,39 @@ def process_total_storage(self): if self.print_header: self.write('\t'.join(header)) - allocations = Allocation.objects.prefetch_related( - 'project', - 'resources', - 'allocationattribute_set', - 'allocationuser_set' - ).filter( + allocations = Allocation.objects.filter( allocationattribute__allocation_attribute_type__name__in=[ XDMOD_STORAGE_GROUP_ATTRIBUTE_NAME, XDMOD_STORAGE_ATTRIBUTE_NAME, ] ) - if self.fetch_expired: - allocations = allocations.filter( - ~Q(status__name='Active'), - ) - else: - allocations = allocations.filter( - status__name='Active', - ) - - if self.filter_user: - allocations = allocations.filter(project__pi__username=self.filter_user) - - if self.filter_account: - allocations = allocations.filter( - Q(allocationattribute__allocation_attribute_type__name=XDMOD_STORAGE_GROUP_ATTRIBUTE_NAME) - & Q(allocationattribute__value=self.filter_account) - ) + allocations = self.filter_allocations(allocations, account_attr_name=XDMOD_STORAGE_GROUP_ATTRIBUTE_NAME) for s in allocations.distinct(): - account_name = s.get_attribute(XDMOD_STORAGE_GROUP_ATTRIBUTE_NAME) - if not account_name: - logger.warning("%s attribute not found for allocation: %s", - XDMOD_STORAGE_GROUP_ATTRIBUTE_NAME, s) - continue - - cpu_hours = s.get_attribute(XDMOD_STORAGE_ATTRIBUTE_NAME) - if cpu_hours is None: - logger.warning("%s attribute not found for allocation: %s", - XDMOD_STORAGE_ATTRIBUTE_NAME, s) + account_name = self.attribute_check(s, XDMOD_STORAGE_GROUP_ATTRIBUTE_NAME) + cpu_hours = self.attribute_check(s, XDMOD_STORAGE_ATTRIBUTE_NAME, num=True) + if None in [account_name, cpu_hours]: continue - 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 - - resources.append(rname) - + resources = self.id_resources(s) if len(resources) == 0: - logger.warning("%s attribute not found on any resouces for allocation: %s", - XDMOD_RESOURCE_ATTRIBUTE_NAME, s) continue + fetcher = XDModFetcher(s.start_date, s.end_date, resources=resources) try: - usage = xdmod_fetch_total_storage( - s.start_date, s.end_date, account_name, - resources=resources, statistics='avg_physical_usage') + usage = fetcher.xdmod_fetch_total_storage( + account_name, statistics='avg_physical_usage' + ) except XdmodNotFoundError: logger.warning( "No data in XDMoD found for allocation %s account %s resources %s", s, account_name, resources) continue - logger.warning("Total GB = %s for allocation %s account %s GB %s resources %s", - usage, s, account_name, cpu_hours, resources) + logger.warning( + "Total GB = %s for allocation %s account %s GB %s resources %s", + usage, s, account_name, cpu_hours, resources) if self.sync: s.set_usage(XDMOD_STORAGE_ATTRIBUTE_NAME, usage) @@ -172,73 +185,29 @@ def process_total_gpu_hours(self): if self.print_header: self.write('\t'.join(header)) - allocations = Allocation.objects.prefetch_related( - 'project', - 'resources', - 'allocationattribute_set', - 'allocationuser_set' - ).filter( + allocations = Allocation.objects.filter( allocationattribute__allocation_attribute_type__name__in=[ XDMOD_ACCOUNT_ATTRIBUTE_NAME, XDMOD_ACC_HOURS_ATTRIBUTE_NAME, ] ) - - if self.fetch_expired: - allocations = allocations.filter( - ~Q(status__name='Active'), - ) - else: - allocations = allocations.filter( - status__name='Active', - ) - - if self.filter_user: - allocations = allocations.filter(project__pi__username=self.filter_user) - - if self.filter_account: - allocations = allocations.filter( - Q(allocationattribute__allocation_attribute_type__name=XDMOD_ACCOUNT_ATTRIBUTE_NAME) - & Q(allocationattribute__value=self.filter_account) - ) + allocations = self.filter_allocations(allocations) for s in allocations.distinct(): - account_name = s.get_attribute(XDMOD_ACCOUNT_ATTRIBUTE_NAME) - if not account_name: - logger.warning("%s attribute not found for allocation: %s", - XDMOD_ACCOUNT_ATTRIBUTE_NAME, s) - continue - - cpu_hours = s.get_attribute(XDMOD_ACC_HOURS_ATTRIBUTE_NAME) - if cpu_hours is None: - logger.warning("%s attribute not found for allocation: %s", - XDMOD_ACC_HOURS_ATTRIBUTE_NAME, s) + account_name = self.attribute_check(s, XDMOD_ACCOUNT_ATTRIBUTE_NAME) + cpu_hours = self.attribute_check(s, XDMOD_ACC_HOURS_ATTRIBUTE_NAME, num=True) + if None in [account_name, cpu_hours]: continue - 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 - - resources.append(rname) - + resources = self.id_resources(s) if len(resources) == 0: - logger.warning("%s attribute not found on any resouces for allocation: %s", - XDMOD_RESOURCE_ATTRIBUTE_NAME, s) continue + fetcher = XDModFetcher(s.start_date, s.end_date, resources=resources) try: - usage = xdmod_fetch_total_cpu_hours( - s.start_date, s.end_date, account_name, - resources=resources, statistics='total_gpu_hours') + usage = fetcher.xdmod_fetch_total_cpu_hours( + account_name, statistics='total_gpu_hours' + ) except XdmodNotFoundError: logger.warning( "No data in XDMoD found for allocation %s account %s resources %s", @@ -273,65 +242,27 @@ def process_total_cpu_hours(self): if self.print_header: self.write('\t'.join(header)) - allocations = Allocation.objects.prefetch_related( - 'project', - 'resources', - 'allocationattribute_set', - 'allocationuser_set' - ).filter( - status__name='Active', - ).filter( + allocations = Allocation.objects.filter( allocationattribute__allocation_attribute_type__name__in=[ XDMOD_ACCOUNT_ATTRIBUTE_NAME, XDMOD_CPU_HOURS_ATTRIBUTE_NAME, ] ) - - if self.filter_user: - allocations = allocations.filter(project__pi__username=self.filter_user) - - if self.filter_account: - allocations = allocations.filter( - Q(allocationattribute__allocation_attribute_type__name=XDMOD_ACCOUNT_ATTRIBUTE_NAME) - & Q(allocationattribute__value=self.filter_account) - ) + allocations = self.filter_allocations(allocations) for s in allocations.distinct(): - account_name = s.get_attribute(XDMOD_ACCOUNT_ATTRIBUTE_NAME) - if not account_name: - logger.warning("%s attribute not found for allocation: %s", - XDMOD_ACCOUNT_ATTRIBUTE_NAME, s) - continue - - cpu_hours = s.get_attribute(XDMOD_CPU_HOURS_ATTRIBUTE_NAME) - if cpu_hours is None: - logger.warning("%s attribute not found for allocation: %s", - XDMOD_CPU_HOURS_ATTRIBUTE_NAME, s) + 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 = [] - 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 - - resources.append(rname) - + resources = self.id_resources(s) if len(resources) == 0: - logger.warning("%s attribute not found on any resouces for allocation: %s", - XDMOD_RESOURCE_ATTRIBUTE_NAME, s) continue + fetcher = XDModFetcher(s.start_date, s.end_date, resources=resources) try: - usage = xdmod_fetch_total_cpu_hours( - s.start_date, s.end_date, account_name, resources=resources) + usage = fetcher.xdmod_fetch_total_cpu_hours(account_name) except XdmodNotFoundError: logger.warning( "No data in XDMoD found for allocation %s account %s resources %s", @@ -370,65 +301,28 @@ def process_cloud_core_time(self): if self.print_header: self.write('\t'.join(header)) - allocations = Allocation.objects.prefetch_related( - 'project', - 'resources', - 'allocationattribute_set', - 'allocationuser_set' - ).filter( - status__name='Active', - ).filter( + allocations = Allocation.objects.filter( allocationattribute__allocation_attribute_type__name__in=[ XDMOD_CLOUD_PROJECT_ATTRIBUTE_NAME, XDMOD_CLOUD_CORE_TIME_ATTRIBUTE_NAME, ] ) - - if self.filter_user: - allocations = allocations.filter(project__pi__username=self.filter_user) - - if self.filter_project: - allocations = allocations.filter( - Q(allocationattribute__allocation_attribute_type__name=XDMOD_CLOUD_PROJECT_ATTRIBUTE_NAME) - & Q(allocationattribute__value=self.filter_project) - ) + allocations = self.filter_allocations(allocations) for s in allocations.distinct(): - project_name = s.get_attribute(XDMOD_CLOUD_PROJECT_ATTRIBUTE_NAME) - if not project_name: - logger.warning("%s attribute not found for allocation: %s", - XDMOD_CLOUD_PROJECT_ATTRIBUTE_NAME, s) - continue - - core_time = s.get_attribute(XDMOD_CLOUD_CORE_TIME_ATTRIBUTE_NAME) - if not core_time: - logger.warning("%s attribute not found for allocation: %s", - XDMOD_CLOUD_CORE_TIME_ATTRIBUTE_NAME, s) + project_name = self.attribute_check( + s, XDMOD_CLOUD_PROJECT_ATTRIBUTE_NAME) + core_time = self.attribute_check( + s, XDMOD_CLOUD_CORE_TIME_ATTRIBUTE_NAME, num=True) + if None in [project_name, core_time]: continue - - 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 - - resources.append(rname) - + resources = self.id_resources(s) if len(resources) == 0: - logger.warning("%s attribute not found on any resouces for allocation: %s", - XDMOD_RESOURCE_ATTRIBUTE_NAME, s) continue + fetcher = XDModFetcher(s.start_date, s.end_date, resources=resources) try: - usage = xdmod_fetch_cloud_core_time( - s.start_date, s.end_date, project_name, resources=resources) + usage = fetcher.xdmod_fetch_cloud_core_time(project_name) except XdmodNotFoundError: logger.warning( "No data in XDMoD found for allocation %s project %s resources %s", @@ -465,7 +359,6 @@ def handle(self, *args, **options): self.sync = True logger.warning("Syncing ColdFront with XDMoD") - filters = { 'username': self.filter_user, 'account': self.filter_account, diff --git a/coldfront/plugins/xdmod/utils.py b/coldfront/plugins/xdmod/utils.py index 433890614..3a2612487 100644 --- a/coldfront/plugins/xdmod/utils.py +++ b/coldfront/plugins/xdmod/utils.py @@ -52,147 +52,82 @@ class XdmodError(Exception): class XdmodNotFoundError(XdmodError): pass -def xdmod_fetch_total_cpu_hours(start, end, account, resources=None, statistics='total_cpu_hours'): - if resources is None: - resources = [] - - url = f'{XDMOD_API_URL}{_ENDPOINT_CORE_HOURS}' - payload = _DEFAULT_PARAMS - payload['pi_filter'] = f'"{account}"' - payload['resource_filter'] = f'"{",".join(resources)}"' - payload['start_date'] = start - payload['end_date'] = end - payload['group_by'] = 'pi' - payload['realm'] = 'Jobs' - payload['operation'] = 'get_data' - payload['statistic'] = statistics - r = requests.get(url, params=payload, auth=HTTPBasicAuth(XDMOD_USER, XDMOD_PASS)) - - logger.info(r.url) - logger.info(r.text) - - try: - error = r.json() - # XXX fix me. Here we assume any json response is bad as we're - # expecting xml but XDMoD should just return json always. - raise XdmodNotFoundError(f'Got json response but expected XML: {error}') - except json.decoder.JSONDecodeError as e: - pass - except requests.exceptions.JSONDecodeError: - pass - - try: - root = ET.fromstring(r.text) - except ET.ParserError as e: - raise XdmodError(f'Invalid XML data returned from XDMoD API: {e}') - - rows = root.find('rows') - if len(rows) != 1: - raise XdmodNotFoundError(f'Rows not found for {account} - {resources}') - - cells = rows.find('row').findall('cell') - if len(cells) != 2: - raise XdmodError('Invalid XML data returned from XDMoD API: Cells not found') - - core_hours = cells[1].find('value').text - - return core_hours - - - -def xdmod_fetch_total_storage(start, end, account, resources=None, statistics='physical_usage'): - if resources is None: - resources = [] - - payload_end = end - if payload_end is None: - payload_end = '2099-01-01' - url = f'{XDMOD_API_URL}{_ENDPOINT_CORE_HOURS}' - payload = _DEFAULT_PARAMS - payload['pi_filter'] = f'"{account}"' - payload['resource_filter'] = f'"{",".join(resources)}"' - payload['start_date'] = start - payload['end_date'] = payload_end - payload['group_by'] = 'pi' - payload['realm'] = 'Storage' - payload['operation'] = 'get_data' - payload['statistic'] = statistics - r = requests.get(url, params=payload, auth=HTTPBasicAuth(XDMOD_USER, XDMOD_PASS)) - - logger.info(r.url) - logger.info(r.text) - if is_json(r.content): - error = r.json() - # XXX fix me. Here we assume any json response is bad as we're - # expecting xml but XDMoD should just return json always. - - # print('XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX\n') - # print(f'XDMOD synchronization error: ({start}, {end}, {account}, {resources}, response: {r})') - # print(r.content) - # print(r.url) - # print(payload) - # print('XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX\n\n') - - raise XdmodNotFoundError(f'Got json response but expected XML: {error}') - - - try: - root = ET.fromstring(r.text) - except ET.ParserError as e: - raise XdmodError(f'Invalid XML data returned from XDMoD API: {e}') - - rows = root.find('rows') - if len(rows) != 1: - raise XdmodNotFoundError(f'Rows not found for {account} - {resources}') - - cells = rows.find('row').findall('cell') - if len(cells) != 2: - raise XdmodError('Invalid XML data returned from XDMoD API: Cells not found') - - physical_usage = float(cells[1].find('value').text) / 1E9 - - return physical_usage - -def xdmod_fetch_cloud_core_time(start, end, project, resources=None): - if resources is None: - resources = [] - - url = f'{XDMOD_API_URL}{_ENDPOINT_CORE_HOURS}' - payload = _DEFAULT_PARAMS - payload['project_filter'] = project - payload['resource_filter'] = f'"{",".join(resources)}"' - payload['start_date'] = start - payload['end_date'] = end - payload['group_by'] = 'project' - payload['realm'] = 'Cloud' - payload['operation'] = 'get_data' - payload['statistic'] = 'cloud_core_time' - r = requests.get(url, params=payload, auth=HTTPBasicAuth(XDMOD_USER, XDMOD_PASS)) - - logger.info(r.url) - logger.info(r.text) - - try: - error = r.json() - # XXXX fix me. Here we assume any json response is bad as we're - # expecting xml but XDMoD should just return json always. - raise XdmodNotFoundError(f'Got json response but expected XML: {error}') - except json.decoder.JSONDecodeError as e: - pass - - try: - root = ET.fromstring(r.text) - except ET.ParserError as e: - raise XdmodError(f'Invalid XML data returned from XDMoD API: {e}') - - rows = root.find('rows') - if len(rows) != 1: - raise XdmodNotFoundError(f'Rows not found for {project} - {resources}') - - cells = rows.find('row').findall('cell') - if len(cells) != 2: - raise XdmodError('Invalid XML data returned from XDMoD API: Cells not found') - - core_hours = cells[1].find('value').text - - return core_hours +class XDModFetcher: + + def __init__(self, start, 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 + payload['resource_filter'] = f'"{",".join(resources)}"' + payload['operation'] = 'get_data' + self.payload = payload + + def fetch_data(self, search_item, payload): + r = requests.get( + self.url, params=payload, auth=HTTPBasicAuth(XDMOD_USER, XDMOD_PASS) + ) + logger.info(r.url) + logger.info(r.text) + + try: + error = r.json() + # XXXX fix me. Here we assume any json response is bad as we're + # expecting xml but XDMoD should just return json always. + raise XdmodNotFoundError(f'Got json response but expected XML: {error}') + except json.decoder.JSONDecodeError as e: + pass + + try: + root = ET.fromstring(r.text) + except ET.ParserError as e: + raise XdmodError(f'Invalid XML data returned from XDMoD API: {e}') from e + + rows = root.find('rows') + if len(rows) != 1: + raise XdmodNotFoundError( + f'Rows not found for {search_item} - {self.payload["resources"]}' + ) + cells = rows.find('row').findall('cell') + if len(cells) != 2: + raise XdmodError('Invalid XML data returned from XDMoD API: Cells not found') + + stats = cells[1].find('value').text + return stats + + def xdmod_fetch_total_cpu_hours(self, account, statistics='total_cpu_hours'): + """fetch total cpu hours.""" + payload = dict(self.payload) + payload['pi_filter'] = f'"{account}"' + payload['group_by'] = 'pi' + payload['realm'] = 'Jobs' + payload['statistic'] = statistics + + core_hours = self.fetch_data(account, payload) + return core_hours + + def xdmod_fetch_total_storage(self, account, statistics='physical_usage'): + """fetch total storage.""" + payload = dict(self.payload) + payload['pi_filter'] = f'"{account}"' + payload['group_by'] = 'pi' + payload['realm'] = 'Storage' + payload['statistic'] = statistics + + stats = self.fetch_data(account, payload) + physical_usage = float(stats) / 1E9 + return physical_usage + + def xdmod_fetch_cloud_core_time(self, project): + """fetch cloud core time.""" + payload = dict(self.payload) + payload['project_filter'] = project + payload['group_by'] = 'project' + payload['realm'] = 'Cloud' + payload['statistic'] = 'cloud_core_time' + + core_hours = self.fetch_data(project, payload) + return core_hours