Allow dynamic quota creation and removal

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),
 ]

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)

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

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)

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 = (

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

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()