Aurora 1.5: interrupted rollout leaves noise accumulation active for later forecasts

Start with aurora/rollout.py, where rollout and rollout_ensemble manage generator cleanup, and review the existing rollout tests alongside the supplied tests/v1p5/test_rollout_cleanup.py reproduction. Run that regression file with the command in the issue, then check that closing, cancelling, and forward errors clear noise accumulation and preserve later seeded forward results.

At main commit 85dd3880683f0ac19aa8fe4386c5dcacfec99892, explicitly closing rollout(..., fine_lead_times=[3.0, 6.0]) after its first yield leaves the stochastic backbone's noise accumulation enabled. Throwing an exception into the generator or a native forward validation error does the same. rollout_ensemble is affected too.

The existing cleanup model.set_noise_accumulation(n=0) runs only after the loop is exhausted. This is distinct from the seeded-noise support in #202: even when the caller restores the same query generator seed, a subsequent standalone forward can see the leftover accumulation configuration/cache and produce a different result.

Native-model reproduction

I used your upstream _make_small_v1p5(stochastic=True) and batch helpers, with the same _unzero_adaptive_layer_norms helper used by your ensemble tests. The latter is necessary because the freshly initialized model otherwise hides conditioning/noise effects. No forward, backbone, RNG, or outputs are mocked. The inputs and reduced randomly initialized weights are constructed test fixtures, not a released pretrained checkpoint.

For both rollout APIs I tested:

1. Explicit .close() after the first yield.
2. .throw(RuntimeError("cancelled")) after the first yield.
3. A real forward ValueError from an invalid empty generator tuple.
4. Normal exhaustion as a control.

Original: 6 failed, 2 passed. The six interrupted/error cases retain _noise_cache_size=2 and _accumulate_noise=True; both exhaustion controls clear them.

The same query batch and a fresh generator seeded 123 were used before and after each sequence, with unchanged model weights. All six interrupted/error paths changed the later standalone forecast; exhaustion controls were exactly equal. These are fixture output differences, not measured weather forecast errors.

Proposed local correction

• Put the rollout computation/yields inside try/finally, retaining the existing normal-completion cleanup in the finally clause.
• Have rollout_ensemble explicitly close its inner generator using contextlib.closing.
• Document that a caller stopping early must explicitly close the generator, e.g. with contextlib.closing. A bare break while retaining a generator reference does not trigger cleanup.

Local candidate: all 8/8 regression cases pass, all caches disabled/empty afterwards, and all later standalone outputs exactly equal their seeded controls. Weights remain unchanged. Related existing tests plus the eight new cases: 30 passed, 1 deselected. The deselected test constructs the full-size plain Aurora solely to reject variable lead times, outside this reduced CPU run.

Scope: generator lifecycle and subsequent forward isolation in the real reduced AuroraV1p5 implementation, CPU float32, not validation of full-size pretrained weights, operational forecasts, GPU behavior, concurrent/interleaved generators sharing a model, or an unclosed retained generator. This is a candidate for maintainer review, not a claim of upstream acceptance.

AI-assisted report and local verification by Massimiliano Brighindi | PHI-OMEGA / M.

Reproducer / regression file

Save as tests/v1p5/test_rollout_cleanup.py and run OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 python -m pytest tests/v1p5/test_rollout_cleanup.py -q from the repository root. Executed on Python 3.12.14, torch 2.5.1+cpu, torchvision 0.20.1+cpu; full dependency versions retained with the audit.

"""Regression checks for state left behind by interrupted forecast generators."""

import pytest
import torch

from ._helpers import _OUTPUT_ONLY_SURF, _SURF_VARS, _make_batch, _make_small_v1p5
from .test_ensemble import _unzero_adaptive_layer_norms
from aurora import rollout, rollout_ensemble


@pytest.mark.parametrize("ending", ["close", "throw", "bad_generator", "exhaust"])
@pytest.mark.parametrize("ensemble", [False, True])
def test_rollout_cleanup_restores_independent_forward(ending, ensemble):
    torch.manual_seed(17)
    model = _make_small_v1p5(stochastic=True)
    # Match upstream ensemble tests: zero-initialised conditioning hides noise effects.
    _unzero_adaptive_layer_norms(model)
    model.eval()
    batch = _make_batch(surf_vars=tuple(v for v in _SURF_VARS if v not in _OUTPUT_ONLY_SURF))
    before = {k: v.clone() for k, v in model.state_dict().items()}

    with torch.inference_mode():
        reference = model.forward(
            batch, lead_times=torch.tensor([6.0]), generator=torch.Generator().manual_seed(123)
        )
        kwargs = dict(
            steps=1,
            fine_lead_times=[3.0, 6.0],
            generator=() if ending == "bad_generator" else torch.Generator().manual_seed(456),
        )
        predictions = (
            rollout_ensemble(model, batch, num_ensemble_members=2, **kwargs)
            if ensemble
            else rollout(model, batch, **kwargs)
        )
        if ending == "bad_generator":
            with pytest.raises(ValueError, match="generator"):
                next(predictions)
        elif ending == "exhaust":
            assert len(list(predictions)) == 2
        else:
            next(predictions)
            assert model.backbone._noise_cache_size == 2
            if ending == "close":
                predictions.close()
            else:
                with pytest.raises(RuntimeError, match="cancelled"):
                    predictions.throw(RuntimeError("cancelled"))

        assert model.backbone._noise_cache_size == 0
        assert not model.backbone._accumulate_noise
        assert not model.backbone._noise_cache
        actual = model.forward(
            batch, lead_times=torch.tensor([6.0]), generator=torch.Generator().manual_seed(123)
        )

    for values, expected in (
        (actual.surf_vars, reference.surf_vars),
        (actual.atmos_vars, reference.atmos_vars),
    ):
        for k in values:
            torch.testing.assert_close(values[k], expected[k], rtol=0, atol=0)
    for k, v in model.state_dict().items():
        assert torch.equal(v, before[k])

Candidate production patch

diff --git a/aurora/rollout.py b/aurora/rollout.py
index e9861e6..b9250ac 100644
--- a/aurora/rollout.py
+++ b/aurora/rollout.py
@@ -2,6 +2,7 @@
 
 import dataclasses
 import math
+from contextlib import closing
 from typing import Generator, Optional, Sequence
 
 import torch
@@ -98,6 +99,11 @@ def rollout(
 
     Yields:
         :class:`aurora.Batch`: The prediction after every (sub-)step.
+
+    Note:
+        If stopping early, explicitly close the generator (for example, with
+        :func:`contextlib.closing`) to disable noise accumulation and clear its cache.
+        A `break` alone does not close a generator that is still referenced.
     """
     # We will need to concatenate data, so ensure that everything is already of the right form.
     batch = model.batch_transform_hook(batch)  # This might modify the available variables.
@@ -123,36 +129,36 @@ def rollout(
     if use_noise_accumulation and fine_lead_times is not None:
         model.set_noise_accumulation(n=len(fine_lead_times))
 
-    # Pre-compute the base lead-time tensor for models with variable lead time support.
-    base_lead_times: Optional[torch.Tensor] = None
-    if model.variable_lead_time:
-        base_lead_times = _make_lead_time_tensor(batch, model.timestep.total_seconds() / 3600.0)
-
-    for _ in range(steps):
-        if fine_lead_times is not None:
-            # Inner loop: iterate over sub-step lead times.
-            for lt_hours in fine_lead_times:
-                sub_lead_times = _make_lead_time_tensor(batch, lt_hours)
-                pred = model.forward(batch, lead_times=sub_lead_times, generator=generator)
+    try:
+        # Pre-compute the base lead-time tensor for models with variable lead time support.
+        base_lead_times: Optional[torch.Tensor] = None
+        if model.variable_lead_time:
+            base_lead_times = _make_lead_time_tensor(batch, model.timestep.total_seconds() / 3600.0)
 
-                yield pred
+        for _ in range(steps):
+            if fine_lead_times is not None:
+                # Inner loop: iterate over sub-step lead times.
+                for lt_hours in fine_lead_times:
+                    sub_lead_times = _make_lead_time_tensor(batch, lt_hours)
+                    pred = model.forward(batch, lead_times=sub_lead_times, generator=generator)
 
-            # If desired, apply clipping before feeding predictions back as inputs.
-            if apply_rollout_input_clipping:
-                pred = model.apply_rollout_input_clipping(pred)
-            batch = _advance_batch(batch, pred)
-        else:
-            pred = model.forward(batch, lead_times=base_lead_times, generator=generator)
+                    yield pred
 
-            yield pred
+                # If desired, apply clipping before feeding predictions back as inputs.
+                if apply_rollout_input_clipping:
+                    pred = model.apply_rollout_input_clipping(pred)
+                batch = _advance_batch(batch, pred)
+            else:
+                pred = model.forward(batch, lead_times=base_lead_times, generator=generator)
 
-            if apply_rollout_input_clipping:
-                pred = model.apply_rollout_input_clipping(pred)
-            batch = _advance_batch(batch, pred)
+                yield pred
 
-    # Disable noise accumulation after roll-out is complete, in case the model will be used for
-    # normal inference or training afterwards.
-    model.set_noise_accumulation(n=0)
+                if apply_rollout_input_clipping:
+                    pred = model.apply_rollout_input_clipping(pred)
+                batch = _advance_batch(batch, pred)
+    finally:
+        # Also reset when explicitly closed, cancelled, or a forward pass raises.
+        model.set_noise_accumulation(n=0)
 
 
 def rollout_ensemble(
@@ -183,6 +189,9 @@ def rollout_ensemble(
 
     Raises:
         ValueError: If `num_ensemble_members` is less than `1` or if `model` is not stochastic.
+
+    Note:
+        As for :func:`rollout`, explicitly close the generator if stopping early.
     """
     if num_ensemble_members < 1:
         raise ValueError(
@@ -192,13 +201,16 @@ def rollout_ensemble(
         raise ValueError("`rollout_ensemble` requires a stochastic model.")
 
     batch = tile_batch(batch, num_ensemble_members)
-    for pred in rollout(
-        model,
-        batch,
-        steps,
-        fine_lead_times=fine_lead_times,
-        use_noise_accumulation=use_noise_accumulation,
-        apply_rollout_input_clipping=apply_rollout_input_clipping,
-        generator=generator,
-    ):
-        yield split_batch(pred, num_ensemble_members)
+    with closing(
+        rollout(
+            model,
+            batch,
+            steps,
+            fine_lead_times=fine_lead_times,
+            use_noise_accumulation=use_noise_accumulation,
+            apply_rollout_input_clipping=apply_rollout_input_clipping,
+            generator=generator,
+        )
+    ) as predictions:
+        for pred in predictions:
+            yield split_batch(pred, num_ensemble_members)