Skip to content

Commit

Permalink
more refactoring
Browse files Browse the repository at this point in the history
  • Loading branch information
@claire-peters
claire-peters committed Nov 11, 2023
1 parent 0e2f8c0 commit 4ab0641
Show file tree
Hide file tree
Showing 2 changed files with 78 additions and 54 deletions.
82 changes: 49 additions & 33 deletions coldfront/plugins/xdmod/management/commands/xdmod_usage.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from django.db.models import Q

from coldfront.core.allocation.models import Allocation
from coldfront.core.resource.models import Resource
from coldfront.plugins.xdmod.utils import (XDMOD_ACCOUNT_ATTRIBUTE_NAME,
XDMOD_CLOUD_CORE_TIME_ATTRIBUTE_NAME,
XDMOD_CLOUD_PROJECT_ATTRIBUTE_NAME,
Expand Down Expand Up @@ -57,23 +58,20 @@ def write(self, data):
os.dup2(devnull, sys.stdout.fileno())
sys.exit(1)

def id_resources(self, s):
def run_resource_checks(self, resource):
rname = resource.get_attribute(XDMOD_RESOURCE_ATTRIBUTE_NAME)
if not rname and resource.parent_resource:
rname = resource.parent_resource.get_attribute(
XDMOD_RESOURCE_ATTRIBUTE_NAME)
if self.filter_resource and self.filter_resource != rname:
return None
return rname

def id_allocation_resources(self, s):
resources = []
for r in s.resources.all():
rname = r.get_attribute(XDMOD_RESOURCE_ATTRIBUTE_NAME)
if not rname and r.parent_resource:
rname = r.parent_resource.get_attribute(
XDMOD_RESOURCE_ATTRIBUTE_NAME)

if not rname:
continue
if self.filter_resource and self.filter_resource != rname:
continue

rname = self.run_resource_checks(r)
resources.append(rname)
if len(resources) == 0:
logger.warning("%s attribute not found on any resouces for allocation: %s",
XDMOD_RESOURCE_ATTRIBUTE_NAME, s)
return resources

def attribute_check(self, s, attr_name, num=False):
Expand All @@ -88,9 +86,18 @@ def attribute_check(self, s, attr_name, num=False):
return attr

def filter_allocations(self, allocations, account_attr_name=XDMOD_ACCOUNT_ATTRIBUTE_NAME):
allocations = allocations.select_related('project').prefetch_related(
'resources', 'allocationattribute_set', 'allocationuser_set'
)
cleared_resources = Resource.objects.filter(
Q(resourceattribute__resource_attribute_type__name=XDMOD_RESOURCE_ATTRIBUTE_NAME) |
Q(parent_resource__resourceattribute__resource_attribute_type__name=XDMOD_RESOURCE_ATTRIBUTE_NAME)
)

allocations = (
allocations.select_related('project')
.prefetch_related(
'resources', 'allocationattribute_set', 'allocationuser_set'
)
.filter(resources__in=cleared_resources)
)
if self.fetch_expired:
allocations = allocations.filter(~Q(status__name='Active'))
else:
Expand Down Expand Up @@ -133,21 +140,20 @@ def process_total_storage(self):
]
)

allocations = self.filter_allocations(allocations, account_attr_name=XDMOD_STORAGE_GROUP_ATTRIBUTE_NAME)
allocations = self.filter_allocations(
allocations, account_attr_name=XDMOD_STORAGE_GROUP_ATTRIBUTE_NAME)

for s in allocations.distinct():
account_name = self.attribute_check(s, XDMOD_STORAGE_GROUP_ATTRIBUTE_NAME)
cpu_hours = self.attribute_check(s, XDMOD_STORAGE_ATTRIBUTE_NAME, num=True)
if None in [account_name, cpu_hours]:
continue

resources = self.id_resources(s)
if len(resources) == 0:
continue
resources = self.id_allocation_resources(s)

fetcher = XDModFetcher(s.start_date, s.end_date, resources=resources)
try:
usage = fetcher.xdmod_fetch_total_storage(
usage = fetcher.xdmod_fetch_storage(
account_name, statistics='avg_physical_usage'
)
except XdmodNotFoundError:
Expand Down Expand Up @@ -199,13 +205,11 @@ def process_total_gpu_hours(self):
if None in [account_name, cpu_hours]:
continue

resources = self.id_resources(s)
if len(resources) == 0:
continue
resources = self.id_allocation_resources(s)

fetcher = XDModFetcher(s.start_date, s.end_date, resources=resources)
try:
usage = fetcher.xdmod_fetch_total_cpu_hours(
usage = fetcher.xdmod_fetch_cpu_hours(
account_name, statistics='total_gpu_hours'
)
except XdmodNotFoundError:
Expand Down Expand Up @@ -250,19 +254,33 @@ def process_total_cpu_hours(self):
)
allocations = self.filter_allocations(allocations)

resource_filter = (
Q(resourceattribute__resource_attribute_type__name=XDMOD_RESOURCE_ATTRIBUTE_NAME) |
Q(parent_resource__resourceattribute__resource_attribute_type__name=XDMOD_RESOURCE_ATTRIBUTE_NAME)
)
resources_all = Resource.objects.filter(
id__in=[r.id for a in allocations for r in a.resources.all()]).filter(
resource_filter
)

# fetcher = XDModFetcher(resources=resources_all)
# try:
# usage = fetcher.xdmod_fetch_all_project_usages('total_cpu_hours')
# except XdmodNotFoundError:
# raise XdmodNotFoundError(
# "No data in XDMoD found for resources %s", resources_all)

for s in allocations.distinct():
account_name = self.attribute_check(s, XDMOD_ACCOUNT_ATTRIBUTE_NAME)
cpu_hours = self.attribute_check(s, XDMOD_CPU_HOURS_ATTRIBUTE_NAME, num=True)
if None in [account_name, cpu_hours]:
continue

resources = self.id_resources(s)
if len(resources) == 0:
continue
resources = self.id_allocation_resources(s)

fetcher = XDModFetcher(s.start_date, s.end_date, resources=resources)
try:
usage = fetcher.xdmod_fetch_total_cpu_hours(account_name)
usage = fetcher.xdmod_fetch_cpu_hours(account_name)
except XdmodNotFoundError:
logger.warning(
"No data in XDMoD found for allocation %s account %s resources %s",
Expand Down Expand Up @@ -316,9 +334,7 @@ def process_cloud_core_time(self):
s, XDMOD_CLOUD_CORE_TIME_ATTRIBUTE_NAME, num=True)
if None in [project_name, core_time]:
continue
resources = self.id_resources(s)
if len(resources) == 0:
continue
resources = self.id_allocation_resources(s)

fetcher = XDModFetcher(s.start_date, s.end_date, resources=resources)
try:
Expand Down
50 changes: 29 additions & 21 deletions coldfront/plugins/xdmod/utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,29 +48,30 @@

logger = logging.getLogger(__name__)

QUARTER_START, QUARTER_END = get_quarter_start_end()


class XdmodError(Exception):
pass

class XdmodNotFoundError(XdmodError):
pass

class XDModFetcher:
def __init__(self, start, end, resources=None,):
def __init__(self, start=QUARTER_START, end=QUARTER_END, resources=None):
self.url = f'{XDMOD_API_URL}{_ENDPOINT_CORE_HOURS}'
if resources is None:
resources = []

payload = _DEFAULT_PARAMS
# payload['start_date'] = start
# payload['end_date'] = end
start_date, end_date = get_quarter_start_end()
payload['start_date'] = start_date
payload['end_date'] = end_date
payload['start_date'] = start
payload['end_date'] = end
payload['resource_filter'] = f'"{",".join(resources)}"'
payload['operation'] = 'get_data'
self.payload = payload
self.group_by = {'total':'pi', 'per-user':'user'}

def fetch_data(self, search_item, payload):
def fetch_data(self, payload, search_item=None):
r = requests.get(
self.url, params=payload, auth=HTTPBasicAuth(XDMOD_USER, XDMOD_PASS)
)
Expand All @@ -93,7 +94,7 @@ def fetch_data(self, search_item, payload):
rows = root.find('rows')
if len(rows) != 1:
raise XdmodNotFoundError(
f'Rows not found for {search_item} - {self.payload["resources"]}'
f'Rows not found for {search_item} - {self.payload["resource_filter"]}'
)
cells = rows.find('row').findall('cell')
if len(cells) != 2:
Expand All @@ -102,26 +103,33 @@ def fetch_data(self, search_item, payload):
stats = cells[1].find('value').text
return stats

def xdmod_fetch_total_cpu_hours(self, account, statistics='total_cpu_hours'):
"""fetch total cpu hours."""
def xdmod_fetch(self, account, statistics, realm, group_by='total'):
"""fetch either total or per-user usage stats"""
payload = dict(self.payload)
payload['pi_filter'] = f'"{account}"'
payload['group_by'] = 'pi'
payload['realm'] = 'Jobs'
payload['group_by'] = self.group_by[group_by]
payload['statistic'] = statistics

core_hours = self.fetch_data(account, payload)
payload['realm'] = realm
core_hours = self.fetch_data(payload, search_item=account)
return core_hours

def xdmod_fetch_total_storage(self, account, statistics='physical_usage'):
"""fetch total storage."""
def xdmod_fetch_all_project_usages(self, statistic):
"""return usage statistics for all projects"""
payload = dict(self.payload)
payload['pi_filter'] = f'"{account}"'
payload['group_by'] = 'pi'
payload['realm'] = 'Storage'
payload['statistic'] = statistics
payload['realm'] = 'Jobs'
payload['statistic'] = statistic
stats = self.fetch_data(payload)
return stats

def xdmod_fetch_cpu_hours(self, account, group_by='total', statistics='total_cpu_hours'):
"""fetch either total or per-user cpu hours"""
core_hours = self.xdmod_fetch(account, statistics, 'Jobs', group_by=group_by)
return core_hours

stats = self.fetch_data(account, payload)
def xdmod_fetch_storage(self, account, group_by='total', statistics='physical_usage'):
"""fetch total or per-user storage stats."""
stats = self.xdmod_fetch(account, statistics, 'Storage', group_by=group_by)
physical_usage = float(stats) / 1E9
return physical_usage

Expand All @@ -133,5 +141,5 @@ def xdmod_fetch_cloud_core_time(self, project):
payload['realm'] = 'Cloud'
payload['statistic'] = 'cloud_core_time'

core_hours = self.fetch_data(project, payload)
core_hours = self.fetch_data(payload, search_item=project)
return core_hours

0 comments on commit 4ab0641

Please sign in to comment.