Integrate nerc-rates for dynamic outage information loading

requirements.txt

@@ -1,4 +1,4 @@
-git+https://github.com/CCI-MOC/nerc-rates@33701ed#egg=nerc_rates
+git+https://github.com/CCI-MOC/nerc-rates@5569bba#egg=nerc_rates
 boto3
 kubernetes
 openshift

src/coldfront_plugin_cloud/management/commands/calculate_storage_gb_hours.py

@@ -20,6 +20,12 @@
 
 _RATES = None
 
+RESOURCE_NAME_TO_NERC_SERVICE = {
+    "NERC": "stack",
+    "NERC-OCP": "ocp-prod",
+    "NERC-OCP-EDU": "academic",
+}
+
 
 def get_rates():
     # nerc-rates doesn't work with Python 3.9, which is what ColdFront is currently
@@ -142,13 +148,6 @@ def add_arguments(self, parser):
             action="store_true",
             help="Upload generated CSV invoice to S3 storage.",
         )
-        parser.add_argument(
-            "--excluded-time-ranges",
-            type=str,
-            default=None,
-            nargs="+",
-            help="List of time ranges excluded from billing, in ISO format.",
-        )
 
     @staticmethod
     def default_start_argument():
@@ -198,8 +197,24 @@ def upload_to_s3(s3_endpoint, s3_bucket, file_location, invoice_month):
         logger.info(f"Uploaded to {secondary_location}.")
 
     def handle(self, *args, **options):
+        def get_outages_for_service(resource_name: str):
+            """Get outages for a service from nerc-rates.
+
+            :param resource_name: Name of the resource to get outages for.
+            :return: List of excluded intervals or None.
+            """
+            service_name = RESOURCE_NAME_TO_NERC_SERVICE.get(resource_name)
+            if service_name:
+                return utils.load_outages_from_nerc_rates(
+                    options["start"], options["end"], service_name
+                )
+            return None
+
         def process_invoice_row(allocation, attrs, su_name, rate):
             """Calculate the value and write the bill using the writer."""
+            resource_name = allocation.resources.first().name
+            excluded_intervals_list = get_outages_for_service(resource_name)
+
             time = 0
             for attribute in attrs:
                 time += utils.calculate_quota_unit_hours(
@@ -234,13 +249,6 @@ def process_invoice_row(allocation, attrs, su_name, rate):
         logger.info(f"Processing invoices for {options['invoice_month']}.")
         logger.info(f"Interval {options['start'] - options['end']}.")
 
-        if options["excluded_time_ranges"]:
-            excluded_intervals_list = utils.load_excluded_intervals(
-                options["excluded_time_ranges"]
-            )
-        else:
-            excluded_intervals_list = None
-
         openstack_resources = Resource.objects.filter(
             resource_type=ResourceType.objects.get(name="OpenStack")
         )
@@ -285,15 +293,13 @@ def process_invoice_row(allocation, attrs, su_name, rate):
             csv_invoice_writer = csv.writer(
                 f, delimiter=",", quotechar="|", quoting=csv.QUOTE_MINIMAL
             )
-            # Write Headers
             csv_invoice_writer.writerow(InvoiceRow.get_headers())
 
             for allocation in openstack_allocations:
                 allocation_str = (
                     f'{allocation.pk} of project "{allocation.project.title}"'
                 )
-                msg = f"Starting billing for allocation {allocation_str}."
-                logger.debug(msg)
+                logger.debug(f"Starting billing for allocation {allocation_str}.")
 
                 process_invoice_row(
                     allocation,
@@ -306,8 +312,7 @@ def process_invoice_row(allocation, attrs, su_name, rate):
                 allocation_str = (
                     f'{allocation.pk} of project "{allocation.project.title}"'
                 )
-                msg = f"Starting billing for allocation {allocation_str}."
-                logger.debug(msg)
+                logger.debug(f"Starting billing for allocation {allocation_str}.")
 
                 process_invoice_row(
                     allocation,

src/coldfront_plugin_cloud/tests/unit/test_calculate_quota_unit_hours.py

@@ -1,6 +1,7 @@
 import datetime
 import pytz
 import tempfile
+from unittest.mock import patch
 
 import freezegun
 
@@ -16,7 +17,11 @@
 
 
 class TestCalculateAllocationQuotaHours(base.TestBase):
-    def test_new_allocation_quota(self):
+    @patch("coldfront_plugin_cloud.utils.load_outages_from_nerc_rates")
+    def test_new_allocation_quota(self, mock_load_outages):
+        """Test quota calculation with nerc-rates outages mocked."""
+        mock_load_outages.return_value = []
+
         self.resource = self.new_openshift_resource(
             name="",
         )
@@ -63,7 +68,7 @@ def test_new_allocation_quota(self):
                 "2020-03",
             )
 
-        # Let's test a complete CLI call including excluded time, while we're at it. This is not for testing
+        # Let's test a complete CLI call. This is not for testing
         # the validity but just the unerrored execution of the complete pipeline.
         # Tests that verify the correct output are further down in the test file.
         with tempfile.NamedTemporaryFile() as fp:
@@ -83,10 +88,12 @@ def test_new_allocation_quota(self):
                 "0.00001",
                 "--invoice-month",
                 "2020-03",
-                "--excluded-time-ranges",
-                "2020-03-02 00:00:00,2020-03-03 05:00:00",
             )
 
+        # Verify that load_outages_from_nerc_rates is not called when resource name
+        # doesn't match NERC service mapping
+        mock_load_outages.assert_not_called()
+
     def test_new_allocation_quota_expired(self):
         """Test that expiration doesn't affect invoicing."""
         self.resource = self.new_openshift_resource(

src/coldfront_plugin_cloud/utils.py

@@ -1,4 +1,5 @@
 import datetime
+import functools
 import math
 import pytz
 import re
@@ -14,6 +15,9 @@
 
 from coldfront_plugin_cloud import attributes
 
+# Load outages data once per program execution
+_OUTAGES_DATA = None
+
 
 def env_safe_name(name):
     return re.sub(r"[^A-Za-z0-9]", "_", str(name)).upper()
@@ -191,6 +195,12 @@ def calculate_quota_unit_hours(
 
 
 def load_excluded_intervals(excluded_interval_arglist):
+    """Parse excluded time ranges from command line arguments.
+
+    :param excluded_interval_arglist: List of time range strings in format "start,end".
+    :return: Sorted list of [start, end] datetime tuples.
+    """
+
     def interval_sort_key(e):
         return e[0]
 
@@ -223,6 +233,27 @@ def check_overlapping_intervals(excluded_intervals_list):
     return excluded_intervals_list
 
 
+@functools.cache
+def load_outages_from_nerc_rates(
+    start: datetime.datetime, end: datetime.datetime, affected_service: str
+) -> list[tuple[datetime.datetime, datetime.datetime]]:
+    """Load outage intervals from nerc-rates for a given time period and service.
+
+    :param start: Start time for outage search.
+    :param end: End time for outage search.
+    :param affected_service: Name of the affected service (e.g., "stack", "ocp-prod").
+    :return: List of [start, end] datetime tuples representing outages.
+    """
+    global _OUTAGES_DATA
+    if _OUTAGES_DATA is None:
+        from nerc_rates import outages
+
+        _OUTAGES_DATA = outages.load_from_url()
+    return _OUTAGES_DATA.get_outages_during(
+        start.isoformat(), end.isoformat(), affected_service
+    )
+
+
 def _clamp_time(time, min_time, max_time):
     if time < min_time:
         time = min_time