From a6a40dc97aea0ccf3a2379ef988f922218349ec0 Mon Sep 17 00:00:00 2001 From: Quan Pham Date: Thu, 15 Jan 2026 13:28:15 -0500 Subject: [PATCH] Allow dynamic quota creation and removal --- src/coldfront_plugin_cloud/attributes.py | 19 +- src/coldfront_plugin_cloud/base.py | 10 + .../commands/add_openshift_resource.py | 13 - .../commands/add_quota_to_resource.py | 118 ++++++ .../commands/calculate_storage_gb_hours.py | 31 +- .../commands/fetch_daily_billable_usage.py | 4 +- .../commands/remove_quota_from_resource.py | 69 ++++ .../commands/validate_allocations.py | 17 +- .../models/quota_models.py | 68 ++++ .../{ => models}/usage_models.py | 0 src/coldfront_plugin_cloud/openshift.py | 20 +- src/coldfront_plugin_cloud/openshift_vm.py | 26 +- src/coldfront_plugin_cloud/openstack.py | 62 ++-- src/coldfront_plugin_cloud/tasks.py | 85 +---- src/coldfront_plugin_cloud/tests/base.py | 4 +- .../functional/openshift/test_allocation.py | 350 ++++++++++-------- .../openshift_vm/test_allocation.py | 117 +++++- .../functional/openstack/test_allocation.py | 6 +- .../tests/unit/openshift/base.py | 20 +- .../tests/unit/openshift/test_rbac.py | 13 +- .../unit/test_calculate_quota_unit_hours.py | 4 +- .../unit/test_fetch_daily_billable_usage.py | 2 +- .../tests/unit/test_usage_models.py | 2 +- 23 files changed, 658 insertions(+), 402 deletions(-) create mode 100644 src/coldfront_plugin_cloud/management/commands/add_quota_to_resource.py create mode 100644 src/coldfront_plugin_cloud/management/commands/remove_quota_from_resource.py create mode 100644 src/coldfront_plugin_cloud/models/quota_models.py rename src/coldfront_plugin_cloud/{ => models}/usage_models.py (100%) diff --git a/src/coldfront_plugin_cloud/attributes.py b/src/coldfront_plugin_cloud/attributes.py index 697d1920..f4590c02 100644 --- a/src/coldfront_plugin_cloud/attributes.py +++ b/src/coldfront_plugin_cloud/attributes.py @@ -24,7 +24,7 @@ class CloudAllocationAttribute: RESOURCE_API_URL = "OpenShift API Endpoint URL" RESOURCE_IDENTITY_NAME = "OpenShift Identity Provider Name" RESOURCE_ROLE = "Role for User in Project" -RESOURCE_IBM_AVAILABLE = "IBM Spectrum Scale Storage Available" +RESOURCE_QUOTA_RESOURCES = "Available Quota Resources" RESOURCE_FEDERATION_PROTOCOL = "OpenStack Federation Protocol" RESOURCE_IDP = "OpenStack Identity Provider" @@ -44,7 +44,7 @@ class CloudAllocationAttribute: CloudResourceAttribute(name=RESOURCE_IDP), CloudResourceAttribute(name=RESOURCE_PROJECT_DOMAIN), CloudResourceAttribute(name=RESOURCE_ROLE), - CloudResourceAttribute(name=RESOURCE_IBM_AVAILABLE), + CloudResourceAttribute(name=RESOURCE_QUOTA_RESOURCES), CloudResourceAttribute(name=RESOURCE_USER_DOMAIN), CloudResourceAttribute(name=RESOURCE_EULA_URL), CloudResourceAttribute(name=RESOURCE_DEFAULT_PUBLIC_NETWORK), @@ -103,17 +103,9 @@ class CloudAllocationAttribute: ########################################################### # OpenShift Quota Attributes -QUOTA_LIMITS_CPU = "OpenShift Limit on CPU Quota" -QUOTA_LIMITS_MEMORY = "OpenShift Limit on RAM Quota (MiB)" QUOTA_LIMITS_EPHEMERAL_STORAGE_GB = "OpenShift Limit on Ephemeral Storage Quota (GiB)" QUOTA_REQUESTS_NESE_STORAGE = "OpenShift Request on NESE Storage Quota (GiB)" QUOTA_REQUESTS_IBM_STORAGE = "OpenShift Request on IBM Storage Quota (GiB)" -QUOTA_REQUESTS_GPU = "OpenShift Request on GPU Quota" -QUOTA_REQUESTS_VM_GPU_A100_SXM4 = "OpenShift Request on GPU A100 SXM4" -QUOTA_REQUESTS_VM_GPU_V100 = "OpenShift Request on GPU V100" -QUOTA_REQUESTS_VM_GPU_H100 = "OpenShift Request on GPU H100" -QUOTA_PVC = "OpenShift Persistent Volume Claims Quota" - ALLOCATION_QUOTA_ATTRIBUTES = [ CloudAllocationAttribute(name=QUOTA_INSTANCES), @@ -125,14 +117,7 @@ class CloudAllocationAttribute: CloudAllocationAttribute(name=QUOTA_FLOATING_IPS), CloudAllocationAttribute(name=QUOTA_OBJECT_GB), CloudAllocationAttribute(name=QUOTA_GPU), - CloudAllocationAttribute(name=QUOTA_LIMITS_CPU), - CloudAllocationAttribute(name=QUOTA_LIMITS_MEMORY), CloudAllocationAttribute(name=QUOTA_LIMITS_EPHEMERAL_STORAGE_GB), CloudAllocationAttribute(name=QUOTA_REQUESTS_NESE_STORAGE), CloudAllocationAttribute(name=QUOTA_REQUESTS_IBM_STORAGE), - CloudAllocationAttribute(name=QUOTA_REQUESTS_GPU), - CloudAllocationAttribute(name=QUOTA_REQUESTS_VM_GPU_A100_SXM4), - CloudAllocationAttribute(name=QUOTA_REQUESTS_VM_GPU_V100), - CloudAllocationAttribute(name=QUOTA_REQUESTS_VM_GPU_H100), - CloudAllocationAttribute(name=QUOTA_PVC), ] diff --git a/src/coldfront_plugin_cloud/base.py b/src/coldfront_plugin_cloud/base.py index e5960d81..8ef37dcf 100644 --- a/src/coldfront_plugin_cloud/base.py +++ b/src/coldfront_plugin_cloud/base.py @@ -1,11 +1,13 @@ import abc import functools +import json from typing import NamedTuple from coldfront.core.allocation import models as allocation_models from coldfront.core.resource import models as resource_models from coldfront_plugin_cloud import attributes +from coldfront_plugin_cloud.models.quota_models import QuotaSpecs class ResourceAllocator(abc.ABC): @@ -25,6 +27,14 @@ def __init__( self.resource = resource self.allocation = allocation + resource_storage_classes_attr = resource_models.ResourceAttribute.objects.get( + resource=resource, + resource_attribute_type__name=attributes.RESOURCE_QUOTA_RESOURCES, + ) + self.resource_quotaspecs = QuotaSpecs.model_validate( + json.loads(resource_storage_classes_attr.value) + ) + def get_or_create_federated_user(self, username): if not (user := self.get_federated_user(username)): user = self.create_federated_user(username) diff --git a/src/coldfront_plugin_cloud/management/commands/add_openshift_resource.py b/src/coldfront_plugin_cloud/management/commands/add_openshift_resource.py index 128c00fc..042b856e 100644 --- a/src/coldfront_plugin_cloud/management/commands/add_openshift_resource.py +++ b/src/coldfront_plugin_cloud/management/commands/add_openshift_resource.py @@ -50,11 +50,6 @@ def add_arguments(self, parser): action="store_true", help="Indicates this is an OpenShift Virtualization resource (default: False)", ) - parser.add_argument( - "--ibm-storage-available", - action="store_true", - help="Indicates that Ibm Scale storage is available in this resource (default: False)", - ) def handle(self, *args, **options): self.validate_role(options["role"]) @@ -97,14 +92,6 @@ def handle(self, *args, **options): resource=openshift, value=options["role"], ) - - ResourceAttribute.objects.get_or_create( - resource_attribute_type=ResourceAttributeType.objects.get( - name=attributes.RESOURCE_IBM_AVAILABLE - ), - resource=openshift, - value="true" if options["ibm_storage_available"] else "false", - ) ResourceAttribute.objects.get_or_create( resource_attribute_type=ResourceAttributeType.objects.get( name=attributes.RESOURCE_CLUSTER_NAME diff --git a/src/coldfront_plugin_cloud/management/commands/add_quota_to_resource.py b/src/coldfront_plugin_cloud/management/commands/add_quota_to_resource.py new file mode 100644 index 00000000..0b657b3d --- /dev/null +++ b/src/coldfront_plugin_cloud/management/commands/add_quota_to_resource.py @@ -0,0 +1,118 @@ +import json +import logging + +from django.core.management.base import BaseCommand +from coldfront.core.resource.models import ( + Resource, + ResourceAttribute, + ResourceAttributeType, +) +from coldfront.core.allocation.models import AllocationAttributeType, AttributeType + +from coldfront_plugin_cloud import attributes +from coldfront_plugin_cloud.models.quota_models import QuotaSpecs, QuotaSpec + +logger = logging.getLogger(__name__) + + +class Command(BaseCommand): + def add_arguments(self, parser): + parser.add_argument( + "--display_name", + type=str, + required=True, + help="The display name for the quota attribute to add to the resource type.", + ) + parser.add_argument( + "--default-quota", + type=int, + required=True, + help="The default quota value for the storage attribute. In GB", + ) + parser.add_argument( + "--resource_name", + type=str, + required=True, + help="The name of the resource to add the storage attribute to.", + ) + parser.add_argument( + "--quota-label", + dest="quota_label", + type=str, + required=True, + help="Human-readable quota_label for this quota (must be unique).", + ) + parser.add_argument( + "--multiplier", + dest="multiplier", + type=int, + default=0, + help="Multiplier applied per SU quantity (int).", + ) + parser.add_argument( + "--static-quota", + dest="static_quota", + type=int, + default=0, + help="Static quota added to every SU quantity (int).", + ) + parser.add_argument( + "--unit-suffix", + dest="unit_suffix", + type=str, + default="", + help='Unit suffix to append to formatted quota values (e.g. "Gi").', + ) + parser.add_argument( + "--is-storage-type", + action="store_true", + help="Indicates if this quota is for a storage type for billing purposes", + ) + parser.add_argument( + "--invoice-name", + type=str, + default="", + help="Name of quota as it appears on invoice. Required if --is-storage-type is set.", + ) + + def handle(self, *args, **options): + if options["is_storage_type"] and not options["invoice_name"]: + logger.error( + "--invoice-name must be provided when --is-storage-type is set." + ) + + resource_name = options["resource_name"] + display_name = options["display_name"] + new_quota_spec = QuotaSpec(**options) + new_quota_dict = {display_name: new_quota_spec.model_dump()} + QuotaSpecs.model_validate(new_quota_dict) + + resource = Resource.objects.get(name=resource_name) + available_quotas_attr, created = ResourceAttribute.objects.get_or_create( + resource=resource, + resource_attribute_type=ResourceAttributeType.objects.get( + name=attributes.RESOURCE_QUOTA_RESOURCES + ), + defaults={"value": json.dumps(new_quota_dict)}, + ) + + # TODO (Quan): Dict update allows migration of existing quotas. This is fine? + if not created: + available_quotas_dict = json.loads(available_quotas_attr.value) + available_quotas_dict.update(new_quota_dict) + QuotaSpecs.model_validate(available_quotas_dict) # Validate uniqueness + available_quotas_attr.value = json.dumps(available_quotas_dict) + available_quotas_attr.save() + + # Now create Allocation Attribute for this quota + AllocationAttributeType.objects.get_or_create( + name=display_name, + defaults={ + "attribute_type": AttributeType.objects.get(name="Int"), + "has_usage": False, + "is_private": False, + "is_changeable": True, + }, + ) + + logger.info("Added quota '%s' to resource '%s'.", display_name, resource_name) diff --git a/src/coldfront_plugin_cloud/management/commands/calculate_storage_gb_hours.py b/src/coldfront_plugin_cloud/management/commands/calculate_storage_gb_hours.py index f2f75fea..749f2685 100644 --- a/src/coldfront_plugin_cloud/management/commands/calculate_storage_gb_hours.py +++ b/src/coldfront_plugin_cloud/management/commands/calculate_storage_gb_hours.py @@ -1,4 +1,5 @@ import csv +import json from decimal import Decimal, ROUND_HALF_UP import dataclasses from datetime import datetime, timedelta, timezone @@ -7,6 +8,7 @@ from coldfront_plugin_cloud import attributes from coldfront_plugin_cloud import utils +from coldfront_plugin_cloud.models.quota_models import QuotaSpecs import boto3 from django.core.management.base import BaseCommand @@ -20,6 +22,10 @@ _RATES = None +QUOTA_LIMITS_EPHEMERAL_STORAGE_GB = "OpenShift Limit on Ephemeral Storage Quota (GiB)" +QUOTA_REQUESTS_NESE_STORAGE = "OpenShift Request on NESE Storage Quota (GiB)" +QUOTA_REQUESTS_IBM_STORAGE = "OpenShift Request on IBM Storage Quota (GiB)" + def get_rates(): # nerc-rates doesn't work with Python 3.9, which is what ColdFront is currently @@ -210,6 +216,16 @@ def upload_to_s3(s3_endpoint, s3_bucket, file_location, invoice_month, end_time) def handle(self, *args, **options): generated_at = datetime.now(tz=timezone.utc).isoformat(timespec="seconds") + def get_storage_quotaspecs(allocation: Allocation): + """Get storage-related quota attributes for an allocation.""" + quotaspecs_dict = json.loads( + allocation.resources.first().get_attribute( + attributes.RESOURCE_QUOTA_RESOURCES + ) + ) + quotaspecs = QuotaSpecs.model_validate(quotaspecs_dict) + return quotaspecs.storage_quotas + def get_outages_for_service(cluster_name: str): """Get outages for a service from nerc-rates. @@ -316,12 +332,15 @@ def process_invoice_row(allocation, attrs, su_name, rate): ) logger.debug(f"Starting billing for allocation {allocation_str}.") - process_invoice_row( - allocation, - [attributes.QUOTA_VOLUMES_GB, attributes.QUOTA_OBJECT_GB], - "OpenStack Storage", - openstack_nese_storage_rate, - ) + # TODO (Quan): An illustration of how billing could be simplified. Shuold I follow with this? + quotaspecs = get_storage_quotaspecs(allocation) + for quota_name, quotaspec in quotaspecs.items(): + process_invoice_row( + allocation, + [quota_name], + quotaspec.invoice_name, + openstack_nese_storage_rate, + ) for allocation in openshift_allocations: allocation_str = ( diff --git a/src/coldfront_plugin_cloud/management/commands/fetch_daily_billable_usage.py b/src/coldfront_plugin_cloud/management/commands/fetch_daily_billable_usage.py index 18d4c9ac..20d33371 100644 --- a/src/coldfront_plugin_cloud/management/commands/fetch_daily_billable_usage.py +++ b/src/coldfront_plugin_cloud/management/commands/fetch_daily_billable_usage.py @@ -9,8 +9,8 @@ from coldfront_plugin_cloud import attributes from coldfront.core.utils.common import import_from_settings -from coldfront_plugin_cloud import usage_models -from coldfront_plugin_cloud.usage_models import UsageInfo, validate_date_str +from coldfront_plugin_cloud.models import usage_models +from coldfront_plugin_cloud.models.usage_models import UsageInfo, validate_date_str from coldfront_plugin_cloud import utils import boto3 diff --git a/src/coldfront_plugin_cloud/management/commands/remove_quota_from_resource.py b/src/coldfront_plugin_cloud/management/commands/remove_quota_from_resource.py new file mode 100644 index 00000000..89ca39f8 --- /dev/null +++ b/src/coldfront_plugin_cloud/management/commands/remove_quota_from_resource.py @@ -0,0 +1,69 @@ +import json +import logging +from django.core.management.base import BaseCommand + +from coldfront.core.resource.models import ( + Resource, + ResourceAttribute, + ResourceAttributeType, +) +from coldfront_plugin_cloud import attributes +from coldfront_plugin_cloud.models.quota_models import QuotaSpecs + +logger = logging.getLogger(__name__) + + +class Command(BaseCommand): + help = "Remove a quota from a resource's available resource quotas. Use --apply to perform the change." + + def add_arguments(self, parser): + parser.add_argument( + "resource_name", + type=str, + help="Name of the Resource to modify.", + ) + parser.add_argument( + "display_name", + type=str, + help="Display name of the quota to remove.", + ) + parser.add_argument( + "--apply", + action="store_true", + dest="apply", + help="If set, apply the removal", + ) + + def handle(self, *args, **options): + resource_name = options["resource_name"] + display_name = options["display_name"] + apply_change = options["apply"] + + resource = Resource.objects.get(name=resource_name) + rat = ResourceAttributeType.objects.get( + name=attributes.RESOURCE_QUOTA_RESOURCES + ) + available_attr = ResourceAttribute.objects.get( + resource=resource, resource_attribute_type=rat + ) + + available_dict = json.loads(available_attr.value or "{}") + + if display_name not in available_dict: + logger.info( + "Display name '%s' not present on resource '%s'. Nothing to remove.", + display_name, + resource_name, + ) + return + + logger.info( + "Removing quota '%s' from resource '%s':", display_name, resource_name + ) + if not apply_change: + return + + del available_dict[display_name] + QuotaSpecs.model_validate(available_dict) + available_attr.value = json.dumps(available_dict) + available_attr.save() diff --git a/src/coldfront_plugin_cloud/management/commands/validate_allocations.py b/src/coldfront_plugin_cloud/management/commands/validate_allocations.py index 9eae7272..b84186d5 100644 --- a/src/coldfront_plugin_cloud/management/commands/validate_allocations.py +++ b/src/coldfront_plugin_cloud/management/commands/validate_allocations.py @@ -85,9 +85,10 @@ def sync_openshift_project_labels(project_id, allocator, apply): @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) + resource_quotaspecs = allocator.resource_quotaspecs + value = resource_quotaspecs.root[coldfront_attr].quota_by_su_quantity( + allocation.quantity + ) utils.set_attribute_on_allocation(allocation, coldfront_attr, value) return value @@ -147,8 +148,8 @@ def handle(self, *args, **options): "keys" ][attributes.QUOTA_OBJECT_GB] - for attr in tasks.get_expected_attributes(allocator): - key = allocator.QUOTA_KEY_MAPPING_ALL_KEYS.get(attr, None) + for attr, quotaspec in allocator.resource_quotaspecs.root.items(): + key = quotaspec.quota_label if not key: # Note(knikolla): Some attributes are only maintained # for bookkeeping purposes and do not have a @@ -251,11 +252,9 @@ def handle(self, *args, **options): project_id, allocator, options["apply"] ) - for attr in tasks.get_expected_attributes(allocator): - key_with_lambda = allocator.QUOTA_KEY_MAPPING.get(attr, None) - + for attr, quotaspec in allocator.resource_quotaspecs.root.items(): # This gives me just the plain key - key = list(key_with_lambda(1).keys())[0] + key = quotaspec.quota_label expected_value = allocation.get_attribute(attr) current_value = quota.get(key, None) diff --git a/src/coldfront_plugin_cloud/models/quota_models.py b/src/coldfront_plugin_cloud/models/quota_models.py new file mode 100644 index 00000000..2b790441 --- /dev/null +++ b/src/coldfront_plugin_cloud/models/quota_models.py @@ -0,0 +1,68 @@ +from typing import Dict +from functools import cached_property + +import pydantic +from pydantic import Field + + +class QuotaSpec(pydantic.BaseModel): + """ + Fields: + - quota_label: human readable label for the quota (must be unique across the dict) + - default_quota: default quota value (int, >= 0) + - multiplier: multiplier applied to the allocation quantity (int, >= 0) + - static_quota: static extra quota added to every project (int, >= 0) + - unit_suffix: textual unit suffix (e.g. "Gi", "Mi", "", etc.) + """ + + quota_label: str + default_quota: int = Field(0, ge=0) + multiplier: int = Field(0, ge=0) + static_quota: int = Field(0, ge=0) + unit_suffix: str = "" + is_storage_type: bool = False + invoice_name: str = "" + + class Config: + model_config = pydantic.ConfigDict(extra="ignore") + + def quota_by_su_quantity(self, quantity: int) -> int: + """ + Compute the quota for a given SU quantity using the formula: + quota = static_quota + multiplier * quantity + """ + return self.static_quota + self.multiplier * int(quantity) + + def formatted_quota(self, quota_value: int) -> int | str: + """ + Return the quota value with the unit_suffix appended as a string when a suffix is set. + If unit_suffix is the empty string, return the original numeric value unchanged. + """ + if self.unit_suffix: + return f"{quota_value}{self.unit_suffix}" + return quota_value + + +class QuotaSpecs(pydantic.RootModel[Dict[str, QuotaSpec]]): + """ + Root model representing a mapping of display_name -> QuotaSpec. + + Validators: + - Ensure quota_label values are unique across all QuotaSpec entries. + """ + + @pydantic.model_validator(mode="after") + def validate_unique_labels(self): + # Ensure quota_label values are unique across the dict + labels = [q.quota_label for q in self.root.values()] + if len(labels) != len(set(labels)): + raise ValueError("Duplicate quota_label values found in QuotaSpecs") + + return self + + @cached_property + def storage_quotas(self) -> dict[str, QuotaSpec]: + """ + Return a list of quota display names that are marked as storage types. + """ + return {name: spec for name, spec in self.root.items() if spec.is_storage_type} diff --git a/src/coldfront_plugin_cloud/usage_models.py b/src/coldfront_plugin_cloud/models/usage_models.py similarity index 100% rename from src/coldfront_plugin_cloud/usage_models.py rename to src/coldfront_plugin_cloud/models/usage_models.py diff --git a/src/coldfront_plugin_cloud/openshift.py b/src/coldfront_plugin_cloud/openshift.py index 44f8c7dd..f58dea29 100644 --- a/src/coldfront_plugin_cloud/openshift.py +++ b/src/coldfront_plugin_cloud/openshift.py @@ -157,22 +157,6 @@ class NotFound(ApiException): class OpenShiftResourceAllocator(base.ResourceAllocator): - QUOTA_KEY_MAPPING = { - attributes.QUOTA_LIMITS_CPU: lambda x: {"limits.cpu": f"{x * 1000}m"}, - attributes.QUOTA_LIMITS_MEMORY: lambda x: {"limits.memory": f"{x}Mi"}, - attributes.QUOTA_LIMITS_EPHEMERAL_STORAGE_GB: lambda x: { - "limits.ephemeral-storage": f"{x}Gi" - }, - attributes.QUOTA_REQUESTS_NESE_STORAGE: lambda x: { - "ocs-external-storagecluster-ceph-rbd.storageclass.storage.k8s.io/requests.storage": f"{x}Gi" - }, - attributes.QUOTA_REQUESTS_IBM_STORAGE: lambda x: { - "ibm-spectrum-scale-fileset.storageclass.storage.k8s.io/requests.storage": f"{x}Gi" - }, - attributes.QUOTA_REQUESTS_GPU: lambda x: {"requests.nvidia.com/gpu": f"{x}"}, - attributes.QUOTA_PVC: lambda x: {"persistentvolumeclaims": f"{x}"}, - } - resource_type = "openshift" project_name_max_length = 45 @@ -286,9 +270,9 @@ def set_quota(self, project_id): object in the project namespace with no extra scopes""" quota_spec = {} - for key, func in self.QUOTA_KEY_MAPPING.items(): + for key, quotaspec in self.resource_quotaspecs.root.items(): if (x := self.allocation.get_attribute(key)) is not None: - quota_spec.update(func(x)) + quota_spec.update({quotaspec.quota_label: quotaspec.formatted_quota(x)}) quota_def = { "metadata": {"name": f"{project_id}-project"}, diff --git a/src/coldfront_plugin_cloud/openshift_vm.py b/src/coldfront_plugin_cloud/openshift_vm.py index 17b47b8a..a475077c 100644 --- a/src/coldfront_plugin_cloud/openshift_vm.py +++ b/src/coldfront_plugin_cloud/openshift_vm.py @@ -1,29 +1,5 @@ -from coldfront_plugin_cloud import attributes, openshift +from coldfront_plugin_cloud import openshift class OpenShiftVMResourceAllocator(openshift.OpenShiftResourceAllocator): - QUOTA_KEY_MAPPING = { - attributes.QUOTA_LIMITS_CPU: lambda x: {"limits.cpu": f"{x * 1000}m"}, - attributes.QUOTA_LIMITS_MEMORY: lambda x: {"limits.memory": f"{x}Mi"}, - attributes.QUOTA_LIMITS_EPHEMERAL_STORAGE_GB: lambda x: { - "limits.ephemeral-storage": f"{x}Gi" - }, - attributes.QUOTA_REQUESTS_NESE_STORAGE: lambda x: { - "ocs-external-storagecluster-ceph-rbd.storageclass.storage.k8s.io/requests.storage": f"{x}Gi" - }, - attributes.QUOTA_REQUESTS_IBM_STORAGE: lambda x: { - "ibm-spectrum-scale-fileset.storageclass.storage.k8s.io/requests.storage": f"{x}Gi" - }, - attributes.QUOTA_REQUESTS_VM_GPU_A100_SXM4: lambda x: { - "requests.nvidia.com/A100_SXM4_40GB": f"{x}" - }, - attributes.QUOTA_REQUESTS_VM_GPU_V100: lambda x: { - "requests.nvidia.com/GV100GL_Tesla_V100": f"{x}" - }, - attributes.QUOTA_REQUESTS_VM_GPU_H100: lambda x: { - "requests.nvidia.com/H100_SXM5_80GB": f"{x}" - }, - attributes.QUOTA_PVC: lambda x: {"persistentvolumeclaims": f"{x}"}, - } - resource_type = "openshift_vm" diff --git a/src/coldfront_plugin_cloud/openstack.py b/src/coldfront_plugin_cloud/openstack.py index 65f03205..53345b50 100644 --- a/src/coldfront_plugin_cloud/openstack.py +++ b/src/coldfront_plugin_cloud/openstack.py @@ -66,36 +66,15 @@ def get_session_for_resource(resource): class OpenStackResourceAllocator(base.ResourceAllocator): # Map the attribute name in ColdFront, to the client of the respective # service, the version of the API, and the key in the payload. - QUOTA_KEY_MAPPING = { - "compute": { - "keys": { - attributes.QUOTA_INSTANCES: "instances", - attributes.QUOTA_VCPU: "cores", - attributes.QUOTA_RAM: "ram", - }, - }, - "network": { - "keys": { - attributes.QUOTA_FLOATING_IPS: "floatingip", - } - }, - "object": { - "keys": { - attributes.QUOTA_OBJECT_GB: "x-account-meta-quota-bytes", - } - }, - "volume": { - "keys": { - attributes.QUOTA_VOLUMES: "volumes", - attributes.QUOTA_VOLUMES_GB: "gigabytes", - } - }, - } - - QUOTA_KEY_MAPPING_ALL_KEYS = { - quota_key: quota_name - for k in QUOTA_KEY_MAPPING.values() - for quota_key, quota_name in k["keys"].items() + SERVICE_QUOTA_MAPPING = { + "compute": [ + attributes.QUOTA_INSTANCES, + attributes.QUOTA_VCPU, + attributes.QUOTA_RAM, + ], + "network": [attributes.QUOTA_FLOATING_IPS], + "object": [attributes.QUOTA_OBJECT_GB], + "volume": [attributes.QUOTA_VOLUMES, attributes.QUOTA_VOLUMES_GB], } resource_type = "openstack" @@ -167,11 +146,14 @@ def set_quota(self, project_id): # If an attribute with the appropriate name is associated with an # allocation, set that as the quota. Otherwise, multiply # the quantity attribute via the mapping table above. - for service_name, service in self.QUOTA_KEY_MAPPING.items(): + for service_name, quotas_list in self.SERVICE_QUOTA_MAPPING.items(): # No need to do any calculations here, just go through each service # and set the value in the attribute. payload = dict() - for coldfront_attr, openstack_key in service["keys"].items(): + for coldfront_attr in quotas_list: + openstack_key = self.resource_quotaspecs.root[ + coldfront_attr + ].quota_label value = self.allocation.get_attribute(coldfront_attr) if value is not None: payload[openstack_key] = value @@ -194,9 +176,9 @@ def _set_object_quota(self, project_id, payload): # Note(knikolla): For consistency with other OpenStack # quotas we're storing this as GB on the attribute and # converting to bytes for Swift. - obj_q_mapping = self.QUOTA_KEY_MAPPING["object"]["keys"][ + obj_q_mapping = self.resource_quotaspecs.root[ attributes.QUOTA_OBJECT_GB - ] + ].quota_label payload[obj_q_mapping] *= GB_IN_BYTES if payload[obj_q_mapping] <= 0: payload[obj_q_mapping] = 1 @@ -245,7 +227,8 @@ def _init_rgw_for_project(self, project_id): def _get_network_quota(self, quotas, project_id): network_quota = self.network.show_quota(project_id)["quota"] - for k in self.QUOTA_KEY_MAPPING["network"]["keys"].values(): + for cf_k in self.SERVICE_QUOTA_MAPPING["network"]: + k = self.resource_quotaspecs.root[cf_k].quota_label quotas[k] = network_quota.get(k) return quotas @@ -254,16 +237,19 @@ def get_quota(self, project_id): quotas = dict() compute_quota = self.compute.quotas.get(project_id) - for k in self.QUOTA_KEY_MAPPING["compute"]["keys"].values(): + for cf_k in self.SERVICE_QUOTA_MAPPING["compute"]: + k = self.resource_quotaspecs.root[cf_k].quota_label quotas[k] = compute_quota.__getattr__(k) volume_quota = self.volume.quotas.get(project_id) - for k in self.QUOTA_KEY_MAPPING["volume"]["keys"].values(): + for cf_k in self.SERVICE_QUOTA_MAPPING["volume"]: + k = self.resource_quotaspecs.root[cf_k].quota_label quotas[k] = volume_quota.__getattr__(k) quotas = self._get_network_quota(quotas, project_id) - key = self.QUOTA_KEY_MAPPING["object"]["keys"][attributes.QUOTA_OBJECT_GB] + # key = self.SERVICE_QUOTA_MAPPING["object"]["keys"][attributes.QUOTA_OBJECT_GB] + key = self.resource_quotaspecs.root[attributes.QUOTA_OBJECT_GB].quota_label try: swift = self.object(project_id).head_account() quotas[key] = int(int(swift.get(key)) / GB_IN_BYTES) diff --git a/src/coldfront_plugin_cloud/tasks.py b/src/coldfront_plugin_cloud/tasks.py index 262f6845..130e56b2 100644 --- a/src/coldfront_plugin_cloud/tasks.py +++ b/src/coldfront_plugin_cloud/tasks.py @@ -17,82 +17,6 @@ logger = logging.getLogger(__name__) -# Map the amount of quota that 1 unit of `quantity` gets you -# This is multiplied to the quantity of that resource allocation. -UNIT_QUOTA_MULTIPLIERS = { - "openstack": { - attributes.QUOTA_INSTANCES: 1, - attributes.QUOTA_VCPU: 1, - attributes.QUOTA_RAM: 4096, - attributes.QUOTA_VOLUMES: 2, - attributes.QUOTA_VOLUMES_GB: 20, - attributes.QUOTA_FLOATING_IPS: 0, - attributes.QUOTA_OBJECT_GB: 1, - attributes.QUOTA_GPU: 0, - }, - "openshift": { - attributes.QUOTA_LIMITS_CPU: 1, - attributes.QUOTA_LIMITS_MEMORY: 4096, - attributes.QUOTA_LIMITS_EPHEMERAL_STORAGE_GB: 5, - attributes.QUOTA_REQUESTS_NESE_STORAGE: 20, - attributes.QUOTA_REQUESTS_IBM_STORAGE: 0, - attributes.QUOTA_REQUESTS_GPU: 0, - attributes.QUOTA_PVC: 2, - }, - "openshift_vm": { - attributes.QUOTA_LIMITS_CPU: 1, - attributes.QUOTA_LIMITS_MEMORY: 4096, - attributes.QUOTA_LIMITS_EPHEMERAL_STORAGE_GB: 5, - attributes.QUOTA_REQUESTS_NESE_STORAGE: 20, - attributes.QUOTA_REQUESTS_IBM_STORAGE: 0, - attributes.QUOTA_REQUESTS_VM_GPU_A100_SXM4: 0, - attributes.QUOTA_REQUESTS_VM_GPU_V100: 0, - attributes.QUOTA_REQUESTS_VM_GPU_H100: 0, - attributes.QUOTA_PVC: 2, - }, - "esi": {attributes.QUOTA_FLOATING_IPS: 0, attributes.QUOTA_NETWORKS: 0}, -} - -# The amount of quota that every projects gets, -# regardless of units of quantity. This is added -# on top of the multiplication. -STATIC_QUOTA = { - "openstack": { - attributes.QUOTA_FLOATING_IPS: 2, - attributes.QUOTA_GPU: 0, - }, - "openshift": { - attributes.QUOTA_REQUESTS_GPU: 0, - }, - "esi": {attributes.QUOTA_FLOATING_IPS: 1, attributes.QUOTA_NETWORKS: 1}, - "openshift_vm": { - attributes.QUOTA_REQUESTS_VM_GPU_A100_SXM4: 0, - attributes.QUOTA_REQUESTS_VM_GPU_V100: 0, - attributes.QUOTA_REQUESTS_VM_GPU_H100: 0, - }, -} - - -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 - resource_expected_quotas = UNIT_QUOTA_MULTIPLIERS[resource_name].copy() - - # If the resource attribute is not set (i.e for OpenStack resources), get_attribute returns None - is_ibm_storage_available = allocator.resource.get_attribute( - attributes.RESOURCE_IBM_AVAILABLE - ) - is_ibm_storage_available = ( - is_ibm_storage_available and is_ibm_storage_available.lower() == "true" - ) - if "openshift" in resource_name and not is_ibm_storage_available: - resource_expected_quotas.pop( - attributes.QUOTA_REQUESTS_IBM_STORAGE, None - ) # The resource may or may not already have this attribute - - return list(resource_expected_quotas.keys()) - - def find_allocator(allocation) -> base.ResourceAllocator: allocators = { "openstack": openstack.OpenStackResourceAllocator, @@ -115,13 +39,10 @@ def set_quota_attributes(): allocation.quantity = 1 # Calculate the quota for the project, and set the attribute for each element - expected_coldfront_attrs = get_expected_attributes(allocator) - for coldfront_attr in expected_coldfront_attrs: + resource_quotaspecs = allocator.resource_quotaspecs + for coldfront_attr, quota_spec in resource_quotaspecs.root.items(): if not allocation.get_attribute(coldfront_attr): - value = allocation.quantity * UNIT_QUOTA_MULTIPLIERS[ - allocator.resource_type - ].get(coldfront_attr, 0) - value += STATIC_QUOTA[allocator.resource_type].get(coldfront_attr, 0) + value = quota_spec.quota_by_su_quantity(allocation.quantity) utils.set_attribute_on_allocation(allocation, coldfront_attr, value) allocation = Allocation.objects.get(pk=allocation_pk) diff --git a/src/coldfront_plugin_cloud/tests/base.py b/src/coldfront_plugin_cloud/tests/base.py index 7cccb1bd..daa4e81e 100644 --- a/src/coldfront_plugin_cloud/tests/base.py +++ b/src/coldfront_plugin_cloud/tests/base.py @@ -87,20 +87,20 @@ def new_openstack_resource( @staticmethod def new_openshift_resource( name=None, + internal_name=None, api_url=None, idp=None, for_virtualization=False, - ibm_storage_available=False, ) -> Resource: resource_name = name or uuid.uuid4().hex call_command( "add_openshift_resource", name=resource_name, + internal_name=internal_name, api_url=api_url or "https://onboarding-onboarding.cluster.local:6443", idp=idp or "developer", for_virtualization=for_virtualization, - ibm_storage_available=ibm_storage_available, ) return Resource.objects.get(name=resource_name) 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 76229e01..fa42db59 100644 --- a/src/coldfront_plugin_cloud/tests/functional/openshift/test_allocation.py +++ b/src/coldfront_plugin_cloud/tests/functional/openshift/test_allocation.py @@ -1,7 +1,6 @@ import os import time import unittest -from unittest import mock import uuid from coldfront_plugin_cloud import attributes, openshift, tasks, utils @@ -11,6 +10,19 @@ import kubernetes.dynamic.exceptions as kexc +# Test OpenShift Quota Attributes +QUOTA_LIMITS_CPU = "OpenShift Limit on CPU Quota" +QUOTA_LIMITS_MEMORY = "OpenShift Limit on RAM Quota (MiB)" +QUOTA_LIMITS_EPHEMERAL_STORAGE_GB = "OpenShift Limit on Ephemeral Storage Quota (GiB)" +QUOTA_REQUESTS_NESE_STORAGE = "OpenShift Request on NESE Storage Quota (GiB)" +QUOTA_REQUESTS_IBM_STORAGE = "OpenShift Request on IBM Storage Quota (GiB)" +QUOTA_REQUESTS_GPU = "OpenShift Request on GPU Quota" +QUOTA_REQUESTS_VM_GPU_A100_SXM4 = "OpenShift Request on GPU A100 SXM4" +QUOTA_REQUESTS_VM_GPU_V100 = "OpenShift Request on GPU V100" +QUOTA_REQUESTS_VM_GPU_H100 = "OpenShift Request on GPU H100" +QUOTA_PVC = "OpenShift Persistent Volume Claims Quota" + + @unittest.skipUnless(os.getenv("FUNCTIONAL_TESTS"), "Functional tests not enabled.") class TestAllocation(base.TestBase): def setUp(self) -> None: @@ -18,7 +30,66 @@ def setUp(self) -> None: self.resource = self.new_openshift_resource( name="Microshift", api_url=os.getenv("OS_API_URL"), - ibm_storage_available=True, + ) + call_command( + "add_quota_to_resource", + display_name=QUOTA_LIMITS_CPU, + default_quota=1, + resource_name=self.resource.name, + quota_label="limits.cpu", + multiplier=1, + static_quota=0, + unit_suffix="", + ) + call_command( + "add_quota_to_resource", + display_name=QUOTA_LIMITS_MEMORY, + default_quota=4096, + resource_name=self.resource.name, + quota_label="limits.memory", + multiplier=4096, + static_quota=0, + unit_suffix="Mi", + ) + call_command( + "add_quota_to_resource", + display_name=QUOTA_LIMITS_EPHEMERAL_STORAGE_GB, + default_quota=5, + resource_name=self.resource.name, + quota_label="limits.ephemeral-storage", + multiplier=5, + static_quota=0, + unit_suffix="Gi", + ) + call_command( + "add_quota_to_resource", + display_name=QUOTA_REQUESTS_NESE_STORAGE, + default_quota=20, + resource_name=self.resource.name, + quota_label="ocs-external-storagecluster-ceph-rbd.storageclass.storage.k8s.io/requests.storage", + multiplier=20, + static_quota=0, + unit_suffix="Gi", + ) + call_command( + "add_quota_to_resource", + display_name=QUOTA_REQUESTS_GPU, + default_quota=0, + resource_name=self.resource.name, + quota_label="requests.nvidia.com/gpu", + multiplier=0, + static_quota=0, + unit_suffix="", + ) + call_command( + "add_quota_to_resource", + display_name=QUOTA_PVC, + default_quota=2, + resource_name=self.resource.name, + quota_label="persistentvolumeclaims", + multiplier=2, + static_quota=0, + unit_suffix="", ) def test_new_allocation(self): @@ -127,19 +198,15 @@ def test_new_allocation_quota(self): project_id = allocation.get_attribute(attributes.ALLOCATION_PROJECT_ID) - self.assertEqual(allocation.get_attribute(attributes.QUOTA_LIMITS_CPU), 2 * 1) - self.assertEqual( - allocation.get_attribute(attributes.QUOTA_LIMITS_MEMORY), 2 * 4096 - ) + self.assertEqual(allocation.get_attribute(QUOTA_LIMITS_CPU), 2 * 1) + self.assertEqual(allocation.get_attribute(QUOTA_LIMITS_MEMORY), 2 * 4096) self.assertEqual( - allocation.get_attribute(attributes.QUOTA_LIMITS_EPHEMERAL_STORAGE_GB), + allocation.get_attribute(QUOTA_LIMITS_EPHEMERAL_STORAGE_GB), 2 * 5, ) - self.assertEqual( - allocation.get_attribute(attributes.QUOTA_REQUESTS_NESE_STORAGE), 2 * 20 - ) - self.assertEqual(allocation.get_attribute(attributes.QUOTA_REQUESTS_GPU), 2 * 0) - self.assertEqual(allocation.get_attribute(attributes.QUOTA_PVC), 2 * 2) + self.assertEqual(allocation.get_attribute(QUOTA_REQUESTS_NESE_STORAGE), 2 * 20) + self.assertEqual(allocation.get_attribute(QUOTA_REQUESTS_GPU), 2 * 0) + self.assertEqual(allocation.get_attribute(QUOTA_PVC), 2 * 2) quota = allocator.get_quota(project_id) # The return value will update to the most relevant unit, so @@ -151,36 +218,29 @@ def test_new_allocation_quota(self): "limits.memory": "8Gi", "limits.ephemeral-storage": "10Gi", "ocs-external-storagecluster-ceph-rbd.storageclass.storage.k8s.io/requests.storage": "40Gi", - "ibm-spectrum-scale-fileset.storageclass.storage.k8s.io/requests.storage": "0", "requests.nvidia.com/gpu": "0", "persistentvolumeclaims": "4", }, ) # change a bunch of attributes - utils.set_attribute_on_allocation(allocation, attributes.QUOTA_LIMITS_CPU, 6) - utils.set_attribute_on_allocation( - allocation, attributes.QUOTA_LIMITS_MEMORY, 8192 - ) - utils.set_attribute_on_allocation( - allocation, attributes.QUOTA_LIMITS_EPHEMERAL_STORAGE_GB, 50 - ) + utils.set_attribute_on_allocation(allocation, QUOTA_LIMITS_CPU, 6) + utils.set_attribute_on_allocation(allocation, QUOTA_LIMITS_MEMORY, 8192) utils.set_attribute_on_allocation( - allocation, attributes.QUOTA_REQUESTS_NESE_STORAGE, 100 + allocation, QUOTA_LIMITS_EPHEMERAL_STORAGE_GB, 50 ) - utils.set_attribute_on_allocation(allocation, attributes.QUOTA_REQUESTS_GPU, 1) - utils.set_attribute_on_allocation(allocation, attributes.QUOTA_PVC, 10) + utils.set_attribute_on_allocation(allocation, QUOTA_REQUESTS_NESE_STORAGE, 100) + utils.set_attribute_on_allocation(allocation, QUOTA_REQUESTS_GPU, 1) + utils.set_attribute_on_allocation(allocation, QUOTA_PVC, 10) - self.assertEqual(allocation.get_attribute(attributes.QUOTA_LIMITS_CPU), 6) - self.assertEqual(allocation.get_attribute(attributes.QUOTA_LIMITS_MEMORY), 8192) - self.assertEqual( - allocation.get_attribute(attributes.QUOTA_LIMITS_EPHEMERAL_STORAGE_GB), 50 - ) + self.assertEqual(allocation.get_attribute(QUOTA_LIMITS_CPU), 6) + self.assertEqual(allocation.get_attribute(QUOTA_LIMITS_MEMORY), 8192) self.assertEqual( - allocation.get_attribute(attributes.QUOTA_REQUESTS_NESE_STORAGE), 100 + allocation.get_attribute(QUOTA_LIMITS_EPHEMERAL_STORAGE_GB), 50 ) - self.assertEqual(allocation.get_attribute(attributes.QUOTA_REQUESTS_GPU), 1) - self.assertEqual(allocation.get_attribute(attributes.QUOTA_PVC), 10) + self.assertEqual(allocation.get_attribute(QUOTA_REQUESTS_NESE_STORAGE), 100) + self.assertEqual(allocation.get_attribute(QUOTA_REQUESTS_GPU), 1) + self.assertEqual(allocation.get_attribute(QUOTA_PVC), 10) # This call should update the openshift quota to match the current attributes call_command("validate_allocations", apply=True) @@ -195,7 +255,6 @@ def test_new_allocation_quota(self): "limits.memory": "8Gi", "limits.ephemeral-storage": "50Gi", "ocs-external-storagecluster-ceph-rbd.storageclass.storage.k8s.io/requests.storage": "100Gi", - "ibm-spectrum-scale-fileset.storageclass.storage.k8s.io/requests.storage": "0", "requests.nvidia.com/gpu": "1", "persistentvolumeclaims": "10", }, @@ -212,7 +271,7 @@ def test_reactivate_allocation(self): project_id = allocation.get_attribute(attributes.ALLOCATION_PROJECT_ID) - self.assertEqual(allocation.get_attribute(attributes.QUOTA_LIMITS_CPU), 2) + self.assertEqual(allocation.get_attribute(QUOTA_LIMITS_CPU), 2) quota = allocator.get_quota(project_id) @@ -225,7 +284,6 @@ def test_reactivate_allocation(self): "limits.memory": "8Gi", "limits.ephemeral-storage": "10Gi", "ocs-external-storagecluster-ceph-rbd.storageclass.storage.k8s.io/requests.storage": "40Gi", - "ibm-spectrum-scale-fileset.storageclass.storage.k8s.io/requests.storage": "0", "requests.nvidia.com/gpu": "0", "persistentvolumeclaims": "4", }, @@ -233,7 +291,7 @@ def test_reactivate_allocation(self): # Simulate an attribute change request and subsequent approval which # triggers a reactivation - utils.set_attribute_on_allocation(allocation, attributes.QUOTA_LIMITS_CPU, 3) + utils.set_attribute_on_allocation(allocation, QUOTA_LIMITS_CPU, 3) tasks.activate_allocation(allocation.pk) allocation.refresh_from_db() @@ -247,7 +305,6 @@ def test_reactivate_allocation(self): "limits.memory": "8Gi", "limits.ephemeral-storage": "10Gi", "ocs-external-storagecluster-ceph-rbd.storageclass.storage.k8s.io/requests.storage": "40Gi", - "ibm-spectrum-scale-fileset.storageclass.storage.k8s.io/requests.storage": "0", "requests.nvidia.com/gpu": "0", "persistentvolumeclaims": "4", }, @@ -330,59 +387,8 @@ def test_create_incomplete(self): ) ) - @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", - }, - ) - def test_migrate_quota_field_names(self): - """When a quota key in QUOTA_KEY_MAPPING changes to a new value, validate_allocations should update the quota.""" + """When a quota changes to a new label name, validate_allocations should update the quota.""" user = self.new_user() project = self.new_project(pi=user) allocation = self.new_allocation(project, self.resource, 1) @@ -401,75 +407,25 @@ def test_migrate_quota_field_names(self): "limits.memory": "4Gi", "limits.ephemeral-storage": "5Gi", "ocs-external-storagecluster-ceph-rbd.storageclass.storage.k8s.io/requests.storage": "20Gi", - "ibm-spectrum-scale-fileset.storageclass.storage.k8s.io/requests.storage": "0", "requests.nvidia.com/gpu": "0", "persistentvolumeclaims": "2", }, ) # Now migrate NESE Storage quota field (ocs-external...) to fake storage quota - with unittest.mock.patch.dict( - openshift.OpenShiftResourceAllocator.QUOTA_KEY_MAPPING, - { - attributes.QUOTA_REQUESTS_NESE_STORAGE: lambda x: { - "fake-storage.storageclass.storage.k8s.io/requests.storage": f"{x}Gi" - } - }, - ): - call_command("validate_allocations", apply=True) - - # Check the quota after migration - quota = allocator.get_quota(project_id) - self.assertEqual( - quota, - { - "limits.cpu": "1", - "limits.memory": "4Gi", - "limits.ephemeral-storage": "5Gi", - "fake-storage.storageclass.storage.k8s.io/requests.storage": "20Gi", # Migrated key - "ibm-spectrum-scale-fileset.storageclass.storage.k8s.io/requests.storage": "0", - "requests.nvidia.com/gpu": "0", - "persistentvolumeclaims": "2", - }, - ) - - def test_ibm_storage_not_available(self): - """If IBM Scale storage is not available, the corresponding quotas should not be set.""" - user = self.new_user() - project = self.new_project(pi=user) - - # Set ibm storage as not available - self.resource.resourceattribute_set.filter( - resource_attribute_type__name=attributes.RESOURCE_IBM_AVAILABLE - ).update(value="false") - allocation = self.new_allocation(project, self.resource, 1) - allocator = openshift.OpenShiftResourceAllocator(self.resource, allocation) - - tasks.activate_allocation(allocation.pk) - allocation.refresh_from_db() - - project_id = allocation.get_attribute(attributes.ALLOCATION_PROJECT_ID) - - quota = allocator.get_quota(project_id) - self.assertEqual( - quota, - { - "limits.cpu": "1", - "limits.memory": "4Gi", - "limits.ephemeral-storage": "5Gi", - "ocs-external-storagecluster-ceph-rbd.storageclass.storage.k8s.io/requests.storage": "20Gi", - "requests.nvidia.com/gpu": "0", - "persistentvolumeclaims": "2", - }, + call_command( + "add_quota_to_resource", + display_name=QUOTA_REQUESTS_NESE_STORAGE, + default_quota=20, + resource_name=self.resource.name, + quota_label="fake-storage.storageclass.storage.k8s.io/requests.storage", + multiplier=20, + static_quota=0, + unit_suffix="Gi", ) - - # Now set IBM Scale storage as available - self.resource.resourceattribute_set.filter( - resource_attribute_type__name=attributes.RESOURCE_IBM_AVAILABLE - ).update(value="true") - call_command("validate_allocations", apply=True) + # Check the quota after migration quota = allocator.get_quota(project_id) self.assertEqual( quota, @@ -477,8 +433,7 @@ def test_ibm_storage_not_available(self): "limits.cpu": "1", "limits.memory": "4Gi", "limits.ephemeral-storage": "5Gi", - "ocs-external-storagecluster-ceph-rbd.storageclass.storage.k8s.io/requests.storage": "20Gi", - "ibm-spectrum-scale-fileset.storageclass.storage.k8s.io/requests.storage": "0", # Newly added IBM key + "fake-storage.storageclass.storage.k8s.io/requests.storage": "20Gi", # Migrated key "requests.nvidia.com/gpu": "0", "persistentvolumeclaims": "2", }, @@ -596,9 +551,102 @@ def test_preexisting_project(self): "limits.memory": "4Gi", "limits.ephemeral-storage": "5Gi", "ocs-external-storagecluster-ceph-rbd.storageclass.storage.k8s.io/requests.storage": "20Gi", - "ibm-spectrum-scale-fileset.storageclass.storage.k8s.io/requests.storage": "0", "requests.nvidia.com/gpu": "0", "persistentvolumeclaims": "2", }, ) assert set([user.username]) == allocator.get_users(project_id) + + +class TestAllocationNewQuota(base.TestBase): + def setUp(self) -> None: + super().setUp() + self.resource = self.new_openshift_resource( + name="Microshift", + api_url=os.getenv("OS_API_URL"), + ) + call_command( + "add_quota_to_resource", + display_name=QUOTA_LIMITS_CPU, + default_quota=1, + resource_name=self.resource.name, + quota_label="limits.cpu", + multiplier=1, + static_quota=0, + unit_suffix="", + ) + + 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(QUOTA_LIMITS_CPU), 2 * 1) + + quota = allocator.get_quota(project_id) + self.assertEqual( + quota, + { + "limits.cpu": "2", + }, + ) + + # Add a new attribute for Openshift + call_command( + "add_quota_to_resource", + display_name=QUOTA_LIMITS_MEMORY, + default_quota=4096, + resource_name=self.resource.name, + quota_label="limits.memory", + multiplier=4096, + static_quota=0, + unit_suffix="Mi", + ) + + call_command("validate_allocations", apply=True) + allocation.refresh_from_db() + + self.assertEqual(allocation.get_attribute(QUOTA_LIMITS_CPU), 2 * 1) + self.assertEqual(allocation.get_attribute(QUOTA_LIMITS_MEMORY), 2 * 4096) + + quota = allocator.get_quota(project_id) + self.assertEqual( + quota, + { + "limits.cpu": "2", + "limits.memory": "8Gi", + }, + ) + + # TODO (Quan): What happens when a quota is removed? Should the attribute be removed from Coldfront? + # Remove the CPU quota + call_command( + "remove_quota_from_resource", + self.resource.name, + QUOTA_LIMITS_CPU, + apply=True, + ) + + call_command("validate_allocations", apply=True) + allocation.refresh_from_db() + + # Check that CPU attribute is removed from ColdFront + self.assertIsNone(allocation.get_attribute(QUOTA_LIMITS_CPU)) + # Memory quota remains + self.assertEqual(allocation.get_attribute(QUOTA_LIMITS_MEMORY), 2 * 4096) + + # Check that OpenShift quota is updated + quota = allocator.get_quota(project_id) + self.assertEqual( + quota, + { + "limits.memory": "8Gi", + }, + ) diff --git a/src/coldfront_plugin_cloud/tests/functional/openshift_vm/test_allocation.py b/src/coldfront_plugin_cloud/tests/functional/openshift_vm/test_allocation.py index 4e8b4c4b..e79d3193 100644 --- a/src/coldfront_plugin_cloud/tests/functional/openshift_vm/test_allocation.py +++ b/src/coldfront_plugin_cloud/tests/functional/openshift_vm/test_allocation.py @@ -4,6 +4,19 @@ from coldfront_plugin_cloud import attributes, openshift_vm, tasks from coldfront_plugin_cloud.tests import base +from django.core.management import call_command + + +# Test OpenShift VM Quota Attributes +QUOTA_LIMITS_CPU = "OpenShift Limit on CPU Quota" +QUOTA_LIMITS_MEMORY = "OpenShift Limit on RAM Quota (MiB)" +QUOTA_LIMITS_EPHEMERAL_STORAGE_GB = "OpenShift Limit on Ephemeral Storage Quota (GiB)" +QUOTA_REQUESTS_NESE_STORAGE = "OpenShift Request on NESE Storage Quota (GiB)" +QUOTA_REQUESTS_VM_GPU_A100_SXM4 = "OpenShift Request on GPU A100 SXM4" +QUOTA_REQUESTS_VM_GPU_V100 = "OpenShift Request on GPU V100" +QUOTA_REQUESTS_VM_GPU_H100 = "OpenShift Request on GPU H100" +QUOTA_PVC = "OpenShift Persistent Volume Claims Quota" + @unittest.skipUnless(os.getenv("FUNCTIONAL_TESTS"), "Functional tests not enabled.") class TestAllocation(base.TestBase): @@ -14,6 +27,86 @@ def setUp(self) -> None: api_url=os.getenv("OS_API_URL"), for_virtualization=True, ) + call_command( + "add_quota_to_resource", + display_name=QUOTA_LIMITS_CPU, + default_quota=1, + resource_name=self.resource.name, + quota_label="limits.cpu", + multiplier=1, + static_quota=0, + unit_suffix="", + ) + call_command( + "add_quota_to_resource", + display_name=QUOTA_LIMITS_MEMORY, + default_quota=4096, + resource_name=self.resource.name, + quota_label="limits.memory", + multiplier=4096, + static_quota=0, + unit_suffix="Mi", + ) + call_command( + "add_quota_to_resource", + display_name=QUOTA_LIMITS_EPHEMERAL_STORAGE_GB, + default_quota=5, + resource_name=self.resource.name, + quota_label="limits.ephemeral-storage", + multiplier=5, + static_quota=0, + unit_suffix="Gi", + ) + call_command( + "add_quota_to_resource", + display_name=QUOTA_REQUESTS_NESE_STORAGE, + default_quota=20, + resource_name=self.resource.name, + quota_label="ocs-external-storagecluster-ceph-rbd.storageclass.storage.k8s.io/requests.storage", + multiplier=20, + static_quota=0, + unit_suffix="Gi", + ) + call_command( + "add_quota_to_resource", + display_name=QUOTA_REQUESTS_VM_GPU_A100_SXM4, + default_quota=0, + resource_name=self.resource.name, + quota_label="requests.nvidia.com/A100_SXM4_40GB", + multiplier=0, + static_quota=0, + unit_suffix="", + ) + call_command( + "add_quota_to_resource", + display_name=QUOTA_REQUESTS_VM_GPU_V100, + default_quota=0, + resource_name=self.resource.name, + quota_label="requests.nvidia.com/GV100GL_Tesla_V100", + multiplier=0, + static_quota=0, + unit_suffix="", + ) + call_command( + "add_quota_to_resource", + display_name=QUOTA_REQUESTS_VM_GPU_H100, + default_quota=0, + resource_name=self.resource.name, + quota_label="requests.nvidia.com/H100_SXM5_80GB", + multiplier=0, + static_quota=0, + unit_suffix="", + ) + call_command( + "add_quota_to_resource", + display_name=QUOTA_PVC, + default_quota=2, + resource_name=self.resource.name, + quota_label="persistentvolumeclaims", + multiplier=2, + static_quota=0, + unit_suffix="", + ) def test_new_allocation(self): # TODO must wait until we know what the quota values for openshift_vm are @@ -27,27 +120,19 @@ def test_new_allocation(self): project_id = allocation.get_attribute(attributes.ALLOCATION_PROJECT_ID) - self.assertEqual(allocation.get_attribute(attributes.QUOTA_LIMITS_CPU), 2 * 1) - self.assertEqual( - allocation.get_attribute(attributes.QUOTA_LIMITS_MEMORY), 2 * 4096 - ) + self.assertEqual(allocation.get_attribute(QUOTA_LIMITS_CPU), 2 * 1) + self.assertEqual(allocation.get_attribute(QUOTA_LIMITS_MEMORY), 2 * 4096) self.assertEqual( - allocation.get_attribute(attributes.QUOTA_LIMITS_EPHEMERAL_STORAGE_GB), + allocation.get_attribute(QUOTA_LIMITS_EPHEMERAL_STORAGE_GB), 2 * 5, ) + self.assertEqual(allocation.get_attribute(QUOTA_REQUESTS_NESE_STORAGE), 2 * 20) self.assertEqual( - allocation.get_attribute(attributes.QUOTA_REQUESTS_NESE_STORAGE), 2 * 20 - ) - self.assertEqual( - allocation.get_attribute(attributes.QUOTA_REQUESTS_VM_GPU_A100_SXM4), 2 * 0 - ) - self.assertEqual( - allocation.get_attribute(attributes.QUOTA_REQUESTS_VM_GPU_V100), 2 * 0 - ) - self.assertEqual( - allocation.get_attribute(attributes.QUOTA_REQUESTS_VM_GPU_H100), 2 * 0 + allocation.get_attribute(QUOTA_REQUESTS_VM_GPU_A100_SXM4), 2 * 0 ) - self.assertEqual(allocation.get_attribute(attributes.QUOTA_PVC), 2 * 2) + self.assertEqual(allocation.get_attribute(QUOTA_REQUESTS_VM_GPU_V100), 2 * 0) + self.assertEqual(allocation.get_attribute(QUOTA_REQUESTS_VM_GPU_H100), 2 * 0) + self.assertEqual(allocation.get_attribute(QUOTA_PVC), 2 * 2) quota = allocator.get_quota(project_id) # The return value will update to the most relevant unit, so diff --git a/src/coldfront_plugin_cloud/tests/functional/openstack/test_allocation.py b/src/coldfront_plugin_cloud/tests/functional/openstack/test_allocation.py index 20a24a61..84ea085a 100644 --- a/src/coldfront_plugin_cloud/tests/functional/openstack/test_allocation.py +++ b/src/coldfront_plugin_cloud/tests/functional/openstack/test_allocation.py @@ -213,9 +213,9 @@ def test_new_allocation_with_quantity(self): # Change allocation attributes for object store quota current_quota = allocator.get_quota(openstack_project.id) - obj_key = openstack.OpenStackResourceAllocator.QUOTA_KEY_MAPPING["object"][ - "keys" - ][attributes.QUOTA_OBJECT_GB] + obj_key = allocator.resource_quotaspecs.root[ + attributes.QUOTA_OBJECT_GB + ].quota_label if obj_key in current_quota.keys(): utils.set_attribute_on_allocation(allocation, attributes.QUOTA_OBJECT_GB, 6) self.assertEqual(allocation.get_attribute(attributes.QUOTA_OBJECT_GB), 6) diff --git a/src/coldfront_plugin_cloud/tests/unit/openshift/base.py b/src/coldfront_plugin_cloud/tests/unit/openshift/base.py index 82662ab8..f8ad60a5 100644 --- a/src/coldfront_plugin_cloud/tests/unit/openshift/base.py +++ b/src/coldfront_plugin_cloud/tests/unit/openshift/base.py @@ -4,10 +4,20 @@ from coldfront_plugin_cloud.openshift import OpenShiftResourceAllocator +class TestOpenshiftResourceAllocator(OpenShiftResourceAllocator): + def __init__(self): + self.resource = mock.Mock() + self.allocation = mock.Mock() + self.resource_quotaspecs = mock.Mock() + self.id_provider = "fake_idp" + self.k8_client = mock.Mock() + + self.verify = False + self.safe_resource_name = "foo" + self.apis = {} + self.member_role_name = "admin" + + class TestUnitOpenshiftBase(base.TestBase): def setUp(self) -> None: - mock_resource = mock.Mock() - mock_allocation = mock.Mock() - self.allocator = OpenShiftResourceAllocator(mock_resource, mock_allocation) - self.allocator.id_provider = "fake_idp" - self.allocator.k8_client = mock.Mock() + self.allocator = TestOpenshiftResourceAllocator() diff --git a/src/coldfront_plugin_cloud/tests/unit/openshift/test_rbac.py b/src/coldfront_plugin_cloud/tests/unit/openshift/test_rbac.py index b7898f7c..54ed32ca 100644 --- a/src/coldfront_plugin_cloud/tests/unit/openshift/test_rbac.py +++ b/src/coldfront_plugin_cloud/tests/unit/openshift/test_rbac.py @@ -2,19 +2,10 @@ import kubernetes.dynamic.exceptions as kexc -from coldfront_plugin_cloud.tests import base -from coldfront_plugin_cloud.openshift import OpenShiftResourceAllocator +from coldfront_plugin_cloud.tests.unit.openshift import base -class TestMocOpenShiftRBAC(base.TestBase): - def setUp(self) -> None: - mock_resource = mock.Mock() - mock_allocation = mock.Mock() - self.allocator = OpenShiftResourceAllocator(mock_resource, mock_allocation) - self.allocator.id_provider = "fake_idp" - self.allocator.k8_client = mock.Mock() - self.allocator.member_role_name = "admin" - +class TestMocOpenShiftRBAC(base.TestUnitOpenshiftBase): def test_user_in_rolebindings_false(self): fake_rb = { "subjects": [ diff --git a/src/coldfront_plugin_cloud/tests/unit/test_calculate_quota_unit_hours.py b/src/coldfront_plugin_cloud/tests/unit/test_calculate_quota_unit_hours.py index f71e12a5..5dad0060 100644 --- a/src/coldfront_plugin_cloud/tests/unit/test_calculate_quota_unit_hours.py +++ b/src/coldfront_plugin_cloud/tests/unit/test_calculate_quota_unit_hours.py @@ -625,14 +625,14 @@ def test_nerc_outages_integration(self, mock_rates_loader): with freezegun.freeze_time("2020-03-01"): user = self.new_user() project = self.new_project(pi=user) - resource = self.new_openstack_resource( + resource = self.new_openshift_resource( name="TEST-RESOURCE", internal_name="test-service" ) allocation = self.new_allocation(project, resource, 100) for attr, val in [ (attributes.ALLOCATION_PROJECT_NAME, "test"), (attributes.ALLOCATION_PROJECT_ID, "123"), - (attributes.QUOTA_VOLUMES_GB, 10), + (attributes.QUOTA_REQUESTS_NESE_STORAGE, 10), ]: utils.set_attribute_on_allocation(allocation, attr, val) diff --git a/src/coldfront_plugin_cloud/tests/unit/test_fetch_daily_billable_usage.py b/src/coldfront_plugin_cloud/tests/unit/test_fetch_daily_billable_usage.py index 9b43217d..a52c2174 100644 --- a/src/coldfront_plugin_cloud/tests/unit/test_fetch_daily_billable_usage.py +++ b/src/coldfront_plugin_cloud/tests/unit/test_fetch_daily_billable_usage.py @@ -8,7 +8,7 @@ Command, ) from coldfront_plugin_cloud import attributes -from coldfront_plugin_cloud import usage_models +from coldfront_plugin_cloud.models import usage_models from coldfront_plugin_cloud.tests import base from coldfront_plugin_cloud import utils diff --git a/src/coldfront_plugin_cloud/tests/unit/test_usage_models.py b/src/coldfront_plugin_cloud/tests/unit/test_usage_models.py index d1f3e81f..bf5f11da 100644 --- a/src/coldfront_plugin_cloud/tests/unit/test_usage_models.py +++ b/src/coldfront_plugin_cloud/tests/unit/test_usage_models.py @@ -1,7 +1,7 @@ from decimal import Decimal from pydantic import ValidationError -from coldfront_plugin_cloud import usage_models +from coldfront_plugin_cloud.models import usage_models from coldfront_plugin_cloud.tests import base