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