fix(waterdata): make parallel_chunks cap a hard ceiling in ChunkPlan._refine

_refine grew the plan with `while self.total < max_chunks`, but with more
than one axis a single split multiplies total by (k+1)/k for the split axis
— it adds the product of the *other* axes, not one — so total could step
*past* the cap (two 8-atom axes at cap 10 landed on 12; three 10-atom axes
at cap 30 landed on 32). That broke the documented "at most n / can't
multiply past it" guarantee the dial exists to provide.

Take a split only when it keeps total within the cap: skip any axis whose
split would overshoot (the increment is per-axis, total // k). The cap is
now a true ceiling — never exceeded — and the plan lands below it when no
whole split hits it exactly (two even axes reach 4, not 5). Single-axis
plans are unchanged (still hit n exactly, or saturate at one atom/chunk).

Docstrings updated (soft cap -> hard ceiling; note the may-undershoot).
Repurpose the many-axes test to a cap the old loop overshot (30, was 32)
and add a parametrized regression over the 2-axis overshoot combos.

Signed-off-by: thodson-usgs <thodson@usgs.gov>

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: thodson-usgs

dataretrieval/ogc/chunking.py

@@ -224,10 +224,13 @@ def parallel_chunks(n: int) -> Iterator[None]:
         integer such as ``2``, ``8``, or ``32``. It caps the plan's *total*
         sub-request count (the cartesian product across every multi-value
         argument combined, not per argument), so several multi-value arguments
-        cannot multiply past it. The actual count is bounded below by what the
-        ~8 KB URL limit already forces and above by the number of values there
-        are to split, so an ``n`` larger than the input allows simply yields one
-        sub-request per value; ``n=1`` asks for no extra fan-out.
+        cannot multiply past it. The cap is a ceiling, never exceeded: the
+        actual count is bounded below by what the ~8 KB URL limit already
+        forces and above by ``n``, so an ``n`` larger than the input allows
+        simply yields one sub-request per value, and with several multi-value
+        arguments the total may land somewhat below ``n`` because splits are
+        whole (the plan can't always divide evenly onto ``n``); ``n=1`` asks
+        for no extra fan-out.
 
         Each sub-request fetches at least one page, so it costs at least one
         request against your hourly rate limit — a larger ``n`` spends more

dataretrieval/ogc/planning.py

@@ -330,15 +330,16 @@ class ChunkPlan:
         Byte budget for the request (URL + body) — a hard ceiling every
         sub-request must fit.
     max_chunks : int, optional
-        Soft cap on the plan's total sub-request count (default ``0`` = off).
+        Hard cap on the plan's total sub-request count (default ``0`` = off).
         ``0`` chunks only as much as ``url_limit`` requires — the most
         conservative plan, fewest sub-requests. A positive cap fans the plan
         out to up to ``max_chunks`` sub-requests overall (the
         cartesian product across axes, never fewer than the byte budget
         already forces) — capped as a whole, not per axis, so several
-        multi-value axes can't multiply past the cap. Set from the
-        :func:`~dataretrieval.ogc.chunking.parallel_chunks` ``n``; see
-        :meth:`_refine`.
+        multi-value axes can't multiply past the cap. The plan never exceeds
+        the cap and may land below it when no whole split lands on it exactly.
+        Set from the :func:`~dataretrieval.ogc.chunking.parallel_chunks` ``n``;
+        see :meth:`_refine`.
 
     Attributes
     ----------
@@ -502,38 +503,58 @@ def _refine(self, max_chunks: int) -> None:
         cartesian product across all axes) at ``max_chunks``, not
         each axis independently — with several multi-value axes, a cap of 32
         still means at most 32 sub-requests overall, not ``32 ** n_axes``.
-        Each split picks the single largest splittable chunk across *every*
-        axis (ties broken by axis-extraction order, then lowest index), so
-        growth is distributed round-robin rather than one axis saturating
-        before another is touched. Purely additive — only ever *splits*
-        existing chunks, so the byte pass's work and the ``url_limit``
-        invariant are both preserved, and it never raises. A no-op at
-        ``max_chunks <= 0``.
+        The cap is a hard ceiling: every split multiplies the plan by
+        ``(k+1)/k`` for the chosen axis (adding ``total // k`` sub-requests,
+        not one), so a split is taken only when it keeps ``total`` within the
+        cap. When no in-budget split remains the plan lands *below* the cap
+        rather than overshooting it — a multi-axis plan may not divide evenly
+        onto ``max_chunks`` (e.g. two even axes can reach 4 but not 5, so a
+        cap of 5 yields 4). Each split picks the single largest splittable
+        chunk among the in-budget axes (ties broken by axis-extraction order,
+        then lowest index), so growth is distributed round-robin rather than
+        one axis saturating before another is touched. Purely additive — only
+        ever *splits* existing chunks, so the byte pass's work and the
+        ``url_limit`` invariant are both preserved, and it never raises. A
+        no-op at ``max_chunks <= 0``.
 
         Parameters
         ----------
         max_chunks : int
-            Soft cap on the plan's total sub-request count (``0`` = off) — the
+            Hard cap on the plan's total sub-request count (``0`` = off) — the
             ``parallel_chunks(n)`` value. The cap is on the whole plan (the
             cartesian product across axes), not per axis, so several multi-value
-            axes can't multiply past it.
+            axes can't multiply past it; the plan never exceeds it and may fall
+            short when no whole split lands on it exactly.
         """
         if max_chunks <= 0:
             return
-        while self.total < max_chunks:
-            # Largest splittable chunk across every axis; a chunk of size 1
-            # can't be split further. ``max`` with a stable input order
-            # breaks ties by axis order, then lowest index within an axis.
+        while True:
+            total = self.total
+            if total >= max_chunks:
+                return
+            # Largest splittable chunk among the axes whose split still fits the
+            # cap. Splitting any chunk of an axis with ``k`` chunks turns that
+            # ``k`` into ``k+1``, so it adds ``total // k`` sub-requests (the
+            # product of the other axes) regardless of which chunk — hence the
+            # budget test is per axis, not per chunk. Skipping an over-budget
+            # axis makes ``max_chunks`` a true ceiling. A chunk of size 1 can't
+            # be split further. ``max`` with a stable input order breaks ties by
+            # axis order, then lowest index within an axis.
             candidate: tuple[_Axis, int] | None = None
             candidate_size = -1
             for axis in self.axes:
-                for idx, chunk in enumerate(self.chunks[axis.arg_key]):
+                axis_chunks = self.chunks[axis.arg_key]
+                if total + total // len(axis_chunks) > max_chunks:
+                    continue  # any split of this axis would overshoot the cap
+                for idx, chunk in enumerate(axis_chunks):
                     if len(chunk) <= 1:
                         continue
                     if len(chunk) > candidate_size:
                         candidate, candidate_size = (axis, idx), len(chunk)
             if candidate is None:
-                return  # every axis saturated at one atom per chunk
+                # Every axis is saturated at one atom per chunk or would
+                # overshoot the cap; stop below it rather than exceed it.
+                return
             axis, idx = candidate
             _split_at(self.chunks[axis.arg_key], idx)
 

tests/waterdata_chunking_test.py

@@ -2237,20 +2237,52 @@ def test_cap_caps_the_total_across_axes():
 
 
 def test_cap_bounds_fan_out_across_many_axes():
-    """The guardrail holds regardless of axis count: three multi-value axes
-    at ``n=32`` still top out at ``n`` sub-requests total, not ``n ** 3`` —
-    the property the single-axis-only cap in the original implementation did
-    not guarantee."""
-    high = 32
+    """The guardrail holds regardless of axis count: three multi-value axes at
+    a cap of 30 fan out to *at most* 30 sub-requests total — never the
+    ``30 ** 3`` a per-axis cap would allow, and never *over* the cap either.
+    30 is deliberately not evenly reachable by these axes: a single split
+    multiplies the plan by more than one, so the naive ``while total < cap``
+    the first refine used stepped past 30 (to 32). The cap is a hard ceiling —
+    the property neither the single-axis-only cap nor that naive loop
+    guaranteed."""
+    cap = 30
     # Three chunkable axes (two list axes + the filter OR-axis), each with 10
-    # atoms — under the old per-axis cap this would have been high**3.
+    # atoms — under the old per-axis cap this would have been cap**3.
     args = {
         "monitoring_location_id": [f"L{i}" for i in range(10)],
         "parameter_code": [f"{i:05d}" for i in range(10)],
         "filter": " OR ".join(f"p='{i}'" for i in range(10)),
     }
-    plan = ChunkPlan(args, _fake_build, url_limit=8000, max_chunks=high)
-    assert plan.total <= high
+    plan = ChunkPlan(args, _fake_build, url_limit=8000, max_chunks=cap)
+    assert 1 < plan.total <= cap  # fanned out, but never past the ceiling
+
+
+@pytest.mark.parametrize(
+    "atoms_per_axis, cap",
+    [
+        (4, 5),  # pre-fix loop overshot 5 -> 6
+        (8, 10),  # pre-fix loop overshot 10 -> 12
+        (10, 7),  # pre-fix loop overshot 7 -> 8
+    ],
+)
+def test_cap_is_a_hard_ceiling_never_overshoots(atoms_per_axis, cap):
+    """The cap is a hard ceiling, not a soft target. With two multi-value axes
+    a single split multiplies the plan by ``(k+1)/k`` for the split axis —
+    adding the product of the *other* axes, not one — so a naive
+    ``while total < cap`` loop steps *past* the cap. These are exactly the
+    (atoms, cap) combos that loop overshot (5->6, 10->12, 7->8). The plan must
+    fan out and cover every atom once, but never exceed the cap, landing below
+    it when no whole split lands on it exactly (two even axes reach 4, not 5)."""
+    args = {
+        "monitoring_location_id": [f"L{i:03d}" for i in range(atoms_per_axis)],
+        "parameter_code": [f"{i:05d}" for i in range(atoms_per_axis)],
+    }
+    plan = ChunkPlan(args, _fake_build, url_limit=8000, max_chunks=cap)
+    assert 1 < plan.total <= cap  # fanned out, but never past the ceiling
+    # Every atom on every axis is still covered exactly once.
+    for key, atoms in args.items():
+        flattened = [a for chunk in plan.chunks[key] for a in chunk]
+        assert sorted(flattened) == sorted(atoms)
 
 
 def test_cap_does_not_mask_unchunkable():