Skip to content

Commit 8703161

Browse files
QuanMPhmQuanMPhm
committed
Allow dynamic quota creation and removal
Allocation quota attributes are no longer hardcoded in `attributes.py` and instead are created dynamically based on the quotas defined for each resource. This eliminates the need to modify code and restart the Coldfront service to update quotas. 3 new commands are added to implement this feature: - `add_quota_to_resource`: Adds a quota to a resource. - `remove_quota_from_resource`: Removes a quota from a resource. - `register_default_quotas`: Registers default quotas for all resources based a predetermined list These commands will interact with a new resource attribute `Available Quota Resources`, which stores a list of all quota resources that can be applied to a resource. When a quota is added or removed using the above commands, this attribute will be updated accordingly. Refactored various files to support dynamic quotas Updated functional and unit tests
1 parent 4bd94ab commit 8703161

27 files changed

Lines changed: 900 additions & 430 deletions

‎src/coldfront_plugin_cloud/attributes.py‎

Lines changed: 2 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ class CloudAllocationAttribute:
2424
RESOURCE_API_URL = "OpenShift API Endpoint URL"
2525
RESOURCE_IDENTITY_NAME = "OpenShift Identity Provider Name"
2626
RESOURCE_ROLE = "Role for User in Project"
27-
RESOURCE_IBM_AVAILABLE = "IBM Spectrum Scale Storage Available"
27+
RESOURCE_QUOTA_RESOURCES = "Available Quota Resources"
2828

2929
RESOURCE_FEDERATION_PROTOCOL = "OpenStack Federation Protocol"
3030
RESOURCE_IDP = "OpenStack Identity Provider"
@@ -44,7 +44,7 @@ class CloudAllocationAttribute:
4444
CloudResourceAttribute(name=RESOURCE_IDP),
4545
CloudResourceAttribute(name=RESOURCE_PROJECT_DOMAIN),
4646
CloudResourceAttribute(name=RESOURCE_ROLE),
47-
CloudResourceAttribute(name=RESOURCE_IBM_AVAILABLE),
47+
CloudResourceAttribute(name=RESOURCE_QUOTA_RESOURCES),
4848
CloudResourceAttribute(name=RESOURCE_USER_DOMAIN),
4949
CloudResourceAttribute(name=RESOURCE_EULA_URL),
5050
CloudResourceAttribute(name=RESOURCE_DEFAULT_PUBLIC_NETWORK),
@@ -116,23 +116,5 @@ class CloudAllocationAttribute:
116116

117117

118118
ALLOCATION_QUOTA_ATTRIBUTES = [
119-
CloudAllocationAttribute(name=QUOTA_INSTANCES),
120-
CloudAllocationAttribute(name=QUOTA_RAM),
121-
CloudAllocationAttribute(name=QUOTA_VCPU),
122-
CloudAllocationAttribute(name=QUOTA_VOLUMES),
123-
CloudAllocationAttribute(name=QUOTA_VOLUMES_GB),
124-
CloudAllocationAttribute(name=QUOTA_NETWORKS),
125-
CloudAllocationAttribute(name=QUOTA_FLOATING_IPS),
126-
CloudAllocationAttribute(name=QUOTA_OBJECT_GB),
127119
CloudAllocationAttribute(name=QUOTA_GPU),
128-
CloudAllocationAttribute(name=QUOTA_LIMITS_CPU),
129-
CloudAllocationAttribute(name=QUOTA_LIMITS_MEMORY),
130-
CloudAllocationAttribute(name=QUOTA_LIMITS_EPHEMERAL_STORAGE_GB),
131-
CloudAllocationAttribute(name=QUOTA_REQUESTS_NESE_STORAGE),
132-
CloudAllocationAttribute(name=QUOTA_REQUESTS_IBM_STORAGE),
133-
CloudAllocationAttribute(name=QUOTA_REQUESTS_GPU),
134-
CloudAllocationAttribute(name=QUOTA_REQUESTS_VM_GPU_A100_SXM4),
135-
CloudAllocationAttribute(name=QUOTA_REQUESTS_VM_GPU_V100),
136-
CloudAllocationAttribute(name=QUOTA_REQUESTS_VM_GPU_H100),
137-
CloudAllocationAttribute(name=QUOTA_PVC),
138120
]

‎src/coldfront_plugin_cloud/base.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
11
import abc
22
import functools
3+
import json
34
from typing import NamedTuple
45

56
from coldfront.core.allocation import models as allocation_models
67
from coldfront.core.resource import models as resource_models
78

89
from coldfront_plugin_cloud import attributes
10+
from coldfront_plugin_cloud.models.quota_models import QuotaSpecs
911

1012

1113
class ResourceAllocator(abc.ABC):
@@ -25,6 +27,19 @@ def __init__(
2527
self.resource = resource
2628
self.allocation = allocation
2729

30+
try:
31+
resource_quota_attr = resource_models.ResourceAttribute.objects.get(
32+
resource=resource,
33+
resource_attribute_type__name=attributes.RESOURCE_QUOTA_RESOURCES,
34+
)
35+
self.resource_quotaspecs = QuotaSpecs.model_validate(
36+
json.loads(resource_quota_attr.value)
37+
)
38+
except resource_models.ResourceAttribute.DoesNotExist as e:
39+
raise ValueError(
40+
f"Resource {resource.name} does not have quota resources defined. Run either register_default_quotas or add_quota_to_resource management command to add quotas to the resource."
41+
) from e
42+
2843
def get_or_create_federated_user(self, username):
2944
if not (user := self.get_federated_user(username)):
3045
user = self.create_federated_user(username)

‎src/coldfront_plugin_cloud/esi.py‎

Lines changed: 0 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,7 @@
1-
from coldfront_plugin_cloud import attributes
21
from coldfront_plugin_cloud.openstack import OpenStackResourceAllocator
32

43

54
class ESIResourceAllocator(OpenStackResourceAllocator):
6-
QUOTA_KEY_MAPPING = {
7-
"network": {
8-
"keys": {
9-
attributes.QUOTA_FLOATING_IPS: "floatingip",
10-
attributes.QUOTA_NETWORKS: "network",
11-
}
12-
}
13-
}
14-
15-
QUOTA_KEY_MAPPING_ALL_KEYS = {
16-
quota_key: quota_name
17-
for k in QUOTA_KEY_MAPPING.values()
18-
for quota_key, quota_name in k["keys"].items()
19-
}
20-
215
resource_type = "esi"
226

237
def get_quota(self, project_id):

‎src/coldfront_plugin_cloud/management/commands/add_openshift_resource.py‎

Lines changed: 0 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -50,11 +50,6 @@ def add_arguments(self, parser):
5050
action="store_true",
5151
help="Indicates this is an OpenShift Virtualization resource (default: False)",
5252
)
53-
parser.add_argument(
54-
"--ibm-storage-available",
55-
action="store_true",
56-
help="Indicates that Ibm Scale storage is available in this resource (default: False)",
57-
)
5853

5954
def handle(self, *args, **options):
6055
self.validate_role(options["role"])
@@ -97,14 +92,6 @@ def handle(self, *args, **options):
9792
resource=openshift,
9893
value=options["role"],
9994
)
100-
101-
ResourceAttribute.objects.get_or_create(
102-
resource_attribute_type=ResourceAttributeType.objects.get(
103-
name=attributes.RESOURCE_IBM_AVAILABLE
104-
),
105-
resource=openshift,
106-
value="true" if options["ibm_storage_available"] else "false",
107-
)
10895
ResourceAttribute.objects.get_or_create(
10996
resource_attribute_type=ResourceAttributeType.objects.get(
11097
name=attributes.RESOURCE_CLUSTER_NAME

‎src/coldfront_plugin_cloud/management/commands/add_quota_to_resource.py‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
import json
2+
import logging
3+
4+
from django.core.management.base import BaseCommand
5+
from coldfront.core.resource.models import (
6+
Resource,
7+
ResourceAttribute,
8+
ResourceAttributeType,
9+
)
10+
from coldfront.core.allocation.models import AllocationAttributeType, AttributeType
11+
12+
from coldfront_plugin_cloud import attributes
13+
from coldfront_plugin_cloud.models.quota_models import QuotaSpecs, QuotaSpec
14+
15+
logger = logging.getLogger(__name__)
16+
17+
18+
class Command(BaseCommand):
19+
def add_arguments(self, parser):
20+
parser.add_argument(
21+
"--display-name",
22+
type=str,
23+
required=True,
24+
help="The display name for the quota attribute to add to the resource type.",
25+
)
26+
parser.add_argument(
27+
"--resource-name",
28+
type=str,
29+
required=True,
30+
help="The name of the resource to add the storage attribute to.",
31+
)
32+
parser.add_argument(
33+
"--quota-label",
34+
type=str,
35+
required=True,
36+
help="The cluster-side label for the quota.",
37+
)
38+
parser.add_argument(
39+
"--multiplier",
40+
type=int,
41+
default=0,
42+
help="Multiplier applied per SU quantity (int).",
43+
)
44+
parser.add_argument(
45+
"--static-quota",
46+
type=int,
47+
default=0,
48+
help="Static quota added to every SU quantity (int).",
49+
)
50+
parser.add_argument(
51+
"--unit-suffix",
52+
type=str,
53+
default="",
54+
help='Unit suffix to append to formatted quota values (e.g. "Gi").',
55+
)
56+
parser.add_argument(
57+
"--resource-type",
58+
type=str,
59+
default="",
60+
help="Indicates which resource type this quota is. Type `storage` is relevant for storage billing",
61+
)
62+
parser.add_argument(
63+
"--invoice-name",
64+
type=str,
65+
default="",
66+
help="Name of quota as it appears on invoice. Required if --resource-type is set to `storage`.",
67+
)
68+
69+
def handle(self, *args, **options):
70+
if options["resource_type"] == "storage" and not options["invoice_name"]:
71+
logger.error(
72+
"--invoice-name must be provided when resource type is `storage`."
73+
)
74+
return
75+
76+
resource_name = options["resource_name"]
77+
display_name = options["display_name"]
78+
new_quota_spec = QuotaSpec(**options)
79+
new_quota_dict = {display_name: new_quota_spec.model_dump()}
80+
QuotaSpecs.model_validate(new_quota_dict)
81+
82+
resource = Resource.objects.get(name=resource_name)
83+
available_quotas_attr, created = ResourceAttribute.objects.get_or_create(
84+
resource=resource,
85+
resource_attribute_type=ResourceAttributeType.objects.get(
86+
name=attributes.RESOURCE_QUOTA_RESOURCES
87+
),
88+
defaults={"value": json.dumps(new_quota_dict)},
89+
)
90+
91+
if not created:
92+
available_quotas_dict = json.loads(available_quotas_attr.value)
93+
available_quotas_dict.update(new_quota_dict)
94+
QuotaSpecs.model_validate(available_quotas_dict) # Validate uniqueness
95+
available_quotas_attr.value = json.dumps(available_quotas_dict)
96+
available_quotas_attr.save()
97+
98+
# Now create Allocation Attribute for this quota
99+
AllocationAttributeType.objects.get_or_create(
100+
name=display_name,
101+
defaults={
102+
"attribute_type": AttributeType.objects.get(name="Int"),
103+
"has_usage": False,
104+
"is_private": False,
105+
"is_changeable": True,
106+
},
107+
)
108+
109+
logger.info("Added quota '%s' to resource '%s'.", display_name, resource_name)

‎src/coldfront_plugin_cloud/management/commands/calculate_storage_gb_hours.py‎

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import csv
2+
import json
23
from decimal import Decimal, ROUND_HALF_UP
34
import dataclasses
45
from datetime import datetime, timedelta, timezone
@@ -7,6 +8,7 @@
78

89
from coldfront_plugin_cloud import attributes
910
from coldfront_plugin_cloud import utils
11+
from coldfront_plugin_cloud.models.quota_models import QuotaSpecs
1012

1113
import boto3
1214
from django.core.management.base import BaseCommand
@@ -19,6 +21,7 @@
1921
logger = logging.getLogger(__name__)
2022

2123
_RATES = None
24+
STORAGE_RESOURCE_TYPE_NAME = "storage"
2225

2326

2427
def get_rates():
@@ -210,6 +213,16 @@ def upload_to_s3(s3_endpoint, s3_bucket, file_location, invoice_month, end_time)
210213
def handle(self, *args, **options):
211214
generated_at = datetime.now(tz=timezone.utc).isoformat(timespec="seconds")
212215

216+
def get_storage_quotaspecs(allocation: Allocation):
217+
"""Get storage-related quota attributes for an allocation."""
218+
quotaspecs_dict = json.loads(
219+
allocation.resources.first().get_attribute(
220+
attributes.RESOURCE_QUOTA_RESOURCES
221+
)
222+
)
223+
quotaspecs = QuotaSpecs.model_validate(quotaspecs_dict)
224+
return quotaspecs.get_quotas_by_type(STORAGE_RESOURCE_TYPE_NAME)
225+
213226
def get_outages_for_service(cluster_name: str):
214227
"""Get outages for a service from nerc-rates.
215228
@@ -316,12 +329,14 @@ def process_invoice_row(allocation, attrs, su_name, rate):
316329
)
317330
logger.debug(f"Starting billing for allocation {allocation_str}.")
318331

319-
process_invoice_row(
320-
allocation,
321-
[attributes.QUOTA_VOLUMES_GB, attributes.QUOTA_OBJECT_GB],
322-
"OpenStack Storage",
323-
openstack_nese_storage_rate,
324-
)
332+
quotaspecs = get_storage_quotaspecs(allocation)
333+
for quota_name, quotaspec in quotaspecs.items():
334+
process_invoice_row(
335+
allocation,
336+
[quota_name],
337+
quotaspec.invoice_name,
338+
openstack_nese_storage_rate,
339+
)
325340

326341
for allocation in openshift_allocations:
327342
allocation_str = (

‎src/coldfront_plugin_cloud/management/commands/fetch_daily_billable_usage.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,8 @@
99

1010
from coldfront_plugin_cloud import attributes
1111
from coldfront.core.utils.common import import_from_settings
12-
from coldfront_plugin_cloud import usage_models
13-
from coldfront_plugin_cloud.usage_models import UsageInfo, validate_date_str
12+
from coldfront_plugin_cloud.models import usage_models
13+
from coldfront_plugin_cloud.models.usage_models import UsageInfo, validate_date_str
1414
from coldfront_plugin_cloud import utils
1515

1616
import boto3

0 commit comments

Comments
 (0)