diff --git a/coldfront/plugins/xdmod/management/commands/xdmod_usage.py b/coldfront/plugins/xdmod/management/commands/xdmod_usage.py
index 37314829e..554f72cd7 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,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.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, start_date=s.start_date,
+ end_date=s.end_date, statistic='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 +180,6 @@ def process_total_storage(self):
]))
-
def process_total_gpu_hours(self):
header = [
'allocation_id',
@@ -166,78 +193,36 @@ 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, start_date=s.start_date, end_date=s.end_date,
+ statistic='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 +249,67 @@ 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.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 +334,36 @@ 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, start_date=s.start_date, end_date=s.end_date)
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 +377,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 +412,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..926434a31 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,113 @@ 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, 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':
+ 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, 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, 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'):
+ """fetch total or per-user storage stats."""
+ 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, 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