diff --git a/src/hyperloom/inference_optimizer/breakdown/reporters/_renderers/final.py b/src/hyperloom/inference_optimizer/breakdown/reporters/_renderers/final.py index 694051839..6f2402a0f 100644 --- a/src/hyperloom/inference_optimizer/breakdown/reporters/_renderers/final.py +++ b/src/hyperloom/inference_optimizer/breakdown/reporters/_renderers/final.py @@ -16,6 +16,10 @@ ) from ._invocation import render_invocation_block +# ``geak_pending.status`` values meaning the candidate was measured but its +# revalidation never landed, so the win was abandoned rather than judged. +_GEAK_DROPPED_STATUSES: frozenset[str] = frozenset({"rebench_cancelled", "rebench_unavailable"}) + @register_renderer("final") def render(breakdown: dict[str, Any]) -> RenderedSection: @@ -87,7 +91,23 @@ def render(breakdown: dict[str, Any]) -> RenderedSection: rationale=f"validated at stack_len={val_stack_len} ts={val_ts}", ) ) - if geak_pending and geak_pending.get("status") == "awaiting_rebench": + geak_pending_status = str(geak_pending.get("status") or "") + if geak_pending and geak_pending_status in _GEAK_DROPPED_STATUSES: + self_gain = geak_pending.get("self_reported_gain_pct") + self_gain_str = fmt_pct(self_gain, plus=True) if isinstance(self_gain, (int, float)) else "unknown" + drop_reason = str(geak_pending.get("revalidation_error") or "").strip() or "reason not recorded" + facts.append( + f"GEAK candidate (self-reported {self_gain_str}) was DROPPED without " + f"revalidation (status={geak_pending_status}, {drop_reason})." + ) + warnings.append( + "A measured GEAK e2e candidate was abandoned because its same-harness " + f"revalidation could not land ({drop_reason}). It was never judged on " + "merit, so this session's gain may understate what the optimizer " + "actually found — the candidate's artefacts are on disk but absent " + "from current_best / action_path / the validated gain." + ) + elif geak_pending and geak_pending_status == "awaiting_rebench": self_gain = geak_pending.get("self_reported_gain_pct") self_gain_str = fmt_pct(self_gain, plus=True) if isinstance(self_gain, (int, float)) else "unknown" facts.append( diff --git a/src/hyperloom/inference_optimizer/tests/test_geak_revalidation_dispatch.py b/src/hyperloom/inference_optimizer/tests/test_geak_revalidation_dispatch.py new file mode 100644 index 000000000..b6e16ed72 --- /dev/null +++ b/src/hyperloom/inference_optimizer/tests/test_geak_revalidation_dispatch.py @@ -0,0 +1,808 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT +"""Regression tests for GEAK same-harness revalidation dispatch (L1/L2).""" + +from __future__ import annotations + +import json +from datetime import datetime, timedelta, timezone + +import pytest + +from hyperloom.orchestrator.phases import geak_rebench as gr +from hyperloom.orchestrator.phases import machine_state as ps + + +@pytest.fixture +def coordinator(tmp_path, monkeypatch): + monkeypatch.setenv("USER_DATA_PATH", str(tmp_path)) + from hyperloom.inference_optimizer.session.paths import make_session_dir as _msd + from hyperloom.orchestrator.loop.coordinator import Coordinator + from hyperloom.orchestrator.roles import ( + MockBackend, + MockCriticBackend, + MockRobustnessBackend, + ScriptedPlan, + ) + from .conftest import seed_target_analysis_marker + + sd = _msd() + seed_target_analysis_marker(sd) + backends = { + "orchestration": MockBackend(ScriptedPlan(turns=[]), name="orchestration"), + "critic": MockCriticBackend(), + "robustness": MockRobustnessBackend(), + } + return Coordinator(sd, backends=backends) + + +def _geak_rebench_params(**extra: object) -> dict: + return { + "source": "resume_stack_revalidate", + "geak_fallback": True, + "reason": "geak_e2e_win", + **extra, + } + + +def _arm_kernel_to_sweep(st) -> None: + now = datetime.now(timezone.utc) + st.phase = ps.PHASE_KERNEL_AGENT + st.phase_started_ts = (now - timedelta(minutes=5)).isoformat() + st.phase_started_unix = (now - timedelta(minutes=5)).timestamp() + st.start_ts = (now - timedelta(minutes=10)).isoformat() + st.max_minutes = 96 * 60 + st.kernel_optimizer = "geak" + st.geak_result = {"status": "ok", "accepted_config": {"flags": "--foo", "env": ""}} + st.geak_pending = {} + st.set_pending_escalate_hint(ps.ESCALATE_HINT_SKIP_TO_SWEEP) + + +def test_geak_revalidate_idempotency_key_scopes_by_macro_cycle() -> None: + assert gr.geak_revalidate_idempotency_key(0) == "geak-revalidate-c0" + assert gr.geak_revalidate_idempotency_key(2) == "geak-revalidate-c2" + assert gr.geak_revalidate_idempotency_key(2) != gr.geak_revalidate_idempotency_key(3) + + +@pytest.mark.asyncio +async def test_cancel_queued_not_allowed_spares_geak_rebench_into_sweep(coordinator) -> None: + c = coordinator + geak_task = await c.tasks.create( + kind="explore", + params=_geak_rebench_params(), + idempotency_key="geak-revalidate-c0", + ) + regular = await c.tasks.create( + kind="explore", + params={"source": "normal_explore"}, + idempotency_key="regular-explore", + ) + + cancelled = await c.tasks.cancel_queued_not_allowed( + allowed_kinds=ps.PHASE_ALLOWED_ACTIONS[ps.PHASE_SWEEP], + reason="phase_transition:KERNEL_AGENT->SWEEP", + spare_queued=lambda _tid, kind, params: gr.spare_geak_rebench_on_phase_transition( + target_phase=ps.PHASE_SWEEP, + kind=kind, + params=params, + ), + ) + + assert geak_task.task_id not in cancelled + assert regular.task_id in cancelled + assert (await c.tasks.get(geak_task.task_id)).state == "queued" + assert (await c.tasks.get(regular.task_id)).state == "cancelled" + + +@pytest.mark.asyncio +async def test_cancel_queued_not_allowed_cancels_geak_rebench_into_close(coordinator) -> None: + c = coordinator + geak_task = await c.tasks.create( + kind="explore", + params=_geak_rebench_params(), + idempotency_key="geak-revalidate-c0", + ) + + cancelled = await c.tasks.cancel_queued_not_allowed( + allowed_kinds=ps.PHASE_ALLOWED_ACTIONS[ps.PHASE_CLOSE], + reason="phase_transition:SWEEP->CLOSE", + spare_queued=lambda _tid, kind, params: gr.spare_geak_rebench_on_phase_transition( + target_phase=ps.PHASE_CLOSE, + kind=kind, + params=params, + ), + ) + + assert geak_task.task_id in cancelled + assert (await c.tasks.get(geak_task.task_id)).state == "cancelled" + + +@pytest.mark.asyncio +async def test_cancel_family_cancels_geak_rebench_explore(coordinator) -> None: + """Explore-family prune must cancel GEAK rebench (no implicit spare).""" + c = coordinator + geak_task = await c.tasks.create( + kind="explore", + params=_geak_rebench_params(), + idempotency_key="geak-revalidate-c1", + ) + + cancelled = await c.tasks.cancel_family(["explore"], reason="prune_branch") + + assert geak_task.task_id in cancelled + assert (await c.tasks.get(geak_task.task_id)).state == "cancelled" + + +@pytest.mark.asyncio +async def test_geak_revalidate_idempotency_key_allows_retry_per_macro_cycle( + coordinator, +) -> None: + c = coordinator + cycle0_key = gr.geak_revalidate_idempotency_key(0) + cycle1_key = gr.geak_revalidate_idempotency_key(1) + + settled = await c.tasks.create( + kind="explore", + params=_geak_rebench_params(), + idempotency_key=cycle0_key, + task_id="geak-revalidate-c0-task", + ) + await c.tasks.transition(settled.task_id, "cancelled", evidence={"reason": "test"}) + + fresh, was_existing = await c.tasks.create_or_return_existing( + kind="explore", + params=_geak_rebench_params(), + idempotency_key=cycle1_key, + ) + + assert not was_existing + assert fresh.task_id != settled.task_id + assert fresh.state == "queued" + + +@pytest.mark.asyncio +async def test_enqueue_internal_stack_rebench_uses_macro_cycle_idempotency_key( + coordinator, +) -> None: + """Production enqueue path must stamp cycle-scoped GEAK rebench keys.""" + c = coordinator + st = c.shared_state + st.baseline_tput = 100.0 + st.geak_result = { + "status": "ok", + "accepted_config": {"flags": "--max-num-batched-tokens 8192", "env": ""}, + } + + st.macro_cycle = 0 + first = await c._enqueue_internal_stack_rebench(reason="geak_e2e_win") + row0 = await c.tasks.get(str(first["task_id"])) + assert row0.idempotency_key == "geak-revalidate-c0" + + st.macro_cycle = 1 + second = await c._enqueue_internal_stack_rebench(reason="geak_e2e_win") + row1 = await c.tasks.get(str(second["task_id"])) + assert row1.idempotency_key == "geak-revalidate-c1" + assert row1.task_id != row0.task_id + + +@pytest.mark.asyncio +async def test_rebench_can_be_rebuilt_after_cancel_within_same_cycle(coordinator) -> None: + """A cancelled rebench must not block a fresh one in the same macro-cycle. + + ``create_or_return_existing`` hands back the cancelled row for a reused key, + which KERNEL then reads as ``rebench_unavailable`` and the GEAK win stays + audit-only for the rest of the cycle. + """ + c = coordinator + st = c.shared_state + st.baseline_tput = 100.0 + st.macro_cycle = 0 + st.geak_result = { + "status": "ok", + "accepted_config": {"flags": "--max-num-batched-tokens 8192", "env": ""}, + } + + first = await c._enqueue_internal_stack_rebench(reason="geak_e2e_win") + first_id = str(first["task_id"]) + assert first["task_state"] == "queued" + + # An explore-family prune settles the queued rebench mid-cycle. + assert first_id in await c.tasks.cancel_family(["explore"], reason="prune_branch") + + second = await c._enqueue_internal_stack_rebench(reason="geak_e2e_win") + + assert second["task_id"] != first_id + assert second["task_state"] == "queued" + + +@pytest.mark.asyncio +async def test_prune_settles_geak_pending_when_rebench_cancelled(coordinator) -> None: + """Pruning the explore family must not leave the slot stuck awaiting.""" + c = coordinator + st = c.shared_state + st.kernel_optimizer = "geak" + st.geak_result = {"status": "ok"} + + rebench = await c.tasks.create( + kind="explore", + params=_geak_rebench_params(), + idempotency_key=gr.geak_revalidate_idempotency_key(0), + task_id="pruned-rebench", + ) + st.geak_pending = {"status": "awaiting_rebench", "revalidation_task_id": rebench.task_id} + st.resume_pending_revalidation = True + + from hyperloom.inference_optimizer.protocol.intent import Intent, IntentType + + await c._handle_prune_branch( + "robustness", + Intent(type=IntentType.PRUNE_BRANCH, payload={"family": "explore", "reason": "prune_branch"}), + ) + + assert (await c.tasks.get(rebench.task_id)).state == "cancelled" + assert st.geak_pending["status"] == "rebench_cancelled" + assert st.resume_pending_revalidation is False + + +@pytest.mark.asyncio +async def test_settled_pending_rejects_late_rebench_result(coordinator) -> None: + """A settled slot must not be revived by a late/orphan rebench completion.""" + c = coordinator + st = c.shared_state + st.baseline_tput = 100.0 + st.current_best = {"action": "baseline", "tput": 100.0, "extra_server_args": ""} + st.geak_result = { + "status": "ok", + "accepted_config": {"flags": "--foo", "env": ""}, + "accepted_kernels": ["k1"], + } + st.geak_pending = {"status": "rebench_cancelled", "revalidation_error": "close_sequence"} + st.resume_pending_revalidation = False + + task = await c.tasks.create( + kind="explore", + params=_geak_rebench_params(), + idempotency_key=gr.geak_revalidate_idempotency_key(0), + task_id="late-rebench", + ) + + await c._promote_to_shared_state( + task.kind, + { + "output_throughput": 150.0, + "best_variant": {"fingerprint": "abc"}, + "winners": [], + }, + task=task, + ) + + assert st.current_best["tput"] == 100.0 + assert st.geak_pending["status"] == "rebench_cancelled" + assert not any(e.get("action") == "geak_e2e" for e in st.optimization_stack) + + +@pytest.mark.asyncio +async def test_cleared_pending_without_resume_flag_rejects_result(coordinator) -> None: + """An empty slot is only a resume signal when the resume flag is set.""" + c = coordinator + st = c.shared_state + st.baseline_tput = 100.0 + st.current_best = {"action": "baseline", "tput": 100.0, "extra_server_args": ""} + st.geak_result = { + "status": "ok", + "accepted_config": {"flags": "--foo", "env": ""}, + "accepted_kernels": ["k1"], + } + st.geak_pending = {} + st.resume_pending_revalidation = False + + task = await c.tasks.create( + kind="explore", + params=_geak_rebench_params(), + idempotency_key=gr.geak_revalidate_idempotency_key(0), + task_id="post-clear-rebench", + ) + + await c._promote_to_shared_state( + task.kind, + { + "output_throughput": 150.0, + "best_variant": {"fingerprint": "abc"}, + "winners": [], + }, + task=task, + ) + + assert st.current_best["tput"] == 100.0 + assert not any(e.get("action") == "geak_e2e" for e in st.optimization_stack) + + +@pytest.mark.asyncio +async def test_settle_preserves_candidate_audit_fields(coordinator) -> None: + """Settling must not discard the self-reported numbers the report shows.""" + c = coordinator + st = c.shared_state + st.baseline_tput = 100.0 + st.geak_result = {"status": "ok"} + c.phase_kernel._record_geak_candidate( + { + "status": "ok", + "final_throughput_tok_s": 116.0, + "throughput_speedup": 1.16, + "accepted_config": {"flags": "--foo", "env": ""}, + } + ) + st.geak_pending = {**st.geak_pending, "revalidation_task_id": "gone-task"} + assert st.geak_pending["self_reported_gain_pct"] == pytest.approx(16.0) + + settled = await gr.settle_dangling_geak_pending(c.tasks, st, reason="close_sequence") + + assert settled is True + assert st.geak_pending["status"] == "rebench_cancelled" + assert st.geak_pending["revalidation_error"] == "close_sequence" + # The audit numbers survive so the report can name what was dropped. + assert st.geak_pending["self_reported_gain_pct"] == pytest.approx(16.0) + assert st.geak_pending["self_reported_tput"] == pytest.approx(116.0) + # The id of a task that will never land must not outlive the slot. + assert "revalidation_task_id" not in st.geak_pending + + +@pytest.mark.asyncio +async def test_wall_clock_closing_stops_rebench_and_settles(coordinator) -> None: + """The wall-clock closing path never reaches ``_on_enter_close``. + + It cancels queued work but left the slot at ``awaiting_rebench`` and a + running rebench alive, so the report claimed a rebench was still coming + while the task had already been cancelled. + """ + c = coordinator + st = c.shared_state + st.kernel_optimizer = "geak" + st.geak_result = {"status": "ok"} + + queued = await c.tasks.create( + kind="explore", + params=_geak_rebench_params(), + idempotency_key=gr.geak_revalidate_idempotency_key(0), + task_id="queued-at-timeout", + ) + running = await c.tasks.create( + kind="explore", + params=_geak_rebench_params(), + idempotency_key=gr.geak_revalidate_idempotency_key(0, 1), + task_id="running-at-timeout", + ) + await c.tasks.transition(running.task_id, "running") + st.geak_pending = {"status": "awaiting_rebench", "revalidation_task_id": running.task_id} + st.resume_pending_revalidation = True + + await c._enter_closing_phase(grace_sec=30.0) + + assert (await c.tasks.get(queued.task_id)).state == "cancelled" + assert (await c.tasks.get(running.task_id)).state == "cancelled" + assert st.geak_pending["status"] == "rebench_cancelled" + assert st.resume_pending_revalidation is False + + +def _render_final(geak_pending: dict) -> tuple[list[str], list[str]]: + from hyperloom.inference_optimizer.breakdown.reporters._renderers.final import render + + section = render( + { + "final": { + "throughput_tok_s_per_gpu": 140.0, + "cumulative_gain_pct_validated": 0.0, + "geak_pending": geak_pending, + }, + "baseline": {"throughput_tok_s_per_gpu": 100.0}, + } + ) + return list(section.key_facts), list(section.warnings) + + +def test_final_report_surfaces_cancelled_geak_revalidation() -> None: + """A measured candidate dropped for a missed rebench must be visible.""" + facts, warnings = _render_final( + { + "status": "rebench_cancelled", + "revalidation_error": "close_sequence", + "self_reported_gain_pct": 12.5, + } + ) + + blob = " ".join(facts + warnings).lower() + assert "rebench" in blob + assert any("close_sequence" in w or "could not" in w.lower() for w in warnings) + # The dropped candidate must not be presented as awaiting anything. + assert not any("awaiting" in f.lower() for f in facts) + + +def test_final_report_still_flags_awaiting_geak_revalidation() -> None: + facts, warnings = _render_final( + {"status": "awaiting_rebench", "self_reported_gain_pct": 12.5} + ) + + assert any("AWAITING" in f for f in facts) + assert warnings + + +def test_legacy_placeholder_does_not_match_cycle_scoped_keys() -> None: + """The legacy slot must not absorb a cycle-scoped rebench from another cycle.""" + from hyperloom.orchestrator.state.task_registry import Task + + def _task(key: str) -> Task: + return Task(task_id="t1", kind="explore", state="queued", params={}, idempotency_key=key) + + assert gr.geak_rebench_tracks_pending_task( + gr.LEGACY_GEAK_REVALIDATE_PLACEHOLDER, + _task(gr.LEGACY_GEAK_REVALIDATE_PLACEHOLDER), + macro_cycle=0, + ) + assert not gr.geak_rebench_tracks_pending_task( + gr.LEGACY_GEAK_REVALIDATE_PLACEHOLDER, + _task("geak-revalidate-c3"), + macro_cycle=0, + ) + + +@pytest.mark.asyncio +async def test_geak_rebench_survives_kernel_to_sweep_transition(coordinator) -> None: + c = coordinator + st = c.shared_state + _arm_kernel_to_sweep(st) + + geak_task = await c.tasks.create( + kind="explore", + params=_geak_rebench_params(), + idempotency_key=gr.geak_revalidate_idempotency_key(st.macro_cycle), + task_id="geak-rebench-survives-transition", + ) + regular = await c.tasks.create( + kind="explore", + params={"source": "normal_explore"}, + idempotency_key="regular-explore-transition", + ) + + cancelled = await c.tasks.cancel_queued_not_allowed( + allowed_kinds=ps.PHASE_ALLOWED_ACTIONS[ps.PHASE_SWEEP], + reason="phase_transition:KERNEL_AGENT->SWEEP", + spare_queued=lambda _tid, kind, params: gr.spare_geak_rebench_on_phase_transition( + target_phase=ps.PHASE_SWEEP, + kind=kind, + params=params, + ), + ) + + assert geak_task.task_id not in cancelled + assert regular.task_id in cancelled + assert (await c.tasks.get(geak_task.task_id)).state == "queued" + assert (await c.tasks.get(regular.task_id)).state == "cancelled" + + +@pytest.mark.asyncio +async def test_duplicate_enqueue_skips_while_rebench_in_flight(coordinator, tmp_path) -> None: + c = coordinator + st = c.shared_state + _arm_kernel_to_sweep(st) + geak_dir = tmp_path / "geak" + geak_dir.mkdir() + result = { + "status": "ok", + "final_throughput_tok_s": 116.0, + "accepted_config": {"flags": "--foo", "env": ""}, + } + (geak_dir / "result.json").write_text(json.dumps(result), encoding="utf-8") + + inflight = await c.tasks.create( + kind="explore", + params=_geak_rebench_params(), + idempotency_key=gr.geak_revalidate_idempotency_key(0), + task_id="cycle0-inflight", + ) + + coord = c + coord.phase_kernel._record_geak_kernel_journey = lambda _result: None + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setattr( + "hyperloom.orchestrator.kernel.request_handlers._kernel_agent_tool_path", + lambda _name: (_ for _ in ()).throw(RuntimeError("runner should not run")), + ) + try: + await coord._run_geak_kernel_phase(from_phase="KERNEL") + finally: + monkeypatch.undo() + + assert len([t for t in await c.tasks.queued() if gr.is_geak_same_harness_rebench_task(t.kind, t.params)]) == 1 + assert st.geak_pending["revalidation_task_id"] == inflight.task_id + created = await c.tasks.get(inflight.task_id) + assert created.state == "queued" + + +@pytest.mark.asyncio +async def test_geak_rebench_failure_clears_pending_when_placeholder_tracked(coordinator) -> None: + c = coordinator + st = c.shared_state + st.kernel_optimizer = "geak" + st.geak_result = {"status": "ok"} + placeholder = gr.geak_revalidate_idempotency_key(0) + st.geak_pending = {"status": "awaiting_rebench", "revalidation_task_id": placeholder} + + task = await c.tasks.create( + kind="explore", + params=_geak_rebench_params(), + idempotency_key=placeholder, + task_id="geak-rebench-fail", + ) + + await c._handle_unpromotable_result( + task, + {"status": "failed", "error_class": "subprocess_nonzero", "error": "revalidation failed"}, + ) + + assert not st.geak_pending + assert st.geak_result["revalidation_status"] == "failed" + + +@pytest.mark.asyncio +async def test_resume_revalidation_with_empty_pending_still_promotes(coordinator) -> None: + """Resume stack rebench with cleared geak_pending must still lift validated gain.""" + c = coordinator + st = c.shared_state + st.baseline_tput = 100.0 + st.current_best = {"action": "geak_e2e", "tput": 116.0, "extra_server_args": "--foo"} + st.optimization_stack = [{"action": "geak_e2e", "tput": 116.0}] + st.geak_result = {"status": "ok", "accepted_config": {"flags": "--foo", "env": ""}} + st.geak_pending = {} + st.resume_pending_revalidation = True + + task = await c.tasks.create( + kind="explore", + params=_geak_rebench_params(), + idempotency_key=gr.geak_revalidate_idempotency_key(0), + task_id="resume-reverify", + ) + + await c._promote_to_shared_state( + task.kind, + { + "output_throughput": 140.0, + "best_variant": {"fingerprint": "abc"}, + "winners": [], + }, + task=task, + ) + + assert st.cumulative_gain_validated == pytest.approx(40.0) + assert st.resume_pending_revalidation is False + + +@pytest.mark.asyncio +async def test_orphan_geak_rebench_success_does_not_promote(coordinator) -> None: + c = coordinator + st = c.shared_state + st.baseline_tput = 100.0 + st.current_best = {"action": "baseline", "tput": 100.0, "extra_server_args": ""} + st.kernel_optimizer = "geak" + st.geak_result = { + "status": "ok", + "accepted_config": {"flags": "--foo", "env": ""}, + "accepted_kernels": ["k1"], + } + tracked = await c.tasks.create( + kind="explore", + params=_geak_rebench_params(), + idempotency_key=gr.geak_revalidate_idempotency_key(0), + task_id="tracked-rebench", + ) + orphan = await c.tasks.create( + kind="explore", + params=_geak_rebench_params(), + idempotency_key=gr.geak_revalidate_idempotency_key(1), + task_id="orphan-rebench", + ) + st.geak_pending = {"status": "awaiting_rebench", "revalidation_task_id": tracked.task_id} + + await c._promote_to_shared_state( + orphan.kind, + { + "output_throughput": 150.0, + "best_variant": {"fingerprint": "abc", "gain_pct": 50.0}, + "winners": [{"fingerprint": "abc", "gain_pct": 50.0}], + }, + task=orphan, + ) + + assert st.current_best["tput"] == 100.0 + assert st.geak_pending["revalidation_task_id"] == tracked.task_id + assert not any(e.get("action") == "geak_e2e" for e in st.optimization_stack) + + +@pytest.mark.asyncio +async def test_orphan_geak_rebench_inconclusive_does_not_run_2a(coordinator) -> None: + """An untracked rebench must not trigger the GEAK-harness 2a fallback. + + A successful 2a writes the stack entry; a failed 2a clears the pending slot + of the genuinely tracked rebench. Both bypass the orphan gate. + """ + c = coordinator + st = c.shared_state + st.baseline_tput = 100.0 + st.current_best = {"action": "baseline", "tput": 100.0, "extra_server_args": ""} + st.kernel_optimizer = "geak" + st.geak_result = { + "status": "ok", + "accepted_config": {"flags": "--foo", "env": ""}, + "accepted_kernels": ["k1"], + } + tracked = await c.tasks.create( + kind="explore", + params=_geak_rebench_params(), + idempotency_key=gr.geak_revalidate_idempotency_key(0), + task_id="tracked-rebench", + ) + orphan = await c.tasks.create( + kind="explore", + # Fingerprint mismatch below makes the 2b decision inconclusive. + params=_geak_rebench_params(expected_cfg_hash="expected-hash"), + idempotency_key=gr.geak_revalidate_idempotency_key(1), + task_id="orphan-rebench", + ) + st.geak_pending = {"status": "awaiting_rebench", "revalidation_task_id": tracked.task_id} + + fallback_calls: list[str] = [] + + async def _record_fallback(*, reason: str) -> dict: + fallback_calls.append(reason) + return {"validated": False, "reason": "should not run"} + + c._validate_geak_via_geak_harness = _record_fallback # type: ignore[assignment] + + await c._promote_to_shared_state( + orphan.kind, + { + "output_throughput": 150.0, + "best_variant": {"fingerprint": "mismatched-hash"}, + "winners": [], + }, + task=orphan, + ) + + assert fallback_calls == [] + assert st.geak_pending["revalidation_task_id"] == tracked.task_id + assert st.geak_result.get("revalidation_status") != "fallback_failed" + + +@pytest.mark.asyncio +async def test_close_entry_settles_pending_after_phase_boundary_cancel(coordinator) -> None: + """CLOSE must settle a dangling ``awaiting_rebench`` slot. + + The SWEEP->CLOSE transition already cancels the queued rebench, so the CLOSE + sequencer finds nothing left to cancel and must still settle the slot. + """ + c = coordinator + st = c.shared_state + st.kernel_optimizer = "geak" + st.geak_result = {"status": "ok", "accepted_config": {"flags": "--foo", "env": ""}} + + rebench = await c.tasks.create( + kind="explore", + params=_geak_rebench_params(), + idempotency_key=gr.geak_revalidate_idempotency_key(0), + task_id="cancelled-by-phase-boundary", + ) + st.geak_pending = {"status": "awaiting_rebench", "revalidation_task_id": rebench.task_id} + st.resume_pending_revalidation = True + + cancelled = await c.tasks.cancel_queued_not_allowed( + allowed_kinds=ps.PHASE_ALLOWED_ACTIONS[ps.PHASE_CLOSE], + reason="phase_transition:SWEEP->CLOSE", + spare_queued=lambda _tid, kind, params: gr.spare_geak_rebench_on_phase_transition( + target_phase=ps.PHASE_CLOSE, + kind=kind, + params=params, + ), + ) + assert rebench.task_id in cancelled + + settled = await gr.settle_dangling_geak_pending(c.tasks, st, reason="close_sequence") + + assert settled is True + assert st.geak_pending["status"] == "rebench_cancelled" + assert st.resume_pending_revalidation is False + + +@pytest.mark.asyncio +async def test_advance_into_close_settles_pending_end_to_end(coordinator) -> None: + """Real entry order: the transition cancels the rebench, CLOSE settles the slot.""" + c = coordinator + st = c.shared_state + st.phase = ps.PHASE_EXPLORE + st.kernel_optimizer = "geak" + st.geak_result = {"status": "ok", "accepted_config": {"flags": "--foo", "env": ""}} + st.set_stop_reason("target_reached") + + rebench = await c.tasks.create( + kind="explore", + params=_geak_rebench_params(), + idempotency_key=gr.geak_revalidate_idempotency_key(0), + task_id="rebench-into-close", + ) + st.geak_pending = {"status": "awaiting_rebench", "revalidation_task_id": rebench.task_id} + st.resume_pending_revalidation = True + + await c._advance_phase_if_needed() + + assert st.phase == ps.PHASE_CLOSE + assert (await c.tasks.get(rebench.task_id)).state == "cancelled" + assert st.geak_pending["status"] == "rebench_cancelled" + assert st.resume_pending_revalidation is False + + +@pytest.mark.asyncio +async def test_settle_waits_while_rebench_still_running(coordinator) -> None: + """Settling is state-driven: a running rebench can still deliver, so wait.""" + c = coordinator + st = c.shared_state + st.kernel_optimizer = "geak" + st.geak_result = {"status": "ok"} + + rebench = await c.tasks.create( + kind="explore", + params=_geak_rebench_params(), + idempotency_key=gr.geak_revalidate_idempotency_key(0), + task_id="running-rebench", + ) + await c.tasks.transition(rebench.task_id, "running") + st.geak_pending = {"status": "awaiting_rebench", "revalidation_task_id": rebench.task_id} + + settled = await gr.settle_dangling_geak_pending(c.tasks, st, reason="close_sequence") + + assert settled is False + assert st.geak_pending["status"] == "awaiting_rebench" + + +@pytest.mark.asyncio +async def test_close_drain_cancels_running_rebench_and_settles(coordinator) -> None: + """CLOSE writes reports only, so a running rebench is stopped, not awaited. + + Leaving it running would hold the GPU lane against the post-opt roofline and + could still rewrite current_best after the report was generated. + """ + c = coordinator + st = c.shared_state + st.kernel_optimizer = "geak" + st.geak_result = {"status": "ok"} + + rebench = await c.tasks.create( + kind="explore", + params=_geak_rebench_params(), + idempotency_key=gr.geak_revalidate_idempotency_key(0), + task_id="running-into-close", + ) + await c.tasks.transition(rebench.task_id, "running") + st.geak_pending = {"status": "awaiting_rebench", "revalidation_task_id": rebench.task_id} + st.resume_pending_revalidation = True + + await c._drain_geak_rebench_for_close() + + assert (await c.tasks.get(rebench.task_id)).state == "cancelled" + assert st.geak_pending["status"] == "rebench_cancelled" + assert st.resume_pending_revalidation is False + + +@pytest.mark.asyncio +async def test_prune_drain_leaves_running_rebench_alone(coordinator) -> None: + """A backlog drain only clears queued work; running rebench keeps going.""" + c = coordinator + running = await c.tasks.create( + kind="explore", + params=_geak_rebench_params(), + idempotency_key=gr.geak_revalidate_idempotency_key(0), + task_id="running-during-prune", + ) + await c.tasks.transition(running.task_id, "running") + + cancelled = await gr.cancel_geak_rebench_tasks(c.tasks, reason="prune_branch") + + assert cancelled == [] + assert (await c.tasks.get(running.task_id)).state == "running" diff --git a/src/hyperloom/orchestrator/loop/coordinator.py b/src/hyperloom/orchestrator/loop/coordinator.py index db198d228..2f348ee20 100644 --- a/src/hyperloom/orchestrator/loop/coordinator.py +++ b/src/hyperloom/orchestrator/loop/coordinator.py @@ -948,6 +948,7 @@ def router(self) -> IntentRouter: "_derive_close_stop_reason": "phase_close", "_session_integrated_kernel_patch": "phase_close", "_maybe_run_close_post_opt_roofline": "phase_close", + "_drain_geak_rebench_for_close": "phase_close", "_on_enter_close": "phase_close", "_enqueue_runnable_internal_task": "phase_close", "_enqueue_internal_report_task": "phase_close", diff --git a/src/hyperloom/orchestrator/loop/intent_router.py b/src/hyperloom/orchestrator/loop/intent_router.py index dd6ef0cd2..99c6a851e 100644 --- a/src/hyperloom/orchestrator/loop/intent_router.py +++ b/src/hyperloom/orchestrator/loop/intent_router.py @@ -1091,6 +1091,20 @@ async def _handle_prune_branch(self, source: str, intent: Intent) -> None: cancelled = await self._drain_queued_baselines(reason=reason) else: cancelled = await self.tasks.cancel_family([family], reason=reason) + # A pruned explore family can take the GEAK 2b rebench with it; settle the + # slot so KERNEL is not held open waiting on a task that will never run. + if cancelled: + from ..phases.geak_rebench import settle_dangling_geak_pending + + try: + if await settle_dangling_geak_pending( + self.tasks, + self.shared_state, + reason=f"prune_branch:{family}", + ): + self.shared_state.save(self.session_dir) + except Exception: # noqa: BLE001 — prune must not fail on bookkeeping + log.exception("prune_branch: GEAK pending settle failed") await self.bus.append_and_seq( Message.new( source, diff --git a/src/hyperloom/orchestrator/loop/writeback.py b/src/hyperloom/orchestrator/loop/writeback.py index 1af99e071..34b5f9afa 100644 --- a/src/hyperloom/orchestrator/loop/writeback.py +++ b/src/hyperloom/orchestrator/loop/writeback.py @@ -870,9 +870,13 @@ async def _handle_unpromotable_result( any_changed = False params = task.params or {} if task.kind == "explore" and bool(params.get("geak_fallback")): - pending = getattr(self.shared_state, "geak_pending", None) or {} - pending_task_id = str(pending.get("revalidation_task_id") or "") if isinstance(pending, dict) else "" - if not pending_task_id or pending_task_id == task.task_id: + from ..phases.geak_rebench import geak_rebench_should_apply_result + + if geak_rebench_should_apply_result( + self.shared_state, + task, + macro_cycle=int(getattr(self.shared_state, "macro_cycle", 0) or 0), + ): geak_result = ( dict(self.shared_state.geak_result) if isinstance(getattr(self.shared_state, "geak_result", None), dict) @@ -3523,7 +3527,47 @@ async def _promote_explore( prev_best_envs=cb_now.get("extra_envs") or {}, ): decision = "no_material" - if decision == "validated": + pending = getattr(self.shared_state, "geak_pending", None) or {} + pending_tid = ( + str(pending.get("revalidation_task_id") or "") if isinstance(pending, dict) else "" + ) + from ..phases.geak_rebench import geak_rebench_should_apply_result + + macro_cycle = int(getattr(self.shared_state, "macro_cycle", 0) or 0) + pending_status = str(pending.get("status") or "") if isinstance(pending, dict) else "" + if not geak_rebench_should_apply_result( + self.shared_state, task, macro_cycle=macro_cycle + ): + # The slot either names another task or already carries a + # verdict, so this result is orphaned or late. Record it: + # silently dropping a measured rebench is hard to diagnose. + log.warning( + "geak 2b: ignoring %s result from rebench task %s not tracked by " + "geak_pending (pending_task=%s status=%s)", + decision, + task.task_id, + pending_tid or "", + pending_status or "", + ) + try: + await self._record_observation( + "coordinator", + "observation", + { + "kind": "geak_rebench_result_ignored", + "decision": decision, + "task_id": task.task_id, + "idempotency_key": str(task.idempotency_key or ""), + "pending_task_id": pending_tid, + "pending_status": pending_status, + "measured_tput": ( + float(measured) if isinstance(measured, (int, float)) else None + ), + }, + ) + except Exception: # noqa: BLE001 - observation is best-effort + log.exception("geak orphan rebench: observation emit failed") + elif decision == "validated": # Write the headline from the measured orchestrator-harness # rebench: lift current_best + optimization_stack + the # validated gain and clear geak_pending. @@ -4965,8 +5009,9 @@ async def _enqueue_internal_stack_rebench(self, *, reason: str) -> dict[str, Any delta becomes the validated cumulative gain. Tagged ``source=resume_stack_revalidate`` so ``_promote_to_shared_state`` reconciles ``cumulative_gain_validated_stack_len`` + clears - ``resume_pending_revalidation`` from the measured throughput. Idempotent - via a fixed idempotency key. + ``resume_pending_revalidation`` from the measured throughput. GEAK 2b + revalidations are idempotent per macro-cycle via + ``geak_revalidate_idempotency_key``. Args: reason: Human-readable reason stamped on the task params. @@ -5021,10 +5066,16 @@ async def _enqueue_internal_stack_rebench(self, *, reason: str) -> dict[str, Any } if self.shared_state.baseline_config_path: params_ps["config_path"] = self.shared_state.baseline_config_path + from ..phases.geak_rebench import resolve_geak_revalidate_idempotency_key + + idempotency_key = await resolve_geak_revalidate_idempotency_key( + self.tasks, + int(getattr(self.shared_state, "macro_cycle", 0) or 0), + ) task, existing = await self.tasks.create_or_return_existing( kind="explore", params=params_ps, - idempotency_key="geak-revalidate", + idempotency_key=idempotency_key, ) try: from hyperloom.inference_optimizer.breakdown.recorder import instrument diff --git a/src/hyperloom/orchestrator/phases/close.py b/src/hyperloom/orchestrator/phases/close.py index 9133861d2..51136bf3e 100644 --- a/src/hyperloom/orchestrator/phases/close.py +++ b/src/hyperloom/orchestrator/phases/close.py @@ -11,6 +11,7 @@ from datetime import datetime, timezone from typing import Any import logging as _logging +from . import geak_rebench as _geak_rebench from . import machine_state as _phase_state from ..bus.message_bus import Message from ..state.task_registry import Task @@ -160,6 +161,60 @@ async def _maybe_run_close_post_opt_roofline(self) -> None: state = getattr(result, "state", None) log.info("CLOSE step 0: post-opt roofline finished (state=%s)", state) + async def _drain_geak_rebench_for_close(self, *, reason: str = "close_sequence") -> None: + """Stop any GEAK 2b rebench and close its pending slot as the run winds down. + + Shared by both wind-down paths: the CLOSE sequencer and the wall-clock + closing phase. Neither can still turn a rebench into a headline, and a + running one holds the GPU lane against the post-opt roofline, so the task + is cancelled and the slot settled. + + Args: + reason: Stamped on the cancellations and the settled slot. + """ + try: + dropped = await _geak_rebench.cancel_geak_rebench_tasks( + self.tasks, + reason=reason, + include_running=True, + ) + if dropped: + log.info( + "%s: cancelled %d in-flight GEAK rebench task(s)", + reason, + len(dropped), + ) + settled = await _geak_rebench.settle_dangling_geak_pending( + self.tasks, + self.shared_state, + reason=reason, + ) + if not (dropped or settled): + return + if settled: + log.info("%s: settled a GEAK revalidation slot that can no longer land", reason) + try: + self.shared_state.save(self.session_dir) + except Exception: # noqa: BLE001 + log.exception("%s: geak_pending settle save failed", reason) + await self._record_observation( + "coordinator", + "observation", + { + "kind": "geak_rebench_close_drain", + "reason": reason, + "cancelled_task_ids": dropped, + "pending_settled": bool(settled), + }, + ) + except Exception: # noqa: BLE001 — wind-down must proceed even if this fails + log.exception("%s: GEAK rebench drain failed (non-fatal)", reason) + await self._record_close_step( + "geak_rebench_drain", + status="failed", + detail="see log; geak_pending may remain awaiting_rebench", + ) + async def _on_enter_close(self, *, from_phase: str) -> None: """CLOSE sequencer (fixed order): post-opt roofline → fact_finalize → report → session_breakdown → langfuse flush → artifact_package → ndjson_drain (no-op) → mark close_sequence_done + stop_reason. Best-effort steps; final done step always runs. The ``CLOSE step N`` log labels are non-contiguous for historical reasons. @@ -167,6 +222,7 @@ async def _on_enter_close(self, *, from_phase: str) -> None: from_phase: The phase being left, used only for logging. """ log.info("CLOSE entered (from=%s); starting 7-step close sequence", from_phase or "") + await self._drain_geak_rebench_for_close() await self._record_close_step("sequencer_started", status="running") # stop_reason must persist before step 2's breakdown (collector derives it from state.json); fill only when blank. @@ -728,6 +784,12 @@ async def _enter_closing_phase(self, *, grace_sec: float) -> float: "closing_phase: cancel of queued tasks failed (non-fatal)", ) + # The wall-clock path never reaches ``_on_enter_close``, so it owns the + # same wind-down: a rebench left running would keep writing back during + # the grace window, and an unsettled slot makes the report promise a + # rebench whose task the loop above has already cancelled. + await self._drain_geak_rebench_for_close(reason="closing_phase") + idempotency_key = f"closing-report-{int(closing_started)}-{uuid.uuid4().hex[:6]}" task, _existing = await self.tasks.create_or_return_existing( kind="report", diff --git a/src/hyperloom/orchestrator/phases/geak_rebench.py b/src/hyperloom/orchestrator/phases/geak_rebench.py new file mode 100644 index 000000000..c8ac1ff9a --- /dev/null +++ b/src/hyperloom/orchestrator/phases/geak_rebench.py @@ -0,0 +1,233 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT + +"""GEAK same-harness revalidation task identity and phase-boundary policy.""" + +from __future__ import annotations + +from typing import Any, TYPE_CHECKING + +from .machine_state import PHASE_CLOSE + +if TYPE_CHECKING: + from ..state.task_registry import Task, TaskRegistry + +LEGACY_GEAK_REVALIDATE_PLACEHOLDER = "geak-revalidate" + +# ``geak_pending.status`` values that record a closed verdict. A result arriving +# against one of these is late or orphaned and must not reopen the slot. +SETTLED_PENDING_STATUSES: frozenset[str] = frozenset({"rebench_cancelled", "rebench_unavailable"}) + +# Fresh keys a single macro-cycle may mint. Each cancelled attempt burns one, so +# this bounds how often a prune/cancel loop can re-dispatch the same rebench. +MAX_REBENCH_ATTEMPTS_PER_CYCLE = 4 + + +def geak_revalidate_idempotency_key(macro_cycle: int, attempt: int = 0) -> str: + """Return the idempotency key for a GEAK 2b rebench task. + + ``attempt`` distinguishes retries within one macro-cycle: reusing the key of + a settled row would hand that row back from + ``create_or_return_existing`` and be read as ``rebench_unavailable``. + """ + base = f"geak-revalidate-c{macro_cycle}" + return base if attempt <= 0 else f"{base}-r{attempt}" + + +def geak_revalidation_placeholder_keys(macro_cycle: int) -> frozenset[str]: + """Placeholder ids written to ``geak_pending`` before the task row exists.""" + return frozenset({LEGACY_GEAK_REVALIDATE_PLACEHOLDER, geak_revalidate_idempotency_key(macro_cycle)}) + + +def is_geak_same_harness_rebench_task(kind: str, params: dict[str, Any] | None) -> bool: + """True when a queued/running task is the orchestrator GEAK 2b revalidation explore.""" + payload = params if isinstance(params, dict) else {} + return ( + str(kind or "").strip() == "explore" + and str(payload.get("source") or "") == "resume_stack_revalidate" + and bool(payload.get("geak_fallback")) + ) + + +def spare_geak_rebench_on_phase_transition(*, target_phase: str, kind: str, params: dict[str, Any]) -> bool: + """Return True to leave a queued GEAK rebench alive across a phase boundary. + + Deny-list rather than allow-list: only ``CLOSE`` kills the rebench, every + other target spares it. The window that actually matters is KERNEL through + SWEEP (plus the SWEEP re-loop back into EXPLORE), so an allow-list would be + tighter — but the phase set changes over time and a missing entry silently + reintroduces the audit-only bug this whole path exists to prevent, whereas a + surplus entry costs at most one wasted bench. + + Correctness of the surplus is owned elsewhere: a rebench that outlives its + macro-cycle is refused by :func:`geak_rebench_should_apply_result`, because + the slot it would have to be tracked in has moved on. So this predicate only + decides whether the task keeps running, never whether its result counts. + """ + if (target_phase or "").strip().upper() == PHASE_CLOSE: + return False + return is_geak_same_harness_rebench_task(kind, params) + + +def geak_rebench_tracks_pending_task( + pending_task_id: str, + task: Task, + *, + macro_cycle: int, +) -> bool: + """True when ``geak_pending.revalidation_task_id`` tracks this rebench task. + + The slot normally holds a task id. It holds a key instead only inside the + reservation window — the slot is published before the task row exists so the + phase guard can see a pending revalidation — and, on state written before + keys were cycle-scoped, the bare legacy key. Both are matched through + :func:`geak_revalidation_placeholder_keys`, which requires the task to carry + that same key, so a rebench from another macro-cycle cannot claim the slot. + """ + tracked = str(pending_task_id or "").strip() + if not tracked: + return False + key = str(task.idempotency_key or "") + if tracked in {task.task_id, key}: + return True + placeholders = geak_revalidation_placeholder_keys(macro_cycle) + return tracked in placeholders and key in placeholders + + +def geak_rebench_should_apply_result(state: Any, task: Task, *, macro_cycle: int) -> bool: + """True when a finished 2b task may mutate ``geak_pending`` / ``geak_result``. + + Ordered so the closed verdicts win first: + + * a settled slot rejects everything — the verdict is already recorded; + * a slot naming a task accepts only that task, so orphans are ignored; + * ``awaiting_rebench`` with no id yet is the window between recording the + candidate and publishing the task id, so it accepts; + * an empty slot accepts only while ``resume_pending_revalidation`` marks a + resume revalidation, which owns no candidate slot of its own. + """ + pending = getattr(state, "geak_pending", None) or {} + if not isinstance(pending, dict): + pending = {} + status = str(pending.get("status") or "").strip().lower() + if status in SETTLED_PENDING_STATUSES: + return False + tracked = str(pending.get("revalidation_task_id") or "").strip() + if tracked: + return geak_rebench_tracks_pending_task(tracked, task, macro_cycle=macro_cycle) + if status == "awaiting_rebench": + return True + return bool(getattr(state, "resume_pending_revalidation", False)) + + +async def find_inflight_geak_rebench_task(tasks: TaskRegistry) -> Task | None: + """Return the oldest queued/running GEAK same-harness rebench, if any.""" + queued_fn = getattr(tasks, "queued", None) + running_fn = getattr(tasks, "running", None) + if not callable(queued_fn) or not callable(running_fn): + return None + for pool in (await queued_fn(), await running_fn()): + for task in pool: + if is_geak_same_harness_rebench_task(task.kind, task.params): + return task + return None + + +async def cancel_geak_rebench_tasks( + tasks: TaskRegistry, + *, + reason: str, + include_running: bool = False, +) -> list[str]: + """Cancel in-flight GEAK 2b rebench tasks. + + Args: + tasks: The task registry. + reason: Stamped onto each cancellation's history evidence. + include_running: Also cancel a rebench already executing. CLOSE sets + this: the phase only writes reports, so a running rebench holds the + GPU lane against post-opt roofline and its result can no longer be + consumed. A backlog drain (prune) leaves running work alone. + + Returns: + The cancelled task ids. + """ + queued_fn = getattr(tasks, "queued", None) + running_fn = getattr(tasks, "running", None) + if not callable(queued_fn): + return [] + pools = [await queued_fn()] + if include_running and callable(running_fn): + pools.append(await running_fn()) + cancelled: list[str] = [] + for pool in pools: + for task in pool: + if not is_geak_same_harness_rebench_task(task.kind, task.params): + continue + await tasks.transition(task.task_id, "cancelled", evidence={"reason": reason}) + cancelled.append(task.task_id) + return cancelled + + +async def resolve_geak_revalidate_idempotency_key(tasks: TaskRegistry, macro_cycle: int) -> str: + """Pick the key for the next 2b rebench in ``macro_cycle``. + + Steps past attempts whose row already settled, because reusing their key + returns that terminal row instead of dispatching. Stops at + ``MAX_REBENCH_ATTEMPTS_PER_CYCLE`` so a cancel loop cannot dispatch forever. + """ + lookup = getattr(tasks, "find_by_idempotency_key", None) + if not callable(lookup): + return geak_revalidate_idempotency_key(macro_cycle) + last = geak_revalidate_idempotency_key(macro_cycle) + for attempt in range(MAX_REBENCH_ATTEMPTS_PER_CYCLE): + last = geak_revalidate_idempotency_key(macro_cycle, attempt) + row = await lookup(last) + if row is None or row.state in {"queued", "running"}: + return last + return last + + +async def settle_dangling_geak_pending(tasks: TaskRegistry, state: Any, *, reason: str) -> bool: + """Settle ``geak_pending`` once no rebench can still land. + + Driven by state rather than by what a caller just cancelled: the phase + boundary into CLOSE already cancels the queued rebench, so the CLOSE + sequencer finds nothing left to cancel yet still has to close the slot. A + rebench still in flight is left alone so its result can arrive. + + Only the verdict fields change. The candidate's self-reported numbers are + what the report uses to say *what* was dropped, so they are kept; the id of + a task that will never land is not. + """ + pending = getattr(state, "geak_pending", None) or {} + if not isinstance(pending, dict): + return False + if str(pending.get("status") or "").strip().lower() != "awaiting_rebench": + return False + if await find_inflight_geak_rebench_task(tasks) is not None: + return False + settled = dict(pending) + settled["status"] = "rebench_cancelled" + settled["revalidation_error"] = str(reason)[:500] + settled.pop("revalidation_task_id", None) + state.geak_pending = settled + state.resume_pending_revalidation = False + return True + + +__all__ = [ + "LEGACY_GEAK_REVALIDATE_PLACEHOLDER", + "MAX_REBENCH_ATTEMPTS_PER_CYCLE", + "SETTLED_PENDING_STATUSES", + "cancel_geak_rebench_tasks", + "find_inflight_geak_rebench_task", + "geak_rebench_should_apply_result", + "geak_rebench_tracks_pending_task", + "geak_revalidate_idempotency_key", + "geak_revalidation_placeholder_keys", + "is_geak_same_harness_rebench_task", + "resolve_geak_revalidate_idempotency_key", + "settle_dangling_geak_pending", + "spare_geak_rebench_on_phase_transition", +] diff --git a/src/hyperloom/orchestrator/phases/kernel.py b/src/hyperloom/orchestrator/phases/kernel.py index 3b5532d36..1f734f311 100644 --- a/src/hyperloom/orchestrator/phases/kernel.py +++ b/src/hyperloom/orchestrator/phases/kernel.py @@ -18,6 +18,7 @@ from datetime import datetime, timezone from pathlib import Path from typing import Any, Callable +from . import geak_rebench as _geak_rebench from . import machine_state as _phase_state from ..kernel import collective_recovery as _collective_recovery from ..actions.stop_attribution import stopped_by_the_run_class @@ -46,12 +47,6 @@ # on a developer box that happens to have the real checkout mounted. _CONTAINER_AITER_CONFIG_DIR = Path("/sgl-workspace/aiter/aiter/configs") -# Idempotency key of the same-harness GEAK rebench enqueued by -# ``_enqueue_internal_stack_rebench``. Doubles as the placeholder that reserves -# ``geak_pending`` before the task row exists, so the phase guard already sees a -# pending revalidation while the enqueue is in flight. -_GEAK_REVALIDATE_IDEMPOTENCY_KEY = "geak-revalidate" - def _collective_comm_share(state: Any) -> tuple[float | None, str]: """Return the communication share gating the lane, and its provenance. @@ -810,15 +805,32 @@ def _finish_skip(result: dict[str, Any]) -> None: async def _enqueue_geak_revalidation(*, reason: str) -> bool: """Enqueue and persist the rebench that keeps a GEAK win pending.""" # Reserve the pending slot BEFORE the task exists. The rebench runs - # as an ``explore`` task, a kind no later phase allows, so the phase - # boundary cancels it on sight; ``kernel_work_pending`` is what holds - # KERNEL open until the rebench lands. Publishing the reservation - # after enqueue left a window where the guard saw no revalidation id, - # let KERNEL exit, and the cancel swept the freshly queued task — - # stranding a validated win as audit-only. + # as an ``explore`` task; non-CLOSE phase boundaries may spare it + # via ``spare_geak_rebench_on_phase_transition`` so the rebench can + # finish after KERNEL winds down. Publishing the reservation after + # enqueue left a window where KERNEL could exit and cancel the row. + cycle = int(getattr(state, "macro_cycle", 0) or 0) + placeholder_keys = _geak_rebench.geak_revalidation_placeholder_keys(cycle) + inflight = await _geak_rebench.find_inflight_geak_rebench_task(self.tasks) + if inflight is not None and inflight.state in {"queued", "running"}: + pending = dict(state.geak_pending) if isinstance(state.geak_pending, dict) else {} + pending["status"] = "awaiting_rebench" + pending["revalidation_task_id"] = inflight.task_id + pending.pop("revalidation_error", None) + state.geak_pending = pending + state.save(self.session_dir) + log.info( + "geak: revalidation already in flight (%s); skipping duplicate enqueue", + inflight.task_id, + ) + return True + reserved = dict(state.geak_pending) if isinstance(state.geak_pending, dict) else {} reserved["status"] = "awaiting_rebench" - reserved.setdefault("revalidation_task_id", _GEAK_REVALIDATE_IDEMPOTENCY_KEY) + if not str(reserved.get("revalidation_task_id") or "").strip(): + reserved["revalidation_task_id"] = _geak_rebench.geak_revalidate_idempotency_key( + cycle + ) reserved.pop("revalidation_error", None) state.geak_pending = reserved state.save(self.session_dir) @@ -840,8 +852,8 @@ async def _enqueue_geak_revalidation(*, reason: str) -> bool: return True pending["status"] = "rebench_unavailable" - # Drop the reservation placeholder so no stale id outlives the slot. - if pending.get("revalidation_task_id") == _GEAK_REVALIDATE_IDEMPOTENCY_KEY: + # Drop reservation placeholders (current + legacy) so no stale id outlives the slot. + if pending.get("revalidation_task_id") in placeholder_keys: pending.pop("revalidation_task_id", None) pending["revalidation_error"] = str( (summary or {}).get("reason") or f"task settled before dispatch ({task_state or 'unknown'})" diff --git a/src/hyperloom/orchestrator/phases/machine.py b/src/hyperloom/orchestrator/phases/machine.py index 3e017d405..150498c3b 100644 --- a/src/hyperloom/orchestrator/phases/machine.py +++ b/src/hyperloom/orchestrator/phases/machine.py @@ -7,6 +7,7 @@ from __future__ import annotations import logging as _logging from typing import Any +from . import geak_rebench as _geak_rebench from . import machine_state as _phase_state from ..bus.message_bus import Message from ..prompts import write_prompt_snapshot as _write_prompt_snapshot @@ -312,9 +313,15 @@ async def _advance_phase_if_needed(self) -> None: ): state.no_gain_cycle_streak = int(evidence.get("no_gain_cycle_streak_effective", 0) or 0) allowed_kinds = _phase_state.PHASE_ALLOWED_ACTIONS.get(target, frozenset()) + target_phase = str(target or "").strip().upper() cancelled = await self.tasks.cancel_queued_not_allowed( allowed_kinds=allowed_kinds, reason=f"phase_transition:{str(prior or '').strip().upper()}->{target}", + spare_queued=lambda _task_id, kind, params: _geak_rebench.spare_geak_rebench_on_phase_transition( + target_phase=target_phase, + kind=kind, + params=params, + ), ) if cancelled: log.info( diff --git a/src/hyperloom/orchestrator/state/_shared_state/render.py b/src/hyperloom/orchestrator/state/_shared_state/render.py index 3445954c6..2b7071b1c 100644 --- a/src/hyperloom/orchestrator/state/_shared_state/render.py +++ b/src/hyperloom/orchestrator/state/_shared_state/render.py @@ -97,12 +97,17 @@ def to_mission_summary(self, *, now: datetime | None = None) -> str: if bool(getattr(self, "resume_pending_revalidation", False)) else "" ) - geak_pending_tag = ( - " ⚠ geak candidate awaiting main-flow rebench — NOT in headline until validated" + geak_pending_status = ( + str(self.geak_pending.get("status") or "") if isinstance(getattr(self, "geak_pending", None), dict) - and self.geak_pending.get("status") == "awaiting_rebench" else "" ) + if geak_pending_status == "awaiting_rebench": + geak_pending_tag = " ⚠ geak candidate awaiting main-flow rebench — NOT in headline until validated" + elif geak_pending_status in {"rebench_cancelled", "rebench_unavailable"}: + geak_pending_tag = f" ⚠ geak candidate dropped unvalidated ({geak_pending_status})" + else: + geak_pending_tag = "" from hyperloom.inference_optimizer import framework_registry lines = [ diff --git a/src/hyperloom/orchestrator/state/task_registry.py b/src/hyperloom/orchestrator/state/task_registry.py index 74e987f4a..c4c4ebf37 100644 --- a/src/hyperloom/orchestrator/state/task_registry.py +++ b/src/hyperloom/orchestrator/state/task_registry.py @@ -19,7 +19,7 @@ import json import uuid -from collections.abc import Iterable +from collections.abc import Callable, Iterable from dataclasses import dataclass, field from datetime import datetime, timezone from typing import Any @@ -27,6 +27,7 @@ from hyperloom.common.timeutil import now_iso from hyperloom.orchestrator.bus.storage.connection import SqliteConnection +SpareQueuedFn = Callable[[str, str, dict[str, Any]], bool] TASK_STATES = ( "queued", @@ -424,6 +425,22 @@ async def record_progress( (json.dumps(history), task_id), ) + async def find_by_idempotency_key(self, idempotency_key: str) -> Task | None: + """Return the task registered under ``idempotency_key``, or None. + + Args: + idempotency_key: The UNIQUE key to look up. + + Returns: + Task | None: The matching task in any state, or ``None`` when the + key has never been used. + """ + row = await self.db.fetchone( + "SELECT * FROM tasks WHERE idempotency_key=?", + (idempotency_key,), + ) + return None if row is None else Task.from_row(row) + async def queued(self) -> list[Task]: """Return all queued tasks ordered oldest-first. @@ -680,17 +697,43 @@ async def cancel_queued_not_allowed( *, allowed_kinds: set[str] | frozenset[str], reason: str, + spare_queued: SpareQueuedFn | None = None, ) -> list[str]: - """Bulk-cancel queued tasks whose kind is not allowed at a phase boundary.""" + """Bulk-cancel queued tasks whose kind is not allowed at a phase boundary. + + Args: + allowed_kinds: Task kinds permitted in the phase being entered. + reason: Stamped onto each cancellation's history evidence. + spare_queued: Optional ``(task_id, kind, params) -> bool`` hook. + When it returns ``True`` the queued row is left untouched even + though its kind is outside ``allowed_kinds``. Callers use this + for narrowly scoped cross-phase work (for example GEAK 2b + rebench survives into SWEEP but not into CLOSE). + + Returns: + The task ids that were cancelled (empty when none matched). + """ allowed = {str(kind or "").strip() for kind in allowed_kinds if str(kind or "").strip()} cancelled: list[str] = [] async with self.db.transaction() as cur: - cur.execute("SELECT task_id, kind, history FROM tasks WHERE state='queued'") - rows = [(r["task_id"], r["kind"], r["history"]) for r in cur.fetchall()] + cur.execute("SELECT task_id, kind, params, history FROM tasks WHERE state='queued'") + rows = [(r["task_id"], r["kind"], r["params"], r["history"]) for r in cur.fetchall()] now = _now_iso() - for task_id, kind, history_json in rows: + for task_id, kind, params_json, history_json in rows: if str(kind or "").strip() in allowed: continue + try: + params = json.loads(params_json) if params_json else {} + except json.JSONDecodeError: + params = {} + if not isinstance(params, dict): + params = {} + if spare_queued is not None and spare_queued( + str(task_id or "").strip(), + str(kind or "").strip(), + params, + ): + continue history = json.loads(history_json) history.append( {