From f088fe8955870746a774918952c939d1c5e0b6ab Mon Sep 17 00:00:00 2001 From: Jonathan Brodrick Date: Mon, 16 Feb 2026 21:46:55 +0000 Subject: [PATCH 1/3] chore: log tests to mlflow in CI --- .github/workflows/cpu-tests.yaml | 9 +++++++++ pyproject.toml | 1 + 2 files changed, 10 insertions(+) diff --git a/.github/workflows/cpu-tests.yaml b/.github/workflows/cpu-tests.yaml index c8cfe2c3..0d81182b 100644 --- a/.github/workflows/cpu-tests.yaml +++ b/.github/workflows/cpu-tests.yaml @@ -13,6 +13,15 @@ on: # Allows you to run this workflow manually from the Actions tab workflow_dispatch: +# Environment variables available to all jobs +env: + BRANCH_NAME: ${{ github.head_ref || github.ref_name }} + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID_UPLOAD }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY_UPLOAD }} + MLFLOW_TRACKING_URI: "https://continuum.ergodic.io/experiments/" + MLFLOW_TRACKING_USERNAME: "githubactions-bot" + MLFLOW_TRACKING_PASSWORD: ${{ secrets.MLFLOW_TRACKING_PASSWORD }} + # Cancel in-progress tests on new commits to the same branch concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} diff --git a/pyproject.toml b/pyproject.toml index 5d1a0db0..36c96e8a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,6 +51,7 @@ docs = [ "adept" ] dev = [ + "gitpython", "pre-commit", "pytest", "pytest-cov", From 6c592cf172a8f2bc4fba7420ec84b703375f6143 Mon Sep 17 00:00:00 2001 From: Jonathan Brodrick Date: Tue, 24 Feb 2026 14:21:58 +0000 Subject: [PATCH 2/3] update uv lock and rename some experiments --- tests/test_tf1d/configs/resonance_search.yaml | 4 ++-- tests/test_vfp1d/epp-short-reflective.yaml | 2 +- tests/test_vfp1d/epp-short.yaml | 4 ++-- uv.lock | 2 ++ 4 files changed, 7 insertions(+), 5 deletions(-) diff --git a/tests/test_tf1d/configs/resonance_search.yaml b/tests/test_tf1d/configs/resonance_search.yaml index e760afff..d7a7019c 100644 --- a/tests/test_tf1d/configs/resonance_search.yaml +++ b/tests/test_tf1d/configs/resonance_search.yaml @@ -1,8 +1,8 @@ solver: tf-1d mlflow: - experiment: tf1d-resonance-search - run: test + experiment: test-tf1d + run: resonance-search adjoint: Backsolve diff --git a/tests/test_vfp1d/epp-short-reflective.yaml b/tests/test_vfp1d/epp-short-reflective.yaml index 4e13b67b..67dee891 100644 --- a/tests/test_vfp1d/epp-short-reflective.yaml +++ b/tests/test_vfp1d/epp-short-reflective.yaml @@ -58,7 +58,7 @@ save: solver: vfp-1d mlflow: - experiment: vfp1d + experiment: test-vfp1d-eppshort run: epperlein-short-reflective drivers: diff --git a/tests/test_vfp1d/epp-short.yaml b/tests/test_vfp1d/epp-short.yaml index ea3d31d9..a11fcb0b 100644 --- a/tests/test_vfp1d/epp-short.yaml +++ b/tests/test_vfp1d/epp-short.yaml @@ -57,8 +57,8 @@ save: solver: vfp-1d mlflow: - experiment: vfp1d - run: epperlein-short + experiment: test-vfp1d-eppshort + run: epperlein-short-periodic drivers: ex: {} diff --git a/uv.lock b/uv.lock index 32bf39bc..ac711197 100644 --- a/uv.lock +++ b/uv.lock @@ -45,6 +45,7 @@ dependencies = [ [package.optional-dependencies] dev = [ + { name = "gitpython" }, { name = "pre-commit" }, { name = "pytest" }, { name = "pytest-cov" }, @@ -77,6 +78,7 @@ requires-dist = [ { name = "adept", extras = ["docs"], marker = "extra == 'dev'" }, { name = "boto3" }, { name = "diffrax" }, + { name = "gitpython", marker = "extra == 'dev'" }, { name = "h5netcdf" }, { name = "interpax", git = "https://github.com/f0uriest/interpax.git" }, { name = "jax" }, From fe6d5feec3172108fbad9337ed2f4a6765cc084e Mon Sep 17 00:00:00 2001 From: Jonathan Brodrick Date: Tue, 24 Feb 2026 20:30:13 +0000 Subject: [PATCH 3/3] rename tests and log git tags --- tests/conftest.py | 19 +++++++++++++++++++ tests/test_base/configs/example.yaml | 2 +- tests/test_base/test_ergoExo.py | 7 ++++++- tests/test_lpse2d/configs/epw.yaml | 2 +- .../test_lpse2d/configs/resonance_search.yaml | 2 +- tests/test_lpse2d/configs/tpd.yaml | 2 +- tests/test_lpse2d/test_tpd_threshold.py | 12 ++++++++---- .../configs/landau-damping-test.yaml | 2 +- tests/test_spectrax1d/test_landau_damping.py | 10 +++++++--- tests/test_spectrax1d/test_maxwell.py | 14 +++++++++----- tests/test_tf1d/configs/resonance.yaml | 2 +- tests/test_tf1d/configs/resonance_search.yaml | 2 +- .../test_tf1d/configs/vlasov_comparison.yaml | 2 +- tests/test_tf1d/test_against_vlasov.py | 7 +++++-- tests/test_tf1d/test_landau_damping.py | 5 +++-- tests/test_tf1d/test_resonance.py | 7 +++++-- tests/test_tf1d/test_resonance_search.py | 6 ++++-- tests/test_tf1d/vlasov-reference/config.yaml | 2 +- tests/test_vfp1d/epp-short-reflective.yaml | 2 +- tests/test_vfp1d/epp-short.yaml | 2 +- tests/test_vfp1d/test_kappa_eh.py | 13 ++++++++----- .../configs/fokker_planck_conservation.yaml | 2 +- .../configs/multispecies_ion_acoustic.yaml | 2 +- tests/test_vlasov1d/configs/resonance.yaml | 2 +- .../test_fokker_planck_conservation.py | 9 +++++++-- tests/test_vlasov1d/test_ion_acoustic_wave.py | 13 +++++++++---- tests/test_vlasov1d/test_landau_damping.py | 5 ++++- tests/test_vlasov2d/configs/damping.yaml | 2 +- tests/test_vlasov2d/test_landau_damping.py | 2 +- 29 files changed, 110 insertions(+), 49 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index b6031b08..910d7bd5 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,5 +1,24 @@ +import os + +import git import jax import numpy as np +import pytest jax.config.update("jax_enable_x64", True) np.seterr(divide="ignore") + + +@pytest.fixture +def tags(): + repo = git.Repo() + try: + branch_name = os.environ["BRANCH_NAME"] + except KeyError: + branch_name = repo.active_branch.name + return dict( + branch=branch_name, + commit_id=repo.head.object.hexsha, + commit_author=repo.head.object.author.name, + commit_author_email=repo.head.object.author.email, + ) diff --git a/tests/test_base/configs/example.yaml b/tests/test_base/configs/example.yaml index 38c039cf..e8f1a2e2 100644 --- a/tests/test_base/configs/example.yaml +++ b/tests/test_base/configs/example.yaml @@ -50,7 +50,7 @@ save: solver: vlasov-1d mlflow: - experiment: vlasov1d + experiment: test-adept-vlasov1d run: test drivers: diff --git a/tests/test_base/test_ergoExo.py b/tests/test_base/test_ergoExo.py index 39f1bd13..0497439d 100644 --- a/tests/test_base/test_ergoExo.py +++ b/tests/test_base/test_ergoExo.py @@ -1,14 +1,19 @@ import yaml +import adept.patched_mlflow as mlflow from adept import ergoExo -def test_reuse_config_dict(): +def test_reuse_config_dict(tags): with open("tests/test_base/configs/example.yaml") as file: cfg = yaml.safe_load(file) exo = ergoExo() exo.setup(cfg) + with mlflow.start_run(run_id=exo.mlflow_run_id, nested=exo.mlflow_nested): + mlflow.set_tags(tags) exo = ergoExo() exo.setup(cfg) + with mlflow.start_run(run_id=exo.mlflow_run_id, nested=exo.mlflow_nested): + mlflow.set_tags(tags) diff --git a/tests/test_lpse2d/configs/epw.yaml b/tests/test_lpse2d/configs/epw.yaml index 2edb2163..4c7e9e4b 100644 --- a/tests/test_lpse2d/configs/epw.yaml +++ b/tests/test_lpse2d/configs/epw.yaml @@ -33,7 +33,7 @@ grid: ymin: -1um mlflow: - experiment: test-lpse + experiment: test-adept-lpse2d-epw run: epw-test save: fields: diff --git a/tests/test_lpse2d/configs/resonance_search.yaml b/tests/test_lpse2d/configs/resonance_search.yaml index ade6d2f6..b1bb6937 100644 --- a/tests/test_lpse2d/configs/resonance_search.yaml +++ b/tests/test_lpse2d/configs/resonance_search.yaml @@ -72,7 +72,7 @@ grid: dt: 0.5 mlflow: - experiment: epw + experiment: test-adept-lpse2d-resonance-search run: compare-against-vlasov terms: diff --git a/tests/test_lpse2d/configs/tpd.yaml b/tests/test_lpse2d/configs/tpd.yaml index cde62e05..5f8e94ae 100644 --- a/tests/test_lpse2d/configs/tpd.yaml +++ b/tests/test_lpse2d/configs/tpd.yaml @@ -36,7 +36,7 @@ grid: ymax: 3um ymin: -3um mlflow: - experiment: tpd + experiment: test-adept-lpse2d-tpd run: 1.5e14 save: fields: diff --git a/tests/test_lpse2d/test_tpd_threshold.py b/tests/test_lpse2d/test_tpd_threshold.py index 0deaa35b..e9309fa5 100644 --- a/tests/test_lpse2d/test_tpd_threshold.py +++ b/tests/test_lpse2d/test_tpd_threshold.py @@ -5,9 +5,10 @@ from adept.lpse2d import calc_threshold_intensity -def run_once(L, Te, I0): +def run_once(L, Te, I0, tags=None): import yaml + import adept.patched_mlflow as mlflow from adept import ergoExo with open("tests/test_lpse2d/configs/tpd.yaml") as fi: @@ -17,17 +18,20 @@ def run_once(L, Te, I0): cfg["density"]["gradient scale length"] = f"{L}um" cfg["units"]["reference electron temperature"] = f"{Te}keV" cfg["mlflow"]["run"] = f"{I0}W/cm^2" # I0 - cfg["mlflow"]["experiment"] = f"I2-threshold-L={L}um, Te={Te}keV" + cfg["mlflow"]["experiment"] = f"test-adept-lpse2d-tpd-threshold-L={L}um-Te={Te}keV" exo = ergoExo() modules = exo.setup(cfg) + if tags is not None: + with mlflow.start_run(run_id=exo.mlflow_run_id, nested=exo.mlflow_nested): + mlflow.set_tags(tags) sol, ppo, mlrunid = exo(modules) es = ppo["metrics"]["log10_total_e_sq"] return es -def test_threshold(): +def test_threshold(tags): if not any(["gpu" == device.platform for device in devices()]): pytest.skip("Takes too long without a GPU") else: @@ -44,7 +48,7 @@ def test_threshold(): I_scan = np.round(I_scan, 2) for I0 in I_scan: - es = run_once(L, Te, I0) + es = run_once(L, Te, I0, tags) ess.append(es) # ess = np.array(ess) diff --git a/tests/test_spectrax1d/configs/landau-damping-test.yaml b/tests/test_spectrax1d/configs/landau-damping-test.yaml index 39e16888..a6301e37 100644 --- a/tests/test_spectrax1d/configs/landau-damping-test.yaml +++ b/tests/test_spectrax1d/configs/landau-damping-test.yaml @@ -4,7 +4,7 @@ solver: spectrax-1d mlflow: - experiment: spectrax1d-test + experiment: test-adept-spectrax1d-landau-damping run: landau-damping-k266 units: diff --git a/tests/test_spectrax1d/test_landau_damping.py b/tests/test_spectrax1d/test_landau_damping.py index 5897a3e7..b3cd2841 100644 --- a/tests/test_spectrax1d/test_landau_damping.py +++ b/tests/test_spectrax1d/test_landau_damping.py @@ -26,6 +26,7 @@ import pytest import yaml +import adept.patched_mlflow as mlflow from adept import electrostatic, ergoExo # --------------------------------------------------------------------------- @@ -33,7 +34,7 @@ # --------------------------------------------------------------------------- -def _run_driven_epw(cfg: dict, klambda_D: float) -> tuple[float, float, float, float]: +def _run_driven_epw(cfg: dict, klambda_D: float, tags: dict | None = None) -> tuple[float, float, float, float]: """ Configure klambda_D, run a driven EPW simulation, return measured and expected values. @@ -58,6 +59,9 @@ def _run_driven_epw(cfg: dict, klambda_D: float) -> tuple[float, float, float, f exo = ergoExo() exo.setup(cfg) + if tags is not None: + with mlflow.start_run(run_id=exo.mlflow_run_id, nested=exo.mlflow_nested): + mlflow.set_tags(tags) _, post, _ = exo(None) metrics = post["metrics"] @@ -91,7 +95,7 @@ def base_cfg(): [False, True], ids=["mobile-ions", "static-ions"], ) -def test_driven_epw_dispersion(base_cfg, static_ions): +def test_driven_epw_dispersion(base_cfg, static_ions, tags): """Driven EPW frequency and damping rate match the electrostatic dispersion within 10%.""" klambda_D = np.random.uniform(0.26, 0.34) @@ -102,7 +106,7 @@ def test_driven_epw_dispersion(base_cfg, static_ions): base_cfg["mlflow"]["experiment"] = "test-adept-spectrax1d-epw" base_cfg["mlflow"]["run"] = f"epw1d-exponential-{'static' if static_ions else 'mobile'}-{klambda_D:.3f}" - measured_freq, measured_damp, expected_freq, expected_damp = _run_driven_epw(base_cfg, klambda_D) + measured_freq, measured_damp, expected_freq, expected_damp = _run_driven_epw(base_cfg, klambda_D, tags) print(f"\nklambda_D={klambda_D:.4f} integrator=exponential static_ions={static_ions}") print( diff --git a/tests/test_spectrax1d/test_maxwell.py b/tests/test_spectrax1d/test_maxwell.py index cf04ad28..f7e5f8bb 100644 --- a/tests/test_spectrax1d/test_maxwell.py +++ b/tests/test_spectrax1d/test_maxwell.py @@ -30,6 +30,7 @@ import pytest import yaml +import adept.patched_mlflow as mlflow from adept import ergoExo # --------------------------------------------------------------------------- @@ -37,10 +38,13 @@ # --------------------------------------------------------------------------- -def _run_maxwell(cfg: dict) -> dict: +def _run_maxwell(cfg: dict, tags: dict | None = None) -> dict: """Run a hermite-maxwell-1d simulation and return the post-process metrics dict.""" exo = ergoExo() exo.setup(cfg) + if tags is not None: + with mlflow.start_run(run_id=exo.mlflow_run_id, nested=exo.mlflow_nested): + mlflow.set_tags(tags) _, post, _ = exo(None) return post["metrics"] @@ -67,7 +71,7 @@ def absorption_cfg(): # --------------------------------------------------------------------------- -def test_em_wave_dispersion(dispersion_cfg): +def test_em_wave_dispersion(dispersion_cfg, tags): """ MaxwellExponential propagates vacuum EM waves at the correct light speed ω = c|k|. @@ -90,7 +94,7 @@ def test_em_wave_dispersion(dispersion_cfg): dispersion_cfg["mlflow"]["experiment"] = "test-adept-spectrax1d-wave-dispersion" dispersion_cfg["mlflow"]["run"] = "em-dispersion-exponential" - metrics = _run_maxwell(dispersion_cfg) + metrics = _run_maxwell(dispersion_cfg, tags) measured_freq = metrics["em_avg_frequency_k1"] print("\nEM vacuum dispersion test (Lawson-RK4)") @@ -111,7 +115,7 @@ def test_em_wave_dispersion(dispersion_cfg): # --------------------------------------------------------------------------- -def test_em_wave_absorption(absorption_cfg): +def test_em_wave_absorption(absorption_cfg, tags): """ Sponge on right boundary absorbs > 90% of the wave energy after driver off. @@ -125,7 +129,7 @@ def test_em_wave_absorption(absorption_cfg): absorption_cfg["mlflow"]["experiment"] = "test-adept-spectrax1d-wave-absorption" absorption_cfg["mlflow"]["run"] = "wave-absorption-exponential" - metrics = _run_maxwell(absorption_cfg) + metrics = _run_maxwell(absorption_cfg, tags) absorption_ratio = metrics["em_absorption_ratio"] peak_energy = metrics["em_peak_energy"] final_energy = metrics["em_final_energy"] diff --git a/tests/test_tf1d/configs/resonance.yaml b/tests/test_tf1d/configs/resonance.yaml index 5a9519b8..26d04126 100644 --- a/tests/test_tf1d/configs/resonance.yaml +++ b/tests/test_tf1d/configs/resonance.yaml @@ -1,7 +1,7 @@ solver: tf-1d mlflow: - experiment: tf1d-ions-test + experiment: test-adept-tf1d run: test grid: diff --git a/tests/test_tf1d/configs/resonance_search.yaml b/tests/test_tf1d/configs/resonance_search.yaml index d7a7019c..cf87c36e 100644 --- a/tests/test_tf1d/configs/resonance_search.yaml +++ b/tests/test_tf1d/configs/resonance_search.yaml @@ -1,7 +1,7 @@ solver: tf-1d mlflow: - experiment: test-tf1d + experiment: test-adept-tf1d-resonance-search run: resonance-search adjoint: Backsolve diff --git a/tests/test_tf1d/configs/vlasov_comparison.yaml b/tests/test_tf1d/configs/vlasov_comparison.yaml index 22c4e85d..898afa03 100644 --- a/tests/test_tf1d/configs/vlasov_comparison.yaml +++ b/tests/test_tf1d/configs/vlasov_comparison.yaml @@ -1,7 +1,7 @@ solver: tf-1d mlflow: - experiment: tf1d-ions-test + experiment: test-adept-tf1d run: test models: False diff --git a/tests/test_tf1d/test_against_vlasov.py b/tests/test_tf1d/test_against_vlasov.py index 1b491426..f647a319 100644 --- a/tests/test_tf1d/test_against_vlasov.py +++ b/tests/test_tf1d/test_against_vlasov.py @@ -4,6 +4,7 @@ import xarray as xr import yaml +import adept.patched_mlflow as mlflow from adept import electrostatic, ergoExo @@ -21,12 +22,12 @@ def _modify_defaults_(defaults): defaults["grid"]["xmax"] = xmax defaults["save"]["x"]["xmax"] = xmax defaults["save"]["kx"]["kxmax"] = rand_k0 - defaults["mlflow"]["experiment"] = "test-against-vlasov" + defaults["mlflow"]["experiment"] = "test-adept-tf1d-against-vlasov" return defaults, float(np.imag(root)) -def test_single_resonance(): +def test_single_resonance(tags): with open("tests/test_tf1d/configs/vlasov_comparison.yaml") as file: defaults = yaml.safe_load(file) @@ -35,6 +36,8 @@ def test_single_resonance(): exo = ergoExo() exo.setup(mod_defaults) + with mlflow.start_run(run_id=exo.mlflow_run_id, nested=exo.mlflow_nested): + mlflow.set_tags(tags) result, datasets, run_id = exo(None) result = result["solver result"] vds = xr.open_dataset("tests/test_tf1d/vlasov-reference/all-fields-kx.nc", engine="h5netcdf") diff --git a/tests/test_tf1d/test_landau_damping.py b/tests/test_tf1d/test_landau_damping.py index d8427985..5f6d6504 100644 --- a/tests/test_tf1d/test_landau_damping.py +++ b/tests/test_tf1d/test_landau_damping.py @@ -22,12 +22,12 @@ def _modify_defaults_(defaults, rng): defaults["grid"]["xmax"] = xmax defaults["save"]["x"]["xmax"] = xmax defaults["save"]["kx"]["kxmax"] = rand_k0 - defaults["mlflow"]["experiment"] = "test-landau-damping" + defaults["mlflow"]["experiment"] = "test-adept-tf1d-landau-damping" return defaults, float(np.imag(root)) -def test_single_resonance(): +def test_single_resonance(tags): with open("tests/test_tf1d/configs/resonance.yaml") as file: defaults = yaml.safe_load(file) @@ -41,6 +41,7 @@ def test_single_resonance(): result = result["solver result"] with mlflow.start_run(run_id=run_id, log_system_metrics=True) as mlflow_run: + mlflow.set_tags(tags) kx = ( np.fft.fftfreq( mod_defaults["save"]["x"]["nx"], diff --git a/tests/test_tf1d/test_resonance.py b/tests/test_tf1d/test_resonance.py index 56b2ebb8..9e70e3ec 100644 --- a/tests/test_tf1d/test_resonance.py +++ b/tests/test_tf1d/test_resonance.py @@ -5,6 +5,7 @@ import yaml from jax import numpy as jnp +import adept.patched_mlflow as mlflow from adept import electrostatic, ergoExo @@ -22,13 +23,13 @@ def _modify_defaults_(defaults, rng, gamma): xmax = float(2.0 * np.pi / rand_k0) defaults["grid"]["xmax"] = xmax defaults["save"]["x"]["xmax"] = xmax - defaults["mlflow"]["experiment"] = "test-resonance" + defaults["mlflow"]["experiment"] = "test-adept-tf1d-resonance" return defaults, float(root) @pytest.mark.parametrize("gamma", ["kinetic", 3.0]) -def test_single_resonance(gamma): +def test_single_resonance(gamma, tags): with open("tests/test_tf1d/configs/resonance.yaml") as file: defaults = yaml.safe_load(file) @@ -39,6 +40,8 @@ def test_single_resonance(gamma): # run exo = ergoExo() exo.setup(mod_defaults) + with mlflow.start_run(run_id=exo.mlflow_run_id, nested=exo.mlflow_nested): + mlflow.set_tags(tags) result, datasets, run_id = exo(None) result = result["solver result"] diff --git a/tests/test_tf1d/test_resonance_search.py b/tests/test_tf1d/test_resonance_search.py index 389d7efb..26248b55 100644 --- a/tests/test_tf1d/test_resonance_search.py +++ b/tests/test_tf1d/test_resonance_search.py @@ -59,8 +59,8 @@ def load_cfg(rand_k0, gamma, adjoint): @pytest.mark.parametrize("adjoint", ["Recursive", "Backsolve"]) @pytest.mark.parametrize("gamma", ["kinetic", 3.0]) -def test_resonance_search(gamma, adjoint): - mlflow.set_experiment("tf1d-resonance-search") +def test_resonance_search(gamma, adjoint, tags): + mlflow.set_experiment("test-adept-tf1d-resonance-search") with mlflow.start_run(run_name="res-search-opt", log_system_metrics=True) as mlflow_run: # sim_k0, actual_w0 = init_w0(gamma, adjoint) rng = np.random.default_rng(420) @@ -83,6 +83,8 @@ def test_resonance_search(gamma, adjoint): exo = ergoExo(mlflow_nested=True) mod_defaults, _ = load_cfg(sim_k0, gamma, adjoint) exo.setup(mod_defaults, Resonance) + with mlflow.start_run(run_id=exo.mlflow_run_id, nested=exo.mlflow_nested): + mlflow.set_tags(tags) val, grad, results = exo.val_and_grad(params) updates, opt_state = optimizer.update(grad, opt_state, params) params = optax.apply_updates(params, updates) diff --git a/tests/test_tf1d/vlasov-reference/config.yaml b/tests/test_tf1d/vlasov-reference/config.yaml index d5460b75..31d330a2 100644 --- a/tests/test_tf1d/vlasov-reference/config.yaml +++ b/tests/test_tf1d/vlasov-reference/config.yaml @@ -98,7 +98,7 @@ krook: width: 100000.0 machine: local mlflow: - experiment: landau_damping + experiment: test-adept-vlasov1d-landau-damping run: test nu: time-profile: diff --git a/tests/test_vfp1d/epp-short-reflective.yaml b/tests/test_vfp1d/epp-short-reflective.yaml index 67dee891..53b4f96a 100644 --- a/tests/test_vfp1d/epp-short-reflective.yaml +++ b/tests/test_vfp1d/epp-short-reflective.yaml @@ -58,7 +58,7 @@ save: solver: vfp-1d mlflow: - experiment: test-vfp1d-eppshort + experiment: test-adept-vfp1d-epp-short run: epperlein-short-reflective drivers: diff --git a/tests/test_vfp1d/epp-short.yaml b/tests/test_vfp1d/epp-short.yaml index a11fcb0b..86f17dea 100644 --- a/tests/test_vfp1d/epp-short.yaml +++ b/tests/test_vfp1d/epp-short.yaml @@ -57,7 +57,7 @@ save: solver: vfp-1d mlflow: - experiment: test-vfp1d-eppshort + experiment: test-adept-vfp1d-epp-short run: epperlein-short-periodic drivers: diff --git a/tests/test_vfp1d/test_kappa_eh.py b/tests/test_vfp1d/test_kappa_eh.py index c7082ffc..4fc923aa 100644 --- a/tests/test_vfp1d/test_kappa_eh.py +++ b/tests/test_vfp1d/test_kappa_eh.py @@ -9,7 +9,7 @@ from adept import ergoExo -def _run_(Z, ee, config_name="epp-short"): +def _run_(Z, ee, config_name="epp-short", tags=None): # with open("configs/tf-1d/damping.yaml", "r") as fi: with open(f"{os.path.join(os.getcwd(), 'tests/test_vfp1d', config_name)}.yaml") as fi: cfg = yaml.safe_load(fi) @@ -22,6 +22,9 @@ def _run_(Z, ee, config_name="epp-short"): exo = ergoExo() exo.setup(cfg) + if tags is not None: + with mlflow.start_run(run_id=exo.mlflow_run_id, nested=exo.mlflow_nested): + mlflow.set_tags(tags) sol, datasets, run_id = exo(None) dataT = datasets["fields"]["fields-T keV"].data @@ -40,10 +43,10 @@ def _run_(Z, ee, config_name="epp-short"): @pytest.mark.parametrize("Z", list(range(1, 22, 4)) + [40, 60, 80]) @pytest.mark.parametrize("ee", [True, False]) -def test_kappa_eh(Z, ee): +def test_kappa_eh(Z, ee, tags): if not any(["gpu" == device.platform for device in devices()]): if Z in [1, 21, 80]: - _run_(Z, ee) + _run_(Z, ee, tags=tags) else: pytest.skip(f"Skipping Z={Z} to save time because no GPU is available") @@ -53,10 +56,10 @@ def test_kappa_eh(Z, ee): @pytest.mark.parametrize("Z", list(range(1, 22, 4)) + [40, 60, 80]) @pytest.mark.parametrize("ee", [True, False]) -def test_kappa_eh_reflective(Z, ee): +def test_kappa_eh_reflective(Z, ee, tags): if not any(["gpu" == device.platform for device in devices()]): if Z in [1, 21, 80]: - _run_(Z, ee, config_name="epp-short-reflective") + _run_(Z, ee, config_name="epp-short-reflective", tags=tags) else: pytest.skip(f"Skipping Z={Z} to save time because no GPU is available") diff --git a/tests/test_vlasov1d/configs/fokker_planck_conservation.yaml b/tests/test_vlasov1d/configs/fokker_planck_conservation.yaml index 2baba8c7..d76bedfd 100644 --- a/tests/test_vlasov1d/configs/fokker_planck_conservation.yaml +++ b/tests/test_vlasov1d/configs/fokker_planck_conservation.yaml @@ -49,7 +49,7 @@ save: solver: vlasov-1d mlflow: - experiment: vlasov1d-test-fokker-planck + experiment: test-adept-vlasov1d-fp-conservation run: conservation-test drivers: diff --git a/tests/test_vlasov1d/configs/multispecies_ion_acoustic.yaml b/tests/test_vlasov1d/configs/multispecies_ion_acoustic.yaml index dc4195a3..577ee8c2 100644 --- a/tests/test_vlasov1d/configs/multispecies_ion_acoustic.yaml +++ b/tests/test_vlasov1d/configs/multispecies_ion_acoustic.yaml @@ -47,7 +47,7 @@ save: solver: vlasov-1d mlflow: - experiment: multispecies-test + experiment: test-adept-vlasov1d-ion-acoustic run: ion-acoustic-wave drivers: diff --git a/tests/test_vlasov1d/configs/resonance.yaml b/tests/test_vlasov1d/configs/resonance.yaml index 86d49cbc..2ea0dfcc 100644 --- a/tests/test_vlasov1d/configs/resonance.yaml +++ b/tests/test_vlasov1d/configs/resonance.yaml @@ -58,7 +58,7 @@ save: solver: vlasov-1d mlflow: - experiment: vlasov1d + experiment: test-adept-vlasov1d run: test drivers: diff --git a/tests/test_vlasov1d/test_fokker_planck_conservation.py b/tests/test_vlasov1d/test_fokker_planck_conservation.py index 71829ce1..e62b763f 100644 --- a/tests/test_vlasov1d/test_fokker_planck_conservation.py +++ b/tests/test_vlasov1d/test_fokker_planck_conservation.py @@ -19,11 +19,12 @@ import pytest import yaml +import adept.patched_mlflow as mlflow from adept import ergoExo @pytest.mark.parametrize("operator_type", ["chang_cooper_dougherty", "chang_cooper", "Dougherty", "Lenard_Bernstein"]) -def test_fokker_planck_conservation(operator_type): +def test_fokker_planck_conservation(operator_type, tags): """ Test that Fokker-Planck operators conserve density and energy. @@ -39,6 +40,8 @@ def test_fokker_planck_conservation(operator_type): # Run simulation exo = ergoExo() exo.setup(config) + with mlflow.start_run(run_id=exo.mlflow_run_id, nested=exo.mlflow_nested): + mlflow.set_tags(tags) result, datasets, run_id = exo(None) result = result["solver result"] @@ -96,7 +99,7 @@ def test_fokker_planck_conservation(operator_type): @pytest.mark.parametrize("operator_type", ["chang_cooper", "Lenard_Bernstein", "Dougherty"]) -def test_fokker_planck_thermalization(operator_type): +def test_fokker_planck_thermalization(operator_type, tags): """ Test that Fokker-Planck operators drive distribution toward Maxwellian. @@ -122,6 +125,8 @@ def test_fokker_planck_thermalization(operator_type): # Run simulation exo = ergoExo() exo.setup(config) + with mlflow.start_run(run_id=exo.mlflow_run_id, nested=exo.mlflow_nested): + mlflow.set_tags(tags) result, datasets, run_id = exo(None) result = result["solver result"] diff --git a/tests/test_vlasov1d/test_ion_acoustic_wave.py b/tests/test_vlasov1d/test_ion_acoustic_wave.py index 5394d5e7..ece3e5c5 100644 --- a/tests/test_vlasov1d/test_ion_acoustic_wave.py +++ b/tests/test_vlasov1d/test_ion_acoustic_wave.py @@ -20,6 +20,7 @@ import pytest import yaml +import adept.patched_mlflow as mlflow from adept import ergoExo @@ -96,7 +97,7 @@ def compute_ion_density(f_ion, dv): @pytest.mark.parametrize("time_integrator", ["sixth"]) -def test_ion_acoustic_simulation_runs(time_integrator): +def test_ion_acoustic_simulation_runs(time_integrator, tags): """ Smoke test that multi-species ion acoustic simulation runs end-to-end. @@ -122,12 +123,14 @@ def test_ion_acoustic_simulation_runs(time_integrator): # Modify config for this test config["terms"]["time"] = time_integrator - config["mlflow"]["experiment"] = "vlasov1d-test-ion-acoustic" + config["mlflow"]["experiment"] = "test-adept-vlasov1d-ion-acoustic" config["mlflow"]["run"] = f"ion-acoustic-smoke-{time_integrator}" # Run simulation exo = ergoExo() exo.setup(config) + with mlflow.start_run(run_id=exo.mlflow_run_id, nested=exo.mlflow_nested): + mlflow.set_tags(tags) result, datasets, run_id = exo(None) solver_result = result["solver result"] @@ -148,7 +151,7 @@ def test_ion_acoustic_simulation_runs(time_integrator): @pytest.mark.parametrize("time_integrator", ["sixth"]) -def test_ion_acoustic_dispersion(time_integrator): +def test_ion_acoustic_dispersion(time_integrator, tags): """ Test that ion acoustic wave frequency matches theoretical prediction. @@ -165,7 +168,7 @@ def test_ion_acoustic_dispersion(time_integrator): # Modify config for this test config["grid"]["nx"] = 64 # Override for better resolution config["terms"]["time"] = time_integrator - config["mlflow"]["experiment"] = "vlasov1d-test-ion-acoustic" + config["mlflow"]["experiment"] = "test-adept-vlasov1d-ion-acoustic" config["mlflow"]["run"] = f"ion-acoustic-dispersion-{time_integrator}" # Extract physical parameters @@ -180,6 +183,8 @@ def test_ion_acoustic_dispersion(time_integrator): # Run simulation exo = ergoExo() exo.setup(config) + with mlflow.start_run(run_id=exo.mlflow_run_id, nested=exo.mlflow_nested): + mlflow.set_tags(tags) result, datasets, run_id = exo(None) solver_result = result["solver result"] diff --git a/tests/test_vlasov1d/test_landau_damping.py b/tests/test_vlasov1d/test_landau_damping.py index 82dafd6a..e0208ccb 100644 --- a/tests/test_vlasov1d/test_landau_damping.py +++ b/tests/test_vlasov1d/test_landau_damping.py @@ -6,6 +6,7 @@ import pytest import yaml +import adept.patched_mlflow as mlflow from adept import electrostatic, ergoExo @@ -38,7 +39,7 @@ def _modify_defaults_(defaults, rng, real_or_imag, time, field, edfdv): ["real", "imag"], ["sixth", "leapfrog"], ["poisson", "ampere", "hampere"], ["exponential", "cubic-spline"] ), ) -def test_single_resonance(real_or_imag, time, field, edfdv): +def test_single_resonance(real_or_imag, time, field, edfdv, tags): if (time == "sixth") and (field == "ampere"): print("not implemented - skipping test") elif (time == "sixth") and (field == "hampere"): @@ -56,6 +57,8 @@ def test_single_resonance(real_or_imag, time, field, edfdv): exo = ergoExo() exo.setup(mod_defaults) + with mlflow.start_run(run_id=exo.mlflow_run_id, nested=exo.mlflow_nested): + mlflow.set_tags(tags) result, datasets, run_id = exo(None) result = result["solver result"] efs = result.ys["fields"]["e"] diff --git a/tests/test_vlasov2d/configs/damping.yaml b/tests/test_vlasov2d/configs/damping.yaml index 8d7bce67..11b6a666 100644 --- a/tests/test_vlasov2d/configs/damping.yaml +++ b/tests/test_vlasov2d/configs/damping.yaml @@ -82,7 +82,7 @@ machine: s-t4 # solver: vlasov-2d mlflow: - experiment: ld-2d2v + experiment: test-adept-vlasov2d-landau-damping run: envelope-driver-small-amp drivers: diff --git a/tests/test_vlasov2d/test_landau_damping.py b/tests/test_vlasov2d/test_landau_damping.py index 271312cd..78c407f4 100644 --- a/tests/test_vlasov2d/test_landau_damping.py +++ b/tests/test_vlasov2d/test_landau_damping.py @@ -23,7 +23,7 @@ def _modify_defaults_(defaults, rng, real_or_imag): defaults["drivers"]["ex"]["0"]["w0"] = float(np.real(root)) xmax = float(2.0 * np.pi / rand_k0) defaults["grid"]["xmax"] = xmax - defaults["mlflow"]["experiment"] = "vlasov2d-test" + defaults["mlflow"]["experiment"] = "test-adept-vlasov2d-landau-damping" return defaults, root