From da838a2a9aa12a53bf6004c2094c9e2e4a2c766e Mon Sep 17 00:00:00 2001 From: Ryan Waldheim Date: Wed, 6 May 2026 14:50:59 -0700 Subject: [PATCH 1/7] 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 2/7] 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 3/7] 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 4/7] 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 5/7] 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 6/7] 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 7/7] 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()