Skip to content

Commit 35273aa

Browse files
QuanMPhmQuanMPhm
committed
Allow dynamic quota creation and removal
1 parent 7054761 commit 35273aa

25 files changed

+649
-385
lines changed

‎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: 10 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,14 @@ def __init__(
2527
self.resource = resource
2628
self.allocation = allocation
2729

30+
resource_storage_classes_attr = resource_models.ResourceAttribute.objects.get(
31+
resource=resource,
32+
resource_attribute_type__name=attributes.RESOURCE_QUOTA_RESOURCES,
33+
)
34+
self.resource_quotaspecs = QuotaSpecs.model_validate(
35+
json.loads(resource_storage_classes_attr.value)
36+
)
37+
2838
def get_or_create_federated_user(self, username):
2939
if not (user := self.get_federated_user(username)):
3040
user = self.create_federated_user(username)

‎src/coldfront_plugin_cloud/esi.py‎

Lines changed: 2 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -3,21 +3,9 @@
33

44

55
class ESIResourceAllocator(OpenStackResourceAllocator):
6-
QUOTA_KEY_MAPPING = {
7-
"network": {
8-
"keys": {
9-
attributes.QUOTA_FLOATING_IPS: "floatingip",
10-
attributes.QUOTA_NETWORKS: "network",
11-
}
12-
}
6+
SERVICE_QUOTA_MAPPING = {
7+
"network": [attributes.QUOTA_FLOATING_IPS, attributes.QUOTA_NETWORKS],
138
}
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-
219
resource_type = "esi"
2210

2311
def get_quota(self, project_id):

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

Lines changed: 37 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from django.core.management.base import BaseCommand
2+
from django.core.management import call_command
23

34
from coldfront.core.resource.models import (
45
Resource,
@@ -50,11 +51,6 @@ def add_arguments(self, parser):
5051
action="store_true",
5152
help="Indicates this is an OpenShift Virtualization resource (default: False)",
5253
)
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-
)
5854

5955
def handle(self, *args, **options):
6056
self.validate_role(options["role"])
@@ -97,14 +93,6 @@ def handle(self, *args, **options):
9793
resource=openshift,
9894
value=options["role"],
9995
)
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-
)
10896
ResourceAttribute.objects.get_or_create(
10997
resource_attribute_type=ResourceAttributeType.objects.get(
11098
name=attributes.RESOURCE_CLUSTER_NAME
@@ -114,3 +102,39 @@ def handle(self, *args, **options):
114102
if options["internal_name"]
115103
else options["name"],
116104
)
105+
106+
# Add common Openshift resources (cpu, memory, etc)
107+
call_command(
108+
"add_quota_to_resource",
109+
display_name=attributes.QUOTA_LIMITS_CPU,
110+
default_quota=1,
111+
resource_name=options["name"],
112+
quota_label="limits.cpu",
113+
multiplier=1,
114+
)
115+
call_command(
116+
"add_quota_to_resource",
117+
display_name=attributes.QUOTA_LIMITS_MEMORY,
118+
default_quota=4096,
119+
resource_name=options["name"],
120+
quota_label="limits.memory",
121+
multiplier=4096,
122+
unit_suffix="Mi",
123+
)
124+
call_command(
125+
"add_quota_to_resource",
126+
display_name=attributes.QUOTA_LIMITS_EPHEMERAL_STORAGE_GB,
127+
default_quota=5,
128+
resource_name=options["name"],
129+
quota_label="limits.ephemeral-storage",
130+
multiplier=5,
131+
unit_suffix="Gi",
132+
)
133+
call_command(
134+
"add_quota_to_resource",
135+
display_name=attributes.QUOTA_PVC,
136+
default_quota=2,
137+
resource_name=options["name"],
138+
quota_label="persistentvolumeclaims",
139+
multiplier=2,
140+
)

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

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
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+
"--default-quota",
28+
type=int,
29+
required=True,
30+
help="The default quota value for the storage attribute. In GB",
31+
)
32+
parser.add_argument(
33+
"--resource_name",
34+
type=str,
35+
required=True,
36+
help="The name of the resource to add the storage attribute to.",
37+
)
38+
parser.add_argument(
39+
"--quota-label",
40+
dest="quota_label",
41+
type=str,
42+
required=True,
43+
help="Human-readable quota_label for this quota (must be unique).",
44+
)
45+
parser.add_argument(
46+
"--multiplier",
47+
dest="multiplier",
48+
type=int,
49+
default=0,
50+
help="Multiplier applied per SU quantity (int).",
51+
)
52+
parser.add_argument(
53+
"--static-quota",
54+
dest="static_quota",
55+
type=int,
56+
default=0,
57+
help="Static quota added to every SU quantity (int).",
58+
)
59+
parser.add_argument(
60+
"--unit-suffix",
61+
dest="unit_suffix",
62+
type=str,
63+
default="",
64+
help='Unit suffix to append to formatted quota values (e.g. "Gi").',
65+
)
66+
parser.add_argument(
67+
"--resource-type",
68+
type=str,
69+
default="",
70+
help="Indicates which resource type this quota is. Type `storage` is relevant for storage billing",
71+
)
72+
parser.add_argument(
73+
"--invoice-name",
74+
type=str,
75+
default="",
76+
help="Name of quota as it appears on invoice. Required if --is-storage-type is set.",
77+
)
78+
79+
def handle(self, *args, **options):
80+
if options["resource_type"] == "storage" and not options["invoice_name"]:
81+
logger.error(
82+
"--invoice-name must be provided when storage type is `storage`."
83+
)
84+
85+
resource_name = options["resource_name"]
86+
display_name = options["display_name"]
87+
new_quota_spec = QuotaSpec(**options)
88+
new_quota_dict = {display_name: new_quota_spec.model_dump()}
89+
QuotaSpecs.model_validate(new_quota_dict)
90+
91+
resource = Resource.objects.get(name=resource_name)
92+
available_quotas_attr, created = ResourceAttribute.objects.get_or_create(
93+
resource=resource,
94+
resource_attribute_type=ResourceAttributeType.objects.get(
95+
name=attributes.RESOURCE_QUOTA_RESOURCES
96+
),
97+
defaults={"value": json.dumps(new_quota_dict)},
98+
)
99+
100+
# TODO (Quan): Dict update allows migration of existing quotas. This is fine?
101+
if not created:
102+
available_quotas_dict = json.loads(available_quotas_attr.value)
103+
available_quotas_dict.update(new_quota_dict)
104+
QuotaSpecs.model_validate(available_quotas_dict) # Validate uniqueness
105+
available_quotas_attr.value = json.dumps(available_quotas_dict)
106+
available_quotas_attr.save()
107+
108+
# Now create Allocation Attribute for this quota
109+
AllocationAttributeType.objects.get_or_create(
110+
name=display_name,
111+
defaults={
112+
"attribute_type": AttributeType.objects.get(name="Int"),
113+
"has_usage": False,
114+
"is_private": False,
115+
"is_changeable": True,
116+
},
117+
)
118+
119+
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
@@ -210,6 +212,16 @@ def upload_to_s3(s3_endpoint, s3_bucket, file_location, invoice_month, end_time)
210212
def handle(self, *args, **options):
211213
generated_at = datetime.now(tz=timezone.utc).isoformat(timespec="seconds")
212214

215+
def get_storage_quotaspecs(allocation: Allocation):
216+
"""Get storage-related quota attributes for an allocation."""
217+
quotaspecs_dict = json.loads(
218+
allocation.resources.first().get_attribute(
219+
attributes.RESOURCE_QUOTA_RESOURCES
220+
)
221+
)
222+
quotaspecs = QuotaSpecs.model_validate(quotaspecs_dict)
223+
return quotaspecs.storage_quotas
224+
213225
def get_outages_for_service(cluster_name: str):
214226
"""Get outages for a service from nerc-rates.
215227
@@ -316,12 +328,15 @@ def process_invoice_row(allocation, attrs, su_name, rate):
316328
)
317329
logger.debug(f"Starting billing for allocation {allocation_str}.")
318330

319-
process_invoice_row(
320-
allocation,
321-
[attributes.QUOTA_VOLUMES_GB, attributes.QUOTA_OBJECT_GB],
322-
"OpenStack Storage",
323-
openstack_nese_storage_rate,
324-
)
331+
# TODO (Quan): An illustration of how billing could be simplified. Shuold I follow with this?
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)