feat: Add wait() and join()

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

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()