Console Output
+ diff --exclude=.git --exclude=version_clubb_core.txt --exclude=version_silhs.txt -r clubb clubb_release
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/clubb_jax/README.md clubb_release/clubb_jax/README.md
38,41d37
< Bare `-jax` selects CPU unless the legacy `CLUBB_JAX_ACCELERATOR` environment
< variable is set. The explicit forms `-jax=cpu` and `-jax=gpu` are also
< accepted, case-insensitively.
<
68,69c64
< number of timesteps as well, but reducing iterations makes testing fundamentally
< more permissive, so use it only when a partial run is sufficient:
---
> number of timesteps as well, but reducing iterations makes testing fundamentally more permissive, so changing it should be considered only for rapid smoke tests:
78c73,74
< Use `-jax=gpu` to run on an NVIDIA GPU on Linux or an Apple Silicon GPU on macOS:
---
> To use an NVIDIA GPU with the CUDA 13 JAX packages, select the accelerator when
> initializing the environment and running:
81,88c77
< ./run_scripts/run_scm.py -jax=gpu -stats none -debug 0 arm
< ```
<
< The launcher checks GPU compatibility and prepares a separate environment on
< first use. It reports an error if the requested GPU backend is unavailable.
< Apple Metal runs use float32 because the plugin does not support float64.
< See [Inspect Runtime Support](#inspect-runtime-support) to check your setup
< without starting a run.
---
> CLUBB_JAX_ACCELERATOR=cuda13 ./clubb_jax/run_jax_wrapper.sh --init_env
90,93c79,80
< On NVIDIA systems, select a card using its index from `nvidia-smi`:
<
< ```bash
< CUDA_VISIBLE_DEVICES=1 ./run_scripts/run_scm.py -jax=gpu arm
---
> CLUBB_JAX_ACCELERATOR=cuda13 CUDA_VISIBLE_DEVICES=0 \
> ./run_scripts/run_scm.py -jax -stats none -debug 0 arm
96,98c83,90
< CUDA memory preallocation is off by default unless enabled through the
< `XLA_PYTHON_CLIENT_PREALLOCATE` environment variable. This lets JAX allocate
< memory as needed when sharing a GPU. To enable up-front memory reservation:
---
> The launcher verifies that JAX initialized the CUDA backend and fails rather
> than silently falling back to CPU. `CUDA_VISIBLE_DEVICES` selects the GPU.
>
> By default, JAX reserves 75% of the GPU's memory when the first JAX operation
> runs. This reduces allocation overhead and memory fragmentation when JAX owns
> the device, but it can leave too little memory for a desktop display or other
> processes using a shared GPU and can cause an out-of-memory error at startup.
> Disable preallocation when that reservation causes contention:
101c93,95
< ./run_scripts/run_scm.py -jax=gpu,xla_prealloc arm
---
> CLUBB_JAX_ACCELERATOR=cuda13 CUDA_VISIBLE_DEVICES=0 \
> XLA_PYTHON_CLIENT_PREALLOCATE=false \
> ./run_scripts/run_scm.py -jax arm
104,106c98,101
< This option applies only to NVIDIA CUDA. See
< [Advanced GPU Options](#advanced-gpu-options) for device selection details
< and memory trade-offs.
---
> With preallocation disabled, JAX allocates memory as the run needs it. This
> usually lowers its initial footprint, but it is more vulnerable to memory
> fragmentation; keep the default when the run owns the GPU and needs most of its
> memory.
163,166c158,159
< Apple Silicon uses the isolated [`requirements-metal.txt`](./requirements-metal.txt)
< profile with Python 3.11 or 3.12 and JAX/JAXLIB 0.4.34 because Apple's experimental
< plugin requires its own compatible JAX line. CPU and CUDA use JAX/JAXLIB 0.11.0
< on Python 3.12 or newer; Python 3.11 is supported with JAX/JAXLIB 0.10.0.
---
> Python 3.12 or newer uses JAX/JAXLIB 0.11.0; Python 3.11 is supported with
> JAX/JAXLIB 0.10.0.
173,174c166
< ./clubb_jax/run_jax.py --init_env
< ./clubb_jax/run_jax.py --profile=gpu --init_env
---
> ./clubb_jax/run_jax_wrapper.sh --init_env
183,184c175
< 4. Creates `.venv-jax` for CPU, `.venv-jax-cuda13` for CUDA 13, or
< `.venv-jax-metal` for Apple Metal.
---
> 4. Creates `.venv-jax` for CPU or `.venv-jax-cuda13` for CUDA 13.
197,218c188
< ./clubb_jax/run_jax.py --init_env
< ```
<
< ### Inspect Runtime Support
<
< The launcher can inspect the selected profile without creating an environment,
< installing packages, or initializing JAX:
<
< ```bash
< ./clubb_jax/run_jax.py --profile=cpu --info
< ./clubb_jax/run_jax.py --profile=gpu --info
< ```
<
< The report shows the detected devices, Python and JAX versions, environment
< path, and whether setup is needed. Hardware discovery works before the JAX
< environment is installed. The report does not initialize JAX or test available
< GPU memory; the run itself reports the devices JAX actually uses.
<
< For machine-readable output:
<
< ```bash
< ./clubb_jax/run_jax.py --profile=gpu --info=json
---
> ./clubb_jax/run_jax_wrapper.sh --init_env
221,225d190
< The CUDA 13 profile checks for an NVIDIA driver version of at least 580 and
< compute capability of at least 7.5 on each exposed GPU. JAX's installed packages
< supply the CUDA runtime libraries, so a local CUDA toolkit is not required.
< Metal support remains experimental.
<
240,241c205,206
< Use `requirements-cuda13.txt` on NVIDIA Linux or `requirements-metal.txt` on
< Apple Silicon. In either case, `-jax=gpu` selects the native GPU profile.
---
> Use `requirements-cuda13.txt` and set `CLUBB_JAX_ACCELERATOR=cuda13` for a GPU
> environment.
258,279c223
< requirements if necessary.
<
< ### Advanced GPU Options
<
< `CUDA_VISIBLE_DEVICES` accepts physical indices from `nvidia-smi`, full GPU
< UUIDs, or unique UUID prefixes. Find UUIDs with `nvidia-smi -L`. The launcher
< converts selections to full UUIDs so CUDA's enumeration order cannot change
< which cards are selected. Lists retain their order: `1,0` exposes physical
< GPU 1 as JAX device 0. Invalid, ambiguous, duplicate, or empty selections are
< rejected before setup. Selection supports whole GPUs, not MIG instances;
< all exposed GPUs must meet the CUDA requirements.
<
< Without `xla_prealloc`, the launcher respects an existing
< `XLA_PYTHON_CLIENT_PREALLOCATE` setting and otherwise defaults to `false`.
< The modifier overrides that variable to `true`. Direct launcher users can pass
< `--profile=gpu --xla-prealloc`.
<
< Preallocation can reduce allocation overhead and fragmentation when a run
< has the GPU to itself. Leaving it disabled lowers the initial memory footprint,
< which helps on shared GPUs. JAX can still cache allocated memory; neither
< setting limits total memory use or guarantees that a run will fit. The launch
< report shows the effective setting.
---
> requirements if necessary.
\ No newline at end of file
Only in clubb/clubb_jax: requirements-metal.txt
Only in clubb/clubb_jax: run_jax.py
Only in clubb_release/clubb_jax: run_jax_wrapper.sh
Only in clubb/clubb_jax: runtime_info.py
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/clubb_jax/src/advance_clubb_to_end.py clubb_release/clubb_jax/src/advance_clubb_to_end.py
148c148
< f" -- time = {time_current:10.1f} / {state['time_final']:10.1f}", flush=True)
---
> f" -- time = {time_current:10.1f} / {state['time_final']:10.1f}")
455c455
< l_rad_itime, state['dt_main'], state['day'], state['month'], state['year'],
---
> jnp.asarray(l_rad_itime), state['dt_main'], state['day'], state['month'], state['year'],
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/clubb_jax/src/clubb_standalone.py clubb_release/clubb_jax/src/clubb_standalone.py
8d7
< import os
11,12d9
<
< import jax
29,40d25
<
< if os.environ.get('CLUBB_JAX_PROFILE') == 'gpu':
< devices = jax.devices()
< accelerator = os.environ.get('CLUBB_JAX_ACCELERATOR', 'cuda13')
< expected_backend = 'metal' if accelerator == 'metal' else 'gpu'
< if jax.default_backend().lower() != expected_backend:
< raise RuntimeError(
< f'The GPU profile did not initialize the {expected_backend} backend.'
< )
< print('==> JAX GPU backend initialized; visible devices:', flush=True)
< for device in devices:
< print(f' JAX GPU {device.id}: {device.device_kind}', flush=True)
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/clubb_jax/src/Radiation/radiation_module.py clubb_release/clubb_jax/src/Radiation/radiation_module.py
40,43c40
< static_argnames=(
< "ngrdcol", "hydromet_dim", "pdf_dim", "lh_num_samples",
< "l_rad_itime", "day", "month", "year",
< ),
---
> static_argnames=("ngrdcol", "hydromet_dim", "pdf_dim", "lh_num_samples", "day", "month", "year"),
68,69d64
<
< ``l_rad_itime`` is a static Boolean supplied by the host timestep schedule.
92a88,89
> # JAX adaptation: ``lax.cond`` requires branch callables for the source
> # ``if ( l_rad_itime )`` block and its retained module-output values.
94c91
< if l_rad_itime:
---
> def advance(_):
101,104c98
< (
< stats, err_info, radht, Frad, Frad_SW_up, Frad_LW_up,
< Frad_SW_down, Frad_LW_down, radht_SW, radht_LW, Frad_SW, Frad_LW,
< ) = radiation_driver(
---
> return radiation_driver(
110a105,115
>
> def retain(_):
> return (
> stats, err_info, radht, Frad, Frad_SW_up, Frad_LW_up,
> Frad_SW_down, Frad_LW_down, radht_SW, radht_LW, Frad_SW, Frad_LW,
> )
>
> (
> stats, err_info, radht, Frad, Frad_SW_up, Frad_LW_up,
> Frad_SW_down, Frad_LW_down, radht_SW, radht_LW, Frad_SW, Frad_LW,
> ) = jax.lax.cond(l_rad_itime, advance, retain, operand=None)
Only in clubb/clubb_jax/tests: test_jax_cli_options.py
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/clubb_jax/tests/test_radiation_module.py clubb_release/clubb_jax/tests/test_radiation_module.py
3,10d2
< from dataclasses import replace
<
< import jax
< import jax.numpy as jnp
< import numpy as np
< import pytest
<
< from clubb_jax.src import clubb_case_initalization
12,13d3
< from clubb_jax.src.Radiation.simple_rad_module import simple_rad_lba
< from clubb_jax.src.Radiation.soil_vegetation import advance_soil_veg
14a5
> from clubb_jax.src.clubb_case_initalization import init_clubb_case
17,39c8,10
< RADIATION_FIELDS = (
< "radht", "Frad", "Frad_SW_up", "Frad_LW_up", "Frad_SW_down", "Frad_LW_down",
< "radht_SW", "radht_LW", "Frad_SW", "Frad_LW",
< )
< SOIL_FIELDS = ("deep_soil_T_in_K", "sfc_soil_T_in_K", "veg_T_in_K")
<
<
< @pytest.fixture(scope="module", params=["arm", "bomex", "lba", "dycoms2_rf01"])
< def radiation_state(request):
< read_namelist = clubb_case_initalization.read_namelist
< # Exercise radiation independently of unsupported microphysics/SILHS and
< # avoid opening case output files. Keep the real case's radiation settings.
< with pytest.MonkeyPatch.context() as patch:
< patch.setattr(
< clubb_case_initalization, "read_namelist",
< lambda path: dict(
< read_namelist(path), l_stats=False,
< microphys_scheme="none", lh_microphys_type="disabled",
< ),
< )
< state = clubb_case_initalization.init_clubb_case(
< f"input/case_setups/{request.param}_model.in"
< )
---
> def test_sampled_radiation_stats_do_not_depend_on_radiation_cadence():
> """Fortran updates radiation stats on every sampled timestep."""
> state = init_clubb_case("input/case_setups/bomex_model.in")
42,43c13,14
< names=("radht", "Frad"),
< grids=("zt", "zm"),
---
> names=("radht",),
> grids=("zt",),
45,46c16,17
< max_nlev=state["nzm"],
< grid_nlev=(state["nzt"], state["nzm"], 1, state["nzt"], 1, state["nzt"], state["nzt"]),
---
> max_nlev=state["nzt"],
> grid_nlev=(state["nzt"], state["nzt"], 1, state["nzt"], 1, state["nzt"], state["nzt"]),
48,70d18
< return state
<
<
< @pytest.mark.parametrize("radiation_state", ["lba"], indirect=True)
< def test_radiation_schedule_can_toggle_between_calls(radiation_state):
< """Static schedule variants still consume the current time and array values."""
< state = dict(radiation_state)
< previous = np.asarray(state["radht"]).copy()
< for step, advance in enumerate((True, False, True, False, True), start=1):
< time_current = state["time_initial"] + step * 600.0
< _advance_radiation(state, time_current, advance)
< if advance:
< expected = simple_rad_lba(
< state["gr"], state["ngrdcol"], time_current,
< state["time_initial"], state["radiation_parameters"],
< )
< np.testing.assert_allclose(state["radht"], expected, rtol=2e-6, atol=1e-12)
< assert np.any(np.asarray(state["radht"]) != previous)
< else:
< np.testing.assert_array_equal(state["radht"], previous)
< bank, slot = state["_jax_stats"].name_to_slot["radht"]
< np.testing.assert_array_equal(state["_jax_stats"].nsamples[bank][slot], step)
< previous = np.asarray(state["radht"]).copy()
71a20
> _advance_radiation(state, 0.0, l_rad_itime=False)
73,131c22
< @pytest.mark.parametrize("l_rad_itime", [False, True])
< @pytest.mark.parametrize("l_sample", [False, True])
< def test_radiation_schedule_preserves_outputs_and_stats(radiation_state, l_rad_itime, l_sample):
< """Compiled scheduling preserves retained values and samples every requested step."""
< state = dict(radiation_state)
< state["_jax_stats"] = state["_jax_stats"].begin_timestep(
< l_sample=l_sample, reset_accumulators=True,
< )
< # Nonzero, distinct values catch accidental clearing on a retained step.
< for index, name in enumerate(RADIATION_FIELDS, start=1):
< state[name] = jnp.full_like(state[name], index)
< before = dict(state)
< expected = dict(state)
< time_current = state["time_initial"] + 300.0
< with jax.disable_jit():
< _advance_radiation(expected, time_current, l_rad_itime)
< _advance_radiation(state, time_current, l_rad_itime)
<
< for name in (*RADIATION_FIELDS, *SOIL_FIELDS, "_jax_stats", "err_info"):
< for actual, reference in zip(
< jax.tree_util.tree_leaves(state[name]), jax.tree_util.tree_leaves(expected[name]),
< ):
< # Compiled float32 arithmetic can differ slightly from eager
< # evaluation (including fused operations on CPU).
< tolerance = 5e-6 if actual.dtype == jnp.float32 else 1e-12
< np.testing.assert_allclose(actual, reference, rtol=tolerance, atol=1e-12, equal_nan=False)
< if not l_rad_itime:
< for name in RADIATION_FIELDS:
< np.testing.assert_array_equal(state[name], before[name])
< for name, levels in (("radht", state["nzt"]), ("Frad", state["nzm"])):
< stats = state["_jax_stats"]
< bank, slot = stats.name_to_slot[name]
< np.testing.assert_array_equal(stats.nsamples[bank][slot], int(l_sample))
< np.testing.assert_array_equal(
< stats.buffers[bank][slot], state[name] if l_sample else np.zeros((state["ngrdcol"], levels)),
< )
<
<
< @pytest.mark.parametrize("radiation_state", ["arm"], indirect=True)
< @pytest.mark.parametrize("l_rad_itime", [False, True])
< def test_soil_updates_use_incoming_fluxes_on_every_step(radiation_state, l_rad_itime):
< """Soil advances before radiation, including steps which retain radiation."""
< state = dict(radiation_state)
< state["radiation_parameters"] = replace(state["radiation_parameters"], l_soil_veg=True)
< for name, value in zip(SOIL_FIELDS, (288.58, 295.0, 300.0)):
< state[name] = jnp.full_like(state[name], value)
< for name, value in (("Frad_SW_up", 40.0), ("Frad_SW_down", 200.0), ("Frad_LW_down", 320.0)):
< state[name] = jnp.full_like(state[name], value)
< before = {name: state[name] for name in SOIL_FIELDS}
< _, *expected = advance_soil_veg(
< state["ngrdcol"], state["dt_main"], state["rho_zm"][:, 0],
< state["Frad_SW_up"][:, 0], state["Frad_SW_down"][:, 0], state["Frad_LW_down"][:, 0],
< state["wpthlp_sfc"], state["wprtp_sfc"], state["p_sfc"], state["_jax_stats"],
< *(state[name] for name in SOIL_FIELDS),
< )
< _advance_radiation(state, state["time_initial"], l_rad_itime)
< for name, value in zip(SOIL_FIELDS, expected):
< np.testing.assert_allclose(state[name], value, rtol=1e-7, atol=1e-12)
< assert np.any(np.asarray(state[name]) != np.asarray(before[name]))
---
> assert sum(int(bank.sum()) for bank in state["_jax_stats"].nsamples) == state["nzt"]
Only in clubb/clubb_jax/tests: test_runtime_info.py
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/clubb_python_driver/advance_clubb_to_end.py clubb_release/clubb_python_driver/advance_clubb_to_end.py
90c90
< f" -- time = {time_current:10.1f} / {state['time_final']:10.1f}", flush=True)
---
> f" -- time = {time_current:10.1f} / {state['time_final']:10.1f}")
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/dash_app/assets/09_selected_build_badge.css clubb_release/dash_app/assets/09_selected_build_badge.css
38c38
< font-size: 10px;
---
> font-size: 9px;
50c50
< font-size: 13px;
---
> font-size: 11px;
87c87
< gap: 12px;
---
> gap: 7px;
90c90
< padding: 14px;
---
> padding: 8px;
92c92
< border-radius: 14px;
---
> border-radius: 9px;
100,113c100,104
< gap: 10px;
< min-width: 0;
< }
<
< .compile-selector-toolbar {
< display: flex;
< align-items: center;
< justify-content: space-between;
< }
<
< /* Keep plot-style help above the selector and its click-away backdrop. */
< .compile-selector-help-layer {
< position: relative;
< z-index: 31000;
---
> gap: 7px;
> padding: 8px;
> border: 1px solid #334155;
> border-radius: 7px;
> background: rgba(15, 23, 42, .82);
118c109
< font-size: 11px;
---
> font-size: 9px;
127,128c118
< gap: 4px;
< padding: 4px;
---
> overflow: hidden;
130,136c120
< border-radius: 10px;
< }
<
< .compile-run-jax-profile-choices {
< display: grid;
< grid-template-columns: repeat(2, minmax(0, 1fr));
< gap: 10px;
---
> border-radius: 6px;
140,142c124,125
< min-height: 40px;
< min-width: 0;
< padding: 8px 10px;
---
> min-height: 34px;
> padding: 6px 10px;
144c127
< border-radius: 7px;
---
> border-left: 1px solid #475569;
149c132
< font-size: 14px;
---
> font-size: 11px;
154c137,141
< .compile-run-implementation-choice:not(:disabled):hover {
---
> .compile-run-implementation-choice:first-child {
> border-left: 0;
> }
>
> .compile-run-implementation-choice:hover {
159,160c146
< .compile-run-implementation-choice-selected,
< .compile-run-implementation-choice-selected:not(:disabled):hover {
---
> .compile-run-implementation-choice-selected {
165,278d150
< .compile-run-implementation-choice:disabled {
< background: #111827;
< color: #64748b;
< cursor: not-allowed;
< opacity: .72;
< }
<
< .compile-run-implementation-choice:focus-visible {
< outline: 2px solid #38bdf8;
< outline-offset: 2px;
< }
<
< .compile-run-implementation-choice.compile-profile-card {
< padding: 13px;
< border: 1px solid #475569;
< border-radius: 10px;
< text-align: left;
< letter-spacing: normal;
< }
<
< .compile-profile-card.compile-run-implementation-choice-selected,
< .compile-profile-card.compile-run-implementation-choice-selected:not(:disabled):hover {
< border-color: #60a5fa;
< background: #172e50;
< box-shadow: inset 0 3px 0 #60a5fa;
< }
<
< .compile-profile-card-content {
< display: flex;
< height: 100%;
< flex-direction: column;
< gap: 10px;
< }
<
< .compile-profile-card-header {
< display: flex;
< align-items: center;
< flex-wrap: wrap;
< gap: 6px;
< justify-content: space-between;
< }
<
< .compile-profile-name {
< font-size: 17px;
< font-weight: 800;
< }
<
< .compile-profile-status {
< padding: 3px 7px;
< border: 1px solid currentColor;
< border-radius: 20px;
< font-size: 10px;
< font-weight: 600;
< white-space: nowrap;
< }
<
< .compile-profile-description {
< font-size: 13px;
< font-weight: 400;
< line-height: 1.5;
< overflow-wrap: anywhere;
< }
<
< .compile-jax-prealloc-label {
< display: flex;
< align-items: center;
< gap: 8px;
< padding: 4px 0;
< font-size: 12px;
< color: #94a3b8;
< cursor: pointer;
< }
<
< .compile-jax-prealloc-label:has(input:disabled) {
< opacity: .55;
< cursor: not-allowed;
< }
<
< .compile-jax-prealloc input {
< accent-color: #2563eb;
< }
<
< #app-root.theme-light .compile-jax-prealloc-label {
< color: #475569;
< }
<
< .compile-run-jax-profile-info {
< min-width: 0;
< padding: 12px 0;
< color: #cbd5e1;
< font-size: 13px;
< line-height: 1.4;
< overflow-wrap: anywhere;
< }
<
< .compile-run-jax-profile-info + .compile-run-jax-profile-info {
< border-top: 1px solid #334155;
< }
<
< .compile-run-jax-profile-info-title {
< margin-bottom: 3px;
< color: #f8fafc;
< font-weight: 800;
< }
<
< .compile-run-jax-profile-info-runtime {
< color: #94a3b8;
< }
<
< .compile-run-jax-profile-info-reason {
< margin-top: 4px;
< color: #fca5a5;
< }
<
281c153
< font-size: 12px;
---
> font-size: 10px;
366a239,243
> #app-root.theme-light .compile-run-implementation-panel {
> border-color: #cbd5e1;
> background: #f8fafc;
> }
>
368d244
< #app-root.theme-light .compile-run-jax-profile-choices,
378c254
< #app-root.theme-light .compile-run-implementation-choice:not(:disabled):hover {
---
> #app-root.theme-light .compile-run-implementation-choice:hover {
383,384c259
< #app-root.theme-light .compile-run-implementation-choice-selected,
< #app-root.theme-light .compile-run-implementation-choice-selected:not(:disabled):hover {
---
> #app-root.theme-light .compile-run-implementation-choice-selected {
389,421d263
< #app-root.theme-light .compile-profile-card.compile-run-implementation-choice-selected,
< #app-root.theme-light .compile-profile-card.compile-run-implementation-choice-selected:not(:disabled):hover {
< border-color: #2563eb;
< background: #eff6ff;
< color: #1e40af;
< box-shadow: inset 0 3px 0 #2563eb;
< }
<
< #app-root.theme-light .compile-run-implementation-choice:disabled {
< background: #f1f5f9;
< color: #94a3b8;
< }
<
< #app-root.theme-light .compile-run-jax-profile-info + .compile-run-jax-profile-info {
< border-color: #cbd5e1;
< }
<
< #app-root.theme-light .compile-run-jax-profile-info {
< color: #334155;
< }
<
< #app-root.theme-light .compile-run-jax-profile-info-title {
< color: #0f172a;
< }
<
< #app-root.theme-light .compile-run-jax-profile-info-runtime {
< color: #64748b;
< }
<
< #app-root.theme-light .compile-run-jax-profile-info-reason {
< color: #b91c1c;
< }
<
443,448d284
< }
< }
<
< @media (max-width: 420px) {
< .compile-run-jax-profile-choices {
< grid-template-columns: minmax(0, 1fr);
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/dash_app/assets/11_tab_run_theme.css clubb_release/dash_app/assets/11_tab_run_theme.css
17,171d16
< /* Existing output decision */
< .run-overwrite-modal {
< position: fixed;
< inset: 0;
< z-index: 30000;
< display: grid;
< place-items: center;
< box-sizing: border-box;
< padding: 24px;
< background: rgba(2, 6, 23, 0.68);
< backdrop-filter: blur(4px);
< }
<
< .run-overwrite-modal-hidden { display: none; }
<
< .run-overwrite-panel {
< display: grid;
< width: min(560px, calc(100vw - 32px));
< gap: 18px;
< box-sizing: border-box;
< padding: 24px;
< border: 1px solid;
< border-top: 4px solid #f59e0b;
< border-radius: 8px;
< box-shadow: 0 28px 90px rgba(0, 0, 0, 0.5);
< }
<
< #app-root.theme-dark .run-overwrite-panel {
< border-color: #334155;
< border-top-color: #f59e0b;
< background: #111827;
< color: #e5e7eb;
< }
<
< #app-root.theme-light .run-overwrite-panel {
< border-color: #cbd5e1;
< border-top-color: #d97706;
< background: #ffffff;
< color: #0f172a;
< }
<
< .run-overwrite-heading {
< display: grid;
< grid-template-columns: 44px minmax(0, 1fr);
< align-items: center;
< gap: 13px;
< }
<
< .run-overwrite-icon {
< display: grid;
< width: 42px;
< height: 42px;
< place-items: center;
< border: 1px solid rgba(245, 158, 11, 0.5);
< border-radius: 50%;
< background: rgba(245, 158, 11, 0.13);
< color: #f59e0b;
< font-size: 23px;
< font-weight: 900;
< }
<
< .run-overwrite-title {
< margin-bottom: 4px;
< font-size: 18px;
< font-weight: 800;
< }
<
< .run-overwrite-message {
< color: #94a3b8;
< font-size: 12.5px;
< line-height: 1.45;
< }
<
< #app-root.theme-light .run-overwrite-message { color: #64748b; }
<
< .run-overwrite-details {
< display: grid;
< gap: 7px;
< padding: 12px;
< border: 1px solid;
< border-radius: 6px;
< font-size: 12.5px;
< }
<
< #app-root.theme-dark .run-overwrite-details {
< border-color: #334155;
< background: #0b1220;
< }
<
< #app-root.theme-light .run-overwrite-details {
< border-color: #d9e2ef;
< background: #f8fafc;
< }
<
< .run-overwrite-fact {
< display: grid;
< grid-template-columns: 90px minmax(0, 1fr);
< gap: 10px;
< min-width: 0;
< }
<
< .run-overwrite-fact-label { color: #94a3b8; font-weight: 700; }
< #app-root.theme-light .run-overwrite-fact-label { color: #64748b; }
< .run-overwrite-fact-value { min-width: 0; overflow-wrap: anywhere; }
<
< .run-overwrite-label {
< margin-bottom: -11px;
< color: #94a3b8;
< font-size: 10px;
< font-weight: 800;
< text-transform: uppercase;
< }
<
< #app-root.theme-light .run-overwrite-label { color: #64748b; }
<
< #app-root .run-overwrite-input {
< box-sizing: border-box;
< width: 100%;
< min-height: 43px;
< border: 1px solid #64748b;
< border-radius: 6px;
< padding: 9px 11px;
< background: transparent;
< color: inherit;
< font: inherit;
< }
<
< #app-root .run-overwrite-input:focus {
< border-color: #2563eb;
< box-shadow: 0 0 0 3px rgba(37, 99, 235, 0.18);
< outline: none;
< }
<
< .run-overwrite-actions {
< display: grid;
< grid-template-columns: repeat(3, minmax(0, 1fr));
< gap: 10px;
< }
<
< .run-overwrite-button {
< min-height: 42px;
< border: 1px solid transparent;
< border-radius: 6px;
< color: #ffffff;
< cursor: pointer;
< font-size: 13px;
< font-weight: 800;
< }
<
< .run-overwrite-button:not(:disabled):hover { filter: brightness(1.1); }
< .run-overwrite-button-danger { background: #dc2626; }
< .run-overwrite-button-primary { background: #2563eb; }
< .run-overwrite-button-cancel { border-color: #64748b; background: transparent; color: inherit; }
< .run-overwrite-button:disabled { cursor: not-allowed; opacity: 0.45; }
<
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/dash_app/compile_tab/build_selector.py clubb_release/dash_app/compile_tab/build_selector.py
5d4
< import json
7,8d5
< import platform
< import subprocess
13d9
< from dash_app.shared.jax_device import jax_device_env, normalize_jax_gpu
23,69d18
< JAX_PROFILES = ("cpu", "gpu")
<
<
< def inspect_jax_runtime_profiles(repo_root: Path = REPO_ROOT, *, jax_gpu="", jax_xla_prealloc=False) -> dict[str, dict[str, Any]]:
< """Read profile metadata from the JAX wrapper without preparing an environment."""
< wrapper = Path(repo_root) / "clubb_jax" / "run_jax.py"
< profiles: dict[str, dict[str, Any]] = {}
< if not wrapper.is_file() or not os.access(wrapper, os.X_OK):
< return profiles
< for profile in JAX_PROFILES:
< error = "JAX runtime inspection returned invalid metadata."
< try:
< result = subprocess.run(
< [str(wrapper), f"--profile={profile}", "--info=json"],
< cwd=repo_root,
< env=jax_device_env({"implementation": "jax", "jax_profile": profile,
< "jax_xla_prealloc": jax_xla_prealloc if profile == "gpu" else None,
< "jax_gpu": jax_gpu if profile == "gpu" else ""}),
< capture_output=True,
< text=True,
< timeout=20,
< )
< payload = json.loads(result.stdout) if result.returncode == 0 else None
< if result.returncode != 0:
< detail = next(iter((result.stderr or result.stdout).strip().splitlines()), "")
< error = f"JAX runtime inspection failed (exit {result.returncode}). {detail}".strip()
< except subprocess.TimeoutExpired:
< payload = None
< error = "JAX runtime inspection timed out. Reopen the selector to retry."
< except (OSError, json.JSONDecodeError) as exc:
< payload = None
< error = f"JAX runtime inspection failed: {exc}"
< if (
< isinstance(payload, dict)
< and payload.get("schema_version") == 1
< and payload.get("profile") == profile
< and payload.get("status")
< in {"ready", "setup_required", "unavailable", "unknown"}
< and isinstance(payload.get("selectable"), bool)
< ):
< profiles[profile] = payload
< else:
< profiles[profile] = {
< "schema_version": 1, "profile": profile, "status": "unknown",
< "selectable": False, "reason": error,
< }
< return profiles
158,162d106
< def normalize_jax_profile(value: Any) -> str:
< profile = str(value or "cpu").strip().lower()
< return profile if profile in JAX_PROFILES else "cpu"
<
<
167,168d110
< *,
< jax_profile: Any = "cpu",
170,191c112
< """Return whether the selected implementation has its required runtime."""
< implementation = normalize_run_implementation(implementation)
< if implementation == "jax":
< profile = normalize_jax_profile(jax_profile)
< jax_root = Path(repo_root) / "clubb_jax"
< wrapper = jax_root / "run_jax.py"
< driver = jax_root / "src" / "clubb_standalone.py"
< requirements = jax_root / (
< "requirements-metal.txt"
< if profile == "gpu" and platform.system() == "Darwin"
< else "requirements-cuda13.txt"
< if profile == "gpu"
< else "requirements.txt"
< )
< if not driver.is_file():
< return False, "JAX standalone driver is missing"
< if not wrapper.is_file() or not os.access(wrapper, os.X_OK):
< return False, "JAX environment wrapper is missing or not executable"
< if not requirements.is_file():
< return False, "JAX requirements file is missing"
< return True, ""
<
---
> """Return whether an install can launch one implementation today."""
192a114
> implementation = normalize_run_implementation(implementation)
209a132,133
> if implementation == "jax" and not (Path(repo_root) / "clubb_jax" / "clubb_standalone.py").is_file():
> return False, "JAX standalone driver is missing"
213,220c137
< def selected_launch_target(
< implementation: Any,
< repo_root: Path = REPO_ROOT,
< *,
< jax_profile: Any = "cpu",
< jax_gpu: Any = "",
< jax_xla_prealloc: bool | None = None,
< ) -> dict[str, str]:
---
> def selected_launch_target(implementation: Any, repo_root: Path = REPO_ROOT) -> dict[str, str]:
223,238d139
< if implementation == "jax":
< profile = normalize_jax_profile(jax_profile)
< available, reason = build_implementation_capability(
< "", implementation, repo_root, jax_profile=profile
< )
< if not available:
< raise ValueError(f"JAX cannot start: {reason}")
< return {
< "implementation": implementation,
< "jax_profile": profile,
< **({"jax_xla_prealloc": jax_xla_prealloc} if profile == "gpu" and jax_xla_prealloc is not None else {}),
< **({"jax_gpu": normalize_jax_gpu(jax_gpu)} if profile == "gpu" and jax_gpu else {}),
< "install_dir": "",
< "build_name": f"{profile.upper()} environment",
< }
<
286,294d186
< dcc.Store(
< id="compile-run-jax-profile",
< data="cpu",
< storage_type="local",
< ),
< dcc.Store(id="compile-jax-runtime-info", data={}),
< dcc.Store(id="compile-run-jax-gpu", data="", storage_type="local"),
< dcc.Store(id="compile-run-jax-xla-prealloc", data=False, storage_type="local"),
< html.Div(id="compile-build-selector-help", className="compile-selector-help-layer"),
309c201
< **{"aria-label": "Choose a runtime or rebuild CLUBB"},
---
> **{"aria-label": "Choose or rebuild a CLUBB build"},
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/dash_app/compile_tab/callbacks.py clubb_release/dash_app/compile_tab/callbacks.py
11c11
< from dash import ALL, Input, Output, State, callback_context, dcc, html, no_update
---
> from dash import ALL, Input, Output, State, callback_context, html, no_update
13c13
< from dash_app.shared.notecard import information_body, notecard
---
> from dash_app.shared.notecard import notecard
17d16
< JAX_PROFILES,
21,22d19
< inspect_jax_runtime_profiles,
< normalize_jax_profile,
453,617d449
< def render_jax_profile_info(profile, info, *, source_error="", compact=False, label=None):
< """Render wrapper-owned hardware and runtime metadata for one JAX profile."""
< info = dict(info or {})
< hardware = info.get("hardware") if isinstance(info.get("hardware"), dict) else {}
< runtime = info.get("runtime") if isinstance(info.get("runtime"), dict) else {}
< cpu = hardware.get("cpu") if isinstance(hardware.get("cpu"), dict) else {}
< gpus = hardware.get("gpus") if isinstance(hardware.get("gpus"), list) else []
< gpu = hardware.get("selected_gpu") or (gpus[0] if gpus else {})
< gpu = gpu if isinstance(gpu, dict) else {}
< python = runtime.get("python") if isinstance(runtime.get("python"), dict) else {}
< jax = runtime.get("jax") if isinstance(runtime.get("jax"), dict) else {}
< status = str(info.get("status") or "checking")
< status_label = {
< "ready": "Ready",
< "setup_required": "Setup on first run",
< "unavailable": "Unavailable",
< "unknown": "Unknown",
< "checking": "Inspecting",
< }.get(status, status.replace("_", " ").title())
<
< if profile == "cpu":
< model = str(cpu.get("model") or "Inspecting CPU")
< logical_cpus = cpu.get("logical_cpus")
< hardware_line = model + (
< f" | {logical_cpus} logical CPUs" if logical_cpus else ""
< )
< elif gpu:
< details = []
< memory_mib = gpu.get("memory_mib")
< if isinstance(memory_mib, int):
< details.append(f"{memory_mib / 1024:.1f} GiB")
< if gpu.get("driver_version"):
< details.append(f"driver {gpu['driver_version']}")
< if gpu.get("compute_capability"):
< details.append(f"compute {gpu['compute_capability']}")
< hardware_line = str(gpu.get("name") or "NVIDIA GPU")
< if details:
< hardware_line += " | " + " | ".join(details)
< else:
< if info.get("status") == "unknown":
< hardware_line = "GPU inspection failed"
< else:
< hardware_line = "No usable GPU detected" if info else "Inspecting GPU"
<
< runtime_parts = []
< if python.get("version"):
< runtime_parts.append(
< f"Python {python['version']}"
< + ("" if python.get("installed") else " planned")
< )
< installed_jax = jax.get("installed")
< required_jax = jax.get("required")
< if installed_jax and required_jax and installed_jax != required_jax:
< runtime_parts.append(f"JAX {installed_jax} -> {required_jax}")
< elif installed_jax:
< runtime_parts.append(f"JAX {installed_jax}")
< elif required_jax:
< runtime_parts.append(f"JAX {required_jax} planned")
< if runtime.get("cuda_major"):
< runtime_parts.append(f"CUDA {runtime['cuda_major']}")
< runtime_parts.append(status_label)
< if runtime.get("xla_preallocate") is not None:
< runtime_parts.append(f"XLA preallocation: {runtime['xla_preallocate']}")
< reason = source_error or str(info.get("reason") or "")
< if compact:
< if profile == "gpu" and gpu:
< hardware_line = str(gpu.get("name") or "NVIDIA GPU")
< if not hardware.get("selected_gpu") and len(gpus) > 1:
< hardware_line = f"{len(gpus)} NVIDIA GPUs available"
< elif isinstance(gpu.get("memory_mib"), int):
< hardware_line += f" · {gpu['memory_mib'] / 1024:g} GiB"
< badge = {
< "ready": "Ready", "setup_required": "Auto setup",
< "unavailable": "Unavailable", "unknown": "Unknown", "checking": "Checking",
< "unchecked": "Check on selection",
< }.get(status, status_label)
< if source_error:
< badge = "Unavailable"
< return html.Span(
< [
< html.Span([
< html.Span(label or profile.upper(), className="compile-profile-name"),
< html.Span(badge, className="compile-profile-status"),
< ], className="compile-profile-card-header"),
< html.Span(hardware_line, className="compile-profile-description"),
< ],
< className="compile-profile-card-content",
< )
< return html.Div(
< [
< html.Div(
< "CPU" if profile == "cpu" else "GPU",
< className="compile-run-jax-profile-info-title",
< ),
< html.Div(hardware_line, className="compile-run-jax-profile-info-hardware"),
< html.Div(
< " | ".join(runtime_parts),
< className="compile-run-jax-profile-info-runtime",
< ),
< html.Div(reason, className="compile-run-jax-profile-info-reason")
< if reason
< else None,
< ],
< className=(
< "compile-run-jax-profile-info "
< f"compile-run-jax-profile-info-{status}"
< ),
< )
<
<
< def render_build_selector_help(runtime_info):
< """Use the same explanatory notecard as the plot-card help controls."""
< body = information_body(
< "Choose how CLUBB runs in Run and Profile. Your choice is remembered; "
< "jobs already running won't change.",
< [
< {"heading": "Choose a version", "bullets": [
< "Fortran: use one of your compiled CLUBB builds.",
< "Python: run through Python using a Python-enabled Fortran build.",
< "JAX: choose CPU or GPU below. Required software is set up automatically.",
< ]},
< {"heading": "CPU or GPU?", "paragraphs": [
< "CPU uses your processor. Each GPU tile selects one NVIDIA graphics card; its number matches nvidia-smi. Click a tile to choose it.",
< "Auto setup means a first-run installation is needed. Unavailable means that option can't be used. Ready means setup is complete, but a run still needs enough free memory.",
< ]},
< {"heading": "Preallocate GPU memory", "paragraphs": [
< "Leave this unchecked when sharing the GPU with other apps. Checking it reserves memory up front, which can help performance but leaves less for other apps. It's available only after selecting a compatible GPU and has no effect on CPU runs.",
< ]},
< {"heading": "Using Tune?", "paragraphs": [
< "Tune uses Fortran regardless of this chooser.",
< ]},
< ],
< )
< details = [
< render_jax_profile_info(profile, (runtime_info or {}).get(profile),
< source_error=build_implementation_capability("", "jax", jax_profile=profile)[1])
< for profile in JAX_PROFILES
< ]
< return notecard(
< "Choosing a runtime", [body, html.Details([
< html.Summary("Technical details"),
< html.P("GPU choices are saved by UUID and override CUDA_VISIBLE_DEVICES for each new job. "
< "The memory checkbox overrides XLA_PYTHON_CLIENT_PREALLOCATE. "
< "For command-line runs, -jax=gpu,xla_prealloc enables preallocation; "
< "otherwise an explicit environment setting is preserved, defaulting to false."),
< *details,
< ])],
< {"type": "compile-selector-help-close", "index": "runtime"}, size="medium",
< )
<
<
< def jax_preallocation_available(runtime_info, profile, gpu):
< """Require a compatible selection, not stale metadata for a different GPU."""
< info = (runtime_info or {}).get("gpu") or {}
< selected = ((info.get("hardware") or {}).get("selected_gpu") or {}).get("uuid")
< return (
< profile == "gpu"
< and (info.get("runtime") or {}).get("accelerator") != "metal"
< and info.get("selectable") is True
< and info.get("status") in {"ready", "setup_required"}
< and (not gpu or gpu == selected)
< and build_implementation_capability("", "jax", jax_profile="gpu")[0]
< )
<
<
626,629d457
< jax_profile="cpu",
< jax_runtime_info=None,
< jax_gpu="",
< jax_xla_prealloc=False,
633,634d460
< jax_profile = normalize_jax_profile(jax_profile)
< jax_runtime_info = dict(jax_runtime_info or {})
642,647c468
< html.Div([
< html.Div("Run with", className="compile-build-selector-heading"),
< html.Button("?", id={"type": "compile-selector-help-open", "index": "runtime"},
< type="button", n_clicks=0, className="plots-card-help",
< title="About runtime selection", **{"aria-label": "About runtime selection"}),
< ], className="compile-selector-toolbar"),
---
> html.Div("Implementation", className="compile-build-selector-heading"),
655d475
< **{"aria-pressed": str(name == implementation).lower()},
668c488,490
< "Tune uses the Fortran worker.",
---
> "Tune currently uses its F2PY worker backend; this choice is saved for future Tune support."
> if trigger_id == "tune-selected-build-badge"
> else "Choose the implementation, then its supporting CLUBB build.",
670c492
< ) if trigger_id == "tune-selected-build-badge" else None,
---
> ),
675,742d496
< if implementation == "jax":
< profile_info = {
< profile: dict(jax_runtime_info.get(profile) or {})
< for profile in JAX_PROFILES
< }
< profile_source = {
< profile: build_implementation_capability(
< "", "jax", jax_profile=profile
< )
< for profile in JAX_PROFILES
< }
< tiles = []
<
< def add_tile(profile, index, label, info, selected, title=""):
< available, source_error = profile_source[profile]
< disabled = not available or info.get("selectable") is False
< tiles.append(html.Button(
< render_jax_profile_info(profile, info, source_error=source_error,
< compact=True, label=label),
< id={"type": "compile-run-jax-profile-choice", "index": index},
< type="button", n_clicks=0, disabled=disabled,
< title=(source_error or str(info.get("reason") or "")) if disabled
< else title or f"Use {label}",
< className="compile-profile-card compile-run-implementation-choice" + (
< " compile-run-implementation-choice-selected" if selected else ""),
< **{"aria-pressed": str(selected).lower()},
< ))
<
< add_tile("cpu", "cpu", "CPU", profile_info["cpu"], jax_profile == "cpu")
< hardware = profile_info["gpu"].get("hardware") or {}
< gpus = [gpu for gpu in hardware.get("gpus") or [] if gpu.get("uuid")]
< inspected_uuid = (hardware.get("selected_gpu") or {}).get("uuid")
< effective_uuid = jax_gpu or inspected_uuid
< for gpu in gpus:
< uuid = gpu["uuid"]
< # The wrapper report describes its selected device, not every card.
< # Do not apply one card's compatibility result to another device.
< info = dict(profile_info["gpu"]) if uuid == inspected_uuid else {
< "status": "unchecked", "selectable": True,
< }
< info["hardware"] = {"selected_gpu": gpu, "gpus": [gpu]}
< add_tile("gpu", uuid, f"GPU {gpu.get('index', '?')}", info,
< jax_profile == "gpu" and uuid == effective_uuid,
< f"{uuid} | PCI {gpu.get('pci_bus_id', 'unknown')}")
<
< if jax_gpu and jax_gpu not in {gpu["uuid"] for gpu in gpus}:
< add_tile("gpu", jax_gpu, "Saved GPU", {
< "status": "unavailable", "selectable": False,
< "reason": f"Selected GPU is no longer detected: {jax_gpu}",
< }, jax_profile == "gpu")
< elif not gpus or (jax_profile == "gpu" and not effective_uuid):
< # Keep an inherited, unmapped selection honest until a card is chosen.
< add_tile("gpu", "gpu", "GPU Default" if gpus else "GPU", profile_info["gpu"],
< jax_profile == "gpu", "Use the dashboard environment")
<
< items.append(html.Div([
< html.Div("Compute", className="compile-build-selector-heading"),
< html.Div(tiles, className="compile-run-jax-profile-choices"),
< html.Div(dcc.Checklist(
< id={"type": "compile-jax-prealloc-choice", "index": "setting"},
< options=[{"label": "Preallocate GPU memory", "value": "enabled",
< "disabled": not jax_preallocation_available(jax_runtime_info, jax_profile, jax_gpu)}],
< value=["enabled"] if jax_xla_prealloc else [],
< className="compile-jax-prealloc",
< labelClassName="compile-jax-prealloc-label",
< ), title="Available after selecting a compatible GPU. Has no effect on CPU runs."),
< ], className="compile-run-implementation-panel"))
< return items
1261,1275d1014
< @app.callback(
< Output("compile-build-selector-help", "children"),
< Input({"type": "compile-selector-help-open", "index": ALL}, "n_clicks"),
< Input({"type": "compile-selector-help-close", "index": ALL}, "n_clicks"),
< State("compile-jax-runtime-info", "data"),
< prevent_initial_call=True,
< )
< def toggle_selector_help(_open, _close, runtime_info):
< trigger = clicked_trigger_id()
< if not isinstance(trigger, dict):
< return no_update
< if trigger.get("type") == "compile-selector-help-close":
< return ""
< return render_build_selector_help(runtime_info)
<
1299,1348d1037
< Output("compile-jax-runtime-info", "data"),
< Input("compile-build-selector-anchor", "data"),
< Input("compile-run-implementation", "data"),
< Input("compile-run-jax-gpu", "data"),
< Input("compile-run-jax-xla-prealloc", "data"),
< prevent_initial_call=True,
< )
< def refresh_jax_runtime_info(anchor, implementation, jax_gpu, jax_xla_prealloc):
< if not anchor or normalize_run_implementation(implementation) != "jax":
< return no_update
< return inspect_jax_runtime_profiles(jax_gpu=jax_gpu, jax_xla_prealloc=jax_xla_prealloc)
<
< @app.callback(
< Output("compile-run-jax-xla-prealloc", "data"),
< Input({"type": "compile-jax-prealloc-choice", "index": ALL}, "value"),
< State("compile-jax-runtime-info", "data"),
< State("compile-run-jax-profile", "data"),
< State("compile-run-jax-gpu", "data"),
< State("compile-run-jax-xla-prealloc", "data"),
< prevent_initial_call=True,
< )
< def select_jax_preallocation(values, runtime_info, profile, gpu, current):
< if not values or not jax_preallocation_available(runtime_info, profile, gpu):
< return no_update
< enabled = "enabled" in (values[0] or [])
< return enabled if enabled != current else no_update
<
< @app.callback(
< Output("compile-run-jax-profile", "data"),
< Output("compile-run-jax-gpu", "data"),
< Input({"type": "compile-run-jax-profile-choice", "index": ALL}, "n_clicks"),
< State("compile-jax-runtime-info", "data"),
< prevent_initial_call=True,
< )
< def select_run_jax_compute(_clicks, runtime_info):
< """Select the backend and physical device atomically, even on reselect."""
< trigger = clicked_trigger_id()
< if not isinstance(trigger, dict):
< return no_update, no_update
< choice = trigger.get("index")
< if choice == "cpu":
< return "cpu", no_update
< if choice == "gpu":
< return "gpu", ""
< gpus = ((runtime_info or {}).get("gpu", {}).get("hardware") or {}).get("gpus") or []
< if choice in {gpu.get("uuid") for gpu in gpus}:
< return "gpu", choice
< return no_update, no_update
<
< @app.callback(
1359,1362d1047
< Input("compile-run-jax-profile", "data"),
< Input("compile-jax-runtime-info", "data"),
< Input("compile-run-jax-gpu", "data"),
< Input("compile-run-jax-xla-prealloc", "data"),
1364,1375c1049
< def update_build_selector(
< anchor,
< discovery,
< statuses,
< failures,
< job,
< implementation,
< jax_profile,
< jax_runtime_info,
< jax_gpu,
< jax_xla_prealloc,
< ):
---
> def update_build_selector(anchor, discovery, statuses, failures, job, implementation):
1378d1051
< jax_profile = normalize_jax_profile(jax_profile)
1388,1391d1060
< jax_profile,
< jax_runtime_info,
< jax_gpu,
< jax_xla_prealloc,
1395,1419c1064
< if implementation == "jax":
< available, reason = build_implementation_capability(
< "", "jax", jax_profile=jax_profile
< )
< runtime_info = dict((jax_runtime_info or {}).get(jax_profile) or {})
< runtime_status = str(runtime_info.get("status") or "checking")
< if runtime_info.get("selectable") is False:
< available = False
< reason = str(runtime_info.get("reason") or "JAX profile unavailable")
< selected_name = "CPU" if jax_profile == "cpu" else "GPU"
< if jax_profile == "gpu" and jax_gpu:
< gpus = (runtime_info.get("hardware") or {}).get("gpus") or []
< selected_gpu = next((gpu for gpu in gpus if gpu.get("uuid") == jax_gpu), {})
< selected_name = f"GPU {selected_gpu.get('index', '?')} · {selected_gpu.get('name', 'unavailable')}"
< if not available:
< status_class = "compile-build-card-failed"
< status_label = reason
< elif runtime_status == "ready":
< status_class = "compile-build-card-current"
< status_label = "ready"
< else:
< status_class = "compile-build-card-checking"
< status_label = runtime_status.replace("_", " ")
< elif selected_build:
< selected_name = selected_info["name"]
---
> if selected_build:
1430d1074
< selected_name = selected_info["name"]
1441c1085
< build_selector_trigger_children(selected_name, implementation),
---
> build_selector_trigger_children(selected_info["name"], implementation),
1451c1095
< f"{implementation.title()} using {selected_name} — {status_label}. Click to configure.",
---
> f"{implementation.title()} using {selected_info['name']} — {status_label}. Click to configure.",
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/dash_app/DEVELOPMENT.md clubb_release/dash_app/DEVELOPMENT.md
185,246d184
<
< ## Runtime selection and JAX job settings
<
< The runtime chooser uses the read-only report from
< [`clubb_jax/run_jax.py`](../clubb_jax/run_jax.py). Hardware discovery does not
< require an installed JAX environment. A compatible device whose environment
< needs setup remains selectable.
<
< GPU choices are stored by UUID. Until a specific GPU is chosen, the inherited
< selection retains the server's `CUDA_VISIBLE_DEVICES` setting (or CUDA's
< default); an unmapped active default has its own tile. Explicit selections
< override visibility for that job and all its cases or profiling workers. They
< do not distribute work across GPUs. CPU runs ignore a saved GPU selection.
<
< The preallocation checkbox is enabled only for a compatible CUDA selection,
< including one that needs environment setup. It is disabled during checks and
< for CPU, Metal, or unavailable selections. Both checked and unchecked values
< explicitly override the server's `XLA_PYTHON_CLIENT_PREALLOCATE` setting. Freeze
< the effective GPU and preallocation settings into each submitted Run/Profile
< request and include them in command previews. Enabling preallocation maps to
< `-jax=gpu,xla_prealloc`.
<
< Compiled implementations follow `install/selected`, falling back to
< `install/latest`. Explicit runner `-exe` or `-install_dir` options override
< that default. Tune continues to use its Fortran/F2PY worker independently of
< the chooser's JAX selection.
<
< ## Profile results
<
< The Profile tab is a browser interface to `utilities/time_clubb.py`. Its top
< benchmark panel configures the case, process/per-process-batch-size sweep,
< repetitions, executable, configuration, overrides, and additional
< `run_scm.py` arguments. A direct one-second polling path reads the active
< summary and process rows and renders figures server-side, so results appear
< after each measured repetition while the broker-owned job is running; warmups
< remain hidden. Browser stores retain only compact timer/process choices rather
< than the growing raw timing table. The running row counter and all four figures
< are returned by the same callback response, so visible progress cannot advance
< independently of the plots. Stored profiles can be overlaid, compared with a baseline, or
< viewed as process distributions and exclusive-cost decompositions. The right
< rail has a profile-selection section above a separate set of shared comparison
< controls; plot-specific options remain beside the plot they affect. The
< profile chooser shows the three newest unselected results by default, expands
< to the full library, and displays active comparisons as removable pills. The
< benchmark-label field indicates when its normalized profile name already
< exists. Starting that benchmark asks for confirmation, then replaces the
< existing profile in place instead of creating a timestamped version; any older
< same-label/same-case versions are removed from the active comparison selection.
<
< The selected directory is a collection of compact, directly commit-able
< profile folders. Each benchmark creates `<profile-name>/` containing
< `README.md`, profile-wide `profile.json` provenance, one workload row per
< process-count/batch-size point in `batches.csv`, raw timer observations in
< `timings.csv`, and one representative input/setup/log/native-timing set under
< `logs/<batch-id>/`. Child processes otherwise run in temporary directories,
< which are deleted after each workload is aggregated. Warmups are retained with
< `phase=warmup` but excluded from the default plots. Dash derives statistical
< summaries in memory instead of storing duplicate summary files. **Export
< selected** downloads complete profiles as a ZIP; **Import** accepts those ZIPs
< on another machine or checkout. Provenance includes the effective vertical
< level count, observed model steps, source revision, executable checksum, host,
< timer backend, and time basis so Dash can flag potentially incomparable runs.
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/dash_app/profile_tab/callbacks.py clubb_release/dash_app/profile_tab/callbacks.py
62,65c62
< def collect_profile_settings(
< values: list[Any], implementation: Any = None, jax_profile: Any = None, jax_gpu: Any = None,
< jax_xla_prealloc: bool | None = None,
< ) -> dict[str, Any]:
---
> def collect_profile_settings(values: list[Any], implementation: Any = None) -> dict[str, Any]:
83,86c80
< settings.update(
< selected_launch_target(implementation, jax_profile=jax_profile, jax_gpu=jax_gpu,
< jax_xla_prealloc=jax_xla_prealloc)
< )
---
> settings.update(selected_launch_target(implementation))
441,443d434
< State("compile-run-jax-profile", "data"),
< State("compile-run-jax-gpu", "data"),
< State("compile-run-jax-xla-prealloc", "data"),
457,461c448,449
< setting_values = values[:-6]
< implementation = values[-6]
< jax_profile = values[-5]
< jax_gpu = values[-4]
< jax_xla_prealloc = values[-3]
---
> setting_values = values[:-3]
> implementation = values[-3]
466,468c454
< settings = collect_profile_settings(
< list(setting_values), implementation, jax_profile, jax_gpu, jax_xla_prealloc
< )
---
> settings = collect_profile_settings(list(setting_values), implementation)
604,606d589
< Input("compile-run-jax-profile", "data"),
< Input("compile-run-jax-gpu", "data"),
< Input("compile-run-jax-xla-prealloc", "data"),
611,613c594
< collect_profile_settings(
< list(values[:-4]), values[-4], values[-3], values[-2], values[-1]
< )
---
> collect_profile_settings(list(values[:-1]), values[-1])
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/dash_app/profile_tab/runtime.py clubb_release/dash_app/profile_tab/runtime.py
17d16
< from dash_app.shared.jax_device import jax_device_env, normalize_jax_gpu, jax_command_display
129,135d127
< jax_profile = _clean(settings.get("jax_profile")).lower() or "cpu"
< if jax_profile not in {"cpu", "gpu"}:
< raise ValueError("JAX profile must be CPU or GPU")
< jax_gpu = normalize_jax_gpu(settings.get("jax_gpu"))
< jax_xla_prealloc = settings.get("jax_xla_prealloc")
< jax_device_env({"implementation": implementation, "jax_profile": jax_profile,
< "jax_gpu": jax_gpu, "jax_xla_prealloc": jax_xla_prealloc}, {})
177,179d168
< "jax_profile": jax_profile,
< "jax_gpu": jax_gpu,
< "jax_xla_prealloc": jax_xla_prealloc,
222,228c211,212
< modifier = ",xla_prealloc" if normalized["jax_xla_prealloc"] is True else ""
< command.append(f"-jax={normalized['jax_profile']}{modifier}")
< if (
< normalized["install_dir"]
< and not normalized["executable"]
< and normalized["implementation"] != "jax"
< ):
---
> command.append("-jax")
> if normalized["install_dir"] and not normalized["executable"]:
235c219
< return jax_command_display(profile_command(settings), settings)
---
> return shlex.join(profile_command(settings))
308c292
< env=jax_device_env(normalized),
---
> env=os.environ.copy(),
326c310
< "command_display": jax_command_display(command, normalized),
---
> "command_display": shlex.join(command),
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/dash_app/pytests/test_agent_services.py clubb_release/dash_app/pytests/test_agent_services.py
36,78d35
< def test_jax_build_identity_tracks_repository_runtime_not_fortran_install(tmp_path, monkeypatch):
< import platform
<
< from dash_app.shared import actions
<
< jax_root = tmp_path / "clubb_jax"
< (jax_root / "src").mkdir(parents=True)
< wrapper = jax_root / "run_jax.py"
< driver = jax_root / "src" / "clubb_standalone.py"
< requirements = jax_root / (
< "requirements-metal.txt"
< if platform.system() == "Darwin"
< else "requirements-cuda13.txt"
< )
< wrapper.write_text("#!/usr/bin/env bash\n", encoding="utf-8")
< driver.write_text("def main(): pass\n", encoding="utf-8")
< requirements.write_text("jax GPU requirements\n", encoding="utf-8")
< monkeypatch.setattr(actions, "REPO_ROOT", tmp_path)
< monkeypatch.setenv("CLUBB_JAX_ACCELERATOR", "cpu")
< monkeypatch.setenv("CLUBB_JAX_PRECISION", "single")
<
< identity = actions._scm_build_identity(
< {
< "implementation": "jax",
< "jax_profile": "gpu",
< "install_dir": str(tmp_path / "install" / "selected"),
< }
< )
<
< assert identity["implementation"] == "jax"
< assert identity["runtime"] == "repository-managed"
< assert identity["install_directory"] is None
< assert identity["executable"]["path"] == str(wrapper)
< assert identity["driver"]["path"] == str(driver)
< assert identity["requirements"]["path"] == str(requirements)
< assert identity["precision"] == "single"
< assert identity["profile"] == "gpu"
< assert identity["accelerator"] == (
< "metal" if platform.system() == "Darwin" else "cuda13"
< )
< assert identity["driver"]["sha256"]
<
<
133,145d89
< with pytest.raises(ValidationError):
< ScmRunRequest(
< request_id="request-123",
< case="arm",
< implementation="numpy",
< )
< with pytest.raises(ValidationError):
< ScmRunRequest(
< request_id="request-123",
< case="arm",
< implementation="jax",
< jax_profile="tpu",
< )
168,171c112
< @pytest.mark.parametrize("profile,gpu,prealloc", [("cpu", "", None),
< ("gpu", "GPU-aaaaaaaa-1111-2222-3333-000000000001", False),
< ("gpu", "GPU-aaaaaaaa-1111-2222-3333-000000000001", True)])
< def test_scm_batch_submission_uses_one_flat_output_and_is_idempotent(tmp_path, monkeypatch, profile, gpu, prealloc):
---
> def test_scm_batch_submission_uses_one_flat_output_and_is_idempotent(tmp_path, monkeypatch):
180,183d120
< assert cli_options["implementation"] == "jax"
< assert cli_options["jax_profile"] == profile
< assert cli_options["jax_gpu"] == gpu
< assert cli_options["jax_xla_prealloc"] is prealloc
203,211c140
< request = ScmRunBatchRequest(
< request_id="batch-request-123",
< cases=["arm", "bomex"],
< implementation="jax",
< jax_profile=profile,
< jax_gpu=gpu,
< jax_xla_prealloc=prealloc,
< max_workers=2,
< )
---
> request = ScmRunBatchRequest(request_id="batch-request-123", cases=["arm", "bomex"], max_workers=2)
229,230d157
< assert manifest["job"]["build_identity"]["implementation"] == "jax"
< assert manifest["job"]["build_identity"]["install_directory"] is None
891,896d817
< assert set(schema["$defs"]["ScmRunRequest"]["properties"]["implementation"]["enum"]) == {
< "fortran", "python", "jax"
< }
< assert set(schema["$defs"]["ScmRunRequest"]["properties"]["jax_profile"]["enum"]) == {
< "cpu", "gpu"
< }
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/dash_app/pytests/test_callbacks_runs.py clubb_release/dash_app/pytests/test_callbacks_runs.py
8,10d7
< prepared_run_with_output,
< run_output_rename_available,
< submit_prepared_run,
13c10
< from dash_app.run_tab.runtime import build_case_command, output_directory_details
---
> from dash_app.run_tab.runtime import build_case_command
47,60c44,47
< action_callback = next(
< entry
< for key, entry in app.callback_map.items()
< if "run-action-result.data" in key
< )
< action_inputs = {item["id"] for item in action_callback["inputs"]}
< assert {
< "run-button",
< "run-cancel",
< "run-clear",
< "run-overwrite-button",
< "run-rename-button",
< "run-overwrite-cancel-button",
< } <= action_inputs
---
> action_inputs = {
> item["id"] for item in app.callback_map["run-action-result.data"]["inputs"]
> }
> assert {"run-button", "run-cancel", "run-clear"} <= action_inputs
164,168c151
< {
< "implementation": "jax",
< "jax_profile": "gpu",
< "install_dir": "/tmp/build-two",
< },
---
> {"implementation": "jax", "install_dir": "/tmp/build-two"},
172,258c155
< assert "-jax=gpu bomex" in jax_command
< assert "-install_dir" not in jax_command
<
<
< def test_output_directory_details_count_only_stats_cases(tmp_path):
< output = tmp_path / "results"
< output.mkdir()
< (output / "arm_stats.nc").write_bytes(b"CDF")
< (output / "bomex_stats.nc").write_bytes(b"CDF")
< (output / "run.log").write_text("done", encoding="utf-8")
< (output / "nested").mkdir()
<
< details = output_directory_details(output)
<
< assert details["path"] == str(output.resolve())
< assert details["nonempty"] is True
< assert details["case_count"] == 2
< assert details["created"] != "Not created yet"
< assert details["last_edited"] != "Not created yet"
<
<
< def test_output_rename_changes_only_the_frozen_output_target(tmp_path):
< current = tmp_path / "current"
< renamed = tmp_path / "renamed"
< pending = {
< "cases": ["arm"],
< "output_dir": str(current),
< "cli_options": {"out_dir": str(current), "debug": "0"},
< }
<
< updated = prepared_run_with_output(pending, renamed)
<
< assert run_output_rename_available(renamed, pending) is True
< assert run_output_rename_available(current, pending) is False
< assert updated["output_dir"] == str(renamed)
< assert updated["cli_options"] == {"out_dir": str(renamed), "debug": "0"}
< assert pending["cli_options"]["out_dir"] == str(current)
<
<
< def test_prepared_run_submission_preserves_frozen_settings():
< calls = []
< gpu = "GPU-aaaaaaaa-1111-2222-3333-000000000001"
<
< def perform_action(action, payload, *, internal):
< calls.append((action, payload, internal))
< return {"job_id": "batch-job"}
<
< result = submit_prepared_run(
< {
< "cases": ["arm", "bomex"],
< "stats": "standard_stats.in",
< "config": "default",
< "overrides": {"flags": {"l_uv_nudge": ".true."}},
< "typed_overrides": {"l_uv_nudge": ".true."},
< "cli_options": {
< "implementation": "jax",
< "jax_profile": "gpu",
< "jax_gpu": gpu,
< "jax_xla_prealloc": False,
< "out_dir": "renamed",
< },
< "typed_options": {"max_iters": 10},
< "max_workers": 2,
< "output_dir": "renamed",
< "implementation": "jax",
< "jax_profile": "gpu",
< "jax_gpu": gpu,
< "jax_xla_prealloc": False,
< },
< perform_action,
< )
<
< assert result["job_id"] == "batch-job"
< action, payload, internal = calls[0]
< assert action == "domain_submit_scm_batch"
< assert internal is True
< assert payload["request"]["cases"] == ["arm", "bomex"]
< assert payload["request"]["implementation"] == "jax"
< assert payload["request"]["jax_profile"] == "gpu"
< assert payload["request"]["jax_gpu"] == gpu
< assert payload["request"]["jax_xla_prealloc"] is False
< assert payload["request"]["max_workers"] == 2
< assert payload["native_cli_options"]["implementation"] == "jax"
< assert payload["native_cli_options"]["jax_profile"] == "gpu"
< assert payload["native_cli_options"]["jax_gpu"] == gpu
< assert payload["native_cli_options"]["jax_xla_prealloc"] is False
< assert payload["native_cli_options"]["out_dir"] == "renamed"
---
> assert "-jax -install_dir /tmp/build-two bomex" in jax_command
Only in clubb/dash_app/pytests: test_jax_device.py
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/dash_app/pytests/test_profile_tab.py clubb_release/dash_app/pytests/test_profile_tab.py
200,202c200
< @pytest.mark.parametrize(
< ("implementation", "flag"), (("python", "-python"), ("jax", "-jax=cpu"))
< )
---
> @pytest.mark.parametrize(("implementation", "flag"), (("python", "-python"), ("jax", "-jax")))
212,215c210
< if implementation == "python":
< assert command[command.index("-install_dir") + 1] == str(install.resolve())
< else:
< assert "-install_dir" not in command
---
> assert command[command.index("-install_dir") + 1] == str(install.resolve())
638c633
< {"state": "finished", "run_id": "candidate-run"},
---
> {"state": "finished", "run_id": "rtx3080"},
646c641
< {"state": "running", "run_id": "candidate-run"},
---
> {"state": "running", "run_id": "rtx3080"},
650,651c645,646
< assert preferred == ["imported", "candidate-run"]
< assert replacement == "candidate-run"
---
> assert preferred == ["imported", "rtx3080"]
> assert replacement == "rtx3080"
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/dash_app/pytests/test_run_broker_simplification.py clubb_release/dash_app/pytests/test_run_broker_simplification.py
2d1
< import shlex
12d10
< from dash_app.run_tab.telemetry import scm_run_view
42,75d39
<
<
< @pytest.mark.parametrize("profile,prealloc", [("cpu", None), ("gpu", False), ("gpu", True)])
< def test_api_jax_batch_copy_commands_preserve_launch_target(tmp_path, monkeypatch, profile, prealloc):
< _isolated_activity(tmp_path, monkeypatch)
< store = _isolated_batch_services(tmp_path, monkeypatch)
< gpu = "GPU-aaaaaaaa-1111-2222-3333-000000000001" if profile == "gpu" else ""
< request = ScmRunBatchRequest(
< request_id="jax-copy-command-audit",
< cases=["arm", "bomex"],
< implementation="jax",
< jax_profile=profile,
< jax_gpu=gpu,
< jax_xla_prealloc=prealloc,
< run_options={"max_iters": 1},
< )
< batch = actions.submit_scm_batch(request)
< for child in batch["children"]:
< for state in ("queued", "running", "finished"):
< if state != "queued":
< store.update(child["job_id"], state=state, runtime={"cli_options": {
< "implementation": "jax", "jax_profile": profile,
< "jax_gpu": gpu, "jax_xla_prealloc": prealloc, "max_iters": 1,
< }})
< command = shlex.split(scm_run_view(store.get(child["job_id"]))["command"])
< modifier = ",xla_prealloc" if prealloc else ""
< assert f"-jax={profile}{modifier}" in command
< assert command[command.index("-max_iters") + 1] == "1"
< assert command[-1] == child["case"]
< if gpu:
< assert f"CUDA_VISIBLE_DEVICES={gpu}" in command
< assert f"XLA_PYTHON_CLIENT_PREALLOCATE={str(prealloc).lower()}" in command
< else:
< assert not any(part.startswith("CUDA_VISIBLE_DEVICES=") for part in command)
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/dash_app/pytests/test_selected_build.py clubb_release/dash_app/pytests/test_selected_build.py
1,2d0
< import json
< import subprocess
5,6d2
< import pytest
<
16d11
< inspect_jax_runtime_profiles,
19d13
< selected_launch_target,
51,60d44
< def component_text(component) -> str:
< if isinstance(component, str):
< return component
< if not isinstance(component, Component):
< return ""
< children = getattr(component, "children", None)
< values = children if isinstance(children, (list, tuple)) else [children]
< return " ".join(component_text(child) for child in values if child is not None)
<
<
153d136
< assert component_id_count(app.layout, "compile-build-selector-help") == 1
209,211c192,193
< jax_root = tmp_path / "clubb_jax"
< jax_driver = jax_root / "src" / "clubb_standalone.py"
< jax_driver.parent.mkdir(parents=True)
---
> jax_driver = tmp_path / "clubb_jax" / "clubb_standalone.py"
> jax_driver.parent.mkdir()
213,215d194
< wrapper = jax_root / "run_jax.py"
< wrapper.touch(mode=0o755)
< (jax_root / "requirements.txt").touch()
217,443d195
<
<
< def test_jax_launch_target_does_not_require_a_compiled_install(tmp_path):
< jax_root = tmp_path / "clubb_jax"
< (jax_root / "src").mkdir(parents=True)
< (jax_root / "src" / "clubb_standalone.py").touch()
< (jax_root / "requirements.txt").touch()
< (jax_root / "run_jax.py").touch(mode=0o755)
<
< target = selected_launch_target("jax", tmp_path)
<
< assert target == {
< "implementation": "jax",
< "jax_profile": "cpu",
< "install_dir": "",
< "build_name": "CPU environment",
< }
<
<
< def test_jax_selector_replaces_compiled_build_rows_with_managed_runtime():
< menu = callbacks.render_compact_build_selector(
< {"builds": []},
< implementation="jax",
< )
<
< assert len(menu) == 2
< assert menu[0].children[2] is None
< toolbar = menu[0].children[0]
< assert toolbar.children[0].children == "Run with"
< assert toolbar.children[1].children == "?"
< assert toolbar.children[1].className == "plots-card-help"
< assert menu[1].children[0].children == "Compute"
< buttons = menu[1].children[1].children
< assert [button.id["index"] for button in buttons] == ["cpu", "gpu"]
< assert component_text(buttons[0]).startswith("CPU ")
< assert component_text(buttons[1]).startswith("GPU ")
< assert buttons[0].to_plotly_json()["props"]["aria-pressed"] == "true"
< assert buttons[1].to_plotly_json()["props"]["aria-pressed"] == "false"
< assert "compile-run-implementation-choice-selected" in buttons[0].className
<
<
< def test_jax_selector_displays_wrapper_metadata_and_disables_unavailable_gpu():
< runtime_info = {
< "cpu": {
< "status": "ready",
< "selectable": True,
< "hardware": {"cpu": {"model": "Test CPU", "logical_cpus": 16}},
< "runtime": {
< "python": {"version": "3.12.4", "installed": True},
< "jax": {"required": "0.11.0", "installed": "0.11.0"},
< },
< },
< "gpu": {
< "status": "unavailable",
< "selectable": False,
< "reason": "CUDA 13 requires NVIDIA driver 580 or newer; detected 470.239",
< "hardware": {
< "cpu": {},
< "gpus": [
< {
< "name": "Test GPU",
< "memory_mib": 24576,
< "driver_version": "470.239",
< "compute_capability": "8.0",
< }
< ],
< },
< "runtime": {
< "cuda_major": 13,
< "python": {"version": "3.12.4", "installed": True},
< "jax": {"required": "0.11.0", "installed": "0.11.0"},
< },
< },
< }
<
< menu = callbacks.render_compact_build_selector(
< {"builds": []}, implementation="jax", jax_runtime_info=runtime_info
< )
<
< buttons = menu[1].children[1].children
< assert buttons[0].disabled is False
< assert buttons[1].disabled is True
< text = component_text(menu[1])
< assert "Test CPU | 16 logical CPUs" in text
< assert "Test GPU · 24 GiB" in text
< assert "Unavailable" in text
< assert "Python 3.12.4" not in text
< assert "requires NVIDIA driver 580" in buttons[1].title
< help_text = component_text(callbacks.render_build_selector_help(runtime_info))
< assert "Test GPU | 24.0 GiB | driver 470.239 | compute 8.0" in help_text
< assert "Python 3.12.4 | JAX 0.11.0 | CUDA 13 | Unavailable" in help_text
< assert "requires NVIDIA driver 580" in help_text
< assert "CUDA_VISIBLE_DEVICES" in help_text
<
<
< def test_profile_card_describes_selected_gpu_not_first_inventory_entry():
< info = {
< "status": "ready",
< "hardware": {
< "gpus": [{"name": "Test GPU A"}, {"name": "Test GPU B"}],
< "selected_gpu": {"name": "Test GPU B", "memory_mib": 8192},
< },
< }
< card = callbacks.render_jax_profile_info("gpu", info, compact=True)
< assert component_text(card) == "GPU Ready Test GPU B · 8 GiB"
< del info["hardware"]["selected_gpu"]
< assert "2 NVIDIA GPUs available" in component_text(
< callbacks.render_jax_profile_info("gpu", info, compact=True)
< )
<
<
< def test_runtime_help_ignores_mount_events_and_opens_and_closes(monkeypatch):
< app = Dash(__name__, suppress_callback_exceptions=True)
< callbacks.register_compile_callbacks(app)
< toggle = app.callback_map["compile-build-selector-help.children"]["callback"].__wrapped__
< monkeypatch.setattr(callbacks, "clicked_trigger_id", lambda: None)
< assert toggle([0], [], {}) is callbacks.no_update
< monkeypatch.setattr(callbacks, "clicked_trigger_id", lambda: {
< "type": "compile-selector-help-open", "index": "runtime",
< })
< help_card = toggle([1], [], {})
< assert help_card.className == "shared-notecard-overlay"
< assert "Choosing a runtime" in component_text(help_card)
< assert "Close" in component_text(help_card)
< monkeypatch.setattr(callbacks, "clicked_trigger_id", lambda: {
< "type": "compile-selector-help-close", "index": "runtime",
< })
< assert toggle([1], [1], {}) == ""
<
<
< def test_jax_selector_checks_each_profile_source_independently(monkeypatch):
< monkeypatch.setattr(
< callbacks,
< "build_implementation_capability",
< lambda _install, _implementation, *, jax_profile: (
< (True, "")
< if jax_profile == "cpu"
< else (False, "GPU requirements are missing")
< ),
< )
<
< menu = callbacks.render_compact_build_selector(
< {"builds": []}, implementation="jax", jax_profile="cpu"
< )
<
< buttons = menu[1].children[1].children
< assert buttons[0].disabled is False
< assert buttons[1].disabled is True
< assert buttons[1].title == "GPU requirements are missing"
<
<
< def test_dash_runtime_inspection_accepts_only_wrapper_schema(tmp_path, monkeypatch):
< wrapper = tmp_path / "clubb_jax" / "run_jax.py"
< wrapper.parent.mkdir(parents=True)
< wrapper.touch(mode=0o755)
<
< def fake_run(command, **_kwargs):
< profile = command[1].split("=", 1)[1]
< payload = {
< "schema_version": 1,
< "profile": profile,
< "selectable": profile == "cpu",
< "status": "ready" if profile == "cpu" else "unavailable",
< }
< return subprocess.CompletedProcess(command, 0, json.dumps(payload), "")
<
< monkeypatch.setattr("dash_app.compile_tab.build_selector.subprocess.run", fake_run)
<
< info = inspect_jax_runtime_profiles(tmp_path)
<
< assert info["cpu"]["selectable"] is True
< assert info["gpu"]["selectable"] is False
<
<
< def test_jax_gpu_capability_and_launch_target_use_native_requirements(tmp_path):
< import platform
<
< jax_root = tmp_path / "clubb_jax"
< (jax_root / "src").mkdir(parents=True)
< (jax_root / "src" / "clubb_standalone.py").touch()
< (jax_root / "requirements.txt").touch()
< (jax_root / "run_jax.py").touch(mode=0o755)
<
< available, reason = build_implementation_capability(
< "", "jax", tmp_path, jax_profile="gpu"
< )
< assert available is False
< assert "requirements" in reason.lower()
<
< requirements_name = (
< "requirements-metal.txt"
< if platform.system() == "Darwin"
< else "requirements-cuda13.txt"
< )
< (jax_root / requirements_name).touch()
< target = selected_launch_target("jax", tmp_path, jax_profile="GPU")
< assert target["jax_profile"] == "gpu"
< assert target["build_name"] == "GPU environment"
<
<
< @pytest.mark.parametrize("failure", ["exit", "timeout", "json", "schema"])
< def test_dash_runtime_probe_failure_is_visible(tmp_path, monkeypatch, failure):
< wrapper = tmp_path / "clubb_jax" / "run_jax.py"
< wrapper.parent.mkdir()
< wrapper.touch(mode=0o755)
<
< def fake_run(command, **kwargs):
< if failure == "timeout":
< raise subprocess.TimeoutExpired(command, kwargs["timeout"])
< if failure == "exit":
< return subprocess.CompletedProcess(command, 1, "", "probe failed")
< return subprocess.CompletedProcess(command, 0, "not json" if failure == "json" else "{}", "")
<
< monkeypatch.setattr("dash_app.compile_tab.build_selector.subprocess.run", fake_run)
< for profile, info in inspect_jax_runtime_profiles(tmp_path).items():
< assert info["status"] == "unknown"
< assert info["selectable"] is False
< assert info["reason"]
< card = component_text(callbacks.render_jax_profile_info(profile, info, compact=True))
< assert "Checking" not in card
< assert "Unknown" in card
<
<
< def test_metal_preallocation_is_disabled():
< info = {"gpu": {"status": "setup_required", "selectable": True,
< "runtime": {"accelerator": "metal"}}}
< assert not callbacks.jax_preallocation_available(info, "gpu", "")
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/dash_app/README.md clubb_release/dash_app/README.md
8c8,48
< ## Quick start
---
> The Run, Profile, and Tune action areas each show the effective default CLUBB
> build next to their launch buttons. This read-only badge follows
> `install/selected` (or the same `install/latest` fallback used by
> `run_scm.py`) and updates when the Compile tab selects another build. Hover it
> for the resolved install and CMake paths, Fortran compiler, build type,
> precision, accelerator, OpenMP, and GPTL details. Explicit runner `-exe` or
> `-install_dir` options still override the displayed default.
>
> The Profile tab is a browser interface to `utilities/time_clubb.py`. Its top
> benchmark panel configures the case, process/per-process-batch-size sweep,
> repetitions, executable, configuration, overrides, and additional
> `run_scm.py` arguments. A direct one-second polling path reads the active
> summary and process rows and renders figures server-side, so results appear
> after each measured repetition while the broker-owned job is running; warmups
> remain hidden. Browser stores retain only compact timer/process choices rather
> than the growing raw timing table. The running row counter and all four figures
> are returned by the same callback response, so visible progress cannot advance
> independently of the plots. Stored profiles can be overlaid, compared with a baseline, or
> viewed as process distributions and exclusive-cost decompositions. The right
> rail has a profile-selection section above a separate set of shared comparison
> controls; plot-specific options remain beside the plot they affect. The
> profile chooser shows the three newest unselected results by default, expands
> to the full library, and displays active comparisons as removable pills. The
> benchmark-label field indicates when its normalized profile name already
> exists. Starting that benchmark asks for confirmation, then replaces the
> existing profile in place instead of creating a timestamped version; any older
> same-label/same-case versions are removed from the active comparison selection.
>
> The selected directory is a collection of compact, directly commit-able
> profile folders. Each benchmark creates `<profile-name>/` containing
> `README.md`, profile-wide `profile.json` provenance, one workload row per
> process-count/batch-size point in `batches.csv`, raw timer observations in
> `timings.csv`, and one representative input/setup/log/native-timing set under
> `logs/<batch-id>/`. Child processes otherwise run in temporary directories,
> which are deleted after each workload is aggregated. Warmups are retained with
> `phase=warmup` but excluded from the default plots. Dash derives statistical
> summaries in memory instead of storing duplicate summary files. **Export
> selected** downloads complete profiles as a ZIP; **Import** accepts those ZIPs
> on another machine or checkout. Provenance includes the effective vertical
> level count, observed model steps, source revision, executable checksum, host,
> timer backend, and time basis so Dash can flag potentially incomparable runs.
10c50,53
< From the repository root:
---
> ## Install
>
> The top-level launcher can create the local virtualenv, install dependencies,
> and start the foreground dashboard manager:
16,48c59,60
< The launcher installs the dashboard dependencies and opens the app in your
< browser, normally at port `23404`. Keep the launching terminal open while using
< the app. Running the command again reopens this checkout's existing dashboard.
<
< To run a case:
<
< 1. Open **Run** and click the runtime badge beside the launch button to choose
< Fortran, Python, or JAX.
< 2. Prepare the selected implementation: Fortran needs a build from the
< **Compile** tab or `./compile.py`; Python needs a
< [Python build](#python-runs-and-tuning). JAX sets up its own environment on
< first use and needs no Fortran build.
< 3. Choose a case and output directory, launch the run, and watch its log.
< 4. Open **Plots** to view the output. You can also plot existing NetCDF files
< without compiling or running anything.
<
< ### Choosing CPU or GPU for JAX
<
< In the runtime chooser, select JAX and then the CPU or a compatible GPU tile.
< NVIDIA tiles show the physical GPU index, model, and memory. The chooser
< remembers your preference in this browser. The first run may install the
< selected runtime; a missing or incompatible GPU produces an error.
<
< A selected GPU is used for all cases in that job; selecting it does not spread
< the cases across multiple GPUs. Leave **Preallocate GPU memory** off for a
< shared NVIDIA GPU. Enable it when you want up-front memory reservation; it is
< unavailable for CPU and Metal. The chooser's **?** button explains the options.
< See the [JAX guide](../clubb_jax/README.md#gpu-running) for
< hardware requirements and command-line equivalents.
<
< Fortran and Python use the selected compiled build, which you can change or
< rebuild through the chooser. Tune uses its own Fortran/Python worker regardless
< of the JAX selection.
---
> The manager starts the runtime broker and Dash as child processes. Arguments
> are passed through to `dash_app/app.py`, for example:
50c62,64
< ## Basic Workflows
---
> ```bash
> ./launch_dashboard.sh --port 23404 -debug
> ```
52,66c66,80
< - **Run tab:** choose benchmark cases and settings, launch CLUBB, and watch the
< run output in the browser.
< - **Profile tab:** measure runtime across process counts and batch sizes, then
< compare saved profiles. See [profiling](#profiling) for details.
< - **Plots tab:** load one or more CLUBB output directories and make profile,
< time-height, time-series, budget, and subcolumn plots from the NetCDF files.
< - **Tune tab:** configure and monitor tuning runs. This requires a
< [Python build](#python-runs-and-tuning).
< - **Tutorial tab:** explore CLUBB concepts through interactive lessons,
< including a guide to the model equations and the ADG1 two-Gaussian explorer.
< - **Reports tab:** browse saved investigation reports from `doc/reports/`,
< including their figures, data, and provenance.
< - **Misc tab:** open focused diagnostics such as the SAM w–rₜ neighborhood
< viewer and Mixing Length Trajectories explorer. Setup and implementation
< notes are in [DEVELOPMENT.md](./DEVELOPMENT.md#misc-subtabs).
---
> If Dash crashes or stops reporting its broker heartbeat, the manager retries it
> every 10 seconds for up to 5 minutes. A successful restart is selected without
> opening another browser tab. If Dash does not recover within that window, the
> manager reports the last failure, gracefully stops broker-owned work, stops the
> broker, and exits nonzero. `SIGINT`, `SIGTERM`, and terminal hangup use the same
> ordered shutdown.
>
> The broker also watches a private manager heartbeat. If the manager is killed
> without a chance to clean up, a replacement launcher can adopt the broker for
> 30 seconds. After that grace period, the broker stops the orphaned Dash process
> group and active Compile/Run/Tune work, then exits.
>
> Dash serializes ordinary callbacks by default to protect NetCDF/HDF5 access.
> Explicitly expensive callbacks use isolated background worker processes. Use
> `--threaded` only for short diagnostics on a stack known to be thread-safe.
68c82
< ## Advanced usage
---
> For manual setup, run this from the repository root:
70c84,86
< ### Python runs and tuning
---
> ```bash
> python3 -m pip install -r dash_app/requirements.txt
> ```
72,73c88
< Python runs and Tune jobs require CLUBB's Python/F2PY interface. After the
< launcher has prepared the Dash environment, build it with that same Python:
---
> Run the Dash test suite with the same environment:
76c91
< .venv-dash/bin/python compile.py -python
---
> tests/run_pytests.sh -dash
79,82c94,98
< Use the corresponding `bin/python` path if `CLUBB_DASH_VENV` names a different
< virtual environment. Using the same environment avoids NumPy compatibility
< problems when loading the compiled interface. The interface is not needed just
< to open the dashboard or configure Tune controls.
---
> Compile CLUBB before using the run tab:
>
> ```bash
> ./compile.py
> ```
84c100
< ### Launch options and manual setup
---
> Plotting existing NetCDF output does not require a fresh compile.
86c102,105
< Pass application options through the launcher, for example:
---
> Dash builds the Tune controls from a checked-in, Fortran-validated bound table;
> it does not need the Python/F2PY interface merely to start. Actual Tune jobs
> use CLUBB's in-memory F2PY loss driver, so compile with `-python` before
> running a tuning workflow:
89c108
< ./launch_dashboard.sh --port 23404 -debug
---
> ./compile.py -python
92c111,124
< For manual setup in your chosen Python environment:
---
> When Dash is launched with `./launch_dashboard.sh`, build the extension with
> the same virtualenv Python that runs Dash. This avoids loading an F2PY module
> compiled against NumPy 1.x into a Tune worker's NumPy 2 environment:
>
> ```bash
> .venv-dash/bin/python compile.py -python
> ```
>
> Use the corresponding `bin/python` path if `CLUBB_DASH_VENV` names a different
> virtualenv.
>
> ## Run
>
> From the repository root:
95d126
< python3 -m pip install -r dash_app/requirements.txt
97a129,132
> or
> ```bash
> python3 dash_app/app.py &
> ```
99,119c134,137
< Use `python3 dash_app/app.py --help` for host, port, debug, and threading options.
< Dash serializes ordinary callbacks to protect NetCDF/HDF5 access; use
< `--threaded` only for diagnostics on a stack known to be thread-safe.
<
< The launcher supervises Dash and attempts recovery after a crash. Lifecycle
< and broker details are in the
< [development notes](./DEVELOPMENT.md#local-mcp-endpoint-lifecycle).
<
< ### Profiling
<
< Use **Profile** to choose a case, process counts, batch sizes, and repetitions.
< Results appear after each measured repetition; warmups are excluded from the
< default plots. Saved profiles can be overlaid, compared with a baseline, or
< viewed as process distributions and exclusive-cost decompositions.
<
< Reusing a profile name asks for confirmation before replacing the existing
< profile. **Export selected** downloads complete profiles as a ZIP; **Import**
< loads those ZIPs on another machine or checkout. Profiles include run and build
< metadata so Dash can flag potentially incomparable results. Storage and update
< details are documented in the
< [development notes](./DEVELOPMENT.md#profile-results).
---
> By default the app opens in a browser at port `23404`, or the next available
> port. Starting it again while this checkout's dashboard is already running
> reopens the registered dashboard instead of starting a second process. Use
> `python3 dash_app/app.py --help` for host, port, debug, and threading options.
121c139
< ### JULY_2017 statistics vs. 3-D recreation viewer
---
> ## JULY_2017 statistics vs. 3-D recreation viewer
134c152
< ### Local agent integration
---
> ## Local agent integration
136c154
< #### Runtime boundary
---
> ### Runtime boundary
164c182
< #### Add the running dashboard to a Codex chat
---
> ### Add the running dashboard to a Codex chat
350c368,414
< ### LES Benchmark Overlays
---
> ### ADG1 two-Gaussian explorer
>
> The Tutorial tab includes the active ADG1 two-Gaussian explorer. It visualizes
> the normalized ADG1 diagnosis and a direct-control trivariate comparison using
> the same grid moments. It is a teaching visualization, not a replacement for
> the full Fortran PDF diagnosis.
>
> ## Basic Workflows
>
> - **Run tab:** choose benchmark cases and settings, launch CLUBB, and watch the
> run output in the browser.
> - **Plots tab:** load one or more CLUBB output directories and make profile,
> time-height, time-series, budget, and subcolumn plots from the NetCDF files.
> - **Tune tab:** configure and monitor tuner runs when the branch and local build
> support the tuner workflow. This requires a CLUBB build compiled with
> `./compile.py -python`.
> - **Tutorial tab:** follow short interactive explanations of CLUBB concepts.
> The welcome page suggests a path through the lessons; a vertical page rail
> opens each lesson. The CLUBB Equations page provides a clickable quick
> reference for the core prognostic budgets, PDF transport closures, and
> cloud/buoyancy diagnostics. Its colors preserve the official equation
> document's ownership convention, while a stable inspector explains each
> term's physical role, source, closure path, and implementation relevance.
> The next lesson opens the ADG1 two-Gaussian explorer, where shared moments
> and the supplied moments show how ADG1 component placement and covariance
> allocation change the PDF geometry.
> - **Reports tab:** browse immutable, static investigation bundles from
> `doc/reports/`. Each bundle carries its own HTML, figures, excerpts, data,
> and provenance. Dash polls the published JSON catalog, so an agent can add a
> completed report without editing dashboard source or restarting the app.
> - **Misc tab:** browse living investigations and focused diagnostics from a
> persistent left-side vertical directory, including the SAM w–rₜ time-height
> neighborhood and the Mixing Length Trajectories explorer. The neighborhood
> browses a pre-rendered 5×5 atlas generated by
> `python -m dash_app.misc_tab.sam_w_rt_neighborhood.atlas`; the trajectory
> explorer reconstructs upward and downward parcel-energy paths from a
> compatible CLUBB statistics file and compares them with its stored
> `Lscale_up`, `Lscale_down`, and `Lscale` profiles.
>
> The equation-guide content is curated in
> `dash_app/tutorial_tab/clubb_equations_demo/` from
> `doc/CLUBBeqns.tex` and the corresponding current routines in
> `src/CLUBB_core/`. It deliberately emphasizes the continuous closed equations;
> an on-page caveat distinguishes them from implicit discretization, host
> coupling, surface forcing, clipping, limiters, and other enabled adjustments.
>
> ## LES Benchmark Overlays
370c434
< ### Shared UI components
---
> ## Shared UI components
378,388d441
<
< ### Development and tests
<
< Run the Dash test suite with the dashboard environment:
<
< ```bash
< tests/run_pytests.sh -dash
< ```
<
< See [DEVELOPMENT.md](./DEVELOPMENT.md) for UI conventions, service boundaries,
< and how runtime selections are recorded in jobs.
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/dash_app/run_tab/callbacks_runs.py clubb_release/dash_app/run_tab/callbacks_runs.py
8c8
< from dash import ALL, Input, Output, State, callback_context, html, no_update
---
> from dash import ALL, Input, Output, State, callback_context, no_update
13d12
< output_directory_details,
19d17
< from utilities.output_paths import resolve_output_dir
29,32d26
< RUN_OVERWRITE_OPEN = "run-overwrite-modal"
< RUN_OVERWRITE_CLOSED = "run-overwrite-modal run-overwrite-modal-hidden"
<
<
143,219d136
< def run_output_rename_available(proposed_output, pending):
< """Return whether a pending Run can move to a different output target."""
< proposed = clean_cli_option(proposed_output)
< current = clean_cli_option((pending or {}).get("output_dir"))
< if not proposed or not current:
< return False
< try:
< return resolve_output_dir(proposed).resolve() != resolve_output_dir(current).resolve()
< except (OSError, TypeError, ValueError):
< return False
<
<
< def run_output_detail_components(details):
< """Render the requested output-folder facts for the overwrite dialog."""
< case_count = int(details.get("case_count") or 0)
< cases = f"{case_count} case" if case_count == 1 else f"{case_count} cases"
< return [
< html.Div(
< [
< html.Span(label, className="run-overwrite-fact-label"),
< html.Span(value, className="run-overwrite-fact-value"),
< ],
< className="run-overwrite-fact",
< )
< for label, value in (
< ("Folder", str(details.get("path") or "")),
< ("Created", str(details.get("created") or "Unknown")),
< ("Last edited", str(details.get("last_edited") or "Unknown")),
< ("Cases", f"{cases} with *_stats.nc output"),
< )
< ]
<
<
< def prepared_run_with_output(pending, output_dir):
< """Copy a frozen Run submission while changing only its output target."""
< updated = dict(pending or {})
< updated["output_dir"] = clean_cli_option(output_dir) or "output"
< updated["cli_options"] = dict(updated.get("cli_options") or {})
< if updated["output_dir"] == "output":
< updated["cli_options"].pop("out_dir", None)
< else:
< updated["cli_options"]["out_dir"] = updated["output_dir"]
< return updated
<
<
< def submit_prepared_run(pending, perform_action):
< """Submit one previously validated and frozen Run request."""
< request_material = json.dumps(pending, sort_keys=True, default=str)
< result = perform_action(
< "domain_submit_scm_batch",
< {
< "request": {
< "request_id": fresh_batch_request_id(request_material),
< "cases": list(pending.get("cases") or []),
< "implementation": pending.get("implementation") or "fortran",
< "jax_profile": pending.get("jax_profile") or "cpu",
< "jax_gpu": pending.get("jax_gpu") or "",
< "jax_xla_prealloc": pending.get("jax_xla_prealloc"),
< "stats_file": pending.get("stats") or DEFAULT_STATS_NAME,
< "config": pending.get("config") or "default",
< "overrides": dict(pending.get("typed_overrides") or {}),
< "run_options": dict(pending.get("typed_options") or {}),
< "max_workers": int(pending.get("max_workers") or 1),
< },
< "native_overrides": dict(pending.get("overrides") or {}),
< "native_cli_options": dict(pending.get("cli_options") or {}),
< "submission_origin": "dash",
< },
< internal=True,
< )
< return {
< "action": "run",
< "at": time.time(),
< "job_id": result.get("job_id"),
< }
<
<
225,230d141
< Output("run-pending-request", "data"),
< Output("run-overwrite-modal", "className"),
< Output("run-overwrite-name", "value"),
< Output("run-overwrite-message", "children"),
< Output("run-overwrite-details", "children"),
< Output("run-opt-out-dir", "value"),
234,236d144
< Input("run-overwrite-button", "n_clicks"),
< Input("run-rename-button", "n_clicks"),
< Input("run-overwrite-cancel-button", "n_clicks"),
262,266d169
< State("compile-run-jax-profile", "data"),
< State("compile-run-jax-gpu", "data"),
< State("compile-run-jax-xla-prealloc", "data"),
< State("run-pending-request", "data"),
< State("run-overwrite-name", "value"),
273,275d175
< _overwrite_clicks,
< _rename_clicks,
< _overwrite_cancel_clicks,
301,305d200
< run_jax_profile,
< run_jax_gpu,
< run_jax_xla_prealloc,
< pending_request,
< proposed_output,
312,315c207
< return (
< {"action": "clear", "at": time.time()},
< {}, RUN_OVERWRITE_CLOSED, "", "", [], no_update,
< )
---
> return {"action": "clear", "at": time.time()}
319,377c211
< return (
< {"action": "cancel", "at": time.time(), "result": result},
< {}, RUN_OVERWRITE_CLOSED, "", "", [], no_update,
< )
<
< pending = dict(pending_request or {})
< if trigger == "run-overwrite-button":
< if not pending:
< return (no_update,) * 7
< try:
< action = submit_prepared_run(pending, perform_action)
< return (
< action, {}, RUN_OVERWRITE_CLOSED, "", "", [],
< pending["output_dir"],
< )
< except (OSError, RuntimeError, ValueError) as exc:
< return (
< {"action": "error", "at": time.time(), "message": str(exc)},
< pending,
< RUN_OVERWRITE_OPEN,
< proposed_output,
< str(exc),
< no_update,
< no_update,
< )
<
< if trigger == "run-rename-button":
< if not pending or not run_output_rename_available(proposed_output, pending):
< return (no_update,) * 7
< try:
< renamed = prepared_run_with_output(pending, proposed_output)
< details = output_directory_details(renamed["output_dir"])
< if details["nonempty"]:
< return (
< no_update,
< renamed,
< RUN_OVERWRITE_OPEN,
< proposed_output,
< "That folder also contains files. Choose another folder or overwrite it.",
< run_output_detail_components(details),
< no_update,
< )
< action = submit_prepared_run(renamed, perform_action)
< return (
< action, {}, RUN_OVERWRITE_CLOSED, "", "", [], proposed_output,
< )
< except (OSError, RuntimeError, TypeError, ValueError) as exc:
< return (
< {"action": "error", "at": time.time(), "message": str(exc)},
< pending,
< RUN_OVERWRITE_OPEN,
< proposed_output,
< str(exc),
< no_update,
< no_update,
< )
<
< if trigger == "run-overwrite-cancel-button":
< return no_update, {}, RUN_OVERWRITE_CLOSED, "", "", [], no_update
---
> return {"action": "cancel", "at": time.time(), "result": result}
380c214
< return (no_update,) * 7
---
> return no_update
384c218
< return (no_update,) * 7
---
> return no_update
409,417c243,248
< return (
< {
< "action": "error",
< "at": time.time(),
< "cases": cases_to_run,
< "message": message,
< },
< {}, RUN_OVERWRITE_CLOSED, "", "", [], no_update,
< )
---
> return {
> "action": "error",
> "at": time.time(),
> "cases": cases_to_run,
> "message": message,
> }
426,429c257
< launch_target = selected_launch_target(
< run_implementation, jax_profile=run_jax_profile, jax_gpu=run_jax_gpu,
< jax_xla_prealloc=run_jax_xla_prealloc,
< )
---
> launch_target = selected_launch_target(run_implementation)
434,437d261
< if launch_target["implementation"] == "jax":
< cli_options["jax_profile"] = launch_target["jax_profile"]
< cli_options["jax_gpu"] = launch_target.get("jax_gpu") or ""
< cli_options["jax_xla_prealloc"] = launch_target.get("jax_xla_prealloc")
450,458c274,279
< return (
< {
< "action": "error",
< "at": time.time(),
< "cases": cases_to_run,
< "message": str(exc),
< },
< {}, RUN_OVERWRITE_CLOSED, "", "", [], no_update,
< )
---
> return {
> "action": "error",
> "at": time.time(),
> "cases": cases_to_run,
> "message": str(exc),
> }
482a304,314
> request_material = json.dumps(
> {
> "cases": cases_to_run,
> "stats": stats_name,
> "config": config_name,
> "overrides": overrides,
> "cli_options": cli_options,
> },
> sort_keys=True,
> default=str,
> )
489,503d320
< prepared = {
< "cases": cases_to_run,
< "stats": stats_name,
< "config": config_name,
< "overrides": overrides,
< "typed_overrides": typed_overrides,
< "cli_options": cli_options,
< "typed_options": typed_options,
< "max_workers": max_tasks,
< "output_dir": output_dir or "output",
< "implementation": launch_target["implementation"],
< "jax_profile": launch_target.get("jax_profile") or "cpu",
< "jax_gpu": launch_target.get("jax_gpu") or "",
< "jax_xla_prealloc": launch_target.get("jax_xla_prealloc"),
< }
505,519c322,323
< details = output_directory_details(prepared["output_dir"])
< if details["nonempty"]:
< return (
< no_update,
< prepared,
< RUN_OVERWRITE_OPEN,
< output_dir or "output",
< "Running here may replace matching case output files.",
< run_output_detail_components(details),
< no_update,
< )
< action = submit_prepared_run(prepared, perform_action)
< return action, {}, RUN_OVERWRITE_CLOSED, "", "", [], no_update
< except (OSError, RuntimeError, TypeError, ValueError) as exc:
< return (
---
> result = perform_action(
> "domain_submit_scm_batch",
521,524c325,336
< "action": "error",
< "at": time.time(),
< "cases": cases_to_run,
< "message": str(exc),
---
> "request": {
> "request_id": fresh_batch_request_id(request_material),
> "cases": cases_to_run,
> "stats_file": stats_name,
> "config": config_name,
> "overrides": typed_overrides,
> "run_options": typed_options,
> "max_workers": max_tasks,
> },
> "native_overrides": overrides,
> "native_cli_options": cli_options,
> "submission_origin": "dash",
526c338
< {}, RUN_OVERWRITE_CLOSED, "", "", [], no_update,
---
> internal=True,
528,542c340,351
<
< @app.callback(
< Output("run-rename-button", "disabled"),
< Output("run-rename-button", "title"),
< Input("run-overwrite-name", "value"),
< State("run-pending-request", "data"),
< )
< def update_run_rename_action(proposed_output, pending):
< available = run_output_rename_available(proposed_output, pending)
< return (
< not available,
< "Rename the output folder and start the run."
< if available
< else "Enter a different output folder to rename and run.",
< )
---
> except (OSError, RuntimeError, ValueError) as exc:
> return {
> "action": "error",
> "at": time.time(),
> "cases": cases_to_run,
> "message": str(exc),
> }
> return {
> "action": "run",
> "at": time.time(),
> "job_id": result.get("job_id"),
> }
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/dash_app/run_tab/layout.py clubb_release/dash_app/run_tab/layout.py
472,545d471
< def build_output_overwrite_dialog():
< """Render the Run output-folder collision decision dialog."""
< return html.Div(
< html.Div(
< [
< html.Div(
< [
< html.Div("!", className="run-overwrite-icon", **{"aria-hidden": "true"}),
< html.Div(
< [
< html.Div(
< "Output folder already contains files",
< className="run-overwrite-title",
< ),
< html.Div(
< id="run-overwrite-message",
< className="run-overwrite-message",
< ),
< ]
< ),
< ],
< className="run-overwrite-heading",
< ),
< html.Div(id="run-overwrite-details", className="run-overwrite-details"),
< html.Label(
< "Output folder",
< htmlFor="run-overwrite-name",
< className="run-overwrite-label",
< ),
< dcc.Input(
< id="run-overwrite-name",
< type="text",
< value="",
< debounce=False,
< className="run-overwrite-input",
< ),
< html.Div(
< [
< html.Button(
< "Overwrite",
< id="run-overwrite-button",
< type="button",
< n_clicks=0,
< className="run-overwrite-button run-overwrite-button-danger",
< ),
< html.Button(
< "Rename",
< id="run-rename-button",
< type="button",
< n_clicks=0,
< disabled=True,
< className="run-overwrite-button run-overwrite-button-primary",
< title="Enter a different output folder to rename and run.",
< ),
< html.Button(
< "Cancel",
< id="run-overwrite-cancel-button",
< type="button",
< n_clicks=0,
< className="run-overwrite-button run-overwrite-button-cancel",
< ),
< ],
< className="run-overwrite-actions",
< ),
< ],
< className="run-overwrite-panel",
< role="dialog",
< **{"aria-modal": "true", "aria-labelledby": "run-overwrite-message"},
< ),
< id="run-overwrite-modal",
< className="run-overwrite-modal run-overwrite-modal-hidden",
< )
<
<
927d852
< dcc.Store(id="run-pending-request", data={}),
929d853
< build_output_overwrite_dialog(),
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/dash_app/run_tab/runtime.py clubb_release/dash_app/run_tab/runtime.py
10,11d9
< from datetime import datetime
< from pathlib import Path
30,70d27
< from dash_app.shared.jax_device import jax_device_env, jax_command_display
< from utilities.output_paths import resolve_output_dir
<
<
< def _format_output_timestamp(timestamp):
< """Format one filesystem timestamp in the dashboard server's timezone."""
< return datetime.fromtimestamp(timestamp).astimezone().strftime("%Y-%m-%d %H:%M:%S %Z")
<
<
< def output_directory_details(value):
< """Describe an existing Run output target without creating or changing it."""
< path = resolve_output_dir(value).resolve()
< if path.exists() and not path.is_dir():
< raise ValueError("output path is an existing file")
< if not path.exists():
< return {
< "path": str(path),
< "exists": False,
< "nonempty": False,
< "created": "Not created yet",
< "last_edited": "Not created yet",
< "case_count": 0,
< }
<
< entries = list(path.iterdir())
< path_stat = path.stat()
< created_at = getattr(path_stat, "st_birthtime", path_stat.st_ctime)
< edited_at = max(
< [path_stat.st_mtime]
< + [entry.stat().st_mtime for entry in entries if entry.exists()]
< )
< return {
< "path": str(path),
< "exists": True,
< "nonempty": bool(entries),
< "created": _format_output_timestamp(created_at),
< "last_edited": _format_output_timestamp(edited_at),
< "case_count": sum(
< 1 for entry in entries if entry.is_file() and entry.name.endswith("_stats.nc")
< ),
< }
200,202c157
< jax_profile = clean_cli_option((cli_options or {}).get("jax_profile")).lower()
< modifier = ",xla_prealloc" if (cli_options or {}).get("jax_xla_prealloc") is True else ""
< cmd.append(f"-jax={jax_profile}{modifier}" if jax_profile else "-jax")
---
> cmd.append("-jax")
204c159
< if install_dir and implementation != "jax":
---
> if install_dir:
225c180
< return jax_command_display([str(part) for part in cmd], cli_options)
---
> return " ".join(shlex.quote(str(part)) for part in cmd)
274c229
< env=jax_device_env(cli_options, run_child_env()),
---
> env=run_child_env(),
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/dash_app/run_tab/telemetry.py clubb_release/dash_app/run_tab/telemetry.py
46,54c46
< # Queued API jobs have their launch target in the typed request; older
< # display records may contain only run_options. Once launched, prefer the
< # actual runtime options over either submission snapshot.
< cli_options = dict(request.get("run_options") or {})
< for key in ("implementation", "jax_profile", "jax_gpu", "jax_xla_prealloc"):
< if key in request:
< cli_options[key] = request[key]
< cli_options.update(display.get("cli_options") or {})
< cli_options.update(runtime.get("cli_options") or {})
---
> cli_options = dict(display.get("cli_options") or runtime.get("cli_options") or {})
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/dash_app/services/models.py clubb_release/dash_app/services/models.py
8d7
< from dash_app.shared.jax_device import GPU_UUID_PATTERN
42,45d40
< implementation: Literal["fortran", "python", "jax"] = "fortran"
< jax_profile: Literal["cpu", "gpu"] = "cpu"
< jax_gpu: str = Field(default="", pattern=GPU_UUID_PATTERN, description="Full GPU UUID; empty inherits the server environment.")
< jax_xla_prealloc: bool | None = Field(default=None, strict=True, description="GPU preallocation override; null preserves the launcher default/environment.")
67,70d61
< implementation: Literal["fortran", "python", "jax"] = "fortran"
< jax_profile: Literal["cpu", "gpu"] = "cpu"
< jax_gpu: str = Field(default="", pattern=GPU_UUID_PATTERN, description="Full GPU UUID; empty inherits the server environment.")
< jax_xla_prealloc: bool | None = Field(default=None, strict=True, description="GPU preallocation override; null preserves the launcher default/environment.")
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/dash_app/shared/actions.py clubb_release/dash_app/shared/actions.py
11d10
< import platform
151d149
< from dash_app.shared.jax_device import jax_device_env
1222d1219
< jax_device_env(request.model_dump(), {})
1263c1260
< "tout", "out_dir", "extra_args", "implementation", "install_dir", "jax_profile", "jax_gpu", "jax_xla_prealloc",
---
> "tout", "out_dir", "extra_args", "implementation", "install_dir",
1269,1271c1266
< if raw.get("jax_xla_prealloc") is not None:
< normalized["jax_xla_prealloc"] = raw["jax_xla_prealloc"]
< for key in allowed - {"extra_args", "jax_xla_prealloc"}:
---
> for key in allowed - {"extra_args"}:
1283,1292d1277
< if normalized.get("jax_profile"):
< jax_profile = normalized["jax_profile"].lower()
< if implementation != "jax":
< raise ValueError("jax_profile requires the JAX implementation")
< if jax_profile not in {"cpu", "gpu"}:
< raise ValueError("JAX profile must be cpu or gpu")
< normalized["jax_profile"] = jax_profile
< jax_device_env(normalized, {})
< if implementation == "jax":
< normalized.pop("install_dir", None)
1356c1341
< """Capture the exact implementation selected by ``run_scm.py``.
---
> """Capture the exact compiled executable selected by ``run_scm.py``.
1364,1425d1348
< if implementation == "jax":
< jax_root = REPO_ROOT / "clubb_jax"
< requested_profile = str(options.get("jax_profile") or "").strip().lower()
< if requested_profile:
< profile = requested_profile
< accelerator = (
< "metal" if profile == "gpu" and platform.system() == "Darwin"
< else "cuda13" if profile == "gpu"
< else "cpu"
< )
< else:
< accelerator = os.environ.get("CLUBB_JAX_ACCELERATOR", "cpu").strip().lower()
< profile = "gpu" if accelerator in {"cuda13", "metal"} else "cpu"
< default_precision = "single" if accelerator == "metal" else "double"
< precision_value = os.environ.get(
< "CLUBB_JAX_PRECISION", default_precision
< ).strip().lower()
< precision = (
< "single"
< if precision_value in {"single", "float32", "f32", "32", "real4", "sp"}
< else "double"
< )
< wrapper = jax_root / "run_jax.py"
< driver = jax_root / "src" / "clubb_standalone.py"
< requirements = jax_root / (
< "requirements-cuda13.txt"
< if accelerator == "cuda13"
< else "requirements-metal.txt"
< if accelerator == "metal"
< else "requirements.txt"
< )
<
< def source_identity(path: Path) -> dict[str, Any]:
< return {
< "path": str(path),
< "sha256": sha256_file(path),
< "bytes": path.stat().st_size if path.is_file() else None,
< }
<
< return {
< "install_selector": None,
< "install_directory": None,
< "implementation": "jax",
< "runtime": "repository-managed",
< "profile": profile,
< "executable": source_identity(wrapper),
< "launcher": source_identity(wrapper),
< "driver": source_identity(driver),
< "requirements": source_identity(requirements),
< "environment": {
< "CUDA_VISIBLE_DEVICES": options.get("jax_gpu") or os.environ.get("CUDA_VISIBLE_DEVICES", ""),
< "CLUBB_JAX_ACCELERATOR": accelerator,
< "CLUBB_JAX_PRECISION": precision_value or default_precision,
< "CLUBB_JAX_VENV": os.environ.get("CLUBB_JAX_VENV", ""),
< "CLUBB_JAX_TOOLS_DIR": os.environ.get("CLUBB_JAX_TOOLS_DIR", ""),
< "XLA_PYTHON_CLIENT_PREALLOCATE": jax_device_env(options).get(
< "XLA_PYTHON_CLIENT_PREALLOCATE", "false"),
< },
< "precision": precision,
< "accelerator": accelerator,
< }
<
1566,1573c1489,1493
< cli_options = request.run_options.model_dump(exclude_none=True)
< if native_cli_options is not None:
< cli_options.update(native_cli_options)
< cli_options.setdefault("implementation", request.implementation)
< if request.implementation == "jax":
< cli_options.setdefault("jax_profile", request.jax_profile)
< cli_options.setdefault("jax_gpu", request.jax_gpu)
< cli_options.setdefault("jax_xla_prealloc", request.jax_xla_prealloc)
---
> cli_options = dict(
> native_cli_options
> if native_cli_options is not None
> else request.run_options.model_dump(exclude_none=True)
> )
1634,1637d1553
< implementation=request.implementation,
< jax_profile=request.jax_profile,
< jax_gpu=request.jax_gpu,
< jax_xla_prealloc=request.jax_xla_prealloc,
1660,1663d1575
< implementation=request.implementation,
< jax_profile=request.jax_profile,
< jax_gpu=request.jax_gpu,
< jax_xla_prealloc=request.jax_xla_prealloc,
1811,1819c1723
< build_identity = _scm_build_identity(
< normalized_native_cli_options
< or {
< "implementation": request.implementation,
< "jax_profile": request.jax_profile,
< "jax_gpu": request.jax_gpu,
< "jax_xla_prealloc": request.jax_xla_prealloc,
< }
< )
---
> build_identity = _scm_build_identity(normalized_native_cli_options)
1954,1957d1857
< implementation=request.implementation,
< jax_profile=request.jax_profile,
< jax_gpu=request.jax_gpu,
< jax_xla_prealloc=request.jax_xla_prealloc,
Only in clubb/dash_app/shared: jax_device.py
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/run_scripts/run_scm.py clubb_release/run_scripts/run_scm.py
26,42d25
< def extract_jax_options(argv):
< # argparse's nargs="?" would consume CASE in "-jax CASE". Extract only
< # attached values (-jax=VALUE); the JAX launcher interprets their contents.
< normalized = []
< value = None
< occurrences = 0
< for token in argv:
< option, separator, attached = token.partition("=")
< if option == "-jax":
< occurrences += 1
< value = attached if separator else None
< normalized.append("-jax")
< else:
< normalized.append(token)
< return normalized, value, occurrences
<
<
166c149,152
< jax_launcher = os.path.join(CLUBB_ROOT, "clubb_jax", "run_jax.py")
---
> jax_driver = os.path.join(CLUBB_ROOT, "clubb_jax", "src", "clubb_standalone.py")
> if not os.path.isfile(jax_driver):
> sys.exit(f"JAX standalone driver not found: {jax_driver}")
> jax_launcher = os.path.join(CLUBB_ROOT, "clubb_jax", "run_jax_wrapper.sh")
171,174d156
< if args.jax_options is not None:
< # Keep profile/modifier parsing in the launcher so CLI and Dash
< # share its validation and environment setup rules.
< run_cmd.append(f"--options={args.jax_options}")
293,294c275
< help=("Run through the JAX wrapper. An optional attached -jax=VALUE is "
< "forwarded unchanged; see clubb_jax/run_jax.py --launcher-help."))
---
> help="Run the JAX standalone driver (python -m clubb_jax.src.clubb_standalone)")
341,345c322
< normalized_argv, jax_options, jax_occurrences = extract_jax_options(sys.argv[1:])
< args = parser.parse_args(normalized_argv)
< if jax_occurrences > 1:
< parser.error("-jax may be specified only once.")
< args.jax_options = jax_options
---
> args = parser.parse_args()
diff '--exclude=.git' '--exclude=version_clubb_core.txt' '--exclude=version_silhs.txt' -r clubb/tests/run_jax_vs_fortran_cases.py clubb_release/tests/run_jax_vs_fortran_cases.py
18a19,54
> def _reexec_with_repo_jax_python() -> None:
> """Initialize and use the repository-local JAX environment."""
> repo_root = Path(__file__).resolve().parents[1]
> launcher = repo_root / "clubb_jax" / "run_jax_wrapper.sh"
> if not launcher.is_file():
> return
>
> initialized_env_var = "_CLUBB_JAX_HARNESS_ENV_INITIALIZED"
> if os.environ.get(initialized_env_var) != "1":
> init_env = subprocess.run([str(launcher), "--init_env"], check=False)
> if init_env.returncode != 0:
> raise SystemExit(init_env.returncode)
>
> accelerator = os.environ.get("CLUBB_JAX_ACCELERATOR", "cpu").lower()
> default_venv = ".venv-jax-cuda13" if accelerator == "cuda13" else ".venv-jax"
> venv_dir = Path(os.environ.get("CLUBB_JAX_VENV", repo_root / default_venv))
> if not venv_dir.is_absolute():
> venv_dir = repo_root / venv_dir
> venv_python = venv_dir / "bin" / "python"
> if not venv_python.is_file():
> return
>
> if Path(sys.executable).absolute() == venv_python.absolute():
> return
>
> exec_env = os.environ.copy()
> exec_env[initialized_env_var] = "1"
> os.execve(
> str(venv_python),
> [str(venv_python), str(Path(__file__).resolve()), *sys.argv[1:]],
> exec_env,
> )
>
>
> _reexec_with_repo_jax_python()
>
23d58
< from clubb_jax.run_jax import ensure_environment # noqa: E402
129a165,197
> def _check_jax_runtime() -> str | None:
> """Validate the interpreter before launching every case with it."""
> accelerator = os.environ.get("CLUBB_JAX_ACCELERATOR", "cpu").lower()
> probe = subprocess.run(
> [
> sys.executable,
> "-c",
> (
> "import jax, jaxlib, netCDF4, sys, tabulate; "
> "backend = jax.default_backend(); "
> "expected = 'gpu' if sys.argv[1] == 'cuda13' else 'cpu'; "
> "assert backend == expected, "
> "f'requested {sys.argv[1]} but JAX initialized {backend}: {jax.devices()}'; "
> "print(f'jax={jax.__version__} jaxlib={jaxlib.__version__} "
> "backend={backend} devices={jax.devices()}')"
> ),
> accelerator,
> ],
> text=True,
> capture_output=True,
> )
> if probe.returncode == 0:
> return probe.stdout.strip()
>
> detail = (probe.stderr or probe.stdout).strip()
> print(f"ERROR: JAX runtime check failed with {sys.executable}:\n{detail}")
> print(
> "Create or repair the selected repository JAX environment, or invoke this script with a Python "
> "environment containing compatible jax and jaxlib packages."
> )
> return None
>
>
338d405
< ensure_environment()
340c407
< if accelerator in {"cuda13", "metal"} and args.jobs != 1:
---
> if accelerator == "cuda13" and args.jobs != 1:
347c414,417
< print(f"Python runtime: {sys.executable}")
---
> jax_runtime = _check_jax_runtime()
> if jax_runtime is None:
> return 2
> print(f"Python runtime: {sys.executable} ({jax_runtime})")
