From 3b8927b8c3d46e2ec92548a1580d8eec9c899a2a Mon Sep 17 00:00:00 2001 From: Claire Peters Date: Tue, 19 Dec 2023 16:39:02 -0800 Subject: [PATCH 1/4] refactor xdmod plugin, add testing library --- .../xdmod/management/commands/xdmod_usage.py | 494 ++++++++---------- coldfront/plugins/xdmod/tests.py | 121 +++++ coldfront/plugins/xdmod/utils.py | 283 +++++----- 3 files changed, 467 insertions(+), 431 deletions(-) create mode 100644 coldfront/plugins/xdmod/tests.py diff --git a/coldfront/plugins/xdmod/management/commands/xdmod_usage.py b/coldfront/plugins/xdmod/management/commands/xdmod_usage.py index 37314829e..382dd643e 100644 --- a/coldfront/plugins/xdmod/management/commands/xdmod_usage.py +++ b/coldfront/plugins/xdmod/management/commands/xdmod_usage.py @@ -1,14 +1,13 @@ import logging import os import sys -import tempfile -from django.conf import settings -from django.core.management.base import BaseCommand, CommandError +from django.core.management.base import BaseCommand from django.db.models import Q +from django.contrib.auth import get_user_model -from coldfront.core.allocation.models import Allocation -from coldfront.core.utils.common import import_from_settings +from coldfront.core.allocation.models import Allocation, AllocationUserStatusChoice +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, @@ -17,15 +16,21 @@ XDMOD_RESOURCE_ATTRIBUTE_NAME, XDMOD_STORAGE_ATTRIBUTE_NAME, XDMOD_STORAGE_GROUP_ATTRIBUTE_NAME, - XdmodError, XdmodNotFoundError, - xdmod_fetch_cloud_core_time, - xdmod_fetch_total_cpu_hours, xdmod_fetch_total_storage) + XdmodNotFoundError, + XDModFetcher) logger = logging.getLogger(__name__) class Command(BaseCommand): help = 'Sync usage data from XDMoD to ColdFront' + filter_user = '' + filter_project = '' + filter_resource = '' + filter_account = '' + sync = False + print_header = False + fetch_expired = False def add_arguments(self, parser): parser.add_argument("-r", "--resource", @@ -36,14 +41,15 @@ def add_arguments(self, parser): help="Report usage for specific username") parser.add_argument("-p", "--project", help="Report usage for specific cloud project") - parser.add_argument( - "-s", "--sync", help="Update allocation attributes with latest data from XDMoD", action="store_true") - parser.add_argument( - "-x", "--header", help="Include header in output", action="store_true") - parser.add_argument( - "-m", "--statistic", help="XDMoD statistic (default total_cpu_hours)", required=True) - parser.add_argument( - "--expired", help="XDMoD statistic for archived projects", action="store_true") + parser.add_argument("-s", "--sync", + help="Update allocation attributes with latest data from XDMoD", action="store_true") + parser.add_argument("-x", "--header", + help="Include header in output", action="store_true") + parser.add_argument("-m", "--statistic", + help="XDMoD statistic (default total_cpu_hours)", + default='total_cpu_hours') + parser.add_argument("--expired", + help="XDMoD statistic for archived projects", action="store_true") def write(self, data): try: @@ -53,6 +59,68 @@ def write(self, data): os.dup2(devnull, sys.stdout.fileno()) sys.exit(1) + 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 = self.run_resource_checks(r) + resources.append(rname) + 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): + 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: + 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', @@ -66,79 +134,38 @@ 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.warn("%s attribute not found for allocation: %s", - XDMOD_STORAGE_GROUP_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 - cpu_hours = s.get_attribute(XDMOD_STORAGE_ATTRIBUTE_NAME) - if not cpu_hours: - logger.warn("%s attribute not found for allocation: %s", - XDMOD_STORAGE_ATTRIBUTE_NAME, s) - 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) - - if len(resources) == 0: - logger.warn("%s attribute not found on any resouces for allocation: %s", - XDMOD_RESOURCE_ATTRIBUTE_NAME, s) - continue + resources = self.id_allocation_resources(s) + fetcher = XDModFetcher(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_storage( + account_name, statistics='avg_physical_usage' + ) except XdmodNotFoundError: - logger.warn( - "No data in XDMoD found for allocation %s account %s resources %s", s, account_name, resources) + logger.warning( + "No data in XDMoD found for allocation %s account %s resources %s", + s, account_name, resources) continue - logger.warn("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) @@ -152,7 +179,6 @@ def process_total_storage(self): ])) - def process_total_gpu_hours(self): header = [ 'allocation_id', @@ -166,78 +192,35 @@ 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.warn("%s attribute not found for allocation: %s", - XDMOD_ACCOUNT_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 - cpu_hours = s.get_attribute(XDMOD_ACC_HOURS_ATTRIBUTE_NAME) - if not cpu_hours: - logger.warn("%s attribute not found for allocation: %s", - XDMOD_ACC_HOURS_ATTRIBUTE_NAME, s) - 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) - - if len(resources) == 0: - logger.warn("%s attribute not found on any resouces for allocation: %s", - XDMOD_RESOURCE_ATTRIBUTE_NAME, s) - continue + resources = self.id_allocation_resources(s) + fetcher = XDModFetcher(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_cpu_hours( + account_name, statistics='total_gpu_hours' + ) except XdmodNotFoundError: - logger.warn( - "No data in XDMoD found for allocation %s account %s resources %s", s, account_name, resources) + logger.warning( + "No data in XDMoD found for allocation %s account %s resources %s", + s, account_name, resources) continue - logger.warn("Total Accelerator hours = %s for allocation %s account %s gpu_hours %s resources %s", + logger.warning( + "Total Accelerator hours = %s for allocation %s account %s gpu_hours %s resources %s", usage, s, account_name, cpu_hours, resources) if self.sync: s.set_usage(XDMOD_ACC_HOURS_ATTRIBUTE_NAME, usage) @@ -264,73 +247,84 @@ 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) + + + # # bulk collection + # 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 = s.get_attribute(XDMOD_ACCOUNT_ATTRIBUTE_NAME) - if not account_name: - logger.warn("%s attribute not found for allocation: %s", - XDMOD_ACCOUNT_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 - cpu_hours = s.get_attribute(XDMOD_CPU_HOURS_ATTRIBUTE_NAME) - if not cpu_hours: - logger.warn("%s attribute not found for allocation: %s", - XDMOD_CPU_HOURS_ATTRIBUTE_NAME, s) - 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) - - if len(resources) == 0: - logger.warn("%s attribute not found on any resouces for allocation: %s", - XDMOD_RESOURCE_ATTRIBUTE_NAME, s) - continue + resources = self.id_allocation_resources(s) + fetcher = XDModFetcher(resources=resources) try: - usage = xdmod_fetch_total_cpu_hours( - s.start_date, s.end_date, account_name, resources=resources) - except XdmodNotFoundError: - logger.warn( - "No data in XDMoD found for allocation %s account %s resources %s", s, account_name, resources) + usage = fetcher.xdmod_fetch_cpu_hours(account_name) + except XdmodNotFoundError as e: + logger.warning( + "No data in XDMoD found for allocation %s account %s resources %s: %s", + s, account_name, resources, e) continue - logger.warn("Total CPU hours = %s for allocation %s account %s cpu_hours %s resources %s", + logger.warning( + "Total CPU hours = %s for allocation %s account %s cpu_hours %s resources %s", usage, s, account_name, cpu_hours, resources) + # collect user-level usage and update allocationuser entries with them + auser_status_active = AllocationUserStatusChoice.objects.get(name='Active') + + usage_data = fetcher.xdmod_fetch_cpu_hours(account_name, group_by='per-user') + no_use_allocation_users = s.allocationuser_set.filter( + ~Q(user__username__in=usage_data.keys()) + ) + + for user in no_use_allocation_users: + user.usage = 0 + user.save() + for username, user_usage in usage_data.items(): + try: + user_obj = get_user_model().objects.get(username=username) + except: + # if user not present, add to ifx + logger.warning("user missing from ifx: %s", username) + continue + user, created = s.allocationuser_set.get_or_create( + user=user_obj, + defaults={ + 'usage': user_usage, 'unit': 'CPU Hours', + 'status': auser_status_active + } + ) + if not created: + user.usage = user_usage + user.save() if self.sync: + cpu_hours_attr = s.allocationattribute_set.get( + allocation_attribute_type__name=XDMOD_CPU_HOURS_ATTRIBUTE_NAME) + cpu_hours_attr.value = usage + cpu_hours_attr.save() s.set_usage(XDMOD_CPU_HOURS_ATTRIBUTE_NAME, usage) self.write('\t'.join([ @@ -355,72 +349,35 @@ 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.warn("%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.warn("%s attribute not found for allocation: %s", - XDMOD_CLOUD_CORE_TIME_ATTRIBUTE_NAME, s) - 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) - - if len(resources) == 0: - logger.warn("%s attribute not found on any resouces for allocation: %s", - XDMOD_RESOURCE_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 = self.id_allocation_resources(s) + fetcher = XDModFetcher(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.warn( - "No data in XDMoD found for allocation %s project %s resources %s", s, project_name, resources) + logger.warning( + "No data in XDMoD found for allocation %s project %s resources %s", + s, project_name, resources) continue - logger.warn("Cloud core time = %s for allocation %s project %s core_time %s resources %s", - usage, s, project_name, core_time, resources) + logger.warning( + "Cloud core time = %s for allocation %s project %s core_time %s resources %s", + usage, s, project_name, core_time, resources) if self.sync: s.set_usage(XDMOD_CLOUD_CORE_TIME_ATTRIBUTE_NAME, usage) @@ -434,50 +391,31 @@ def process_cloud_core_time(self): ])) def handle(self, *args, **options): - verbosity = int(options['verbosity']) - root_logger = logging.getLogger('') - if verbosity == 0: - root_logger.setLevel(logging.ERROR) - elif verbosity == 2: - root_logger.setLevel(logging.INFO) - elif verbosity == 3: - root_logger.setLevel(logging.DEBUG) - else: - root_logger.setLevel(logging.WARN) - self.sync = False if options['sync']: self.sync = True - logger.warn("Syncing ColdFront with XDMoD") + logger.warning("Syncing ColdFront with XDMoD") + + filters = { + 'username': self.filter_user, + 'account': self.filter_account, + 'project': self.filter_project, + 'resource': self.filter_resource, + } + for filter_name, filter_attr in filters.items(): + if options[filter_name]: + logger.info("Filtering output by %s: %s", filter_name, options[filter_name]) + filter_attr = options[filter_name] + + bool_opts = { + 'header':self.print_header, + 'expired':self.fetch_expired, + } + for opt, attribute in bool_opts.items(): + if options[opt]: + attribute = True statistic = 'total_cpu_hours' - self.filter_user = '' - self.filter_project = '' - self.filter_resource = '' - self.filter_account = '' - self.print_header = False - self.fetch_expired = False - - if options['username']: - logger.info("Filtering output by username: %s", - options['username']) - self.filter_user = options['username'] - if options['account']: - logger.info("Filtering output by account: %s", options['account']) - self.filter_account = options['account'] - if options['project']: - logger.info("Filtering output by project: %s", options['project']) - self.filter_project = options['project'] - if options['resource']: - logger.info("Filtering output by resource: %s", - options['resource']) - self.filter_resource = options['resource'] - if options['header']: - self.print_header = True - - if options['expired']: - self.fetch_expired = True - if options['statistic']: statistic = options['statistic'] @@ -488,7 +426,7 @@ def handle(self, *args, **options): elif statistic == 'total_acc_hours': self.process_total_gpu_hours() elif statistic == 'total_storage': - self.process_total_storage() + self.process_total_storage() else: logger.error("Unsupported XDMoD statistic") sys.exit(1) diff --git a/coldfront/plugins/xdmod/tests.py b/coldfront/plugins/xdmod/tests.py new file mode 100644 index 000000000..d81b7faaa --- /dev/null +++ b/coldfront/plugins/xdmod/tests.py @@ -0,0 +1,121 @@ +"""Testing suite for the xdmod plugin.""" +import json +import xml.etree.ElementTree as ET +from unittest import mock +from unittest.mock import patch, Mock + +from django.test import TestCase + +from coldfront.plugins.xdmod.utils import XDModFetcher, XdmodNotFoundError, XdmodError + +XDMOD_XML_SIMPLE = """ +
+CPU Hours: Total: by PI + + +PI + smith_lab + + +2023-09-01 +2023-12-01 + +PI +CPU Hours: Total + +
+ + + +smith_lab + + +720.1442 + + + +""" + +def read_in_xml(xml): + """Read in the xml.""" + root = ET.fromstring(xml) + return root.find('rows') + + +class XDModFetcherTestCase(TestCase): + """Tests for the xdmod plugin.""" + + def setUp(self): + """Set up the test case.""" + self.xdmodfetcher = XDModFetcher() + # if the return value is valid xml, return the xml + mock_response = Mock(text=XDMOD_XML_SIMPLE) + mock_response.json.side_effect = json.decoder.JSONDecodeError('JSON decode error', '', 0) + self.correct_return_value = mock_response + + @mock.patch("coldfront.plugins.xdmod.utils.requests.get") + def test_fetch_data(self, mock_get): + """Test that the fetcher works.""" + # if the return value is not xml, raise an exception + mock_get.return_value = Mock(text="Hello") + with self.assertRaises(XdmodError): + self.xdmodfetcher.fetch_data(None) + # if the return value is empty xml, raise an exception + mock_get.return_value = Mock(text="") + with self.assertRaises(XdmodNotFoundError): + self.xdmodfetcher.fetch_data(None) + mock_get.return_value = self.correct_return_value + actual_result = ET.tostring(self.xdmodfetcher.fetch_data(None)).decode() + self.assertEqual(actual_result, ET.tostring(read_in_xml(XDMOD_XML_SIMPLE)).decode()) + + @mock.patch("coldfront.plugins.xdmod.utils.requests.get") + def test_fetch_value(self, mock_get): + """Test that fetch_value works.""" + mock_get.return_value = self.correct_return_value + self.assertEqual(self.xdmodfetcher.fetch_value(None), '720.1442') + + @mock.patch("coldfront.plugins.xdmod.utils.requests.get") + def test_fetch_value_error(self, mock_get): + """Test that fetch_value raises an exception if the xml is invalid.""" + mock_get.return_value = Mock(status_code=200, text="no_second_value") + with self.assertRaises(XdmodError): + self.xdmodfetcher.fetch_value(None) + + @mock.patch("coldfront.plugins.xdmod.utils.requests.get") + def test_fetch_table(self, mock_get): + """Test that fetch_table works.""" + mock_get.return_value = self.correct_return_value + self.assertEqual(self.xdmodfetcher.fetch_table(None), {"smith_lab": '720.1442'}) + + @mock.patch("coldfront.plugins.xdmod.utils.requests.get") + def test_fetch_table_error(self, mock_get): + """Test that fetch_table raises an exception if the xml is invalid.""" + mock_get.return_value = Mock(status_code=200, text="no_second_value") + with self.assertRaises(XdmodError): + self.xdmodfetcher.fetch_table(None) + + @mock.patch("coldfront.plugins.xdmod.utils.requests.get") + def test_xdmod_fetch(self, mock_get): + """Test that xdmod_fetch works.""" + mock_get.return_value = self.correct_return_value + self.assertEqual(self.xdmodfetcher.xdmod_fetch("account", "statistic", "realm"), '720.1442') + + @patch("coldfront.plugins.xdmod.utils.requests.get") + def test_xdmod_fetch_all_project_usages(self, mock_get): + mock_get.return_value = self.correct_return_value + self.assertEqual(self.xdmodfetcher.xdmod_fetch_all_project_usages("statistic"), {"smith_lab": '720.1442'}) + + @patch("coldfront.plugins.xdmod.utils.requests.get") + def test_xdmod_fetch_cpu_hours(self, mock_get): + mock_get.return_value = self.correct_return_value + self.assertEqual(self.xdmodfetcher.xdmod_fetch_cpu_hours("account", group_by='total', statistics='total_cpu_hours'), '720.1442') + + @patch("coldfront.plugins.xdmod.utils.requests.get") + def test_xdmod_fetch_storage(self, mock_get): + mock_get.return_value = self.correct_return_value + self.assertEqual(self.xdmodfetcher.xdmod_fetch_storage("account", group_by='total', statistic='physical_usage'), 720.1442 / 1E9) + + @patch("coldfront.plugins.xdmod.utils.requests.get") + def test_xdmod_fetch_cloud_core_time(self, mock_get): + mock_get.return_value = self.correct_return_value + self.assertEqual(self.xdmodfetcher.xdmod_fetch_cloud_core_time("project"), '720.1442') diff --git a/coldfront/plugins/xdmod/utils.py b/coldfront/plugins/xdmod/utils.py index 85f0c338e..28770c8ed 100644 --- a/coldfront/plugins/xdmod/utils.py +++ b/coldfront/plugins/xdmod/utils.py @@ -1,22 +1,32 @@ import logging -import requests import json import xml.etree.ElementTree as ET +from datetime import datetime + +import requests from coldfront.core.utils.common import import_from_settings -XDMOD_CLOUD_PROJECT_ATTRIBUTE_NAME = import_from_settings('XDMOD_CLOUD_PROJECT_ATTRIBUTE_NAME', 'Cloud Account Name') -XDMOD_CLOUD_CORE_TIME_ATTRIBUTE_NAME = import_from_settings('XDMOD_CLOUD_CORE_TIME_ATTRIBUTE_NAME', 'Core Usage (Hours)') +XDMOD_CLOUD_PROJECT_ATTRIBUTE_NAME = import_from_settings( + 'XDMOD_CLOUD_PROJECT_ATTRIBUTE_NAME', 'Cloud Account Name') +XDMOD_CLOUD_CORE_TIME_ATTRIBUTE_NAME = import_from_settings( + 'XDMOD_CLOUD_CORE_TIME_ATTRIBUTE_NAME', 'Core Usage (Hours)') -XDMOD_ACCOUNT_ATTRIBUTE_NAME = import_from_settings('XDMOD_ACCOUNT_ATTRIBUTE_NAME', 'slurm_account_name') +XDMOD_ACCOUNT_ATTRIBUTE_NAME = import_from_settings( + 'XDMOD_ACCOUNT_ATTRIBUTE_NAME', 'slurm_account_name') -XDMOD_RESOURCE_ATTRIBUTE_NAME = import_from_settings('XDMOD_RESOURCE_ATTRIBUTE_NAME', 'xdmod_resource') +XDMOD_RESOURCE_ATTRIBUTE_NAME = import_from_settings( + 'XDMOD_RESOURCE_ATTRIBUTE_NAME', 'xdmod_resource') -XDMOD_CPU_HOURS_ATTRIBUTE_NAME = import_from_settings('XDMOD_CPU_HOURS_ATTRIBUTE_NAME', 'Core Usage (Hours)') -XDMOD_ACC_HOURS_ATTRIBUTE_NAME = import_from_settings('XDMOD_ACC_HOURS_ATTRIBUTE_NAME', 'Accelerator Usage (Hours)') +XDMOD_CPU_HOURS_ATTRIBUTE_NAME = import_from_settings( + 'XDMOD_CPU_HOURS_ATTRIBUTE_NAME', 'Core Usage (Hours)') +XDMOD_ACC_HOURS_ATTRIBUTE_NAME = import_from_settings( + 'XDMOD_ACC_HOURS_ATTRIBUTE_NAME', 'Accelerator Usage (Hours)') -XDMOD_STORAGE_ATTRIBUTE_NAME = import_from_settings('XDMOD_STORAGE_ATTRIBUTE_NAME', 'Storage Quota (GB)') -XDMOD_STORAGE_GROUP_ATTRIBUTE_NAME = import_from_settings('XDMOD_STORAGE_GROUP_ATTRIBUTE_NAME', 'Storage_Group_Name') +XDMOD_STORAGE_ATTRIBUTE_NAME = import_from_settings( + 'XDMOD_STORAGE_ATTRIBUTE_NAME', 'Storage Quota (GB)') +XDMOD_STORAGE_GROUP_ATTRIBUTE_NAME = import_from_settings( + 'XDMOD_STORAGE_GROUP_ATTRIBUTE_NAME', 'Storage_Group_Name') XDMOD_API_URL = import_from_settings('XDMOD_API_URL') @@ -31,6 +41,15 @@ 'query_group': 'tg_usage', } +def get_quarter_start_end(): + y = datetime.today().year + quarter_starts = [f'{y}-01-01', f'{y}-04-01', f'{y}-07-01', f'{y}-10-01'] + quarter_ends = [f'{y}-03-31', f'{y}-06-30', f'{y}-09-30', f'{y}-12-31'] + quarter = (datetime.today().month-1)//3 + return (quarter_starts[quarter], quarter_ends[quarter]) + +QUARTER_START, QUARTER_END = get_quarter_start_end() + logger = logging.getLogger(__name__) class XdmodError(Exception): @@ -39,147 +58,105 @@ 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 = '{}{}'.format(XDMOD_API_URL, _ENDPOINT_CORE_HOURS) - payload = _DEFAULT_PARAMS - payload['pi_filter'] = '"{}"'.format(account) - payload['resource_filter'] = '"{}"'.format(','.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) - - 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('Got json response but expected XML: {}'.format(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('Invalid XML data returned from XDMoD API: {}'.format(e)) - - rows = root.find('rows') - if len(rows) != 1: - raise XdmodNotFoundError('Rows not found for {} - {}'.format(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 = '{}{}'.format(XDMOD_API_URL, _ENDPOINT_CORE_HOURS) - payload = _DEFAULT_PARAMS - payload['pi_filter'] = '"{}"'.format(account) - payload['resource_filter'] = '{}'.format(','.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) - - 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('XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX\n') - # print(f'XDMOD synchronization error: ({start}, {end}, {account}, {resources}, response: {r})') - # print(r.content) - # print(r.url) - # print(payload) - # print('XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX\n\n') - - raise XdmodNotFoundError('Got json response but expected XML: {}'.format(error)) - - - try: - root = ET.fromstring(r.text) - except ET.ParserError as e: - raise XdmodError('Invalid XML data returned from XDMoD API: {}'.format(e)) - - rows = root.find('rows') - if len(rows) != 1: - raise XdmodNotFoundError('Rows not found for {} - {}'.format(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 = '{}{}'.format(XDMOD_API_URL, _ENDPOINT_CORE_HOURS) - payload = _DEFAULT_PARAMS - payload['project_filter'] = project - payload['resource_filter'] = '"{}"'.format(','.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) - - 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('Got json response but expected XML: {}'.format(error)) - except json.decoder.JSONDecodeError as e: - pass - - try: - root = ET.fromstring(r.text) - except ET.ParserError as e: - raise XdmodError('Invalid XML data returned from XDMoD API: {}'.format(e)) - - rows = root.find('rows') - if len(rows) != 1: - raise XdmodNotFoundError('Rows not found for {} - {}'.format(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=QUARTER_START, end=QUARTER_END, resources=None): + self.url = f'{XDMOD_API_URL}{_ENDPOINT_CORE_HOURS}' + self.resources = resources + payload = _DEFAULT_PARAMS + payload['start_date'] = start + payload['end_date'] = end + if resources: + payload['resource_filter'] = f'"{",".join(resources)}"' + self.payload = payload + self.group_by = {'total':'pi', 'per-user':'person'} + + def fetch_data(self, payload, search_item=None): + r = requests.get(self.url, params=payload) + logger.info(r.url) + logger.info(r.text) + + try: + error = r.json() + 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["resource_filter"]}' + ) + return rows + + def fetch_value(self, payload, search_item=None): + rows = self.fetch_data(payload, search_item=search_item) + 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 fetch_table(self, payload, search_item=None): + """make a dictionary of usernames and their associated core hours from + XML data. + """ + # return rows extracted from XML data + rows = self.fetch_data(payload, search_item=search_item) + # Produce a dict of usernames and their associated core hours from those rows + stats = {} + for row in rows: + cells = row.findall('cell') + username = cells[0].find('value').text + stats[username] = cells[1].find('value').text + return stats + + def xdmod_fetch(self, account, statistic, realm, group_by='total'): + """fetch either total or per-user usage stats for specified project""" + payload = dict(self.payload) + payload['pi_filter'] = f'"{account}"' + payload['group_by'] = self.group_by[group_by] + payload['statistic'] = statistic + payload['realm'] = realm + if group_by == 'total': + core_hours = self.fetch_value(payload, search_item=account) + elif group_by == 'per-user': + core_hours = self.fetch_table(payload, search_item=account) + else: + raise Exception('unrecognized group_by value') + return core_hours + + def xdmod_fetch_all_project_usages(self, statistic): + """return usage statistics for all projects""" + payload = dict(self.payload) + payload['group_by'] = 'pi' + payload['realm'] = 'Jobs' + payload['statistic'] = statistic + stats = self.fetch_table(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 + + def xdmod_fetch_storage(self, account, group_by='total', statistic='physical_usage'): + """fetch total or per-user storage stats.""" + stats = self.xdmod_fetch(account, statistic, 'Storage', group_by=group_by) + 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_value(payload, search_item=project) + return core_hours From d8a5c7981b05f34f666b50671902a457d7eb7c45 Mon Sep 17 00:00:00 2001 From: Claire Peters Date: Tue, 19 Dec 2023 17:06:05 -0800 Subject: [PATCH 2/4] add start_date and end_date back to xdmod_usage --- .../xdmod/management/commands/xdmod_usage.py | 9 ++++++--- coldfront/plugins/xdmod/utils.py | 20 +++++++++++++------ 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/coldfront/plugins/xdmod/management/commands/xdmod_usage.py b/coldfront/plugins/xdmod/management/commands/xdmod_usage.py index 382dd643e..9cf2f96c5 100644 --- a/coldfront/plugins/xdmod/management/commands/xdmod_usage.py +++ b/coldfront/plugins/xdmod/management/commands/xdmod_usage.py @@ -155,7 +155,8 @@ def process_total_storage(self): fetcher = XDModFetcher(resources=resources) try: usage = fetcher.xdmod_fetch_storage( - account_name, statistics='avg_physical_usage' + account_name, start_date=s.start_date, + end_date=s.end_date, statistic='avg_physical_usage' ) except XdmodNotFoundError: logger.warning( @@ -211,7 +212,8 @@ def process_total_gpu_hours(self): fetcher = XDModFetcher(resources=resources) try: usage = fetcher.xdmod_fetch_cpu_hours( - account_name, statistics='total_gpu_hours' + account_name, start_date=s.start_date, end_date=s.end_date, + statistics='total_gpu_hours' ) except XdmodNotFoundError: logger.warning( @@ -368,7 +370,8 @@ def process_cloud_core_time(self): fetcher = XDModFetcher(resources=resources) try: - usage = fetcher.xdmod_fetch_cloud_core_time(project_name) + usage = fetcher.xdmod_fetch_cloud_core_time( + project_name, start_date=s.start_date, end_date=s.end_date) except XdmodNotFoundError: logger.warning( "No data in XDMoD found for allocation %s project %s resources %s", diff --git a/coldfront/plugins/xdmod/utils.py b/coldfront/plugins/xdmod/utils.py index 28770c8ed..1bc86964d 100644 --- a/coldfront/plugins/xdmod/utils.py +++ b/coldfront/plugins/xdmod/utils.py @@ -115,13 +115,17 @@ def fetch_table(self, payload, search_item=None): stats[username] = cells[1].find('value').text return stats - def xdmod_fetch(self, account, statistic, realm, group_by='total'): + def xdmod_fetch(self, account, statistic, realm, start_date=None, end_date=None, group_by='total'): """fetch either total or per-user usage stats for specified project""" payload = dict(self.payload) payload['pi_filter'] = f'"{account}"' payload['group_by'] = self.group_by[group_by] payload['statistic'] = statistic payload['realm'] = realm + if start_date: + payload['start_date'] = start_date + if end_date: + payload['end_date'] = end_date if group_by == 'total': core_hours = self.fetch_value(payload, search_item=account) elif group_by == 'per-user': @@ -139,24 +143,28 @@ def xdmod_fetch_all_project_usages(self, statistic): stats = self.fetch_table(payload) return stats - def xdmod_fetch_cpu_hours(self, account, group_by='total', statistics='total_cpu_hours'): + def xdmod_fetch_cpu_hours(self, account, start_date=None, end_date=None, 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) + core_hours = self.xdmod_fetch(account, statistics, 'Jobs', start_date=start_date, end_date=end_date, group_by=group_by) return core_hours - def xdmod_fetch_storage(self, account, group_by='total', statistic='physical_usage'): + def xdmod_fetch_storage(self, account, start_date=None, end_date=None, group_by='total', statistic='physical_usage'): """fetch total or per-user storage stats.""" - stats = self.xdmod_fetch(account, statistic, 'Storage', group_by=group_by) + stats = self.xdmod_fetch(account, statistic, 'Storage', start_date=start_date, end_date=end_date, group_by=group_by) physical_usage = float(stats) / 1E9 return physical_usage - def xdmod_fetch_cloud_core_time(self, project): + def xdmod_fetch_cloud_core_time(self, project, start_date=None, end_date=None): """fetch cloud core time.""" payload = dict(self.payload) payload['project_filter'] = project payload['group_by'] = 'project' payload['realm'] = 'Cloud' payload['statistic'] = 'cloud_core_time' + if start_date: + payload['start_date'] = start_date + if end_date: + payload['end_date'] = end_date core_hours = self.fetch_value(payload, search_item=project) return core_hours From 09aa2247f80bc86be6841d6b9afb8321ecfb6aab Mon Sep 17 00:00:00 2001 From: Claire Peters Date: Tue, 19 Dec 2023 17:08:51 -0800 Subject: [PATCH 3/4] statistic, not statistics --- coldfront/plugins/xdmod/management/commands/xdmod_usage.py | 2 +- coldfront/plugins/xdmod/utils.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/coldfront/plugins/xdmod/management/commands/xdmod_usage.py b/coldfront/plugins/xdmod/management/commands/xdmod_usage.py index 9cf2f96c5..c932a75b0 100644 --- a/coldfront/plugins/xdmod/management/commands/xdmod_usage.py +++ b/coldfront/plugins/xdmod/management/commands/xdmod_usage.py @@ -213,7 +213,7 @@ def process_total_gpu_hours(self): try: usage = fetcher.xdmod_fetch_cpu_hours( account_name, start_date=s.start_date, end_date=s.end_date, - statistics='total_gpu_hours' + statistic='total_gpu_hours' ) except XdmodNotFoundError: logger.warning( diff --git a/coldfront/plugins/xdmod/utils.py b/coldfront/plugins/xdmod/utils.py index 1bc86964d..926434a31 100644 --- a/coldfront/plugins/xdmod/utils.py +++ b/coldfront/plugins/xdmod/utils.py @@ -143,9 +143,9 @@ def xdmod_fetch_all_project_usages(self, statistic): stats = self.fetch_table(payload) return stats - def xdmod_fetch_cpu_hours(self, account, start_date=None, end_date=None, group_by='total', statistics='total_cpu_hours'): + def xdmod_fetch_cpu_hours(self, account, start_date=None, end_date=None, group_by='total', statistic='total_cpu_hours'): """fetch either total or per-user cpu hours""" - core_hours = self.xdmod_fetch(account, statistics, 'Jobs', start_date=start_date, end_date=end_date, group_by=group_by) + core_hours = self.xdmod_fetch(account, statistic, 'Jobs', start_date=start_date, end_date=end_date, group_by=group_by) return core_hours def xdmod_fetch_storage(self, account, start_date=None, end_date=None, group_by='total', statistic='physical_usage'): From 710d9fee2c52b4038d24bba143d2a0de03de4d43 Mon Sep 17 00:00:00 2001 From: Claire Peters Date: Tue, 19 Dec 2023 17:09:46 -0800 Subject: [PATCH 4/4] remove commented out rewrite start --- .../xdmod/management/commands/xdmod_usage.py | 17 ----------------- 1 file changed, 17 deletions(-) diff --git a/coldfront/plugins/xdmod/management/commands/xdmod_usage.py b/coldfront/plugins/xdmod/management/commands/xdmod_usage.py index c932a75b0..554f72cd7 100644 --- a/coldfront/plugins/xdmod/management/commands/xdmod_usage.py +++ b/coldfront/plugins/xdmod/management/commands/xdmod_usage.py @@ -257,23 +257,6 @@ def process_total_cpu_hours(self): ) allocations = self.filter_allocations(allocations) - - # # bulk collection - # 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)