diff --git a/statek/executors/utils.py b/statek/executors/utils.py index 1147cbf..865ed7f 100644 --- a/statek/executors/utils.py +++ b/statek/executors/utils.py @@ -1001,6 +1001,20 @@ async def wrapped(job: Job, *args: Any, **kwargs: Any) -> Any: return wrapped +def _release_ready_continuation(job: Job) -> bool: + """Return whether a job may run, releasing a satisfied continuation once.""" + awaited_result = job.awaited_result + if awaited_result is None: + return True + if not awaited_result.check_condition(): + return False + + job.awaited_result = None + if job.status == JobStatus.SUSPENDED: + job.set_status(JobStatus.STARTED) + return True + + @_with_statek_job_context async def run_job_step(job: Job, provider: str = None) -> bool: """ @@ -1028,6 +1042,10 @@ async def run_job_step(job: Job, provider: str = None) -> bool: if job.status == JobStatus.DONE: return True + # Keep waiting continuations dormant until their recorded future is ready. + if not _release_ready_continuation(job): + return False + # Build shared local_context for all context-aware functions in this step local_context = dict(job.py_env.local_state) if job.py_env.local_state else {} @@ -1378,23 +1396,17 @@ def process_agent_events( def unsuspend_jobs(): """ - Review continuation conditions of suspended jobs and change their status to STARTED - where the continuation criteria are satisfied. + Review waiting jobs and release continuations whose criteria are satisfied. - This function finds all jobs with status SUSPENDED, checks if their awaited_result - condition is satisfied, and changes their status to STARTED if the condition is met. + Ordinary suspended jobs transition to STARTED when released. Warmup jobs retain + WARMING_UP so execution can continue in the active warmup block. - Note: This implementation might be slow for a large number of suspended jobs. + Note: This implementation might be slow for a large number of waiting jobs. Future versions will introduce a more robust Notifier engine and job expiration conditions. """ - # Find all suspended jobs - suspended_jobs = db0.find(Job, JobStatus.SUSPENDED) - for job in suspended_jobs: - # Check if the job has an awaited_result and if its condition is satisfied - condition_met = job.awaited_result.check_condition() - if job.awaited_result is not None and condition_met: - # Change status from SUSPENDED to STARTED - job.set_status(JobStatus.STARTED) + waiting_jobs = db0.find(Job, [JobStatus.SUSPENDED, JobStatus.WARMING_UP]) + for job in waiting_jobs: + _release_ready_continuation(job) def handle_critical_error(error: Optional[Exception] = None) -> None: @@ -1477,8 +1489,16 @@ async def run_jobs_loop(max_concurrency: int = 100, provider: str = None, available_capacity = max_concurrency - len(pending_tasks) start_jobs_func(available_capacity) - ready_or_started_jobs = db0.filter(lambda found_job: found_job not in pending_tasks, - db0.find(Job, [JobStatus.READY, JobStatus.WARMING_UP, JobStatus.STARTED])) + ready_or_started_jobs = db0.filter( + lambda found_job: ( + found_job not in pending_tasks + and not ( + found_job.status == JobStatus.WARMING_UP + and found_job.awaited_result is not None + ) + ), + db0.find(Job, [JobStatus.READY, JobStatus.WARMING_UP, JobStatus.STARTED]), + ) # Make sure not to exceed max_concurrency if len(pending_tasks) < max_concurrency: for job in ready_or_started_jobs: diff --git a/tests/test_run_job_step.py b/tests/test_run_job_step.py index 7d889e3..248abe9 100644 --- a/tests/test_run_job_step.py +++ b/tests/test_run_job_step.py @@ -216,9 +216,12 @@ async def test_resumes_from_future_error_line( assert job.next_instr_num == 2 # suspended at idx 2 assert job.warmup_block_num is None # block did not advance (None means block 0) - # Make the future ready and replace it in env - future_ready = create_future_ready(99) - job.py_env.local_state['future_val'] = future_ready + # Resolve the exact future recorded on the job. + future_not_ready.deps.value = 99 + future_not_ready.set_complement_functions( + complement=_fetch_result_from_deps, + condition=_check_condition_true, + ) # Second run: resumes from idx 2, skipping idx 0 and 1. # counter must stay 1 (not incremented again). diff --git a/tests/test_warmup_blocks.py b/tests/test_warmup_blocks.py index cbe5d8b..9e399d3 100644 --- a/tests/test_warmup_blocks.py +++ b/tests/test_warmup_blocks.py @@ -308,9 +308,12 @@ async def test_run_job_step_multiple_blocks_continuation_after_future_error( # warmup_block_num should NOT advance since block didn't complete assert job.warmup_block_num == 1 - # Make the future ready - future_ready = create_future_ready(42) - job.py_env.local_state['future_val'] = future_ready + # Resolve the exact future recorded on the job. + future_not_ready.deps.value = 42 + future_not_ready.set_complement_functions( + complement=_fetch_result_from_deps, + condition=_check_condition_true, + ) # Resume second block - should complete and advance result3 = await run_job_step(job) diff --git a/tests/test_warmup_wait_scheduling.py b/tests/test_warmup_wait_scheduling.py new file mode 100644 index 0000000..0f49ac4 --- /dev/null +++ b/tests/test_warmup_wait_scheduling.py @@ -0,0 +1,170 @@ +"""Regression tests for scheduling warmups that await a FutureResult.""" +# pylint: disable=no-member + +import asyncio +from contextlib import suppress +from dataclasses import dataclass + +import dbzero as db0 +import pytest + +from statek.exceptions import FutureError +from statek.executors.chat_log_item import WarmupLogItem +from statek.executors.job import Job, JobStatus +from statek.executors.utils import run_job_step, run_jobs_loop, unsuspend_jobs +from statek.future import FutureResult + + +@db0.memo +@dataclass +class WaitState: + """Mutable dependency used to control and count future readiness checks.""" + + ready: bool = False + value: str = "done" + condition_checks: int = 0 + + +def _check_wait_state(future_result: FutureResult) -> bool: + future_result.deps.condition_checks += 1 + return future_result.deps.ready + + +def _fetch_wait_state(future_result: FutureResult) -> str: + if not future_result.deps.ready: + raise FutureError(future_result=future_result) + return future_result.deps.value + + +def _controlled_future(wait_state: WaitState) -> FutureResult: + future = FutureResult(deps=wait_state, state_num=0) + future.set_complement_functions( + complement=_fetch_wait_state, + condition=_check_wait_state, + ) + return future + + +def _warmup_items(job: Job) -> list[WarmupLogItem]: + return [item for item in job.chat_log if isinstance(item, WarmupLogItem)] + + +async def _wait_until_done(job: Job) -> None: + while job.status != JobStatus.DONE: + await asyncio.sleep(0.01) + + +class TestWarmupWaitReadiness: + """Waiting warmups remain dormant until their recorded future is ready.""" + + @pytest.mark.asyncio + async def test_direct_calls_do_not_reexecute_unresolved_warmup( + self, job_def_factory, db0_fixture + ): # pylint: disable=unused-argument + """Direct callers cannot bypass the continuation readiness gate.""" + wait_state = WaitState() + future = _controlled_future(wait_state) + job = Job( + job_def=job_def_factory( + warmup_code="attempts = attempts + 1\nresult = future_value" + ), + job_status=JobStatus.READY, + ) + job.py_env.local_state["attempts"] = 0 + job.py_env.local_state["future_value"] = future + + assert await run_job_step(job) is False + initial_history = list(job.chat_log) + initial_console = list(job.py_env.console or []) + + assert await run_job_step(job) is False + assert await run_job_step(job) is False + + assert job.status == JobStatus.WARMING_UP + assert job.awaited_result is future + assert job.next_instr_num == 1 + assert job.py_env.local_state["attempts"] == 1 + assert list(job.chat_log) == initial_history + assert list(job.py_env.console or []) == initial_console + assert wait_state.condition_checks == 2 + + @pytest.mark.asyncio + async def test_scheduler_skips_blocked_warmup_and_checks_release_once( + self, job_def_factory, db0_fixture + ): # pylint: disable=unused-argument + """A blocked warmup does not consume a worker or starve another ready job.""" + wait_state = WaitState() + future = _controlled_future(wait_state) + blocked_job = Job( + job_def=job_def_factory( + warmup_code="attempts = attempts + 1\nresult = future_value\nexit(result)" + ), + job_status=JobStatus.READY, + ) + blocked_job.py_env.local_state["attempts"] = 0 + blocked_job.py_env.local_state["future_value"] = future + assert await run_job_step(blocked_job) is False + + ready_job = Job( + job_def=job_def_factory(warmup_code='exit("ready job completed")'), + job_status=JobStatus.READY, + ) + initial_history_length = len(_warmup_items(blocked_job)) + scheduler_cycles = 0 + repeated_wait_observed = asyncio.Event() + + def record_scheduler_cycle(_available_capacity: int) -> None: + nonlocal scheduler_cycles + scheduler_cycles += 1 + if scheduler_cycles >= 3: + repeated_wait_observed.set() + + loop_task = asyncio.create_task( + run_jobs_loop( + max_concurrency=1, + queue_prefixes=[db0.get_current_prefix().name], + start_jobs_func=record_scheduler_cycle, + ) + ) + try: + await asyncio.wait_for(repeated_wait_observed.wait(), timeout=2) + + assert ready_job.status == JobStatus.DONE + assert blocked_job.status == JobStatus.WARMING_UP + assert blocked_job.awaited_result is future + assert blocked_job.py_env.local_state["attempts"] == 1 + assert len(_warmup_items(blocked_job)) == initial_history_length + + checks_before_release = wait_state.condition_checks + wait_state.ready = True + await asyncio.wait_for(_wait_until_done(blocked_job), timeout=2) + + assert blocked_job.py_env.exit_status == "done" + assert blocked_job.py_env.local_state["attempts"] == 1 + assert blocked_job.awaited_result is None + assert blocked_job.next_instr_num is None + assert wait_state.condition_checks == checks_before_release + 1 + finally: + loop_task.cancel() + with suppress(asyncio.CancelledError): + await loop_task + + def test_ready_suspended_job_is_released_and_clears_awaited_result( + self, job_def_factory, db0_fixture + ): # pylint: disable=unused-argument + """Ordinary suspended jobs retain their existing transition to STARTED.""" + wait_state = WaitState(ready=True) + future = _controlled_future(wait_state) + job = Job( + job_def=job_def_factory(warmup_code=None), + job_status=JobStatus.SUSPENDED, + ) + job.awaited_result = future + job.next_instr_num = 0 + + unsuspend_jobs() + + assert job.status == JobStatus.STARTED + assert job.awaited_result is None + assert job.next_instr_num == 0 + assert wait_state.condition_checks == 1