Environment Variables
wrf_gpu is configured mainly through your WRF namelist.input and the
gpuwrf run command line. A small set of GPUWRF_* environment
variables tune paths, the compile cache, the runtime path, output, and GPU-memory
preflight. Most users only need the runtime-table path for their physics suite
(typically GPUWRF_WRF_ROOT) plus ordinary CLI flags such as
--input-dir and --output-dir.
Everywhere a variable overlaps a command-line argument, precedence is
explicit CLI argument > environment variable > built-in default. Booleans
accept 1/true/yes/on (and their
0/false/off/no opposites), case-insensitive.
Paths expand ~. See Running the Model for the CLI,
Performance & GPU Modes for the cache/precision levers, and
Output for the history-file knobs.
Paths & data¶
These point the model at its runtime tables and sample data without a hardcoded workstation layout. In an editable checkout they default to repository-relative paths, but those paths may be empty in a source-only clone; set them to your own WRF/corpus trees when the selected physics or validation path needs them.
| Variable | Default | Effect |
|---|---|---|
GPUWRF_WRF_ROOT | <repo>/data/wrf_pristine/WRF | Root of a pristine WRF v4 source/run tree. wrf_gpu parses a few unmodified WRF runtime tables from it at startup — Noah-MP run/MPTABLE.TBL/SOILPARM.TBL/GENPARM.TBL (for sf_surface_physics=4) and the classic RRTM longwave run/RRTM_DATA (for ra_lw_physics=1). See Installation. |
GPUWRF_CANAIRY_ROOT | <repo>/data/canairy_meteo | Root of the bundled sample / CPU-WRF corpus. Also the base for the default L2/L3 run roots. Set it to keep an existing corpus layout working. |
GPUWRF_AEMET_ROOT | <canairy_root>/artifacts/datasets/aemet_stations | Station-observation data. Only needed for the optional gpuwrf run --score verification against AEMET stations. |
GPUWRF_SCRATCH | <output-dir>/.scratch | A real (non-tmpfs) NVMe scratch directory with a few GiB free — never a RAM disk. Precedence: --scratch-dir > $GPUWRF_SCRATCH > <output-dir>/.scratch. Cleaned up on exit. |
GPUWRF_TMPDIR | ~/.cache/gpuwrf | Scratch root for replay traces / cached intermediate outputs, so a clean clone never writes into a hardcoded home path. The CLI also points this at the resolved run scratch directory for the current run. |
The history-file destination is set with gpuwrf run --output-dir runs/<case>,
not an environment variable. Proof JSON defaults to <output-dir>/proofs
(override with --proof-dir).
Compile cache & warm-start¶
The first run of a given geometry pays a one-time cold XLA compile; later runs reuse it. These variables place and control the caches. Full background is on Performance & GPU Modes.
| Variable | Default | Effect |
|---|---|---|
GPUWRF_JAX_CACHE_DIR(alias JAX_COMPILATION_CACHE_DIR) | <repo>/.gpuwrf-cache/jax | XLA JIT compilation-cache directory. The two names are equivalent (the standard JAX name is honoured too); setting either keeps the multi-minute cold compile off subsequent runs. A version mismatch under an explicit directory is warned about, not silently reused. |
GPUWRF_CACHE | unset | Parent cache root. When set, the JIT cache lives under <root>/jit/<version-tag> and the autotune cache under <root>/autotune — one root for both, version-keyed so different builds never share a directory. |
GPUWRF_NESTED_AOT | 1 (on) | Ahead-of-time cheap-key cross-process warm-start for the nested path: a fresh process computes a metadata-only key without lowering and loads a compatible serialized executable without retracing. Readiness depends on the program and cache. Identity-preserving and fail-open (any error falls back to a normal JIT compile). Set =0 to disable. |
GPUWRF_AOT_VERIFY | off | On the first cheap-key hit per key per process, do one confirmatory lower() and assert the HLO digest matches the loaded blob; fail-closed on mismatch. A diagnostic to run on for the first warm-start of a new version, then off in production (it costs one lower per distinct key). |
GPUWRF_NESTED_PARALLEL_COMPILE | unset (auto) | Worker count for the de-fuse cold-compile pre-warm pool. 0 opts out; N>0 overrides the RAM/CPU-derived worker count; unset uses the auto default. Only relevant on the de-fuse compile path below. |
GPUWRF_XLA_AUTOTUNE_CACHE | on (with the compile cache) | Caches XLA GPU autotuning results alongside the JIT cache. Rides with the compile cache by default; a truthy value forces it on, a falsey value (0/off) disables it. The unknown-flag abort guard is unaffected. |
GPUWRF_NESTED_DEFUSE_COMPILE | off | =1 compiles each domain's advance as its own executable instead of fusing the parent + leaves into one XLA module, cutting peak compile-RAM to roughly the single largest domain. Numerically inert (same eager path as the identity switches); it trades some runtime throughput for compile-RAM, so it is an opt-in low-host-RAM lever, not the runtime default. |
Set GPUWRF_JAX_CACHE=0 (or off) to disable the JIT compile cache
entirely for a run — useful for a clean cold-compile measurement.
Runtime path & precision¶
These select which numerical/execution path runs. The v0.3.0 defaults are the validated fused fp32 path, checked against original CPU-WRF at frozen per-variable limits (agreement with WRF is not bitwise); the opt-ins are for speed, capability, or debugging.
| Variable | Default | Effect |
|---|---|---|
GPUWRF_FAST_DEFAULTS | 1 | v0.3.0 release defaults: the validated fused fp32 fast paths are on. 0 turns all of them off (legacy path); a single fast-path flag can be opted out with an explicit =0. |
GPUWRF_ACOUSTIC_PRECISION_MODE | fp64_default | Acoustic-step precision of the legacy (pre-v0.3.0) path; the v0.3.0 default runs the fused fp32 kernels enabled by GPUWRF_FAST_DEFAULTS. On the legacy path fp64_default is the default; mixed_perturb_fp32_v020 opts into the validated mixed-precision fp32 perturbation path (fp32 mode). An unrecognized value fails closed with the list of accepted modes. |
GPUWRF_BATCH_ENSEMBLE | unset (=1) | Set to a positive integer B to run B independent, same-geometry forecasts concurrently under one outer jax.vmap — the v0.23 batched ensemble that fills a launch-bound small-grid GPU. The default (unset / 1) is byte-identical; a non-integer or <1 value is rejected. |
GPUWRF_BATCH_INPUT_DIRS | unset | A comma- (or :-) separated list of exactly B input directories, one per ensemble member — for a same-date, distinct-IC batch (e.g. a perturbation / initial-condition ensemble). Used only when GPUWRF_BATCH_ENSEMBLE>1; if unset, all members reuse the single --input-dir (a replicated batch). Every dir must be homogeneous — same grid + physics + dt + forecast date (only the initial/boundary data values differ); a count mismatch or a mismatched member fails closed with a named reason. Different forecast dates are not batchable in one batch (the date is part of the compiled program key); to run several cases or dates at once, use independent processes with scripts/run_parallel_cases.sh, which shares one GPU. See Batched ensemble → distinct-init. |
GPUWRF_NESTED_FUSE | on (fused) | The nested runtime default is the fused cascade (measured higher-throughput executable). =0 selects the eager per-domain path, which is bit-identical to the fused path — used for identity/debug comparisons. |
GPUWRF_BITWISE | off | =1 forces the eager, non-fused per-domain path for bitwise/debug comparison against a reference. Same eager path as GPUWRF_NESTED_FUSE=0. |
GPUWRF_FINITE_CHECK | on | Default-on fail-fast finite guard: on the first non-finite prognostic value the run aborts with the exact {domain, field, level, step, sim-time, index} instead of propagating a NaN. Set =0 to opt out. There is no masking/clamping — see Validation. |
GPUWRF_K2_EXPERIMENTAL | off | =1 (alias GPUWRF_K2_ENABLED=1) dispatches to the opt-in sharded, multi-GPU runner. Experimental lab-only; the default single-GPU forecast path never touches it. Not a supported production mode. |
Output¶
Knobs for what the wrfout history file contains and how it is named/written. All
of these keep the default output byte-identical unless you opt in. See
Output.
| Variable | Default | Effect |
|---|---|---|
GPUWRF_TRAINING_OUTPUT_SUBSET | off | =1 writes a focused 39-variable training subset plus mandatory coordinates instead of the full variable set — smaller files for ML training. Lossless subset; default output is byte-identical. |
GPUWRF_FULL_WRFOUT_VARIABLES (alias GPUWRF_FULL_WRFOUT) | ON | Full WRF history by default; canonical value wins. Set =0 for the reduced stream; inspect actual file counts. |
GPUWRF_COLONFREE_OUTPUT | off | =1 writes colon-free HH-MM-SS wrfout timestamps (for S3 / network drains that dislike :); the default stays WRF-standard HH:MM:SS. |
GPUWRF_NESTED_ASYNC_OUTPUT | on | Writes history files on a background writer thread so the GPU is not stalled by disk I/O. Set =0 to force the legacy synchronous write (for A/B measurement or fully-on-thread byte-identity proofs). Written bytes are identical either way. |
GPUWRF_NEST_OUTPUT_PIPELINE | off | =1 enables the structural async output pipeline: the finite check, M9 surface re-solve, and device→host pull run on a background thread that overlaps the next compute burst instead of inline at the output boundary. Byte-identical to the off path; opt-in pending its VRAM gate. (Companion GPUWRF_NEST_OUTPUT_PIPELINE_DEPTH, default 1, bounds in-flight frames.) |
GPUWRF_NESTED_M9_RADIATION_FROM_CARRY | off | =1 backs nested M9 surface-output diagnostics from the resident Noah-MP radiation carry (held SOLDN/LWDN/COSZ). Not byte-identical: it omits the instantaneous upwelling / TOA flux slices (SWUPB/LWUPB) rather than fabricating them, so it is an explicit opt-in shortcut. |
Memory / VRAM preflight¶
Nested runs (--max-dom > 1) run a CPU-side free-VRAM preflight that fails closed
before the long compile if the card is too full, and select the GPU allocator. See
Performance & GPU Modes.
| Variable | Default | Effect |
|---|---|---|
GPUWRF_MIN_FREE_VRAM_FRACTION | 0.50 | Card-relative free-VRAM floor for the nested preflight. After explicit GiB and C-auto plan precedence, the cold fallback is max(24 GiB, fraction × total VRAM), so it scales with the card. Clamped to [0, 1]. |
GPUWRF_MIN_FREE_VRAM_GIB | unset | Explicit absolute free-VRAM threshold in GiB. When set it overrides the fractional rule (the absolute floor otherwise defaults to 24 GiB). |
GPUWRF_FORCE_GPU_RUN | off | =1 bypasses the nested GPU lock / free-VRAM preflight for a deliberate operator override. Equivalent to the --force-gpu-run CLI flag. Prefer the threshold variables above for tuning; use this only to force a run. |
GPUWRF_ALLOCATOR | cuda_async (nested) | Selects the XLA GPU allocator for nested runs: cuda_async (stream-ordered pool, default), platform (no-fragment cudaMalloc fallback for fragmentation OOMs), or bfc (XLA's default arena). An explicit XLA_PYTHON_CLIENT_ALLOCATOR always wins and is left untouched. |
Nesting¶
| Variable | Default | Effect |
|---|---|---|
GPUWRF_GWD_NESTED | 1 (on) | Gravity-wave drag (gwd_opt=1) is default-on on the nested 1 km path. Set =0 to force it off. See Nesting and Physics Options. |
The source defines many more GPUWRF_* flags — per-scheme kernel-tiling and
unroll knobs, cache-locking and GPU-lock bookkeeping injected by helper scripts, and internal
probes. Those are advanced/internal and are not part of the supported
user surface; they may change or be removed without notice. The variables documented
above are the ones intended for users.