From da838a2a9aa12a53bf6004c2094c9e2e4a2c766e Mon Sep 17 00:00:00 2001 From: Ryan Waldheim Date: Wed, 6 May 2026 14:50:59 -0700 Subject: [PATCH 01/11] feat: add facility management commands and compute allocation handling --- ansible-runner/project | 2 +- modules/coact.py | 218 ++++++++++++++++++++++ modules/coactd.py | 170 ++++++++++++++++- tests/conftest.py | 15 ++ tests/test_facility_compute_allocation.py | 97 ++++++++++ 5 files changed, 499 insertions(+), 3 deletions(-) create mode 100644 tests/conftest.py create mode 100644 tests/test_facility_compute_allocation.py diff --git a/ansible-runner/project b/ansible-runner/project index 7f18b37..01bd93a 160000 --- a/ansible-runner/project +++ b/ansible-runner/project @@ -1 +1 @@ -Subproject commit 7f18b37e90139968561017baf0debb0a18615fe5 +Subproject commit 01bd93a20b87022c6acf525cd0a85cb7ef2e274b diff --git a/modules/coact.py b/modules/coact.py index d20b3bf..bd44cfd 100644 --- a/modules/coact.py +++ b/modules/coact.py @@ -1077,6 +1077,224 @@ def overaged(self, data: dict, threshold: float = 100.0) -> Iterator[OveragePoin +# ============================================================================ +# Facility Management Commands +# ============================================================================ + +@click.group(name='facility', help="Facility compute management commands", context_settings=CONTEXT_SETTINGS) +@click.pass_context +def facility(ctx): + """Facility management command group for compute allocation operations.""" + ctx.ensure_object(dict) + + +class FacilityManager(GraphQlMixin): + """Handles facility compute management operations.""" + + FACILITY_UPDATE_PURCHASED_GQL = gql(""" + mutation facilityUpdatePurchased( + $facilityName: String!, + $clusterName: String!, + $purchased: Int! + ) { + facilityUpdateComputePurchase( + facility: $facilityName, + cluster: $clusterName, + purchased: $purchased + ) { + name + computepurchases { + clustername + purchased + } + } + } + """) + + FACILITY_QUERY_GQL = gql(""" + query facility($facilityName: String!) { + facility(filter: {name: $facilityName}) { + name + computepurchases { + clustername + purchased + } + } + } + """) + + REPOS_WITH_ALLOCATIONS_GQL = gql(""" + query reposWithAllocations($facilityName: String!, $clusterName: String!) { + repos(filter: {facility: $facilityName}) { + Id + name + facility + currentComputeAllocations { + Id + clustername + percentOfFacility + allocatedNodesCount + } + } + } + """) + + def __init__(self, username: str, password_file: str): + self.username = username + self.password_file = password_file + + def update_purchased_nodes(self, facility: str, cluster: str, nodes: int, dry_run: bool = False) -> bool: + """Update the purchased node count for a facility/cluster.""" + self.back_channel = self.connect_graph_ql( + username=self.username, + password_file=self.password_file, + timeout=60 + ) + + logger.info(f"Updating {facility}@{cluster} to {nodes} purchased nodes (dry_run={dry_run})") + + if not dry_run: + try: + result = self.back_channel.execute( + self.FACILITY_UPDATE_PURCHASED_GQL, + { + 'facilityName': facility, + 'clusterName': cluster, + 'purchased': nodes + } + ) + logger.info(f"Successfully updated facility: {result}") + return True + except Exception as e: + logger.error(f"Failed to update facility: {e}") + return False + else: + logger.info(f"DRY RUN: Would update {facility}@{cluster} to {nodes} nodes") + return True + + def simulate_change(self, facility: str, cluster: str, nodes: int) -> bool: + """Simulate facility change and show what repos would be affected.""" + self.back_channel = self.connect_graph_ql( + username=self.username, + password_file=self.password_file, + timeout=60 + ) + + logger.info(f"Simulating facility change: {facility}@{cluster} -> {nodes} nodes") + + try: + # Get current facility state + current_result = self.back_channel.execute( + self.FACILITY_QUERY_GQL, + {'facilityName': facility} + ) + + current_facility = current_result.get('facility') + if not current_facility: + logger.error(f"Facility '{facility}' not found") + return False + + current_purchased = None + for purchase in current_facility.get('computepurchases', []): + if purchase['clustername'].lower() == cluster.lower(): + current_purchased = purchase['purchased'] + break + + if current_purchased is None: + logger.error(f"Cluster '{cluster}' not found in facility '{facility}'") + return False + + logger.info(f"Current purchased nodes: {current_purchased}") + logger.info(f"Proposed purchased nodes: {nodes}") + + if current_purchased == nodes: + logger.info("No change detected - no repos would be affected") + return True + + # Get affected repos + repos_result = self.back_channel.execute( + self.REPOS_WITH_ALLOCATIONS_GQL, + {'facilityName': facility, 'clusterName': cluster} + ) + + affected_count = 0 + for repo in repos_result['repos']: + for allocation in repo['currentComputeAllocations']: + if allocation['clustername'].lower() == cluster.lower(): + current_nodes = allocation['allocatedNodesCount'] + percent = allocation['percentOfFacility'] + new_nodes = (percent / 100.0) * nodes + + logger.info( + f" {repo['facility']}:{repo['name']} - " + f"{percent}% -> {current_nodes} nodes would become {new_nodes:.2f} nodes" + ) + affected_count += 1 + + logger.info(f"Total affected repo allocations: {affected_count}") + return True + + except Exception as e: + logger.error(f"Simulation failed: {e}") + return False + + +@facility.command(name='update-purchased') +@common_options +@graphql_options +@click.option('--facility', required=True, help='Facility name (e.g., shared, lcls)') +@click.option('--cluster', required=True, help='Cluster name (e.g., roma, ampere)') +@click.option('--nodes', required=True, type=int, help='New number of purchased nodes') +@click.option('--dry-run', is_flag=True, default=False, help='Show what would be done without making changes') +@click.pass_context +def facility_update_purchased(ctx, verbose, username, password_file, facility, cluster, nodes, dry_run): + """Update the number of purchased nodes for a facility/cluster. + + This will trigger automatic cascade updates of all repo allocations on this facility/cluster + if facility monitoring is enabled. + """ + configure_logging_from_verbose(verbose) + ctx.obj['verbose'] = verbose + + manager = FacilityManager( + username=username, + password_file=password_file + ) + + success = manager.update_purchased_nodes(facility, cluster, nodes, dry_run) + if not success: + raise click.ClickException("Failed to update facility purchased nodes") + + +@facility.command(name='simulate-change') +@common_options +@graphql_options +@click.option('--facility', required=True, help='Facility name (e.g., shared, lcls)') +@click.option('--cluster', required=True, help='Cluster name (e.g., roma, ampere)') +@click.option('--nodes', required=True, type=int, help='Proposed number of purchased nodes') +@click.pass_context +def facility_simulate_change(ctx, verbose, username, password_file, facility, cluster, nodes): + """Simulate a facility compute change and show what repo allocations would be affected. + + This is useful for understanding the impact of facility changes before applying them. + """ + configure_logging_from_verbose(verbose) + ctx.obj['verbose'] = verbose + + manager = FacilityManager( + username=username, + password_file=password_file + ) + + success = manager.simulate_change(facility, cluster, nodes) + if not success: + raise click.ClickException("Simulation failed") + + +# Add facility to main coact group +coact.add_command(facility) + + # For backwards compatibility, allow running this module directly if __name__ == '__main__': coact(obj={}) diff --git a/modules/coactd.py b/modules/coactd.py index 94f226c..84e9875 100644 --- a/modules/coactd.py +++ b/modules/coactd.py @@ -148,6 +148,9 @@ class Registration(GraphQlSubscriber, AnsibleRunner): end percentOfFacility allocated + oldPurchased + newPurchased + updateStrategy } operationType } @@ -404,7 +407,8 @@ class RepoRegistration(Registration): 'RepoRemoveUser', 'RepoChangeComputeRequirement', 'RepoComputeAllocation', - 'RepoUpdateFeature' + 'RepoUpdateFeature', + 'FacilityComputeAllocation' ] REPO_USERS_GQL = gql(""" @@ -450,6 +454,24 @@ class RepoRegistration(Registration): } """) + REPOS_WITH_ALLOCATIONS_GQL = gql(""" + query reposWithAllocations($facilityName: String!) { + repos(filter: {facility: $facilityName}) { + Id + name + facility + currentComputeAllocations { + Id + clustername + percentOfFacility + start + end + allocatedNodesCount + } + } + } + """) + FACILITY_CURRENT_COMPUTE_CGL = gql(""" query facility( $facility: String! ) { facility(filter: {name: $facility}) { @@ -515,6 +537,16 @@ 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) + old_purchased = req.get('oldPurchased', None) + new_purchased = req.get('newPurchased', None) + update_strategy = req.get('updateStrategy', 'proportional') + assert facility and clustername and old_purchased is not None and new_purchased is not None + return self.do_facility_compute_allocation_cascade( + facility, clustername, old_purchased, new_purchased, update_strategy, dry_run=dry_run + ) + return None def do_new_repo( @@ -902,6 +934,140 @@ 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, + old_purchased: int, + new_purchased: int, + update_strategy: str = 'proportional', + dry_run: bool = False + ) -> bool: + """ + Handle facility-level compute allocation changes by updating all affected repo allocations. + + Args: + facility: Name of the facility (e.g., 'shared', 'lcls') + clustername: Name of the cluster (e.g., 'roma', 'ampere') + old_purchased: Previous number of purchased nodes + new_purchased: New number of purchased nodes + update_strategy: 'proportional' (maintain percentages) or 'manual' (no auto-update) + dry_run: If True, only log what would be done without making changes + + Returns: + True if successful, False otherwise + """ + self.logger.info( + f"Processing facility compute allocation cascade: {facility}@{clustername} " + f"from {old_purchased} to {new_purchased} nodes (strategy: {update_strategy})" + ) + + if update_strategy != 'proportional': + self.logger.info(f"Update strategy '{update_strategy}' - no automatic updates performed") + return True + + if new_purchased <= 0: + self.logger.warning(f"Invalid new_purchased nodes: {new_purchased} - skipping cascade updates") + return True + + try: + # 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}") + + # Process each affected repository allocation + update_count = 0 + for item in affected_repos: + repo = item['repo'] + allocation = item['allocation'] + + # Calculate new node allocation maintaining the same percentage + percent_of_facility = allocation['percentOfFacility'] + new_allocated_nodes = (percent_of_facility / 100.0) * new_purchased + + self.logger.info( + f"Updating {repo['facility']}:{repo['name']} on {clustername}: " + f"{percent_of_facility}% -> {new_allocated_nodes:.2f} nodes (was {allocation['allocatedNodesCount']})" + ) + + if not dry_run: + try: + # Update the allocation using existing upsert method + start_time = pdl.parse(allocation['start'], timezone='UTC') + end_time = allocation['end'] + + self.upsert_repo_compute_allocation( + repo_id=repo['Id'], + cluster=clustername, + percent=percent_of_facility, + allocated_resource=new_allocated_nodes, + start=start_time, + end=end_time, + dry_run=dry_run + ) + + # Re-run the SLURM configuration to apply the new limits + # Get updated allocation info + repo_req = {'repo': {'facility': repo['facility'], 'name': repo['name']}} + updated_resp = self.back_channel.execute(self.REPO_CURRENT_COMPUTE_REQUIREMENT_GQL, repo_req) + updated_repo = updated_resp['repo'] + + # Find the updated allocation for this cluster + updated_alloc = None + for alloc in updated_repo['currentComputeAllocations']: + if alloc['clustername'].lower() == clustername.lower(): + updated_alloc = alloc + break + + if updated_alloc: + # Update SLURM with the new resource limits + self.run_playbook( + 'coact/slurm/ensure-repo.yaml', + facility=repo['facility'], + repo=repo['name'], + partition=clustername, + cpus=int(updated_alloc['cpus']), + memory=int(updated_alloc['memory']) * 1024, + nodes=int(ceil(updated_alloc['nodes'])), + gpus=int(updated_alloc['gpus']), + state='present', + dry_run=dry_run + ) + + update_count += 1 + self.logger.info(f"Successfully updated {repo['facility']}:{repo['name']} allocation") + + except Exception as e: + self.logger.error(f"Failed to update {repo['facility']}:{repo['name']}: {e}") + # Continue with other repos even if one fails + continue + else: + self.logger.info(f"DRY RUN: Would update {repo['facility']}:{repo['name']} allocation") + update_count += 1 + + self.logger.info( + f"Facility cascade update completed: {update_count}/{len(affected_repos)} " + f"repo allocations updated on {facility}@{clustername}" + ) + return True + + except Exception as e: + self.logger.error(f"Facility compute allocation cascade failed: {e}") + return False + @coactd.command(name='reporegistration') @common_options @@ -911,7 +1077,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 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..c553324 --- /dev/null +++ b/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() diff --git a/tests/test_facility_compute_allocation.py b/tests/test_facility_compute_allocation.py new file mode 100644 index 0000000..de78d4f --- /dev/null +++ b/tests/test_facility_compute_allocation.py @@ -0,0 +1,97 @@ +""" +Behavioral tests for FacilityComputeAllocation handling in the RepoRegistration daemon. +""" +import pytest +from unittest.mock import MagicMock + +from modules.coactd import RepoRegistration, RequestStatus + + +START = '2026-01-01T00:00:00Z' +END = '2031-01-01T00:00:00Z' + + +def make_handler(): + handler = RepoRegistration.__new__(RepoRegistration) + handler.logger = MagicMock() + handler.username = 'sdf-bot' + handler.password_file = '/tmp/fake' + handler.client_name = 'test-client' + handler.dry_run = False + handler.back_channel = MagicMock() + handler.ident = 'test-req-id' + return handler + + +def test_approved_request_dispatches_cascade_with_payload_fields(): + """ + An approved FacilityComputeAllocation request routes to + do_facility_compute_allocation_cascade with facility, cluster, + old/new purchased, and strategy taken directly from the request dict. + """ + handler = make_handler() + handler.do_facility_compute_allocation_cascade = MagicMock(return_value=True) + + req = { + 'facilityname': 'lcls', + 'clustername': 'ada', + 'oldPurchased': 100, + 'newPurchased': 200, + 'updateStrategy': 'proportional', + } + result = handler.do('req1', 'INSERT', 'FacilityComputeAllocation', RequestStatus.APPROVED, req, dry_run=False) + + assert result is True + handler.do_facility_compute_allocation_cascade.assert_called_once_with( + 'lcls', 'ada', 100, 200, 'proportional', dry_run=False + ) + + +def test_cascade_recalculates_every_repo_allocation_by_percentage(): + """ + When purchased nodes change, every repo on that cluster receives a new + absolute allocation of (percentOfFacility / 100) * new_purchased, preserving + each repo's percentage share of the facility. + """ + handler = make_handler() + handler.upsert_repo_compute_allocation = MagicMock(return_value={}) + handler.run_playbook = MagicMock(return_value=None) + + repos = [ + { + 'Id': 'repo-a', 'name': 'alpha', 'facility': 'lcls', + 'currentComputeAllocations': [{ + 'Id': 'alloc-a', 'clustername': 'ada', + 'percentOfFacility': 25.0, 'allocatedNodesCount': 25.0, + 'start': START, 'end': END, + }], + }, + { + 'Id': 'repo-b', 'name': 'beta', 'facility': 'lcls', + 'currentComputeAllocations': [{ + 'Id': 'alloc-b', 'clustername': 'ada', + 'percentOfFacility': 50.0, 'allocatedNodesCount': 50.0, + 'start': START, 'end': END, + }], + }, + ] + handler.back_channel.execute.side_effect = [ + {'repos': repos}, + {'repo': {'currentComputeAllocations': []}}, # SLURM re-query for repo-a + {'repo': {'currentComputeAllocations': []}}, # SLURM re-query for repo-b + ] + + result = handler.do_facility_compute_allocation_cascade( + 'lcls', 'ada', old_purchased=100, new_purchased=200, + update_strategy='proportional', dry_run=False + ) + + assert result is True + assert handler.upsert_repo_compute_allocation.call_count == 2 + + by_repo = { + c.kwargs['repo_id']: c.kwargs['allocated_resource'] + for c in handler.upsert_repo_compute_allocation.call_args_list + } + assert by_repo['repo-a'] == 50.0 # 25% of 200 + assert by_repo['repo-b'] == 100.0 # 50% of 200 From 6187e006e9d85d6ed1e9dcf7839225edc40d39b4 Mon Sep 17 00:00:00 2001 From: Ryan Waldheim Date: Wed, 6 May 2026 14:56:16 -0700 Subject: [PATCH 02/11] feat: update GraphQL mutations for facility and cluster inputs in FacilityManager --- modules/coact.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/modules/coact.py b/modules/coact.py index bd44cfd..570fa10 100644 --- a/modules/coact.py +++ b/modules/coact.py @@ -1093,14 +1093,14 @@ class FacilityManager(GraphQlMixin): FACILITY_UPDATE_PURCHASED_GQL = gql(""" mutation facilityUpdatePurchased( - $facilityName: String!, - $clusterName: String!, - $purchased: Int! + $facility: FacilityInput!, + $cluster: ClusterInput!, + $purchase: Float! ) { - facilityUpdateComputePurchase( - facility: $facilityName, - cluster: $clusterName, - purchased: $purchased + facilityAddUpdateComputePurchase( + facility: $facility, + cluster: $cluster, + purchase: $purchase ) { name computepurchases { @@ -1124,7 +1124,7 @@ class FacilityManager(GraphQlMixin): """) REPOS_WITH_ALLOCATIONS_GQL = gql(""" - query reposWithAllocations($facilityName: String!, $clusterName: String!) { + query reposWithAllocations($facilityName: String!) { repos(filter: {facility: $facilityName}) { Id name @@ -1158,9 +1158,9 @@ def update_purchased_nodes(self, facility: str, cluster: str, nodes: int, dry_ru result = self.back_channel.execute( self.FACILITY_UPDATE_PURCHASED_GQL, { - 'facilityName': facility, - 'clusterName': cluster, - 'purchased': nodes + 'facility': {'name': facility}, + 'cluster': {'name': cluster}, + 'purchase': float(nodes) } ) logger.info(f"Successfully updated facility: {result}") @@ -1214,7 +1214,7 @@ def simulate_change(self, facility: str, cluster: str, nodes: int) -> bool: # Get affected repos repos_result = self.back_channel.execute( self.REPOS_WITH_ALLOCATIONS_GQL, - {'facilityName': facility, 'clusterName': cluster} + {'facilityName': facility} ) affected_count = 0 From 6bf240c244af28a59ee827dff165c0dd5ecbc5aa Mon Sep 17 00:00:00 2001 From: Ryan Waldheim Date: Wed, 6 May 2026 16:12:13 -0700 Subject: [PATCH 03/11] query purchased nodes directly Co-authored-by: Copilot --- modules/coactd.py | 51 ++++++++++------------- tests/test_facility_compute_allocation.py | 15 +++---- 2 files changed, 28 insertions(+), 38 deletions(-) diff --git a/modules/coactd.py b/modules/coactd.py index 84e9875..0388be9 100644 --- a/modules/coactd.py +++ b/modules/coactd.py @@ -148,9 +148,6 @@ class Registration(GraphQlSubscriber, AnsibleRunner): end percentOfFacility allocated - oldPurchased - newPurchased - updateStrategy } operationType } @@ -539,12 +536,9 @@ def do(self, req_id, op_type, req_type, approval, req, dry_run): elif req_type == 'FacilityComputeAllocation': clustername = req.get('clustername', None) - old_purchased = req.get('oldPurchased', None) - new_purchased = req.get('newPurchased', None) - update_strategy = req.get('updateStrategy', 'proportional') - assert facility and clustername and old_purchased is not None and new_purchased is not None + assert facility and clustername return self.do_facility_compute_allocation_cascade( - facility, clustername, old_purchased, new_purchased, update_strategy, dry_run=dry_run + facility, clustername, dry_run=dry_run ) return None @@ -938,38 +932,37 @@ def do_facility_compute_allocation_cascade( self, facility: str, clustername: str, - old_purchased: int, - new_purchased: int, - update_strategy: str = 'proportional', dry_run: bool = False ) -> bool: """ Handle facility-level compute allocation changes by updating all affected repo allocations. - - Args: - facility: Name of the facility (e.g., 'shared', 'lcls') - clustername: Name of the cluster (e.g., 'roma', 'ampere') - old_purchased: Previous number of purchased nodes - new_purchased: New number of purchased nodes - update_strategy: 'proportional' (maintain percentages) or 'manual' (no auto-update) - dry_run: If True, only log what would be done without making changes - - Returns: - True if successful, False otherwise + Queries the facility record directly for the current purchased node count. """ - self.logger.info( - f"Processing facility compute allocation cascade: {facility}@{clustername} " - f"from {old_purchased} to {new_purchased} nodes (strategy: {update_strategy})" + # 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 update_strategy != 'proportional': - self.logger.info(f"Update strategy '{update_strategy}' - no automatic updates performed") - return True + if new_purchased is None: + self.logger.error(f"No purchase record found for {facility}@{clustername} - cannot cascade") + return False if new_purchased <= 0: - self.logger.warning(f"Invalid new_purchased nodes: {new_purchased} - skipping cascade updates") + self.logger.warning(f"Invalid purchased nodes: {new_purchased} - skipping cascade updates") return True + self.logger.info( + f"Processing facility compute allocation cascade: {facility}@{clustername} " + f"-> {new_purchased} purchased nodes" + ) + try: # Get all repositories with allocations on this facility/cluster repos_resp = self.back_channel.execute( diff --git a/tests/test_facility_compute_allocation.py b/tests/test_facility_compute_allocation.py index de78d4f..e347462 100644 --- a/tests/test_facility_compute_allocation.py +++ b/tests/test_facility_compute_allocation.py @@ -26,8 +26,8 @@ def make_handler(): def test_approved_request_dispatches_cascade_with_payload_fields(): """ An approved FacilityComputeAllocation request routes to - do_facility_compute_allocation_cascade with facility, cluster, - old/new purchased, and strategy taken directly from the request dict. + do_facility_compute_allocation_cascade with facility and cluster + extracted from the request dict. """ handler = make_handler() handler.do_facility_compute_allocation_cascade = MagicMock(return_value=True) @@ -35,22 +35,19 @@ def test_approved_request_dispatches_cascade_with_payload_fields(): req = { 'facilityname': 'lcls', 'clustername': 'ada', - 'oldPurchased': 100, - 'newPurchased': 200, - 'updateStrategy': 'proportional', } result = handler.do('req1', 'INSERT', 'FacilityComputeAllocation', RequestStatus.APPROVED, req, dry_run=False) assert result is True handler.do_facility_compute_allocation_cascade.assert_called_once_with( - 'lcls', 'ada', 100, 200, 'proportional', dry_run=False + 'lcls', 'ada', dry_run=False ) def test_cascade_recalculates_every_repo_allocation_by_percentage(): """ When purchased nodes change, every repo on that cluster receives a new - absolute allocation of (percentOfFacility / 100) * new_purchased, preserving + absolute allocation of (percentOfFacility / 100) * purchased, preserving each repo's percentage share of the facility. """ handler = make_handler() @@ -76,14 +73,14 @@ def test_cascade_recalculates_every_repo_allocation_by_percentage(): }, ] handler.back_channel.execute.side_effect = [ + {'facility': {'computepurchases': [{'clustername': 'ada', 'purchased': 200}]}}, # facility query {'repos': repos}, {'repo': {'currentComputeAllocations': []}}, # SLURM re-query for repo-a {'repo': {'currentComputeAllocations': []}}, # SLURM re-query for repo-b ] result = handler.do_facility_compute_allocation_cascade( - 'lcls', 'ada', old_purchased=100, new_purchased=200, - update_strategy='proportional', dry_run=False + 'lcls', 'ada', dry_run=False ) assert result is True From 064b2716e2822ed8bf12db85c655ba6fcf6b1732 Mon Sep 17 00:00:00 2001 From: Ryan Waldheim Date: Thu, 7 May 2026 16:18:28 -0700 Subject: [PATCH 04/11] consolidate usage of do_repo_compute_allocation --- ansible-runner/project | 2 +- modules/coactd.py | 74 ++++++----------------- tests/test_facility_compute_allocation.py | 20 +++--- 3 files changed, 31 insertions(+), 65 deletions(-) diff --git a/ansible-runner/project b/ansible-runner/project index 01bd93a..7b57733 160000 --- a/ansible-runner/project +++ b/ansible-runner/project @@ -1 +1 @@ -Subproject commit 01bd93a20b87022c6acf525cd0a85cb7ef2e274b +Subproject commit 7b577332d8aabcd218d2042b33f87c95aa14d640 diff --git a/modules/coactd.py b/modules/coactd.py index 0388be9..d61eac2 100644 --- a/modules/coactd.py +++ b/modules/coactd.py @@ -955,8 +955,8 @@ def do_facility_compute_allocation_cascade( return False if new_purchased <= 0: - self.logger.warning(f"Invalid purchased nodes: {new_purchased} - skipping cascade updates") - return True + self.logger.error(f"Invalid purchased nodes: {new_purchased} for {facility}@{clustername} - cannot cascade") + return False self.logger.info( f"Processing facility compute allocation cascade: {facility}@{clustername} " @@ -996,60 +996,24 @@ def do_facility_compute_allocation_cascade( f"{percent_of_facility}% -> {new_allocated_nodes:.2f} nodes (was {allocation['allocatedNodesCount']})" ) - if not dry_run: - try: - # Update the allocation using existing upsert method - start_time = pdl.parse(allocation['start'], timezone='UTC') - end_time = allocation['end'] - - self.upsert_repo_compute_allocation( - repo_id=repo['Id'], - cluster=clustername, - percent=percent_of_facility, - allocated_resource=new_allocated_nodes, - start=start_time, - end=end_time, - dry_run=dry_run - ) - - # Re-run the SLURM configuration to apply the new limits - # Get updated allocation info - repo_req = {'repo': {'facility': repo['facility'], 'name': repo['name']}} - updated_resp = self.back_channel.execute(self.REPO_CURRENT_COMPUTE_REQUIREMENT_GQL, repo_req) - updated_repo = updated_resp['repo'] - - # Find the updated allocation for this cluster - updated_alloc = None - for alloc in updated_repo['currentComputeAllocations']: - if alloc['clustername'].lower() == clustername.lower(): - updated_alloc = alloc - break - - if updated_alloc: - # Update SLURM with the new resource limits - self.run_playbook( - 'coact/slurm/ensure-repo.yaml', - facility=repo['facility'], - repo=repo['name'], - partition=clustername, - cpus=int(updated_alloc['cpus']), - memory=int(updated_alloc['memory']) * 1024, - nodes=int(ceil(updated_alloc['nodes'])), - gpus=int(updated_alloc['gpus']), - state='present', - dry_run=dry_run - ) - - update_count += 1 - self.logger.info(f"Successfully updated {repo['facility']}:{repo['name']} allocation") - - except Exception as e: - self.logger.error(f"Failed to update {repo['facility']}:{repo['name']}: {e}") - # Continue with other repos even if one fails - continue - else: - self.logger.info(f"DRY RUN: Would update {repo['facility']}:{repo['name']} allocation") + try: + self.do_repo_compute_allocation( + repo['name'], + repo['facility'], + clustername, + percent_of_facility, + new_allocated_nodes, + pdl.parse(allocation['start'], timezone='UTC'), + allocation['end'], + dry_run=dry_run, + ) update_count += 1 + self.logger.info(f"Successfully updated {repo['facility']}:{repo['name']} allocation") + + except Exception as e: + self.logger.error(f"Failed to update {repo['facility']}:{repo['name']}: {e}") + # Continue with other repos even if one fails + continue self.logger.info( f"Facility cascade update completed: {update_count}/{len(affected_repos)} " diff --git a/tests/test_facility_compute_allocation.py b/tests/test_facility_compute_allocation.py index e347462..37ae2a1 100644 --- a/tests/test_facility_compute_allocation.py +++ b/tests/test_facility_compute_allocation.py @@ -49,10 +49,13 @@ def test_cascade_recalculates_every_repo_allocation_by_percentage(): When purchased nodes change, every repo on that cluster receives a new absolute allocation of (percentOfFacility / 100) * purchased, preserving each repo's percentage share of the facility. + + do_repo_compute_allocation is the single delegate for each repo; it owns + the SLURM feature-flag check, the DB upsert, the SLURM playbook call, and + the user sync. """ handler = make_handler() - handler.upsert_repo_compute_allocation = MagicMock(return_value={}) - handler.run_playbook = MagicMock(return_value=None) + handler.do_repo_compute_allocation = MagicMock(return_value=True) repos = [ { @@ -75,8 +78,6 @@ def test_cascade_recalculates_every_repo_allocation_by_percentage(): handler.back_channel.execute.side_effect = [ {'facility': {'computepurchases': [{'clustername': 'ada', 'purchased': 200}]}}, # facility query {'repos': repos}, - {'repo': {'currentComputeAllocations': []}}, # SLURM re-query for repo-a - {'repo': {'currentComputeAllocations': []}}, # SLURM re-query for repo-b ] result = handler.do_facility_compute_allocation_cascade( @@ -84,11 +85,12 @@ def test_cascade_recalculates_every_repo_allocation_by_percentage(): ) assert result is True - assert handler.upsert_repo_compute_allocation.call_count == 2 + assert handler.do_repo_compute_allocation.call_count == 2 + # args: (repo_name, facility, cluster, percent, allocated_resource, start, end) by_repo = { - c.kwargs['repo_id']: c.kwargs['allocated_resource'] - for c in handler.upsert_repo_compute_allocation.call_args_list + c.args[0]: c.args[4] + for c in handler.do_repo_compute_allocation.call_args_list } - assert by_repo['repo-a'] == 50.0 # 25% of 200 - assert by_repo['repo-b'] == 100.0 # 50% of 200 + assert by_repo['alpha'] == 50.0 # 25% of 200 + assert by_repo['beta'] == 100.0 # 50% of 200 From d6d3e75c647a8ffe108d53cc7f4a6cfc5dedf44f Mon Sep 17 00:00:00 2001 From: Ryan Waldheim Date: Wed, 13 May 2026 15:27:00 -0700 Subject: [PATCH 05/11] rm unnecessary test CLI endpoints --- ansible-runner/project | 2 +- modules/coact.py | 219 ----------------------------------------- 2 files changed, 1 insertion(+), 220 deletions(-) diff --git a/ansible-runner/project b/ansible-runner/project index 7b57733..6b58301 160000 --- a/ansible-runner/project +++ b/ansible-runner/project @@ -1 +1 @@ -Subproject commit 7b577332d8aabcd218d2042b33f87c95aa14d640 +Subproject commit 6b5830182271ebdf3e8ac21666784de3cf4458d4 diff --git a/modules/coact.py b/modules/coact.py index 570fa10..2e04345 100644 --- a/modules/coact.py +++ b/modules/coact.py @@ -1076,225 +1076,6 @@ def overaged(self, data: dict, threshold: float = 100.0) -> Iterator[OveragePoin ) - -# ============================================================================ -# Facility Management Commands -# ============================================================================ - -@click.group(name='facility', help="Facility compute management commands", context_settings=CONTEXT_SETTINGS) -@click.pass_context -def facility(ctx): - """Facility management command group for compute allocation operations.""" - ctx.ensure_object(dict) - - -class FacilityManager(GraphQlMixin): - """Handles facility compute management operations.""" - - FACILITY_UPDATE_PURCHASED_GQL = gql(""" - mutation facilityUpdatePurchased( - $facility: FacilityInput!, - $cluster: ClusterInput!, - $purchase: Float! - ) { - facilityAddUpdateComputePurchase( - facility: $facility, - cluster: $cluster, - purchase: $purchase - ) { - name - computepurchases { - clustername - purchased - } - } - } - """) - - FACILITY_QUERY_GQL = gql(""" - query facility($facilityName: String!) { - facility(filter: {name: $facilityName}) { - name - computepurchases { - clustername - purchased - } - } - } - """) - - REPOS_WITH_ALLOCATIONS_GQL = gql(""" - query reposWithAllocations($facilityName: String!) { - repos(filter: {facility: $facilityName}) { - Id - name - facility - currentComputeAllocations { - Id - clustername - percentOfFacility - allocatedNodesCount - } - } - } - """) - - def __init__(self, username: str, password_file: str): - self.username = username - self.password_file = password_file - - def update_purchased_nodes(self, facility: str, cluster: str, nodes: int, dry_run: bool = False) -> bool: - """Update the purchased node count for a facility/cluster.""" - self.back_channel = self.connect_graph_ql( - username=self.username, - password_file=self.password_file, - timeout=60 - ) - - logger.info(f"Updating {facility}@{cluster} to {nodes} purchased nodes (dry_run={dry_run})") - - if not dry_run: - try: - result = self.back_channel.execute( - self.FACILITY_UPDATE_PURCHASED_GQL, - { - 'facility': {'name': facility}, - 'cluster': {'name': cluster}, - 'purchase': float(nodes) - } - ) - logger.info(f"Successfully updated facility: {result}") - return True - except Exception as e: - logger.error(f"Failed to update facility: {e}") - return False - else: - logger.info(f"DRY RUN: Would update {facility}@{cluster} to {nodes} nodes") - return True - - def simulate_change(self, facility: str, cluster: str, nodes: int) -> bool: - """Simulate facility change and show what repos would be affected.""" - self.back_channel = self.connect_graph_ql( - username=self.username, - password_file=self.password_file, - timeout=60 - ) - - logger.info(f"Simulating facility change: {facility}@{cluster} -> {nodes} nodes") - - try: - # Get current facility state - current_result = self.back_channel.execute( - self.FACILITY_QUERY_GQL, - {'facilityName': facility} - ) - - current_facility = current_result.get('facility') - if not current_facility: - logger.error(f"Facility '{facility}' not found") - return False - - current_purchased = None - for purchase in current_facility.get('computepurchases', []): - if purchase['clustername'].lower() == cluster.lower(): - current_purchased = purchase['purchased'] - break - - if current_purchased is None: - logger.error(f"Cluster '{cluster}' not found in facility '{facility}'") - return False - - logger.info(f"Current purchased nodes: {current_purchased}") - logger.info(f"Proposed purchased nodes: {nodes}") - - if current_purchased == nodes: - logger.info("No change detected - no repos would be affected") - return True - - # Get affected repos - repos_result = self.back_channel.execute( - self.REPOS_WITH_ALLOCATIONS_GQL, - {'facilityName': facility} - ) - - affected_count = 0 - for repo in repos_result['repos']: - for allocation in repo['currentComputeAllocations']: - if allocation['clustername'].lower() == cluster.lower(): - current_nodes = allocation['allocatedNodesCount'] - percent = allocation['percentOfFacility'] - new_nodes = (percent / 100.0) * nodes - - logger.info( - f" {repo['facility']}:{repo['name']} - " - f"{percent}% -> {current_nodes} nodes would become {new_nodes:.2f} nodes" - ) - affected_count += 1 - - logger.info(f"Total affected repo allocations: {affected_count}") - return True - - except Exception as e: - logger.error(f"Simulation failed: {e}") - return False - - -@facility.command(name='update-purchased') -@common_options -@graphql_options -@click.option('--facility', required=True, help='Facility name (e.g., shared, lcls)') -@click.option('--cluster', required=True, help='Cluster name (e.g., roma, ampere)') -@click.option('--nodes', required=True, type=int, help='New number of purchased nodes') -@click.option('--dry-run', is_flag=True, default=False, help='Show what would be done without making changes') -@click.pass_context -def facility_update_purchased(ctx, verbose, username, password_file, facility, cluster, nodes, dry_run): - """Update the number of purchased nodes for a facility/cluster. - - This will trigger automatic cascade updates of all repo allocations on this facility/cluster - if facility monitoring is enabled. - """ - configure_logging_from_verbose(verbose) - ctx.obj['verbose'] = verbose - - manager = FacilityManager( - username=username, - password_file=password_file - ) - - success = manager.update_purchased_nodes(facility, cluster, nodes, dry_run) - if not success: - raise click.ClickException("Failed to update facility purchased nodes") - - -@facility.command(name='simulate-change') -@common_options -@graphql_options -@click.option('--facility', required=True, help='Facility name (e.g., shared, lcls)') -@click.option('--cluster', required=True, help='Cluster name (e.g., roma, ampere)') -@click.option('--nodes', required=True, type=int, help='Proposed number of purchased nodes') -@click.pass_context -def facility_simulate_change(ctx, verbose, username, password_file, facility, cluster, nodes): - """Simulate a facility compute change and show what repo allocations would be affected. - - This is useful for understanding the impact of facility changes before applying them. - """ - configure_logging_from_verbose(verbose) - ctx.obj['verbose'] = verbose - - manager = FacilityManager( - username=username, - password_file=password_file - ) - - success = manager.simulate_change(facility, cluster, nodes) - if not success: - raise click.ClickException("Simulation failed") - - -# Add facility to main coact group -coact.add_command(facility) - - # For backwards compatibility, allow running this module directly if __name__ == '__main__': coact(obj={}) From 7ee556ff31f2e21cc499e00eff8218127753b35e Mon Sep 17 00:00:00 2001 From: Ryan Waldheim Date: Wed, 13 May 2026 17:02:08 -0700 Subject: [PATCH 06/11] consistent handling of end timestamp --- modules/coactd.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/coactd.py b/modules/coactd.py index d61eac2..c1afbf1 100644 --- a/modules/coactd.py +++ b/modules/coactd.py @@ -1004,7 +1004,7 @@ def do_facility_compute_allocation_cascade( percent_of_facility, new_allocated_nodes, pdl.parse(allocation['start'], timezone='UTC'), - allocation['end'], + pdl.parse(allocation['end'], timezone='UTC') if allocation['end'] else None, dry_run=dry_run, ) update_count += 1 From 8d861ea261a78462d0705cc38f01ece2f1abe2fd Mon Sep 17 00:00:00 2001 From: Ryan Waldheim Date: Wed, 13 May 2026 17:02:13 -0700 Subject: [PATCH 07/11] additional tests --- tests/test_facility_compute_allocation.py | 78 ++++++++++++++++++++++- 1 file changed, 77 insertions(+), 1 deletion(-) diff --git a/tests/test_facility_compute_allocation.py b/tests/test_facility_compute_allocation.py index 37ae2a1..8765daa 100644 --- a/tests/test_facility_compute_allocation.py +++ b/tests/test_facility_compute_allocation.py @@ -1,7 +1,6 @@ """ Behavioral tests for FacilityComputeAllocation handling in the RepoRegistration daemon. """ -import pytest from unittest.mock import MagicMock from modules.coactd import RepoRegistration, RequestStatus @@ -94,3 +93,80 @@ def test_cascade_recalculates_every_repo_allocation_by_percentage(): } assert by_repo['alpha'] == 50.0 # 25% of 200 assert by_repo['beta'] == 100.0 # 50% of 200 + + +def test_cascade_returns_false_when_no_purchase_record(): + """ + When the facility has no computepurchases entry for the requested cluster, + the cascade logs an error and returns False without touching any repos. + """ + handler = make_handler() + handler.do_repo_compute_allocation = MagicMock(return_value=True) + + handler.back_channel.execute.return_value = { + 'facility': {'computepurchases': [{'clustername': 'other-cluster', 'purchased': 100}]} + } + + result = handler.do_facility_compute_allocation_cascade('lcls', 'ada', dry_run=False) + + assert result is False + handler.do_repo_compute_allocation.assert_not_called() + handler.logger.error.assert_called_once() + + +def test_cascade_returns_false_when_purchased_nodes_is_zero(): + """ + A purchase record with zero nodes is rejected before any repo is updated. + """ + handler = make_handler() + handler.do_repo_compute_allocation = MagicMock(return_value=True) + + handler.back_channel.execute.return_value = { + 'facility': {'computepurchases': [{'clustername': 'ada', 'purchased': 0}]} + } + + result = handler.do_facility_compute_allocation_cascade('lcls', 'ada', dry_run=False) + + assert result is False + handler.do_repo_compute_allocation.assert_not_called() + handler.logger.error.assert_called_once() + + +def test_cascade_continues_after_per_repo_failure(): + """ + If do_repo_compute_allocation raises for one repo, the cascade logs the + error and continues processing the remaining repos, returning True overall. + """ + handler = make_handler() + + repos = [ + { + 'Id': 'repo-a', 'name': 'alpha', 'facility': 'lcls', + 'currentComputeAllocations': [{ + 'Id': 'alloc-a', 'clustername': 'ada', + 'percentOfFacility': 25.0, 'allocatedNodesCount': 50.0, + 'start': START, 'end': END, + }], + }, + { + 'Id': 'repo-b', 'name': 'beta', 'facility': 'lcls', + 'currentComputeAllocations': [{ + 'Id': 'alloc-b', 'clustername': 'ada', + 'percentOfFacility': 50.0, 'allocatedNodesCount': 100.0, + 'start': START, 'end': END, + }], + }, + ] + handler.back_channel.execute.side_effect = [ + {'facility': {'computepurchases': [{'clustername': 'ada', 'purchased': 200}]}}, + {'repos': repos}, + ] + handler.do_repo_compute_allocation = MagicMock( + side_effect=[Exception("slurm playbook failed"), True] + ) + + result = handler.do_facility_compute_allocation_cascade('lcls', 'ada', dry_run=False) + + assert result is True + assert handler.do_repo_compute_allocation.call_count == 2 + handler.logger.error.assert_called_once() From 71f0b981a477a98f976ad0c5ec6175c225f83be5 Mon Sep 17 00:00:00 2001 From: Ryan Waldheim Date: Mon, 31 Aug 2026 10:23:44 -0700 Subject: [PATCH 08/11] fix dated tests --- tests/test_slurm_node_memory.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/test_slurm_node_memory.py b/tests/test_slurm_node_memory.py index 1946f22..73a4239 100644 --- a/tests/test_slurm_node_memory.py +++ b/tests/test_slurm_node_memory.py @@ -22,35 +22,35 @@ def setup_method(self): def test_parse_slurm_nodelist_single(self): """Test parsing a single node name.""" - result = self.importer.parse_slurm_nodelist("sdfmilan0271") - assert result == ["sdfmilan0271"] + result = self.importer.parse_slurm_nodelist("sdfmilan271") + assert result == ["sdfmilan271"] def test_parse_slurm_nodelist_range(self): """Test parsing a SLURM node range.""" result = self.importer.parse_slurm_nodelist("sdfmilan[269-272]") - expected = ["sdfmilan0269", "sdfmilan0270", "sdfmilan0271", "sdfmilan0272"] + expected = ["sdfmilan269", "sdfmilan270", "sdfmilan271", "sdfmilan272"] assert result == expected def test_parse_slurm_nodelist_list(self): """Test parsing a comma-separated list of nodes.""" result = self.importer.parse_slurm_nodelist("sdfmilan[006,011,027]") - expected = ["sdfmilan0006", "sdfmilan0011", "sdfmilan0027"] + expected = ["sdfmilan006", "sdfmilan011", "sdfmilan027"] assert result == expected def test_parse_slurm_nodelist_mixed(self): """Test parsing a mixed range and list.""" result = self.importer.parse_slurm_nodelist("sdfmilan[001-003,010,020-022]") expected = [ - "sdfmilan0001", "sdfmilan0002", "sdfmilan0003", - "sdfmilan0010", - "sdfmilan0020", "sdfmilan0021", "sdfmilan0022" + "sdfmilan001", "sdfmilan002", "sdfmilan003", + "sdfmilan010", + "sdfmilan020", "sdfmilan021", "sdfmilan022" ] assert result == expected def test_parse_slurm_nodelist_different_prefix(self): """Test parsing with different node prefix.""" result = self.importer.parse_slurm_nodelist("sdfrome[001-003]") - expected = ["sdfrome0001", "sdfrome0002", "sdfrome0003"] + expected = ["sdfrome001", "sdfrome002", "sdfrome003"] assert result == expected def test_parse_slurm_nodelist_unparseable(self): From d93f081d0e61d1129c4c02febdd4d456ab37854f Mon Sep 17 00:00:00 2001 From: Ryan Waldheim Date: Mon, 31 Aug 2026 10:25:06 -0700 Subject: [PATCH 09/11] update CI/CD workflow to use specific versions of actions --- .github/workflows/ci-cd.yaml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci-cd.yaml b/.github/workflows/ci-cd.yaml index 21db93e..acaff56 100644 --- a/.github/workflows/ci-cd.yaml +++ b/.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" From fed3cccde8ceb2146538b59a52dd214653b0d7f7 Mon Sep 17 00:00:00 2001 From: Ryan Waldheim Date: Mon, 31 Aug 2026 10:32:47 -0700 Subject: [PATCH 10/11] sdf-ansible not available in CI --- tests/test_s3df_posixgroup_ldif.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/test_s3df_posixgroup_ldif.py b/tests/test_s3df_posixgroup_ldif.py index 45a32a6..2de1de9 100644 --- a/tests/test_s3df_posixgroup_ldif.py +++ b/tests/test_s3df_posixgroup_ldif.py @@ -7,6 +7,8 @@ import os from unittest.mock import Mock +import pytest + # Mock the ansible modules imported by s3df_posixgroup sys.modules['ansible.module_utils.basic'] = Mock() sys.modules['ansible.module_utils.common.text.converters'] = Mock() @@ -17,6 +19,13 @@ if LIBRARY_PATH not in sys.path: sys.path.insert(0, LIBRARY_PATH) +# sdf-ansible not available in CI +if not os.path.isfile(os.path.join(LIBRARY_PATH, 's3df_posixgroup.py')): + pytest.skip( + "sdf-ansible submodule not checked out; run `git submodule update --init`", + allow_module_level=True, + ) + from s3df_posixgroup import _change_membership_template # noqa: E402 DN = "cn=sdf-cryoem-cd10,ou=Group,dc=sdf,dc=slac,dc=stanford,dc=edu" From c07a38f4e5cccfd69038ffc2f1a9dd50d8bc219a Mon Sep 17 00:00:00 2001 From: Ryan Waldheim Date: Tue, 1 Sep 2026 14:35:15 -0700 Subject: [PATCH 11/11] Enhance RepoRegistration to support burst allocation parameters and update tests accordingly --- modules/coactd.py | 205 +++++++++----- tests/test_facility_compute_allocation.py | 327 ++++++++++++++++++---- 2 files changed, 400 insertions(+), 132 deletions(-) diff --git a/modules/coactd.py b/modules/coactd.py index 50ba875..8c0d462 100644 --- a/modules/coactd.py +++ b/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 } @@ -416,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) { @@ -443,6 +437,8 @@ class RepoRegistration(Registration): start end percentOfFacility + burstPercentOfFacility + burstAllocated cpus: allocatedCpusCount memory: allocatedMemGb nodes: allocatedNodesCount @@ -462,6 +458,8 @@ class RepoRegistration(Registration): Id clustername percentOfFacility + burstPercentOfFacility + burstAllocated start end allocatedNodesCount @@ -474,9 +472,6 @@ class RepoRegistration(Registration): query facility( $facility: String! ) { facility(filter: {name: $facility}) { name - computeallocations { - clustername - } computepurchases { clustername purchased @@ -524,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': @@ -778,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 ): @@ -802,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! ) { @@ -843,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(): @@ -884,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 @@ -1111,6 +1140,10 @@ def do_facility_compute_allocation_cascade( """ 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( @@ -1125,79 +1158,107 @@ def do_facility_compute_allocation_cascade( break if new_purchased is None: - self.logger.error(f"No purchase record found for {facility}@{clustername} - cannot cascade") - return False + raise RuntimeError(f"No purchase record found for {facility}@{clustername} - cannot cascade") - if new_purchased <= 0: - self.logger.error(f"Invalid purchased nodes: {new_purchased} for {facility}@{clustername} - cannot cascade") - return False + 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" ) - try: - # 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 - }) + # Get all repositories with allocations on this facility/cluster + repos_resp = self.back_channel.execute( + self.REPOS_WITH_ALLOCATIONS_GQL, + {'facilityName': facility} + ) - self.logger.info(f"Found {len(affected_repos)} repo allocations to update on {facility}@{clustername}") + 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." + ) - # Process each affected repository allocation - update_count = 0 - for item in affected_repos: - repo = item['repo'] - allocation = item['allocation'] + 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 - # Calculate new node allocation maintaining the same percentage - percent_of_facility = allocation['percentOfFacility'] - new_allocated_nodes = (percent_of_facility / 100.0) * new_purchased + 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" + ) - self.logger.info( - f"Updating {repo['facility']}:{repo['name']} on {clustername}: " - f"{percent_of_facility}% -> {new_allocated_nodes:.2f} nodes (was {allocation['allocatedNodesCount']})" + 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})") - 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, - dry_run=dry_run, - ) - update_count += 1 - self.logger.info(f"Successfully updated {repo['facility']}:{repo['name']} allocation") - - except Exception as e: - self.logger.error(f"Failed to update {repo['facility']}:{repo['name']}: {e}") - # Continue with other repos even if one fails - continue + 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" + ) - self.logger.info( - f"Facility cascade update completed: {update_count}/{len(affected_repos)} " - f"repo allocations updated on {facility}@{clustername}" + 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 - except Exception as e: - self.logger.error(f"Facility compute allocation cascade failed: {e}") - return False + return True @coactd.command(name='reporegistration') diff --git a/tests/test_facility_compute_allocation.py b/tests/test_facility_compute_allocation.py index 8765daa..fe6fa51 100644 --- a/tests/test_facility_compute_allocation.py +++ b/tests/test_facility_compute_allocation.py @@ -3,6 +3,8 @@ """ from unittest.mock import MagicMock +import pytest + from modules.coactd import RepoRegistration, RequestStatus @@ -22,6 +24,32 @@ def make_handler(): return handler +def make_allocation(alloc_id, percent, allocated, burst_percent=0.0, burst_allocated=0.0): + return { + 'Id': alloc_id, 'clustername': 'ada', + 'percentOfFacility': percent, 'allocatedNodesCount': allocated, + 'burstPercentOfFacility': burst_percent, 'burstAllocated': burst_allocated, + 'start': START, 'end': END, + } + + +def make_repos(*named_allocations): + return [ + { + 'Id': f'repo-{i}', 'name': name, 'facility': 'lcls', + 'currentComputeAllocations': [alloc], + } + for i, (name, alloc) in enumerate(named_allocations) + ] + + +def set_responses(handler, purchased, repos): + handler.back_channel.execute.side_effect = [ + {'facility': {'computepurchases': [{'clustername': 'ada', 'purchased': purchased}]}}, + {'repos': repos}, + ] + + def test_approved_request_dispatches_cascade_with_payload_fields(): """ An approved FacilityComputeAllocation request routes to @@ -56,28 +84,10 @@ def test_cascade_recalculates_every_repo_allocation_by_percentage(): handler = make_handler() handler.do_repo_compute_allocation = MagicMock(return_value=True) - repos = [ - { - 'Id': 'repo-a', 'name': 'alpha', 'facility': 'lcls', - 'currentComputeAllocations': [{ - 'Id': 'alloc-a', 'clustername': 'ada', - 'percentOfFacility': 25.0, 'allocatedNodesCount': 25.0, - 'start': START, 'end': END, - }], - }, - { - 'Id': 'repo-b', 'name': 'beta', 'facility': 'lcls', - 'currentComputeAllocations': [{ - 'Id': 'alloc-b', 'clustername': 'ada', - 'percentOfFacility': 50.0, 'allocatedNodesCount': 50.0, - 'start': START, 'end': END, - }], - }, - ] - handler.back_channel.execute.side_effect = [ - {'facility': {'computepurchases': [{'clustername': 'ada', 'purchased': 200}]}}, # facility query - {'repos': repos}, - ] + set_responses(handler, 200, make_repos( + ('alpha', make_allocation('alloc-a', 25.0, 25.0)), + ('beta', make_allocation('alloc-b', 50.0, 50.0)), + )) result = handler.do_facility_compute_allocation_cascade( 'lcls', 'ada', dry_run=False @@ -95,10 +105,48 @@ def test_cascade_recalculates_every_repo_allocation_by_percentage(): assert by_repo['beta'] == 100.0 # 50% of 200 -def test_cascade_returns_false_when_no_purchase_record(): +def test_cascade_recalculates_burst_allocation_by_percentage(): + """ + Burst is a second percentage of the same facility purchase and must be + recomputed alongside the base allocation. The upsert replaces the whole + allocation document, so omitting burst would silently reset it to zero. + """ + handler = make_handler() + handler.do_repo_compute_allocation = MagicMock(return_value=True) + + set_responses(handler, 200, make_repos( + ('alpha', make_allocation('alloc-a', 25.0, 25.0, burst_percent=10.0, burst_allocated=10.0)), + )) + + handler.do_facility_compute_allocation_cascade('lcls', 'ada', dry_run=False) + + kwargs = handler.do_repo_compute_allocation.call_args.kwargs + assert kwargs['burst_percent'] == 10.0 + assert kwargs['burst_allocated'] == 20.0 # 10% of 200 + + +def test_cascade_keeps_burst_at_zero_when_repo_has_none(): + """A repo without burst must not acquire one, and a null field must not blow up.""" + handler = make_handler() + handler.do_repo_compute_allocation = MagicMock(return_value=True) + + alloc = make_allocation('alloc-a', 25.0, 25.0) + alloc['burstPercentOfFacility'] = None + alloc['burstAllocated'] = None + set_responses(handler, 200, make_repos(('alpha', alloc))) + + handler.do_facility_compute_allocation_cascade('lcls', 'ada', dry_run=False) + + kwargs = handler.do_repo_compute_allocation.call_args.kwargs + assert kwargs['burst_percent'] == 0.0 + assert kwargs['burst_allocated'] == 0.0 + + +def test_cascade_raises_when_no_purchase_record(): """ When the facility has no computepurchases entry for the requested cluster, - the cascade logs an error and returns False without touching any repos. + the cascade raises so the daemon marks the request incomplete rather than + leaving it stuck in Approved. """ handler = make_handler() handler.do_repo_compute_allocation = MagicMock(return_value=True) @@ -107,66 +155,225 @@ def test_cascade_returns_false_when_no_purchase_record(): 'facility': {'computepurchases': [{'clustername': 'other-cluster', 'purchased': 100}]} } - result = handler.do_facility_compute_allocation_cascade('lcls', 'ada', dry_run=False) + with pytest.raises(RuntimeError, match='No purchase record'): + handler.do_facility_compute_allocation_cascade('lcls', 'ada', dry_run=False) - assert result is False handler.do_repo_compute_allocation.assert_not_called() - handler.logger.error.assert_called_once() -def test_cascade_returns_false_when_purchased_nodes_is_zero(): - """ - A purchase record with zero nodes is rejected before any repo is updated. - """ +def test_cascade_raises_on_negative_purchased_nodes(): + """A negative purchase is nonsense and must not reach any repo.""" handler = make_handler() handler.do_repo_compute_allocation = MagicMock(return_value=True) handler.back_channel.execute.return_value = { - 'facility': {'computepurchases': [{'clustername': 'ada', 'purchased': 0}]} + 'facility': {'computepurchases': [{'clustername': 'ada', 'purchased': -1}]} } - result = handler.do_facility_compute_allocation_cascade('lcls', 'ada', dry_run=False) + with pytest.raises(RuntimeError, match='Invalid purchased nodes'): + handler.do_facility_compute_allocation_cascade('lcls', 'ada', dry_run=False) - assert result is False handler.do_repo_compute_allocation.assert_not_called() - handler.logger.error.assert_called_once() -def test_cascade_continues_after_per_repo_failure(): +def test_cascade_zeroes_repo_allocations_when_purchase_is_zero(): """ - If do_repo_compute_allocation raises for one repo, the cascade logs the - error and continues processing the remaining repos, returning True overall. + A facility relinquishing all its nodes is legitimate: repos are driven to + zero rather than the cascade failing and leaving them allocated against + hardware that no longer exists. """ handler = make_handler() + handler.do_repo_compute_allocation = MagicMock(return_value=True) - repos = [ - { - 'Id': 'repo-a', 'name': 'alpha', 'facility': 'lcls', + set_responses(handler, 0, make_repos( + ('alpha', make_allocation('alloc-a', 25.0, 50.0, burst_percent=10.0, burst_allocated=20.0)), + )) + + result = handler.do_facility_compute_allocation_cascade('lcls', 'ada', dry_run=False) + + assert result is True + call = handler.do_repo_compute_allocation.call_args + assert call.args[4] == 0.0 + assert call.kwargs['burst_allocated'] == 0.0 + + +def test_cascade_processes_all_repos_then_raises_on_partial_failure(): + """ + A failure on one repo must not strand the others, but the request must not be + reported complete either - the daemon only marks Incomplete on an exception. + """ + handler = make_handler() + + set_responses(handler, 200, make_repos( + ('alpha', make_allocation('alloc-a', 25.0, 10.0)), + ('beta', make_allocation('alloc-b', 50.0, 10.0)), + )) + handler.do_repo_compute_allocation = MagicMock( + side_effect=[Exception("slurm playbook failed"), True] + ) + + with pytest.raises(RuntimeError, match='lcls:alpha'): + handler.do_facility_compute_allocation_cascade('lcls', 'ada', dry_run=False) + + assert handler.do_repo_compute_allocation.call_count == 2 + + +def test_upsert_sends_burst_and_existing_allocation_id(): + """ + repoComputeAllocationUpsert rebuilds the document, so every field it owns has to + be sent. Passing the existing allocation id makes the API replace that row instead + of matching on (repoid, clustername, start) and inserting a duplicate. + """ + handler = make_handler() + + handler.upsert_repo_compute_allocation( + 'repo-a', 'ada', 25.0, 50.0, START, END, + burst_percent=10.0, burst_allocated=20.0, allocation_id='alloc-a', + ) + + repocompute = handler.back_channel.execute.call_args.args[1]['repocompute'] + assert repocompute['Id'] == 'alloc-a' + assert repocompute['burstPercentOfFacility'] == 10.0 + assert repocompute['burstAllocated'] == 20.0 + + +def test_upsert_omits_allocation_id_when_creating(): + """A repo with no allocation yet on this cluster must fall back to upsert-by-key.""" + handler = make_handler() + + handler.upsert_repo_compute_allocation('repo-a', 'ada', 25.0, 50.0, START, END) + + repocompute = handler.back_channel.execute.call_args.args[1]['repocompute'] + assert 'Id' not in repocompute + + +def test_cascade_targets_the_allocation_row_it_read(): + """ + The cascade updates an existing allocation in place, so it must name the row. + Matching on (repoid, clustername, start) instead risks opening a second period + that jobs.allocationId would not follow. + """ + handler = make_handler() + handler.do_repo_compute_allocation = MagicMock(return_value=True) + + set_responses(handler, 200, make_repos( + ('alpha', make_allocation('alloc-a', 25.0, 10.0)), + )) + + handler.do_facility_compute_allocation_cascade('lcls', 'ada', dry_run=False) + + assert handler.do_repo_compute_allocation.call_args.kwargs['allocation_id'] == 'alloc-a' + + +def _repo_with_slurm(burst_percent=10.0, burst_allocated=20.0): + return { + 'repo': { + 'Id': 'repo-a', 'name': 'alpha', 'facility': 'lcls', 'users': ['someone'], + 'features': [{'name': 'slurm', 'state': True, 'options': []}], + 'computerequirement': 'Normal', 'currentComputeAllocations': [{ 'Id': 'alloc-a', 'clustername': 'ada', - 'percentOfFacility': 25.0, 'allocatedNodesCount': 50.0, + 'percentOfFacility': 25.0, + 'burstPercentOfFacility': burst_percent, 'burstAllocated': burst_allocated, + 'cpus': 10, 'memory': 8, 'nodes': 50, 'gpus': 0, 'start': START, 'end': END, }], - }, - { - 'Id': 'repo-b', 'name': 'beta', 'facility': 'lcls', - 'currentComputeAllocations': [{ - 'Id': 'alloc-b', 'clustername': 'ada', - 'percentOfFacility': 50.0, 'allocatedNodesCount': 100.0, - 'start': START, 'end': END, - }], - }, - ] - handler.back_channel.execute.side_effect = [ - {'facility': {'computepurchases': [{'clustername': 'ada', 'purchased': 200}]}}, - {'repos': repos}, - ] - handler.do_repo_compute_allocation = MagicMock( - side_effect=[Exception("slurm playbook failed"), True] + } + } + + +def test_repo_allocation_preserves_stored_burst_when_not_supplied(): + """ + A caller that knows nothing about burst must not wipe it. The upsert replaces the + whole document, so the stored value has to be read back and re-sent. + """ + handler = make_handler() + handler.run_playbook = MagicMock() + handler.upsert_repo_compute_allocation = MagicMock() + handler.back_channel.execute.side_effect = [_repo_with_slurm(), _repo_with_slurm()] + + handler.do_repo_compute_allocation('alpha', 'lcls', 'ada', 25.0, 50.0, START, END) + + kwargs = handler.upsert_repo_compute_allocation.call_args.kwargs + assert kwargs['burst_percent'] == 10.0 + assert kwargs['burst_allocated'] == 20.0 + + +def test_repo_allocation_uses_supplied_burst_over_stored(): + """An explicit burst from the request or cascade must win over the stored value.""" + handler = make_handler() + handler.run_playbook = MagicMock() + handler.upsert_repo_compute_allocation = MagicMock() + handler.back_channel.execute.side_effect = [_repo_with_slurm(), _repo_with_slurm()] + + handler.do_repo_compute_allocation( + 'alpha', 'lcls', 'ada', 25.0, 50.0, START, END, + burst_percent=30.0, burst_allocated=60.0, ) + kwargs = handler.upsert_repo_compute_allocation.call_args.kwargs + assert kwargs['burst_percent'] == 30.0 + assert kwargs['burst_allocated'] == 60.0 + + +def test_repo_allocation_does_not_target_a_row_unless_told_to(): + """ + A RepoComputeAllocation request may legitimately open a new allocation period, + so the id of the current row must not be assumed. + """ + handler = make_handler() + handler.run_playbook = MagicMock() + handler.upsert_repo_compute_allocation = MagicMock() + handler.back_channel.execute.side_effect = [_repo_with_slurm(), _repo_with_slurm()] + + handler.do_repo_compute_allocation('alpha', 'lcls', 'ada', 25.0, 50.0, START, END) + + assert handler.upsert_repo_compute_allocation.call_args.kwargs['allocation_id'] is None + + +def test_cascade_skips_repos_whose_allocation_is_unchanged(): + """ + Each repo update runs two Ansible playbooks serially inside the subscription + callback, so repos already at the correct node count must be left alone. + """ + handler = make_handler() + handler.do_repo_compute_allocation = MagicMock(return_value=True) + + set_responses(handler, 200, make_repos( + ('alpha', make_allocation('alloc-a', 25.0, 50.0)), # already 25% of 200 + ('beta', make_allocation('alloc-b', 50.0, 10.0)), # stale + )) + result = handler.do_facility_compute_allocation_cascade('lcls', 'ada', dry_run=False) assert result is True - assert handler.do_repo_compute_allocation.call_count == 2 - handler.logger.error.assert_called_once() + assert [c.args[0] for c in handler.do_repo_compute_allocation.call_args_list] == ['beta'] + + +def test_cascade_updates_repo_whose_only_change_is_burst(): + """A stale burst allocation alone is enough to warrant an update.""" + handler = make_handler() + handler.do_repo_compute_allocation = MagicMock(return_value=True) + + set_responses(handler, 200, make_repos( + ('alpha', make_allocation('alloc-a', 25.0, 50.0, burst_percent=10.0, burst_allocated=5.0)), + )) + + handler.do_facility_compute_allocation_cascade('lcls', 'ada', dry_run=False) + + assert handler.do_repo_compute_allocation.call_args.kwargs['burst_allocated'] == 20.0 + + +def test_cascade_warns_when_percentages_exceed_100(): + """Oversubscription is not blocked, but it must be visible in the log.""" + handler = make_handler() + handler.do_repo_compute_allocation = MagicMock(return_value=True) + + set_responses(handler, 200, make_repos( + ('alpha', make_allocation('alloc-a', 70.0, 10.0)), + ('beta', make_allocation('alloc-b', 50.0, 10.0)), + )) + + handler.do_facility_compute_allocation_cascade('lcls', 'ada', dry_run=False) + + assert 'oversubscribed' in handler.logger.warning.call_args.args[0]