Facility-level change to repo cascade

.github/workflows/ci-cd.yaml

@@ -15,16 +15,16 @@ jobs:
 
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v6
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
 
       - name: Install uv
-        uses: astral-sh/setup-uv@v7
+        uses: astral-sh/setup-uv@11f9893b081a58869d3b5fccaea48c9e9e46f990 # v8.3.2
         with:
           version: "0.11.7"
           enable-cache: true
 
       - name: Set up Python
-        uses: actions/setup-python@v6
+        uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
         with:
           python-version-file: "pyproject.toml"
 

modules/coact.py

@@ -1120,7 +1120,6 @@ def overaged(self, data: dict, threshold: float = 100.0) -> Iterator[OveragePoin
                         )
 
 
-
 # For backwards compatibility, allow running this module directly
 if __name__ == '__main__':
     coact(obj={})

modules/coactd.py

@@ -14,7 +14,7 @@
 from loguru import logger
 from enum import Enum
 from typing import Any, Optional, List
-from math import ceil
+from math import ceil, isclose
 from timeit import default_timer as timer
 from pathlib import Path
 
@@ -148,6 +148,8 @@ class Registration(GraphQlSubscriber, AnsibleRunner):
                     end
                     percentOfFacility
                     allocated
+                    burstPercentOfFacility
+                    burstAllocated
                 }
                 operationType
             }
@@ -405,7 +407,8 @@ class RepoRegistration(Registration):
         'RepoRemoveUser',
         'RepoChangeComputeRequirement',
         'RepoComputeAllocation',
-        'RepoUpdateFeature'
+        'RepoUpdateFeature',
+        'FacilityComputeAllocation'
     ]
 
     REPO_USERS_GQL = gql("""
@@ -415,14 +418,6 @@ class RepoRegistration(Registration):
         }
       }""")
 
-    COMPUTE_ALLOCATION_UPSERT_GQL = gql("""
-        mutation repoComputeAllocationUpsert($repo: RepoInput!, $repocompute: RepoComputeAllocationInput!, $qosinputs: [QosInput!]!) {
-            repoComputeAllocationUpsert(repo: $repo, repocompute: $repocompute, qosinputs: $qosinputs) {
-                Id
-            }
-        }
-        """)
-
     REPO_CURRENT_COMPUTE_REQUIREMENT_GQL = gql("""
         query repo( $repo: RepoInput! ) {
           repo(filter: $repo) {
@@ -442,6 +437,8 @@ class RepoRegistration(Registration):
               start
               end
               percentOfFacility
+              burstPercentOfFacility
+              burstAllocated
               cpus: allocatedCpusCount
               memory: allocatedMemGb
               nodes: allocatedNodesCount
@@ -451,13 +448,30 @@ class RepoRegistration(Registration):
         }
         """)
 
+    REPOS_WITH_ALLOCATIONS_GQL = gql("""
+        query reposWithAllocations($facilityName: String!) {
+            repos(filter: {facility: $facilityName}) {
+                Id
+                name
+                facility
+                currentComputeAllocations {
+                    Id
+                    clustername
+                    percentOfFacility
+                    burstPercentOfFacility
+                    burstAllocated
+                    start
+                    end
+                    allocatedNodesCount
+                }
+            }
+        }
+    """)
+
     FACILITY_CURRENT_COMPUTE_CGL = gql("""
         query facility( $facility: String! ) {
           facility(filter: {name: $facility}) {
             name
-            computeallocations {
-              clustername
-            }
             computepurchases {
               clustername
               purchased
@@ -505,7 +519,10 @@ def do(self, req_id, op_type, req_type, approval, req, dry_run):
                 if end is not None:
                     end = pdl.parse(end, timezone='UTC')
                 return self.do_repo_compute_allocation(
-                    repo, facility, clustername, percent, allocated, start, end, dry_run=dry_run
+                    repo, facility, clustername, percent, allocated, start, end,
+                    burst_percent=req.get('burstPercentOfFacility', None),
+                    burst_allocated=req.get('burstAllocated', None),
+                    dry_run=dry_run
                 )
 
             elif req_type == 'RepoChangeComputeRequirement':
@@ -516,6 +533,13 @@ def do(self, req_id, op_type, req_type, approval, req, dry_run):
             elif req_type == 'RepoUpdateFeature':
                 return self.do_feature(repo, facility)
 
+            elif req_type == 'FacilityComputeAllocation':
+                clustername = req.get('clustername', None)
+                assert facility and clustername
+                return self.do_facility_compute_allocation_cascade(
+                    facility, clustername, dry_run=dry_run
+                )
+
         return None
 
     def do_new_repo(
@@ -752,6 +776,9 @@ def upsert_repo_compute_allocation(
         allocated_resource: float,
         start: pdl.DateTime,
         end: Optional[str],
+        burst_percent: float = 0.0,
+        burst_allocated: float = 0.0,
+        allocation_id: Optional[str] = None,
         default_end_delta=None,
         dry_run: bool = False
     ):
@@ -776,10 +803,16 @@ def format_datetime(iso, round_off=None):
                 'clustername': cluster,
                 'percentOfFacility': percent,
                 'allocated': allocated_resource,
+                'burstPercentOfFacility': burst_percent,
+                'burstAllocated': burst_allocated,
                 'start': format_datetime(start),
                 'end': format_datetime(end)
             },
         }
+        # repoComputeAllocationUpsert replaces the whole document; without the id it matches
+        # on (repoid, clustername, start) and would insert a duplicate on any datetime skew.
+        if allocation_id:
+            compute_allocation_req['repocompute']['Id'] = allocation_id
         self.logger.info(f'upserting {compute_allocation_req}')
         REPO_COMPUTE_ALLOCATION_UPSERT_GQL = gql("""
             mutation repo( $repo: RepoInput!, $repocompute: RepoComputeAllocationInput! ) {
@@ -817,9 +850,18 @@ def do_repo_compute_allocation(
         allocated_resource: float,
         start: pdl.DateTime,
         end: Optional[str],
+        burst_percent: Optional[float] = None,
+        burst_allocated: Optional[float] = None,
+        allocation_id: Optional[str] = None,
         dry_run: bool = False
     ):
-        """Does all the necessary tasks to setup a new or existing Repo."""
+        """Does all the necessary tasks to setup a new or existing Repo.
+
+        burst_percent/burst_allocated default to None meaning "keep whatever is stored";
+        the upsert replaces the whole document, so they must always be sent explicitly.
+        allocation_id targets an existing allocation row; without it the upsert matches on
+        (repoid, clustername, start) and so may open a new allocation period.
+        """
         self.logger.info(f"set repo compute allocation {facility}:{repo} at {cluster} to {percent}% ({allocated_resource} nodes) between {start} - {end}")
 
         def _get_allocation_info():
@@ -858,9 +900,22 @@ def _get_allocation_info():
             )
             return True
         else:
+            if burst_percent is None or burst_allocated is None:
+                existing = next(
+                    (a for a in repo_obj['currentComputeAllocations'] if a.get('clustername') == cluster),
+                    {}
+                )
+                if burst_percent is None:
+                    burst_percent = existing.get('burstPercentOfFacility') or 0.0
+                if burst_allocated is None:
+                    burst_allocated = existing.get('burstAllocated') or 0.0
+
             # upsert the record
             resp = self.upsert_repo_compute_allocation(
-                repo_obj['Id'], cluster, percent, allocated_resource, start, end
+                repo_obj['Id'], cluster, percent, allocated_resource, start, end,
+                burst_percent=burst_percent,
+                burst_allocated=burst_allocated,
+                allocation_id=allocation_id,
             )
 
             # fetch it again to obtain the correct resources with the new percentage
@@ -1076,6 +1131,135 @@ def do_compute_requirement(self, repo: str, facility: str, requirement: str, pla
     def do_feature(self, repo, facility, dry_run: bool = False) -> bool:
         raise NotImplementedError("do_feature not yet implemented")
 
+    def do_facility_compute_allocation_cascade(
+        self,
+        facility: str,
+        clustername: str,
+        dry_run: bool = False
+    ) -> bool:
+        """
+        Handle facility-level compute allocation changes by updating all affected repo allocations.
+        Queries the facility record directly for the current purchased node count.
+
+        Each repo keeps its percentage share of the facility, so absolute node counts move
+        with the purchase. Only allocations that are current (start <= now < end) are
+        cascaded; future-dated allocations keep the values they were created with.
+        """
+        # Fetch current purchased nodes from the facility record
+        facility_resp = self.back_channel.execute(
+            self.FACILITY_CURRENT_COMPUTE_CGL,
+            {'facility': facility}
+        )
+        fac_data = facility_resp.get('facility', {})
+        new_purchased = None
+        for cp in fac_data.get('computepurchases', []):
+            if cp['clustername'].lower() == clustername.lower():
+                new_purchased = cp['purchased']
+                break
+
+        if new_purchased is None:
+            raise RuntimeError(f"No purchase record found for {facility}@{clustername} - cannot cascade")
+
+        if new_purchased < 0:
+            raise RuntimeError(f"Invalid purchased nodes: {new_purchased} for {facility}@{clustername} - cannot cascade")
+
+        self.logger.info(
+            f"Processing facility compute allocation cascade: {facility}@{clustername} "
+            f"-> {new_purchased} purchased nodes"
+        )
+
+        # Get all repositories with allocations on this facility/cluster
+        repos_resp = self.back_channel.execute(
+            self.REPOS_WITH_ALLOCATIONS_GQL,
+            {'facilityName': facility}
+        )
+
+        affected_repos = []
+        for repo in repos_resp['repos']:
+            for allocation in repo['currentComputeAllocations']:
+                if allocation['clustername'].lower() == clustername.lower():
+                    affected_repos.append({
+                        'repo': repo,
+                        'allocation': allocation
+                    })
+
+        self.logger.info(f"Found {len(affected_repos)} repo allocations to update on {facility}@{clustername}")
+
+        total_percent = sum(item['allocation']['percentOfFacility'] or 0.0 for item in affected_repos)
+        if total_percent > 100.0:
+            self.logger.warning(
+                f"{facility}@{clustername} is oversubscribed: repo allocations total "
+                f"{total_percent}% of the facility. Node counts are rounded up per repo, "
+                f"so the sum may exceed the {new_purchased} purchased nodes."
+            )
+
+        update_count = 0
+        skipped_count = 0
+        failures = []
+        for item in affected_repos:
+            repo = item['repo']
+            allocation = item['allocation']
+
+            # Calculate new node allocation maintaining the same percentage
+            percent_of_facility = allocation['percentOfFacility']
+            burst_percent = allocation.get('burstPercentOfFacility') or 0.0
+            new_allocated_nodes = (percent_of_facility / 100.0) * new_purchased
+            new_burst_allocated = (burst_percent / 100.0) * new_purchased
+
+            current_allocated = allocation.get('allocatedNodesCount') or 0.0
+            current_burst = allocation.get('burstAllocated') or 0.0
+            if (isclose(current_allocated, new_allocated_nodes)
+                    and isclose(current_burst, new_burst_allocated)):
+                # Each update runs two Ansible playbooks, so do not touch untouched repos.
+                self.logger.debug(
+                    f"Skipping {repo['facility']}:{repo['name']} on {clustername}: "
+                    f"already at {new_allocated_nodes:.2f} nodes"
+                )
+                skipped_count += 1
+                continue
+
+            self.logger.info(
+                f"Updating {repo['facility']}:{repo['name']} on {clustername}: "
+                f"{percent_of_facility}% -> {new_allocated_nodes:.2f} nodes (was {current_allocated}), "
+                f"burst {burst_percent}% -> {new_burst_allocated:.2f} nodes"
+            )
+
+            try:
+                self.do_repo_compute_allocation(
+                    repo['name'],
+                    repo['facility'],
+                    clustername,
+                    percent_of_facility,
+                    new_allocated_nodes,
+                    pdl.parse(allocation['start'], timezone='UTC'),
+                    pdl.parse(allocation['end'], timezone='UTC') if allocation['end'] else None,
+                    burst_percent=burst_percent,
+                    burst_allocated=new_burst_allocated,
+                    allocation_id=allocation['Id'],
+                    dry_run=dry_run,
+                )
+                update_count += 1
+                self.logger.info(f"Successfully updated {repo['facility']}:{repo['name']} allocation")
+
+            except Exception as e:
+                # Keep going so one bad repo does not strand the rest, but fail the request below.
+                self.logger.error(f"Failed to update {repo['facility']}:{repo['name']}: {e}")
+                failures.append(f"{repo['facility']}:{repo['name']} ({e})")
+
+        self.logger.info(
+            f"Facility cascade update completed on {facility}@{clustername}: "
+            f"{update_count} updated, {skipped_count} unchanged, {len(failures)} failed "
+            f"of {len(affected_repos)} repo allocations"
+        )
+
+        if failures:
+            raise RuntimeError(
+                f"Facility compute allocation cascade for {facility}@{clustername} failed for "
+                f"{len(failures)}/{len(affected_repos)} repos: {'; '.join(failures)}"
+            )
+
+        return True
+
 
 @coactd.command(name='reporegistration')
 @common_options
@@ -1091,7 +1275,7 @@ def repo_registration(ctx, verbose, username, password_file, client_name, dry_ru
     """Workflow for repository maintenance.
 
     Handles NewRepo, RepoMembership, RepoRemoveUser, RepoChangeComputeRequirement,
-    RepoComputeAllocation, and RepoUpdateFeature request types.
+    RepoComputeAllocation, RepoUpdateFeature, and FacilityComputeAllocation request types.
     """
     configure_logging_from_verbose(verbose)
     ctx.obj['verbose'] = verbose

tests/conftest.py

@@ -0,0 +1,15 @@
+"""
+pytest configuration for the CLI test suite.
+
+ansible_runner imports pkg_resources at the top level, which is a
+setuptools utility not present in the uv-managed test environment. We stub it
+out here, before any test module imports modules.coactd, so the module loads
+cleanly without requiring the full ansible/setuptools stack at test time.
+"""
+import sys
+from unittest.mock import MagicMock
+
+if "pkg_resources" not in sys.modules:
+    sys.modules["pkg_resources"] = MagicMock()
+if "ansible_runner" not in sys.modules:
+    sys.modules["ansible_runner"] = MagicMock()