diff --git a/src/coldfront_plugin_cloud/management/commands/validate_allocations.py b/src/coldfront_plugin_cloud/management/commands/validate_allocations.py index 01561bff..2c6b5388 100644 --- a/src/coldfront_plugin_cloud/management/commands/validate_allocations.py +++ b/src/coldfront_plugin_cloud/management/commands/validate_allocations.py @@ -52,6 +52,16 @@ def sync_users(project_id, allocation, allocator, apply): return failed_validation + @staticmethod + def set_default_quota_on_allocation(allocation, allocator, coldfront_attr): + uqm = tasks.UNIT_QUOTA_MULTIPLIERS[allocator.resource_type] + value = allocation.quantity * uqm.get(coldfront_attr, 0) + value += tasks.STATIC_QUOTA[allocator.resource_type].get(coldfront_attr, 0) + utils.set_attribute_on_allocation( + allocation, coldfront_attr, value + ) + return value + def check_institution_specific_code(self, allocation, apply): attr = attributes.ALLOCATION_INSTITUTION_SPECIFIC_CODE isc = allocation.get_attribute(attr) @@ -103,39 +113,38 @@ def handle(self, *args, **options): obj_key = openstack.OpenStackResourceAllocator.QUOTA_KEY_MAPPING['object']['keys'][attributes.QUOTA_OBJECT_GB] - for attr in attributes.ALLOCATION_QUOTA_ATTRIBUTES: - if 'OpenStack' in attr.name: - key = allocator.QUOTA_KEY_MAPPING_ALL_KEYS.get(attr.name, None) - if not key: - # Note(knikolla): Some attributes are only maintained - # for bookkeeping purposes and do not have a - # corresponding quota set on the service. - continue - - expected_value = allocation.get_attribute(attr.name) - current_value = quota.get(key, None) - if key == obj_key and expected_value <= 0: - expected_obj_value = 1 - current_value = int(allocator.object(project_id).head_account().get(obj_key)) - if current_value != expected_obj_value: - failed_validation = True - msg = (f'Value for quota for {attr.name} = {current_value} does not match expected' - f' value of {expected_obj_value} on allocation {allocation_str}') - logger.warning(msg) - elif expected_value is None and current_value: - msg = (f'Attribute "{attr.name}" expected on allocation {allocation_str} but not set.' - f' Current quota is {current_value}.') - if options['apply']: - utils.set_attribute_on_allocation( - allocation, attr.name, current_value - ) - msg = f'{msg} Attribute set to match current quota.' - logger.warning(msg) - elif not current_value == expected_value: + for attr in tasks.get_expected_attributes(allocator): + key = allocator.QUOTA_KEY_MAPPING_ALL_KEYS.get(attr, None) + if not key: + # Note(knikolla): Some attributes are only maintained + # for bookkeeping purposes and do not have a + # corresponding quota set on the service. + continue + + expected_value = allocation.get_attribute(attr) + current_value = quota.get(key, None) + if key == obj_key and expected_value <= 0: + expected_obj_value = 1 + current_value = int(allocator.object(project_id).head_account().get(obj_key)) + if current_value != expected_obj_value: failed_validation = True - msg = (f'Value for quota for {attr.name} = {current_value} does not match expected' - f' value of {expected_value} on allocation {allocation_str}') + msg = (f'Value for quota for {attr} = {current_value} does not match expected' + f' value of {expected_obj_value} on allocation {allocation_str}') logger.warning(msg) + elif expected_value is None and current_value: + msg = (f'Attribute "{attr}" expected on allocation {allocation_str} but not set.' + f' Current quota is {current_value}.') + if options['apply']: + utils.set_attribute_on_allocation( + allocation, attr, current_value + ) + msg = f'{msg} Attribute set to match current quota.' + logger.warning(msg) + elif not current_value == expected_value: + failed_validation = True + msg = (f'Value for quota for {attr} = {current_value} does not match expected' + f' value of {expected_value} on allocation {allocation_str}') + logger.warning(msg) if failed_validation and options['apply']: try: @@ -185,79 +194,92 @@ def handle(self, *args, **options): failed_validation = Command.sync_users(project_id, allocation, allocator, options["apply"]) - for attr in attributes.ALLOCATION_QUOTA_ATTRIBUTES: - if "OpenShift" in attr.name: - key_with_lambda = openshift.OpenShiftResourceAllocator.QUOTA_KEY_MAPPING.get(attr.name, None) - - # Openshift and Openshift VM allocations have different quota attributes - if not key_with_lambda: - continue - - # This gives me just the plain key - key = list(key_with_lambda(1).keys())[0] - - expected_value = allocation.get_attribute(attr.name) - current_value = quota.get(key, None) - - PATTERN = r"([0-9]+)(m|Ki|Mi|Gi|Ti|Pi|Ei|K|M|G|T|P|E)?" - - suffix = { - "Ki": 2**10, - "Mi": 2**20, - "Gi": 2**30, - "Ti": 2**40, - "Pi": 2**50, - "Ei": 2**60, - "m": 10**-3, - "K": 10**3, - "M": 10**6, - "G": 10**9, - "T": 10**12, - "P": 10**15, - "E": 10**18, - } - - if current_value and current_value != "0": - result = re.search(PATTERN, current_value) - - if result is None: - raise CommandError( - f"Unable to parse current_value = '{current_value}' for {attr.name}" - ) - - value = int(result.groups()[0]) - unit = result.groups()[1] - - # Convert to number i.e. without any unit suffix - - if unit is not None: - current_value = value * suffix[unit] - else: - current_value = value - - # Convert some attributes to units that coldfront uses - - if "RAM" in attr.name: - current_value = round(current_value / suffix["Mi"]) - elif "Storage" in attr.name: - current_value = round(current_value / suffix["Gi"]) - elif current_value and current_value == "0": - current_value = 0 - - if expected_value is None and current_value is not None: + for attr in tasks.get_expected_attributes(allocator): + key_with_lambda = allocator.QUOTA_KEY_MAPPING.get(attr, None) + + # This gives me just the plain key + key = list(key_with_lambda(1).keys())[0] + + expected_value = allocation.get_attribute(attr) + current_value = quota.get(key, None) + + PATTERN = r"([0-9]+)(m|Ki|Mi|Gi|Ti|Pi|Ei|K|M|G|T|P|E)?" + + suffix = { + "Ki": 2**10, + "Mi": 2**20, + "Gi": 2**30, + "Ti": 2**40, + "Pi": 2**50, + "Ei": 2**60, + "m": 10**-3, + "K": 10**3, + "M": 10**6, + "G": 10**9, + "T": 10**12, + "P": 10**15, + "E": 10**18, + } + + if current_value and current_value != "0": + result = re.search(PATTERN, current_value) + + if result is None: + raise CommandError( + f"Unable to parse current_value = '{current_value}' for {attr}" + ) + + value = int(result.groups()[0]) + unit = result.groups()[1] + + # Convert to number i.e. without any unit suffix + + if unit is not None: + current_value = value * suffix[unit] + else: + current_value = value + + # Convert some attributes to units that coldfront uses + + if "RAM" in attr: + current_value = round(current_value / suffix["Mi"]) + elif "Storage" in attr: + current_value = round(current_value / suffix["Gi"]) + elif current_value and current_value == "0": + current_value = 0 + + if expected_value is None and current_value is not None: + msg = ( + f'Attribute "{attr}" expected on allocation {allocation_str} but not set.' + f" Current quota is {current_value}." + ) + if options["apply"]: + utils.set_attribute_on_allocation( + allocation, attr, current_value + ) + msg = f"{msg} Attribute set to match current quota." + logger.warning(msg) + else: + # We just checked the case where the quota value is set in the cluster + # but not in coldfront. This is the only case the cluster value is the + # "source of truth" for the quota value + # If the coldfront value is set, it is always the source of truth. + # But first, we need to check if the quota value is set anywhere at all. + # TODO (Quan): Refactor these if statements so that we can remove this comment block + if current_value is None and expected_value is None: msg = ( - f'Attribute "{attr.name}" expected on allocation {allocation_str} but not set.' - f" Current quota is {current_value}." + f"Value for quota for {attr} is not set anywhere" + f" on allocation {allocation_str}" ) - if options["apply"]: - utils.set_attribute_on_allocation( - allocation, attr.name, current_value - ) - msg = f"{msg} Attribute set to match current quota." logger.warning(msg) - elif not (current_value == expected_value): + + if options["apply"]: + expected_value = self.set_default_quota_on_allocation(allocation, allocator, attr) + logger.warning(f"Added default quota for {attr} to allocation {allocation_str} to {expected_value}") + + if not (current_value == expected_value): msg = ( - f"Value for quota for {attr.name} = {current_value} does not match expected" + f"Value for quota for {attr} = {current_value} does not match expected" f" value of {expected_value} on allocation {allocation_str}" ) logger.warning(msg) diff --git a/src/coldfront_plugin_cloud/tasks.py b/src/coldfront_plugin_cloud/tasks.py index f0f42fde..4fc43c6c 100644 --- a/src/coldfront_plugin_cloud/tasks.py +++ b/src/coldfront_plugin_cloud/tasks.py @@ -76,6 +76,11 @@ }, } +def get_expected_attributes(allocator: base.ResourceAllocator): + """Based on the allocator's resource type, return the expected quotas attributes the allocation should have""" + resource_name = allocator.resource_type + return list(UNIT_QUOTA_MULTIPLIERS[resource_name].keys()) + def find_allocator(allocation) -> base.ResourceAllocator: allocators = { @@ -99,10 +104,10 @@ def set_quota_attributes(): allocation.quantity = 1 # Calculate the quota for the project, and set the attribute for each element - uqm = UNIT_QUOTA_MULTIPLIERS[allocator.resource_type] - for coldfront_attr in uqm.keys(): + expected_coldfront_attrs = get_expected_attributes(allocator) + for coldfront_attr in expected_coldfront_attrs: if not allocation.get_attribute(coldfront_attr): - value = allocation.quantity * uqm.get(coldfront_attr, 0) + value = allocation.quantity * UNIT_QUOTA_MULTIPLIERS[allocator.resource_type].get(coldfront_attr, 0) value += STATIC_QUOTA[allocator.resource_type].get(coldfront_attr, 0) utils.set_attribute_on_allocation(allocation, coldfront_attr, diff --git a/src/coldfront_plugin_cloud/tests/functional/openshift/test_allocation.py b/src/coldfront_plugin_cloud/tests/functional/openshift/test_allocation.py index 8e4c10bc..a25dc790 100644 --- a/src/coldfront_plugin_cloud/tests/functional/openshift/test_allocation.py +++ b/src/coldfront_plugin_cloud/tests/functional/openshift/test_allocation.py @@ -1,6 +1,7 @@ import os import time import unittest +from unittest import mock import uuid from coldfront_plugin_cloud import attributes, openshift, tasks, utils @@ -214,3 +215,47 @@ def test_reactivate_allocation(self): }) allocator._get_role(user.username, project_id) + + @mock.patch.object( + tasks, + "UNIT_QUOTA_MULTIPLIERS", + { + 'openshift': { + attributes.QUOTA_LIMITS_CPU: 1, + } + } + ) + def test_allocation_new_attribute(self): + """When a new attribute is introduced, but pre-existing allocations don't have it""" + user = self.new_user() + project = self.new_project(pi=user) + allocation = self.new_allocation(project, self.resource, 2) + allocator = openshift.OpenShiftResourceAllocator(self.resource, + allocation) + + tasks.activate_allocation(allocation.pk) + allocation.refresh_from_db() + + project_id = allocation.get_attribute(attributes.ALLOCATION_PROJECT_ID) + + self.assertEqual(allocation.get_attribute(attributes.QUOTA_LIMITS_CPU), 2 * 1) + + quota = allocator.get_quota(project_id) + self.assertEqual(quota, { + "limits.cpu": "2", + }) + + # Add a new attribute for Openshift + tasks.UNIT_QUOTA_MULTIPLIERS['openshift'][attributes.QUOTA_LIMITS_MEMORY] = 4096 + + call_command('validate_allocations', apply=True) + allocation.refresh_from_db() + + self.assertEqual(allocation.get_attribute(attributes.QUOTA_LIMITS_CPU), 2 * 1) + self.assertEqual(allocation.get_attribute(attributes.QUOTA_LIMITS_MEMORY), 2 * 4096) + + quota = allocator.get_quota(project_id) + self.assertEqual(quota, { + "limits.cpu": "2", + "limits.memory": "8Gi", + })