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),
@@ -116,23 +116,5 @@ class CloudAllocationAttribute:
 
 
 ALLOCATION_QUOTA_ATTRIBUTES = [
-    CloudAllocationAttribute(name=QUOTA_INSTANCES),
-    CloudAllocationAttribute(name=QUOTA_RAM),
-    CloudAllocationAttribute(name=QUOTA_VCPU),
-    CloudAllocationAttribute(name=QUOTA_VOLUMES),
-    CloudAllocationAttribute(name=QUOTA_VOLUMES_GB),
-    CloudAllocationAttribute(name=QUOTA_NETWORKS),
-    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/esi.py

@@ -3,21 +3,9 @@
 
 
 class ESIResourceAllocator(OpenStackResourceAllocator):
-    QUOTA_KEY_MAPPING = {
-        "network": {
-            "keys": {
-                attributes.QUOTA_FLOATING_IPS: "floatingip",
-                attributes.QUOTA_NETWORKS: "network",
-            }
-        }
+    SERVICE_QUOTA_MAPPING = {
+        "network": [attributes.QUOTA_FLOATING_IPS, attributes.QUOTA_NETWORKS],
     }
-
-    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()
-    }
-
     resource_type = "esi"
 
     def get_quota(self, project_id):

src/coldfront_plugin_cloud/management/commands/add_openshift_resource.py

@@ -1,4 +1,5 @@
 from django.core.management.base import BaseCommand
+from django.core.management import call_command
 
 from coldfront.core.resource.models import (
     Resource,
@@ -50,11 +51,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 +93,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
@@ -114,3 +102,39 @@ def handle(self, *args, **options):
             if options["internal_name"]
             else options["name"],
         )
+
+        # Add common Openshift resources (cpu, memory, etc)
+        call_command(
+            "add_quota_to_resource",
+            display_name=attributes.QUOTA_LIMITS_CPU,
+            default_quota=1,
+            resource_name=options["name"],
+            quota_label="limits.cpu",
+            multiplier=1,
+        )
+        call_command(
+            "add_quota_to_resource",
+            display_name=attributes.QUOTA_LIMITS_MEMORY,
+            default_quota=4096,
+            resource_name=options["name"],
+            quota_label="limits.memory",
+            multiplier=4096,
+            unit_suffix="Mi",
+        )
+        call_command(
+            "add_quota_to_resource",
+            display_name=attributes.QUOTA_LIMITS_EPHEMERAL_STORAGE_GB,
+            default_quota=5,
+            resource_name=options["name"],
+            quota_label="limits.ephemeral-storage",
+            multiplier=5,
+            unit_suffix="Gi",
+        )
+        call_command(
+            "add_quota_to_resource",
+            display_name=attributes.QUOTA_PVC,
+            default_quota=2,
+            resource_name=options["name"],
+            quota_label="persistentvolumeclaims",
+            multiplier=2,
+        )

src/coldfront_plugin_cloud/management/commands/add_quota_to_resource.py

@@ -0,0 +1,119 @@
+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(
+            "--resource-type",
+            type=str,
+            default="",
+            help="Indicates which resource type this quota is. Type `storage` is relevant for storage billing",
+        )
+        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["resource_type"] == "storage" and not options["invoice_name"]:
+            logger.error(
+                "--invoice-name must be provided when storage type is  `storage`."
+            )
+
+        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
@@ -210,6 +212,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 +328,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
@@ -192,7 +192,9 @@ def s3_client(self):
     def load_csv(location) -> DataFrameGroupBy:
         df = pandas.read_csv(
             location,
-            dtype={INVOICE_COLUMN_COST: pandas.ArrowDtype(pyarrow.decimal128(12, 2))},
+        )
+        df = df.astype(
+            {INVOICE_COLUMN_COST: pandas.ArrowDtype(pyarrow.decimal128(12, 2))}
         )
         return df.groupby(INVOICE_COLUMN_ALLOCATION_ID)