diff --git a/alphatrion/experiment/base.py b/alphatrion/experiment/base.py index 0089b5d..60e9d3f 100644 --- a/alphatrion/experiment/base.py +++ b/alphatrion/experiment/base.py @@ -168,6 +168,9 @@ class Experiment(ABC): "_total_runs_counter", # The end status, None, Err or Cancelled. "_end_status", + # True once wait() is called; the experiment auto-completes when all + # runs are finished. + "_waiting", "_stopped", "_received_signal", "_signal_task", @@ -184,10 +187,9 @@ def __init__(self, config: ExperimentConfig | None = None): self._early_stopping_counter = 0 self._total_runs_counter = 0 self._end_status = None - # if experiment starts to wait, it will auto stop when the runs + # if wait() is called, the experiment will auto stop when the runs # are all finished. - self._start_waiting = False - self._end_status = None + self._waiting = False self._stopped = asyncio.Event() self._received_signal: int | None = None self._signal_task: asyncio.Task | None = None @@ -382,12 +384,21 @@ def _timeout(self) -> int | None: return timeout - # Make sure you have termination condition, either by timeout or by calling cancel() - # Before we have logic like once all the tasks are done, we'll call the cancel() - # automatically, however, this is unpredictable because some tasks may wait for - # external events, so we leave it to the user to decide when to stop the experiment. + # wait blocks until all the runs are finished, then the experiment is + # auto stopped. Use this when you have launched all the runs and want to + # wait for them to complete. async def wait(self): - self._start_waiting = True + self._waiting = True + if len(self._runs) == 0: + self.done() + await self._context.wait() + + # wait_until_done blocks until the experiment is terminated, either by + # timeout or by calling done()/cancel() (e.g. from a signal handler). + # Unlike wait(), it does NOT auto stop when all runs are finished, so the + # experiment keeps running even with no active runs. Use this for + # long-running experiments where the termination condition is external. + async def wait_until_done(self): await self._context.wait() def is_done(self) -> bool: @@ -488,9 +499,9 @@ def _post_run(self, run: Run): ): self.done() - # If the experiment starts to wait and all runs are finished, + # If the experiment is waiting and all runs are finished, # we can stop the experiment. - if self._start_waiting and len(self._runs) == 0: + if self._waiting and len(self._runs) == 0: self.done() @classmethod diff --git a/tests/unit/experiment/test_experiment.py b/tests/unit/experiment/test_experiment.py index 667ca6a..a544737 100644 --- a/tests/unit/experiment/test_experiment.py +++ b/tests/unit/experiment/test_experiment.py @@ -227,7 +227,7 @@ async def test_experiment_with_resume(): @pytest.mark.asyncio -async def test_experiment_with_wait(): +async def test_experiment_with_join(): init( team_id=uuid.uuid4(), user_id=uuid.uuid4(), @@ -252,6 +252,63 @@ async def fake_work(): assert exp_obj.status == Status.COMPLETED +@pytest.mark.asyncio +async def test_experiment_join_with_no_runs(): + """join() must auto-complete immediately when there are no active runs. + Without any runs, no _post_run callback fires, so join() would block + forever if it did not complete on its own.""" + init( + team_id=uuid.uuid4(), + user_id=uuid.uuid4(), + org_id=uuid.uuid4(), + ) + + exp_id = None + async with CraftExperiment.start(name="first-experiment") as exp: + exp_id = current_exp_id.get() + + # No runs launched; join() must return promptly instead of hanging. + await asyncio.wait_for(exp.wait(), timeout=3) + assert exp.is_done() + + exp_obj = exp._runtime.metadb.get_experiment(experiment_id=exp_id) + assert exp_obj.status == Status.COMPLETED + + +@pytest.mark.asyncio +async def test_experiment_with_wait_until_done(): + """wait_until_done() must NOT auto-complete when all runs finish; it blocks until the + experiment is terminated externally (here, by the timeout).""" + init( + team_id=uuid.uuid4(), + user_id=uuid.uuid4(), + org_id=uuid.uuid4(), + ) + + async def fake_work(): + await asyncio.sleep(1) + + exp_id = None + async with CraftExperiment.start( + name="first-experiment", + config=experiment.ExperimentConfig(max_execution_seconds=3), + ) as exp: + exp_id = current_exp_id.get() + start_time = datetime.now() + + exp.run(fake_work) + + await exp.wait_until_done() + # The run finishes after ~1s, but wait_until_done() keeps blocking + # until the timeout at ~3s instead of auto-completing when the run + # drains. + assert datetime.now() - start_time >= timedelta(seconds=3) + assert len(exp._runs) == 0 + + exp_obj = exp._runtime.metadb.get_experiment(experiment_id=exp_id) + assert exp_obj.status == Status.COMPLETED + + @pytest.mark.asyncio async def test_create_experiment_with_run(): team_id = uuid.uuid4() @@ -338,7 +395,10 @@ async def test_create_experiment_with_max_execution_seconds(): name="first-experiment", config=experiment.ExperimentConfig(max_execution_seconds=2), ) as exp: - await exp.wait() + # No runs launched; the experiment must terminate on the timeout, so + # use wait_until_done() rather than wait() (which would auto-complete + # immediately with zero active runs). + await exp.wait_until_done() assert exp.is_done() exp_obj = exp._get_obj()