diff --git a/src/openjd/cli/_run/_local_session/_session_manager.py b/src/openjd/cli/_run/_local_session/_session_manager.py index 2413505..d7c6c22 100644 --- a/src/openjd/cli/_run/_local_session/_session_manager.py +++ b/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: diff --git a/test/openjd/cli/test_adaptive_chunk_size.py b/test/openjd/cli/test_adaptive_chunk_size.py new file mode 100644 index 0000000..3d43114 --- /dev/null +++ b/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)}"