fix: Defer adaptive chunk adjustment without a measurable sample (#232)

Adaptive chunking computed

    duration_per_task   = completed_task_duration / completed_task_count
    adaptive_chunk_size = target_runtime_seconds / duration_per_task

with neither divisor guarded. Implements step 3 of the adaptive chunking
algorithm in the CLI specification, added upstream in openjd-rs #282: "If
the cumulative duration is zero or non-finite, keep the current chunk size
and wait for a measurable sample."

Be clear about severity: this is spec conformance and hardening, NOT a fix
for a reachable crash. I could not reach either divisor through a normal
`openjd run`:

- completed_task_count cannot be zero. It accumulates
  len(IntRangeExpr.from_str(...)), and an IntRangeExpr cannot be empty --
  "1-0" and "5-1" raise ValueError, "0-0" and "1-1" have len 1.
- completed_task_duration cannot be 0.0 in practice. run_task() always
  spawns and waits on a real subprocess, and raises before the
  accumulation line on failure, so the measured delta would have to fit
  inside one perf_counter tick (~42 ns on this host, against a ~10 ms
  floor for the cheapest possible subprocess).

Correcting an earlier claim of mine: I previously recorded this as a
reproduced defect. It was not. The "reproduction" replayed the arithmetic
with a hardcoded 0.0, which demonstrates nothing about the code path.

The guard is still worth having, because the two implementations fail
differently if it is ever reached: Rust produces inf and saturates to a
huge chunk size, while Python raises ZeroDivisionError out of the run.

Extracted as a module-level _calculate_adaptive_chunk_size returning
Optional[int], mirroring openjd-rs's calculate_adaptive_chunk_size, so the
deferral is unit-testable without provoking an unreachable timing
condition.

Mutation-checked, 10 mutants, 0 survivors: removing the guard or any one
of its three arms, making it always defer, ignoring the sentinel at the
call site, and removing the blend or the clamp are each caught by name.
Removing the sentinel check is caught by the pre-existing end-to-end
test_openjd_run_on_chunked_job_adaptive_chunking, which is the evidence
the extraction is behaviour-preserving on the real path.

One finding from that exercise worth recording. My first draft used
`continue` on the deferral path, which also skips the maximum-task
countdown further down the loop and so silently breaks --maximum-tasks. I
caught it by reading, and then confirmed the suite would NOT have: the
mutant survived. The existing
test_openjd_run_on_chunked_job_maximum_task_count[1] covers adaptive
chunking with --maximum-tasks, but with real durations the estimate never
defers, so the interaction was unpinned. Adds
test_maximum_task_count_is_honoured_when_every_estimate_defers, which
forces every estimate to defer; it now catches that mutant.

Raised as its own PR rather than added to #230: that PR is the RFC
0007/0008 feature work, is already MERGEABLE and waiting only on
approval, and its merge unblocks the specifications conformance suite.
This change is unrelated to its scope and not urgent enough to reset its
review.

Verified: 301 passed / 2 skipped (289 before), ruff clean, black clean,
mypy clean.

Signed-off-by: David Leong <116610336+leongdl@users.noreply.github.com>

Author: leongdl

src/openjd/cli/_run/_local_session/_session_manager.py

@@ -8,6 +8,7 @@
 from signal import signal, SIGINT, SIGTERM, SIG_DFL
 from itertools import islice
 from datetime import datetime, timedelta, timezone
+from math import isfinite
 
 from ._actions import (
     EnterEnvironmentAction,
@@ -49,6 +50,50 @@ def __init__(self, failed_action: SessionAction):
         super().__init__(f"Action failed: {failed_action}")
 
 
+def _calculate_adaptive_chunk_size(
+    *,
+    current_chunk_size: int,
+    completed_task_count: int,
+    completed_task_duration: float,
+    target_runtime_seconds: float,
+) -> Optional[int]:
+    """The adaptive-chunking size estimate, or ``None`` to defer adjusting.
+
+    Implements steps 3-7 of the adaptive chunking algorithm in the Open Job
+    Description CLI specification (``specs/cli/run.md`` in openjd-rs): if the
+    cumulative duration is zero or non-finite, keep the current chunk size and
+    wait for a measurable sample.
+
+    Deferring matters because a zero ``duration_per_task`` makes the ideal size
+    unrepresentable. The two implementations fail differently without this
+    guard: the Rust CLI produces ``inf``, which saturates to a huge chunk size,
+    while Python raises ``ZeroDivisionError`` out of the run. Neither is
+    reachable through a normal ``openjd run`` today -- ``run_task`` always
+    spawns and waits on a real subprocess, so a measured chunk duration of
+    exactly ``0.0`` would need that round trip to complete inside one
+    ``perf_counter`` tick (~42 ns here) -- so this is spec conformance and
+    hardening rather than a fix for an observed failure.
+
+    Extracted as a module-level function, mirroring openjd-rs's
+    ``calculate_adaptive_chunk_size``, so the deferral is unit-testable without
+    needing to provoke an unreachable timing condition end to end.
+    """
+    if (
+        completed_task_count <= 0
+        or completed_task_duration <= 0.0
+        or not isfinite(completed_task_duration)
+    ):
+        return None
+
+    duration_per_task = completed_task_duration / completed_task_count
+    adaptive_chunk_size = target_runtime_seconds / duration_per_task
+    if completed_task_count < 10 and adaptive_chunk_size > current_chunk_size:
+        # When we have data about only a few tasks, gradually blend in the new
+        # estimate instead of cutting over immediately.
+        adaptive_chunk_size = 0.75 * current_chunk_size + 0.25 * adaptive_chunk_size
+    return max(int(adaptive_chunk_size), 1)
+
+
 class LocalSession:
     """
     A class to manage a `Session` object from the `sessions` module,
@@ -300,27 +345,28 @@ def _run_tasks_adaptive_chunking(
             # Estimate a chunk size based on the statistics, and update the iterator. Note that this
             # logic is very simple, providing a good starting point that behaves reasonably for other implementations
             # to follow.
-            duration_per_task = completed_task_duration / completed_task_count
-            adaptive_chunk_size = target_runtime_seconds / duration_per_task
+            new_chunk_size = _calculate_adaptive_chunk_size(
+                current_chunk_size=task_parameters.chunks_default_task_count,  # type: ignore
+                completed_task_count=completed_task_count,
+                completed_task_duration=completed_task_duration,
+                target_runtime_seconds=target_runtime_seconds,
+            )
+            # `None` means there is no measurable sample yet: keep the current chunk
+            # size and wait for one. Deliberately not `continue` -- the maximum-task
+            # countdown below must still run for this completed chunk.
             if (
-                completed_task_count < 10
-                and adaptive_chunk_size > task_parameters.chunks_default_task_count  # type: ignore
+                new_chunk_size is not None
+                and new_chunk_size != task_parameters.chunks_default_task_count
             ):
-                # When we have data about only a few tasks, gradually blend in the new estimate instead of cutting over immediately
-                adaptive_chunk_size = (
-                    0.75 * task_parameters.chunks_default_task_count + 0.25 * adaptive_chunk_size  # type: ignore
-                )
-            adaptive_chunk_size = max(int(adaptive_chunk_size), 1)
-            if adaptive_chunk_size != task_parameters.chunks_default_task_count:
                 LOG.info(
                     msg=f"Open Job Description CLI: Ran {completed_task_count} tasks in {timedelta(seconds=completed_task_duration)}, average {timedelta(seconds=completed_task_duration / completed_task_count)}",
                     extra={"session_id": self.session_id},
                 )
                 LOG.info(
-                    msg=f"Open Job Description CLI: Adjusting chunk size from {task_parameters.chunks_default_task_count} to {adaptive_chunk_size}",
+                    msg=f"Open Job Description CLI: Adjusting chunk size from {task_parameters.chunks_default_task_count} to {new_chunk_size}",
                     extra={"session_id": self.session_id},
                 )
-                task_parameters.chunks_default_task_count = adaptive_chunk_size
+                task_parameters.chunks_default_task_count = new_chunk_size
 
             # If a maximum task count was specified, count them down
             if maximum_tasks and maximum_tasks > 0:

test/openjd/cli/test_adaptive_chunk_size.py

@@ -0,0 +1,167 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+"""Adaptive chunking must defer, not divide, when there is no measurable sample.
+
+Step 3 of the adaptive chunking algorithm in the CLI specification
+(``specs/cli/run.md`` in openjd-rs): "If the cumulative duration is zero or
+non-finite, keep the current chunk size and wait for a measurable sample."
+
+Without it, a zero ``duration_per_task`` raises ``ZeroDivisionError`` out of the
+run in this implementation, and produces a saturating ``inf`` chunk size in the
+Rust one. These tests pin the deferral and the negative controls around it.
+"""
+
+from math import inf, nan
+from pathlib import Path
+import re
+from unittest.mock import patch
+
+import pytest
+
+from openjd.cli._run._local_session._session_manager import _calculate_adaptive_chunk_size
+
+from . import run_openjd_cli_main, format_capsys_outerr
+
+CHUNKED_JOB_TEMPLATE_FILE = str(Path(__file__).parent / "templates" / "chunked_job.yaml")
+
+
+class TestDefersWithoutAMeasurableSample:
+    @pytest.mark.parametrize(
+        "completed_task_duration",
+        [
+            pytest.param(0.0, id="zero-duration"),
+            pytest.param(-0.0, id="negative-zero-duration"),
+            pytest.param(-1.0, id="negative-duration"),
+            pytest.param(inf, id="infinite-duration"),
+            pytest.param(nan, id="nan-duration"),
+        ],
+    )
+    def test_unusable_duration_defers(self, completed_task_duration: float) -> None:
+        # GIVEN a completed chunk whose cumulative duration is not a usable
+        # sample, WHEN the estimate is calculated
+        result = _calculate_adaptive_chunk_size(
+            current_chunk_size=1,
+            completed_task_count=1,
+            completed_task_duration=completed_task_duration,
+            target_runtime_seconds=60.0,
+        )
+
+        # THEN adjustment is deferred rather than dividing by it. Before the
+        # guard, the 0.0 cases raised ZeroDivisionError.
+        assert result is None
+
+    @pytest.mark.parametrize(
+        "completed_task_count", [pytest.param(0, id="zero"), pytest.param(-1, id="negative")]
+    )
+    def test_non_positive_task_count_defers(self, completed_task_count: int) -> None:
+        # GIVEN no counted tasks -- the other divisor in the same expression
+        result = _calculate_adaptive_chunk_size(
+            current_chunk_size=1,
+            completed_task_count=completed_task_count,
+            completed_task_duration=5.0,
+            target_runtime_seconds=60.0,
+        )
+
+        # THEN
+        assert result is None
+
+
+class TestStillEstimatesNormally:
+    """Negative controls: the guard must not swallow the usable cases."""
+
+    def test_blends_while_the_sample_is_small(self) -> None:
+        # GIVEN 1 task in 1s against a 60s target, so the ideal size is 60, and
+        # fewer than 10 completed tasks means the conservative ramp applies:
+        # 0.75 * 4 + 0.25 * 60 == 18.
+        result = _calculate_adaptive_chunk_size(
+            current_chunk_size=4,
+            completed_task_count=1,
+            completed_task_duration=1.0,
+            target_runtime_seconds=60.0,
+        )
+
+        # THEN
+        assert result == 18
+
+    def test_uses_the_ideal_size_once_the_sample_is_large(self) -> None:
+        # GIVEN 10 or more completed tasks, the blend no longer applies:
+        # 10 tasks in 10s is 1s/task, so a 60s target is 60 tasks.
+        result = _calculate_adaptive_chunk_size(
+            current_chunk_size=4,
+            completed_task_count=10,
+            completed_task_duration=10.0,
+            target_runtime_seconds=60.0,
+        )
+
+        # THEN
+        assert result == 60
+
+    def test_does_not_blend_when_the_ideal_size_shrinks(self) -> None:
+        # GIVEN slow tasks, so the ideal size (2) is below the current size (50).
+        # The ramp only applies when the estimate grows, so it is used directly.
+        result = _calculate_adaptive_chunk_size(
+            current_chunk_size=50,
+            completed_task_count=1,
+            completed_task_duration=30.0,
+            target_runtime_seconds=60.0,
+        )
+
+        # THEN
+        assert result == 2
+
+    def test_clamps_to_at_least_one(self) -> None:
+        # GIVEN tasks far slower than the target, so the ideal size rounds to 0
+        result = _calculate_adaptive_chunk_size(
+            current_chunk_size=1,
+            completed_task_count=1,
+            completed_task_duration=1000.0,
+            target_runtime_seconds=1.0,
+        )
+
+        # THEN a chunk of zero tasks would stall the run
+        assert result == 1
+
+
+class TestDeferringDoesNotSkipTheRestOfTheLoop:
+    def test_maximum_task_count_is_honoured_when_every_estimate_defers(self, capsys) -> None:
+        """Deferring must skip only the adjustment, not the loop body.
+
+        Regression for a bug in the first draft of this fix: returning early with
+        `continue` when the estimate deferred also skipped the maximum-task
+        countdown further down the loop, so ``--maximum-tasks`` was ignored and
+        the entire parameter space ran.
+
+        ``test_openjd_run_on_chunked_job_maximum_task_count[1]`` in
+        ``test_chunked_job.py`` does not catch this: with real subprocess
+        durations the estimate never defers, so the deferral path is never taken
+        there. Forcing every estimate to defer is what exercises it.
+        """
+        # GIVEN adaptive chunking (TargetRuntime=1) where the estimate always
+        # defers, and a maximum of 3 tasks over a larger parameter space
+        with patch(
+            "openjd.cli._run._local_session._session_manager._calculate_adaptive_chunk_size",
+            return_value=None,
+        ) as mock_estimate:
+            # WHEN
+            outerr = run_openjd_cli_main(
+                capsys,
+                args=[
+                    "run",
+                    CHUNKED_JOB_TEMPLATE_FILE,
+                    "--step",
+                    "Chunked Step",
+                    "-p",
+                    "ChunkSize=3",
+                    "-p",
+                    "TargetRuntime=1",
+                    "--maximum-tasks",
+                    "3",
+                ],
+                expected_exit_code=0,
+            )
+
+        # THEN the run still stopped at the limit, and the deferral path really
+        # was the one taken.
+        assert mock_estimate.called, "the adaptive estimate was never consulted"
+        assert re.search(
+            "Chunks run: 3$", outerr.out, re.MULTILINE
+        ), f"Regex 'Chunks run: 3$' not matched in:\n{format_capsys_outerr(outerr)}"