On-the-fly FLARE learning for all cluster backends#426
Open
mesonepigreco wants to merge 7 commits into
Open
Conversation
…n split() Move the backend-agnostic FLARE on-the-fly helpers (_predict_with_model, _compute_properties, _write_model, _update_gp, _train_gp, _clean_runs) from AiiDAEnsemble up to the base Ensemble class, so any cluster backend can use them. Add _otf_setup_defaults, _otf_predict, _otf_update_from_structure and _otf_maybe_train_and_write as the public entry points the generic cluster driver will call, plus a _import_flare_learners() lazy importer that keeps flare optional (G9). Use Rydberg/Bohr from the Ensemble module units (G4) and add the 'deepcopy' import (G5). Propagate the 16 OTF state attributes through Ensemble.split() via the new OTF_STATE_ATTRIBUTES class constant. Without this, the sub-ensemble built by compute_ensemble via get_noncomputed() would have gp_model=None and OTF would be silently disabled on any backend going through the compute_ensemble dispatcher (G1). The AiiDA path never splits with OTF active, so this is behavior-preserving.
Delete _predict_with_model, _compute_properties, _write_model, _update_gp, _train_gp and _clean_runs from AiiDAEnsemble (now inherited unchanged from Ensemble). Replace the inline default setup + empty-model guard and the TRAIN SECTION in compute_ensemble with calls to _otf_setup_defaults, _otf_predict and _otf_maybe_train_and_write. Keep the 'AIIDA CALCULATIONS' summary header via _clean_runs(..., label='AIIDA CALCULATIONS'). Drop the now-unused flare imports (compute_mae, get_env_indices, is_std_in_bound, deepcopy); keep FLARE_Atoms for the workchain monitoring loop. Behavior-preserving: the existing tests/aiida_ensemble/ suite (which calls the inherited helpers on an AiiDAEnsemble) passes unchanged.
…kend Add Modules/BaseCluster.py containing the generic cluster interface: BaseCluster (attribute-locking machinery moved verbatim from Cluster, _check_calculator_interface, _pre_compute_hook template hook, compute_jobarray abstract method, the generic ensemble driver compute_ensemble/compute_ensemble_batch with the on-the-fly FLARE hooks, _compute_jobarray_thread, _ingest_result) and DirectCluster (in-process backend consuming DirectCalculator objects). The OTF loop is written ONCE here and inherited by every cluster flavor. Refactor Modules/Cluster.py: Cluster now inherits from BaseCluster and keeps its entire public API (constructor signature, scheduler/transport methods, namelist constants). The generic driver (compute_ensemble, compute_ensemble_batch, compute_single_jobarray) and the attribute locking (__setattr__, __getstate__, __setstate__) are dropped (inherited from BaseCluster); batch_size/job_number/max_recalc/lock are set by BaseCluster.__init__. Cluster implements compute_jobarray (the template hook, with collect_results now outside the lock — safe parallelization improvement with identical results) and _pre_compute_hook (mkdir the local workdir). De-QE-hardcode prepare_input_file: extensions come from calc.input_extension/output_extension (default .pwi/.pwo), the prefix print is guarded with getattr, and extra_input_files are shipped to the cluster (G2). compute_jobarray uses the calc extensions too (G3). get_output_path takes an out_extension argument. Register the new modules in meson.build (ASEClusterRunner, BaseCluster, ClusterCalculators). ClusterCalculators.py and ASEClusterRunner.py are stubs for now (filled in the next commit). Compatibility verified: OpticalQECluster (sets attrs before super().__init__), Relax.py namelist path, test_save_binary pickling, and the full non-release suite pass (the 3 failing tests are pre-existing on master and unrelated).
Add Modules/ClusterCalculators.py with the calculator hierarchy (Axis B): ClusterCalculator (doc-level base), FileIOCalculator (mid-layer with input_extension/output_extension/extra_input_files/get_execution_command), EspressoCalculator (thin formalization of CC.calculators.Espresso), ASEFileCalculator (bridge any picklable ASE calculator to the scheduler clusters via pickle + JSON + the runner), DirectCalculator (mid-layer with compute(structure)), ASEDirectCalculator (wrap any ASE calculator for in-process use with DirectCluster). Future codes are one class each, or use the ASE bridges with no dedicated code at all. Add Modules/ASEClusterRunner.py: the dependency-light runner invoked on the cluster by ASEFileCalculator (python -m sscha.ASEClusterRunner --calc ... --input ... --output ...). Loads the pickled ASE calculator and the JSON-serialized Atoms, computes energy/forces/stress, dumps a JSON result. No sscha submodule imports (sscha/__init__.py is empty, so no Julia bootstrap — G10). LocalCluster.copy_file now uses shutil.copy (instead of the scp shell-out) for a cleaner local-to-local copy; returns the dest path (truthy). Only affects LocalCluster; strictly more robust (G8).
Add tests/test_otf/ with three test modules: - test_otf_base.py: backend-agnostic OTF helpers on the base Ensemble (split propagation, setup defaults, update+predict cycle, clean_runs label). Uses Au+EMT with a FLARE SGP. Skip without flare (G9). - test_cluster_interface.py: the layered Cluster/calculator refactor with no flare needed. Covers the DirectCluster generic driver with a mock calculator, the pairing rules (DirectCluster rejects file calculators, Cluster rejects direct ones), QE backward compatibility of prepare_input_file, the EspressoCalculator adapter, the ASE file bridge (prepare_input + local roundtrip), the ASE direct calculator, and the non-picklable calculator early failure. - test_cluster_emt_otf.py: end-to-end OTF on gold+EMT through both DirectCluster and LocalCluster (parametrized), proving the driver is shared. Validates the full chain (calculator -> raw results -> ingestion -> unit conversions) against bare EMT values, and that a second ensemble reuses the trained model with zero new ab-initio calls. Add the 'flare' marker to pytest.ini. Workaround for a flare C extension bug: SGP_Wrapper does not keep a Python reference to the kernel objects (only to the descriptor calculators), so the kernel's pybind11 wrapper is garbage-collected when the helper function returns, leaving the C++ SparseGP with a dangling pointer → segfault in add_training_structure. The tests keep kernels alive in a module-level list (_KERNEL_REFS).
…ter) Two end-to-end example scripts from the plan (section 9): - run_gold_otf_directcluster.py: the simplest possible OTF run, in-process ASE calculator (no files, no scheduler, no HPC). - run_gold_otf_localcluster.py: exercises the full scheduler/file machinery (input files, tar, submission script, result retrieval) mocked locally with bash + shutil (no ssh/scp, no SLURM). Both use a FLARE SGP that learns EMT on the fly. The learning cycle is 2 structures (batch_size=2, job_number=1); after the first cycle trains the GP, subsequent structures are predicted by the model.
Collaborator
Author
|
This merge will fix and close issue #425 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Enable on-the-fly (OTF) FLARE active learning for ensembles computed through the standard
Cluster/LocalClustersubmission machinery, not only throughAiiDAEnsemble. The refactor hoists the FLARE helpers into the baseEnsembleclass (Phase 1), propagates OTF state throughEnsemble.split()(Phase 2), introduces a layered cluster architecture (BaseCluster+Cluster/DirectCluster) and a calculator hierarchy (ClusterCalculators) so that OTF works with any backend (Phase 3), and wires the OTF loop into the generic ensemble driver (Phase 4). Adds end-to-end tests with the EMT calculator.