diff --git a/src/dara/hardware.py b/src/dara/hardware.py new file mode 100644 index 0000000..cef2f81 --- /dev/null +++ b/src/dara/hardware.py @@ -0,0 +1,96 @@ +"""Portable CPU core-count detection for Ray + BGMN concurrency sizing. + +Priority, most authoritative first: + +1. An explicit SLURM allocation (``SLURM_CPUS_PER_TASK``, then + ``SLURM_CPUS_ON_NODE``) -- trusts a scheduler-assigned budget over + anything detected locally. +2. OS-level CPU affinity (``os.sched_getaffinity``, Linux only), narrowed by + a cgroup CPU quota if one is set and implies fewer cores than affinity -- + this matters on shared workstations, containers, and HPC nodes. +3. ``os.cpu_count()`` -- universal fallback, also the only option on macOS, + where ``sched_getaffinity`` does not exist. +""" + +from __future__ import annotations + +import math +import os +from pathlib import Path + + +def _slurm_cpu_count() -> int | None: + for var in ("SLURM_CPUS_PER_TASK", "SLURM_CPUS_ON_NODE"): + value = os.environ.get(var) + if not value: + continue + try: + n = int(value) + except ValueError: + continue + if n > 0: + return n + return None + + +def _cgroup_quota_cpu_count() -> int | None: + """Return the core count implied by a cgroup CPU quota, if one is set.""" + cgroup_v2_path = Path("/sys/fs/cgroup/cpu.max") + if cgroup_v2_path.exists(): + try: + quota_str, period_str = cgroup_v2_path.read_text().split() + except (OSError, ValueError): + quota_str = period_str = None + if quota_str is not None and quota_str != "max": + try: + quota, period = int(quota_str), int(period_str) + if quota > 0 and period > 0: + return max(1, math.floor(quota / period)) + except ValueError: + pass + + quota_path = Path("/sys/fs/cgroup/cpu/cpu.cfs_quota_us") + period_path = Path("/sys/fs/cgroup/cpu/cpu.cfs_period_us") + if quota_path.exists() and period_path.exists(): + try: + quota = int(quota_path.read_text().strip()) + period = int(period_path.read_text().strip()) + if quota > 0 and period > 0: + return max(1, math.floor(quota / period)) + except (OSError, ValueError): + pass + + return None + + +def _affinity_cpu_count() -> int | None: + sched_getaffinity = getattr(os, "sched_getaffinity", None) + if sched_getaffinity is None: + return None + try: + affinity = len(sched_getaffinity(0)) + except OSError: + return None + return affinity if affinity > 0 else None + + +def detect_available_cores() -> int: + """ + Detect the number of CPU cores actually available to this process. + + Returns + ------- + the detected core count, always >= 1. Never returns None or 0. + """ + slurm = _slurm_cpu_count() + if slurm is not None: + return slurm + + affinity = _affinity_cpu_count() + if affinity is not None: + cgroup = _cgroup_quota_cpu_count() + if cgroup is not None: + return max(1, min(affinity, cgroup)) + return affinity + + return max(1, os.cpu_count() or 1) diff --git a/src/dara/search/core.py b/src/dara/search/core.py index fc2167b..d1d1362 100644 --- a/src/dara/search/core.py +++ b/src/dara/search/core.py @@ -11,6 +11,7 @@ from dara.search.data_model import PeakMatchingStrategy from dara.search.tree import BaseSearchTree, SearchTree +from dara.settings import DaraSettings if TYPE_CHECKING: from pathlib import Path @@ -26,7 +27,21 @@ "b1": "0_0^0.005", "rp": 4, } -DEFAULT_REFINEMENT_PARAMS = {"n_threads": 8, "eps1": 0, "eps2": "0_-0.05^0.05"} + +# Portable CPU/thread budget: BGMN threads default to 1 (see +# DaraSettings.BGMN_N_THREADS) so worker *concurrency*, not per-task +# threading, scales with available cores (see dara.hardware and the +# num_cpus=... passed to ray.init() below and to each refinement task's +# .options(num_cpus=...) in dara.search.tree.batch_refinement). Both are +# overridable via DaraSettings (~/.dara.yaml or the DARA_RAY_NUM_CPUS / +# DARA_BGMN_N_THREADS env vars). +_settings = DaraSettings() + +DEFAULT_REFINEMENT_PARAMS = { + "n_threads": _settings.BGMN_N_THREADS, + "eps1": 0, + "eps2": "0_-0.05^0.05", +} DEFAULT_PEAK_MATCHING_STRATEGY = PeakMatchingStrategy.default() @@ -100,7 +115,12 @@ def search_phases( refinement_params = {} if not ray.is_initialized(): - ray.init(runtime_env={"working_dir": None}) + # num_cpus is set explicitly to the detected/allocated core count + # (DaraSettings.RAY_NUM_CPUS) rather than left to Ray's own + # auto-detection, which just calls os.cpu_count() and would ignore + # an explicit SLURM allocation or cgroup/container CPU affinity + # restriction (see dara.hardware.detect_available_cores). + ray.init(runtime_env={"working_dir": None}, num_cpus=_settings.RAY_NUM_CPUS) phase_params = {**DEFAULT_PHASE_PARAMS, **phase_params} refinement_params = {**DEFAULT_REFINEMENT_PARAMS, **refinement_params} diff --git a/src/dara/search/tree.py b/src/dara/search/tree.py index e55e91e..73f994e 100644 --- a/src/dara/search/tree.py +++ b/src/dara/search/tree.py @@ -20,6 +20,7 @@ from dara.refine import RefinementPhase from dara.search.data_model import PeakMatchingStrategy, SearchNodeData, SearchResult from dara.search.peak_matcher import PeakMatcher +from dara.settings import DaraSettings from dara.utils import ( estimate_rpb_threshold, find_optimal_intensity_threshold, @@ -39,8 +40,20 @@ logger = get_logger(__name__, level="INFO") +# Defensive fallback for callers that build a SearchTree directly (bypassing +# search_phases()'s merge of DEFAULT_REFINEMENT_PARAMS, which always sets +# n_threads) without an explicit n_threads in refinement_params. See +# dara.hardware and dara.search.core for the portable core/thread sizing +# this is part of. +_DEFAULT_BGMN_N_THREADS = DaraSettings().BGMN_N_THREADS -@ray.remote(num_cpus=1) + +# num_cpus is intentionally not fixed here -- it is set per-submission via +# .options(num_cpus=...) in batch_refinement() to match the actual n_threads +# a given task's BGMN subprocess will use, so Ray's scheduler throttles +# concurrency to roughly floor(cores / n_threads) instead of assuming a +# hardcoded num_cpus=1 regardless of how many OS threads the task spawns. +@ray.remote def remote_do_refinement_no_saving( pattern_path: Path, cif_paths: list[Path], @@ -129,8 +142,9 @@ def batch_refinement( phase_params: dict[str, ...] | None = None, refinement_params: dict[str, float] | None = None, ) -> list[RefinementResult]: + n_threads = (refinement_params or {}).get("n_threads", _DEFAULT_BGMN_N_THREADS) handles = [ - remote_do_refinement_no_saving.remote( + remote_do_refinement_no_saving.options(num_cpus=n_threads).remote( pattern_path, cif_paths, wavelength=wavelength, @@ -324,6 +338,64 @@ def remove_unnecessary_phases( return new_phases +# Minimum relative Rwp improvement over the parent node required to keep an +# out-of-order-intensity branch instead of pruning it (see +# `should_prune_low_weight_fraction`). Lower Rwp is better, so relative +# improvement is (parent_rwp - child_rwp) / parent_rwp; 0.08 means the child +# must fit at least 8% better than its parent, relatively, to survive. +LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT = 0.08 + + +def should_prune_low_weight_fraction( + intensity_out_of_order: bool, + parent_rwp: float | None, + child_rwp: float, + material_improvement_threshold: float = LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT, +) -> bool: + """ + Decide whether an out-of-order-intensity branch should be pruned. + + Phases are expected to be discovered in roughly decreasing order of + abundance (the naive peak-match search tends to find the biggest + contributor to the pattern first), so a newly-added phase that turns out + to have *more* calculated peak intensity than a phase added earlier in + the same branch (`intensity_out_of_order`) is treated as suspicious by + default -- it can mean the refinement is using this phase to absorb + residual intensity that doesn't belong to it (overfitting), rather than + genuinely explaining new pattern features. + + But that's only actually suspicious if the phase isn't earning its keep. + If adding it materially improved the fit (Rwp) over the parent node, the + improvement is real signal, not an ordering artifact, regardless of what + order the search happened to try phases in -- so the branch is kept. + + Args: + intensity_out_of_order: whether the newly-added phase has more + calculated peak intensity than an earlier phase in the same + branch (the raw ordering signal; computed by the caller). + parent_rwp: the Rwp of the parent node, or None if the parent is the + (unrefined) root -- with no parent to compare against, the + ordering signal alone determines the outcome. + child_rwp: the Rwp of the new (child) node being evaluated. + material_improvement_threshold: the minimum relative Rwp + improvement required to treat an out-of-order branch as + materially improved rather than pruning it. + + Returns + ------- + True if the branch should be pruned as a low-weight-fraction + (suspicious, non-material) out-of-order addition; False if it + should be kept. + """ + if not intensity_out_of_order: + return False + if parent_rwp is None or parent_rwp <= 0: + return True + + relative_improvement = (parent_rwp - child_rwp) / parent_rwp + return relative_improvement < material_improvement_threshold + + def get_natural_break_results( results: list[SearchResult], sorting: bool = True ) -> list[SearchResult]: @@ -491,9 +563,13 @@ def expand_node(self, nid: str) -> list[str]: ]["intensity"].sum(), reverse=True, ) - # make sure the newly added phase has the lowest peak intensity - is_low_weight_fraction = ( - sorted_searched_phases[-1] != searched_phases[-1] + # the newly added phase does not have the lowest peak intensity + intensity_out_of_order = sorted_searched_phases[-1] != searched_phases[-1] + parent_rwp = node.data.current_result.lst_data.rwp if node.data.current_result is not None else None + is_low_weight_fraction = should_prune_low_weight_fraction( + intensity_out_of_order, + parent_rwp, + new_result.lst_data.rwp, ) else: is_low_weight_fraction = False diff --git a/src/dara/settings.py b/src/dara/settings.py index 85a58bb..ecee36f 100644 --- a/src/dara/settings.py +++ b/src/dara/settings.py @@ -6,6 +6,8 @@ from pydantic import Field, model_validator from pydantic_settings import BaseSettings, SettingsConfigDict +from dara.hardware import detect_available_cores + _DEFAULT_CONFIG_FILE_PATH = "~/.dara.yaml" @@ -23,6 +25,29 @@ class DaraSettings(BaseSettings): PATH_TO_ICSD: Path = Field(Path("~/ICSD_2024/ICSD_2024_experimental_inorganic/experimental_inorganic").expanduser()) PATH_TO_COD: Path = Field(Path("~/COD_2024").expanduser()) + RAY_NUM_CPUS: int = Field( + default_factory=detect_available_cores, + description=( + "Number of CPUs to give Ray's own resource accounting (passed as " + "num_cpus= to ray.init()). Auto-detected from an explicit SLURM " + "allocation, then OS-level CPU affinity/cgroup quota, then the " + "total CPU count -- see dara.hardware.detect_available_cores. " + "Override via the DARA_RAY_NUM_CPUS env var or ~/.dara.yaml." + ), + ) + BGMN_N_THREADS: int = Field( + 1, + description=( + "Threads per BGMN refinement subprocess (the n_threads refinement " + "param). Kept low (default 1) so worker concurrency -- not " + "per-task threading -- scales with available cores: refinement " + "Ray tasks are submitted with num_cpus=BGMN_N_THREADS, so " + "concurrency stays near floor(cores / BGMN_N_THREADS) " + "automatically even if this is overridden upward. Override via " + "the DARA_BGMN_N_THREADS env var or ~/.dara.yaml." + ), + ) + model_config = SettingsConfigDict(env_prefix="dara_") # prepend dara_ to env vars @model_validator(mode="before") diff --git a/tests/test_search_tree.py b/tests/test_search_tree.py new file mode 100644 index 0000000..3508c0d --- /dev/null +++ b/tests/test_search_tree.py @@ -0,0 +1,264 @@ +import unittest +from pathlib import Path +from unittest.mock import patch + +import pandas as pd +from treelib import Node, Tree + +from dara.refine import RefinementPhase +from dara.result import DiaResult, LstResult, PhaseResult, RefinementMetrics, RefinementResult +from dara.search.data_model import SearchNodeData +from dara.search.tree import ( + LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT, + BaseSearchTree, + should_prune_low_weight_fraction, +) + + +def _phase(name: str) -> RefinementPhase: + return RefinementPhase.make(Path(f"/fake/cifs/{name}.cif")) + + +def _phase_result() -> PhaseResult: + """A structurally-valid but otherwise-meaningless PhaseResult. + + `calculate_fom_and_strain` (the only code that reads these fields) is + mocked out in the integration test below, so content doesn't matter -- + it only needs to satisfy LstResult.phases_results' pydantic validation. + """ + return PhaseResult( + SpacegroupNo=None, + HermannMauguin=None, + XrayDensity=None, + Rphase=None, + UNIT="nm", + GEWICHT=1.0, + GEWICHT_NAME=None, + ) + + +def _make_refinement_result( + rwp: float, + phase_intensities: dict[str, float], +) -> RefinementResult: + """Build a minimal but genuinely valid (pydantic-validated) RefinementResult. + + `phase_intensities` maps phase stem -> total calculated peak intensity + for that phase; both `peak_data` (read by the intensity-ordering check) + and `plot_data.structs` (read by `remove_unnecessary_phases`) are built + from it so removing any one phase changes the calculated pattern enough + for `remove_unnecessary_phases` to consider it necessary. + + `refinement_metrics` is a real `RefinementMetrics` instance, matching how + production code builds one (`RefinementMetrics(rwp=lst_data.rwp)` in + `get_result`/`refine.py`) -- its only required field is `rwp`, so this + passes the same `rwp` used for `lst_data`. `missing_peaks`/`extra_peaks`/ + `intensity_mismatch_peaks` are irrelevant to the pruning-decision logic + under test and are left at their `None` defaults. + """ + peak_rows = [ + {"phase": phase, "2theta": 10.0 + i, "intensity": intensity} + for i, (phase, intensity) in enumerate(phase_intensities.items()) + ] + structs = {phase: [intensity] for phase, intensity in phase_intensities.items()} + y_calc_total = sum(phase_intensities.values()) + + return RefinementResult( + lst_data=LstResult( + raw_lst="", + pattern_name="fake", + num_steps=1, + Rp=0.0, + Rpb=rwp, + R=0.0, + Rwp=rwp, + Rexp=1.0, + d=0.0, + **{"1-rho": rwp}, + phases_results={phase: _phase_result() for phase in phase_intensities}, + ), + plot_data=DiaResult( + x=[0.0], + y_obs=[y_calc_total], + y_calc=[y_calc_total], + y_bkg=[0.0], + structs=structs, + ), + peak_data=pd.DataFrame(peak_rows), + refinement_metrics=RefinementMetrics(rwp=rwp), + ) + + +class TestShouldPruneLowWeightFraction(unittest.TestCase): + """Direct unit tests for the standalone pruning-decision helper.""" + + def test_in_order_never_pruned(self): + """Not out-of-order: the parent/child fit is irrelevant, never prune.""" + self.assertFalse( + should_prune_low_weight_fraction(intensity_out_of_order=False, parent_rwp=40.0, child_rwp=39.9) + ) + self.assertFalse( + should_prune_low_weight_fraction(intensity_out_of_order=False, parent_rwp=40.0, child_rwp=100.0) + ) + self.assertFalse( + should_prune_low_weight_fraction(intensity_out_of_order=False, parent_rwp=None, child_rwp=39.9) + ) + + def test_out_of_order_no_parent_falls_back_to_ordering_only(self): + """With no parent to compare against (root's first child), the + ordering signal alone determines the outcome -- prune.""" + self.assertTrue(should_prune_low_weight_fraction(intensity_out_of_order=True, parent_rwp=None, child_rwp=10.0)) + + def test_out_of_order_zero_or_negative_parent_rwp_is_pruned(self): + """Degenerate parent_rwp <= 0 can't support a relative-improvement + calculation (division by zero); fall back to the ordering-only rule.""" + self.assertTrue(should_prune_low_weight_fraction(intensity_out_of_order=True, parent_rwp=0.0, child_rwp=5.0)) + self.assertTrue(should_prune_low_weight_fraction(intensity_out_of_order=True, parent_rwp=-1.0, child_rwp=5.0)) + + def test_out_of_order_marginal_improvement_still_pruned(self): + """Out-of-order + improvement below the material threshold: pruned.""" + parent_rwp = 40.0 + # 5% relative improvement, below the 8% default threshold + child_rwp = 38.0 + self.assertTrue( + should_prune_low_weight_fraction(intensity_out_of_order=True, parent_rwp=parent_rwp, child_rwp=child_rwp) + ) + + def test_out_of_order_material_improvement_kept(self): + """Out-of-order + material improvement: this is the bug fix -- a + branch that fits clearly better than its parent must survive even + though the newest phase isn't the smallest.""" + parent_rwp = 40.0 + # 50% relative improvement + child_rwp = 20.0 + self.assertFalse( + should_prune_low_weight_fraction(intensity_out_of_order=True, parent_rwp=parent_rwp, child_rwp=child_rwp) + ) + + def test_improvement_exactly_at_threshold_is_kept(self): + """Exactly LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT relative improvement + counts as material (the check is `< threshold`, not `<= threshold`). + + Nudged by a tiny epsilon so the comparison lands unambiguously on + the "kept" side of the boundary regardless of floating-point + rounding in the parent_rwp * (1 - threshold) computation. + """ + parent_rwp = 40.0 + child_rwp = parent_rwp * (1 - LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT) - 1e-9 + self.assertFalse( + should_prune_low_weight_fraction(intensity_out_of_order=True, parent_rwp=parent_rwp, child_rwp=child_rwp) + ) + + def test_improvement_just_under_threshold_is_pruned(self): + parent_rwp = 40.0 + child_rwp = parent_rwp * (1 - LOW_WEIGHT_FRACTION_RWP_IMPROVEMENT) + 1e-6 + self.assertTrue( + should_prune_low_weight_fraction(intensity_out_of_order=True, parent_rwp=parent_rwp, child_rwp=child_rwp) + ) + + def test_worse_fit_is_pruned(self): + """Out-of-order + the child fits WORSE than the parent (negative + 'improvement'): definitely pruned.""" + self.assertTrue(should_prune_low_weight_fraction(intensity_out_of_order=True, parent_rwp=20.0, child_rwp=30.0)) + + def test_custom_threshold_is_respected(self): + parent_rwp, child_rwp = 40.0, 30.0 # 25% relative improvement + self.assertTrue( + should_prune_low_weight_fraction( + intensity_out_of_order=True, + parent_rwp=parent_rwp, + child_rwp=child_rwp, + material_improvement_threshold=0.5, + ) + ) + self.assertFalse( + should_prune_low_weight_fraction( + intensity_out_of_order=True, + parent_rwp=parent_rwp, + child_rwp=child_rwp, + material_improvement_threshold=0.2, + ) + ) + + +def _make_test_tree() -> BaseSearchTree: + """A minimal, real (treelib-backed) BaseSearchTree for driving the actual + `expand_node()` code path without needing a real pattern file or BGMN. + """ + tree = BaseSearchTree.__new__(BaseSearchTree) + Tree.__init__(tree) + tree.pinned_phases = [] + tree.all_phases_result = {} + tree.max_phases = 5 + tree.rpb_threshold = 0.5 + tree.record_peak_matcher_scores = False + tree.express_mode = False + tree.maximum_grouping_distance = 0.1 + tree.peak_obs = pd.DataFrame({"2theta": [10.0, 11.0], "intensity": [100.0, 50.0]})[["2theta", "intensity"]].values + return tree + + +@patch("dara.search.tree.calculate_fom_and_strain", return_value=(0.5, 0.0)) +class TestExpandNodeLowWeightFractionIntegration(unittest.TestCase): + """Integration-level tests exercising the real `expand_node()` method + (real `remove_unnecessary_phases`, real intensity-ordering computation, + real status-decision branching) with `score_phases`/`refine_phases` + mocked to avoid needing BGMN, and `calculate_fom_and_strain` mocked + since it reads a phase's own CIF file (irrelevant to this decision). + """ + + def _run_expand(self, parent_rwp, child_rwp, phase_a_intensity, phase_b_intensity): + """Set up a root -> phase_a chain, then expand with phase_b added. + + phase_b's calculated intensity is set higher than phase_a's, so the + addition is out-of-order by construction; parent_rwp/child_rwp + control whether it's a material improvement. + """ + tree = _make_test_tree() + phase_a, phase_b = _phase("A"), _phase("B") + + parent_result = _make_refinement_result(parent_rwp, {"A": phase_a_intensity}) + root = Node( + data=SearchNodeData( + current_result=parent_result, + current_phases=[phase_a], + status="pending", + ) + ) + tree.add_node(root) + + child_result = _make_refinement_result(child_rwp, {"A": phase_a_intensity, "B": phase_b_intensity}) + + with ( + patch.object(tree, "score_phases", return_value=([phase_b], {}, 0)), + patch.object(tree, "refine_phases", return_value={phase_b: child_result}), + ): + tree.expand_node(root.identifier) + + children = tree.children(root.identifier) + self.assertEqual(len(children), 1) + return children[0].data.status + + def test_material_improvement_is_retained(self, mock_fom): + """Out-of-order (B's intensity > A's) but a clearly better fit + (50% relative Rwp improvement): must NOT be pruned as low-weight-fraction.""" + status = self._run_expand(parent_rwp=40.0, child_rwp=20.0, phase_a_intensity=60.0, phase_b_intensity=200.0) + self.assertNotEqual(status, "low_weight_fraction") + self.assertEqual(status, "pending") + + def test_marginal_improvement_is_still_pruned(self, mock_fom): + """Out-of-order (B's intensity > A's) with only a marginal (5%) + Rwp improvement: still pruned as low-weight-fraction, as before.""" + status = self._run_expand(parent_rwp=40.0, child_rwp=38.0, phase_a_intensity=60.0, phase_b_intensity=200.0) + self.assertEqual(status, "low_weight_fraction") + + def test_in_order_addition_never_flagged_regardless_of_fit(self, mock_fom): + """B's intensity is LOWER than A's (properly ordered): never flagged, + even with a poor Rwp improvement -- the ordering check simply + doesn't fire, independent of the material-improvement carve-out.""" + status = self._run_expand(parent_rwp=40.0, child_rwp=39.9, phase_a_intensity=200.0, phase_b_intensity=60.0) + self.assertNotEqual(status, "low_weight_fraction") + + +if __name__ == "__main__": + unittest.main()