From 6bc45f6605304969e788fc9cce3e8ee8676eca79 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Tue, 21 Jul 2026 18:49:10 +0100 Subject: [PATCH 1/5] Add wait() Signed-off-by: kerthcet --- README.md | 4 +- alphatrion/experiment/base.py | 28 ++++++++---- .../integration/server/test_graphql_query.py | 4 +- tests/integration/test_craft_experiment.py | 2 +- tests/integration/test_log.py | 8 ++-- tests/integration/test_run_hooks.py | 18 ++++---- tests/unit/experiment/test_experiment.py | 45 ++++++++++++++++--- 7 files changed, 76 insertions(+), 33 deletions(-) diff --git a/README.md b/README.md index f349e74b..eab4477b 100644 --- a/README.md +++ b/README.md @@ -94,7 +94,7 @@ async def my_task(): async with CraftExperiment.start(name="my_experiment") as exp: run = exp.run(my_task) - await exp.wait() + await exp.join() ``` ### 4. Launch Dashboard @@ -146,7 +146,7 @@ async with CraftExperiment.start("training") as exp: train_model, post_run_hooks=[PostRunHookFn.sync_metadata, PostRunHookFn.sync_status] ) - await exp.wait() + await exp.join() ``` ### 7. Cleanup diff --git a/alphatrion/experiment/base.py b/alphatrion/experiment/base.py index 0089b5d8..69c357c9 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 join() is called; the experiment auto-completes when all + # runs are finished. + "_joining", "_stopped", "_received_signal", "_signal_task", @@ -184,9 +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 join() is called, the experiment will auto stop when the runs # are all finished. - self._start_waiting = False + self._joining = False self._end_status = None self._stopped = asyncio.Event() self._received_signal: int | None = None @@ -382,12 +385,19 @@ 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. + # join 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 join(self): + self._joining = True + await self._context.wait() + + # wait blocks until the experiment is terminated, either by timeout or by + # calling done()/cancel() (e.g. from a signal handler). Unlike join(), 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(self): - self._start_waiting = True await self._context.wait() def is_done(self) -> bool: @@ -488,9 +498,9 @@ def _post_run(self, run: Run): ): self.done() - # If the experiment starts to wait and all runs are finished, + # If the experiment is joining and all runs are finished, # we can stop the experiment. - if self._start_waiting and len(self._runs) == 0: + if self._joining and len(self._runs) == 0: self.done() @classmethod diff --git a/tests/integration/server/test_graphql_query.py b/tests/integration/server/test_graphql_query.py index cf1f7a9f..56c8a050 100644 --- a/tests/integration/server/test_graphql_query.py +++ b/tests/integration/server/test_graphql_query.py @@ -340,7 +340,7 @@ async def test_query_single_run( run = exp.run(create_joke) run_id = run.id exp_id = exp.id - await exp.wait() + await exp.join() query = f""" query {{ @@ -510,7 +510,7 @@ async def test_query_experiment_with_usage( exp.run(create_joke) exp._on_signal(signal.SIGTERM) # Simulate sending a signal to trigger resume - await exp.wait() + await exp.join() query = f""" query {{ diff --git a/tests/integration/test_craft_experiment.py b/tests/integration/test_craft_experiment.py index 93c63fb1..5b16a5f7 100644 --- a/tests/integration/test_craft_experiment.py +++ b/tests/integration/test_craft_experiment.py @@ -29,7 +29,7 @@ async def fake_work(duration: int): exp.run(lambda: fake_work(5)) exp.run(lambda: fake_work(6)) - await exp.wait() + await exp.join() runtime = global_runtime() diff --git a/tests/integration/test_log.py b/tests/integration/test_log.py index d4e1c748..44264c9e 100644 --- a/tests/integration/test_log.py +++ b/tests/integration/test_log.py @@ -504,7 +504,7 @@ async def fake_sleep(value: float): exp.run(lambda: fake_work(0.1)) exp.run(lambda: fake_work(0.2)) # trigger early stopping - await exp.wait() + await exp.join() assert ( len( @@ -542,7 +542,7 @@ async def fake_sleep(value: float): exp.run(lambda: fake_work(1)) exp.run(lambda: fake_sleep(5)) # running in parallel. - await exp.wait() + await exp.join() assert ( len(exp._runtime.metadb.list_metrics_by_experiment_id(experiment_id=exp.id)) @@ -607,7 +607,7 @@ async def fake_sleep(value: float): exp.run(lambda: fake_work(0.3)) exp.run(lambda: fake_sleep(0.4)) exp.run(lambda: fake_work(0.9)) - await exp.wait() + await exp.join() assert ( len( @@ -644,7 +644,7 @@ async def fake_sleep(value: float): exp.run(lambda: fake_work(0.3)) exp.run(lambda: fake_sleep(0.4)) exp.run(lambda: fake_work(0.2)) - await exp.wait() + await exp.join() assert ( len( diff --git a/tests/integration/test_run_hooks.py b/tests/integration/test_run_hooks.py index b0d5b9aa..826eac60 100644 --- a/tests/integration/test_run_hooks.py +++ b/tests/integration/test_run_hooks.py @@ -47,7 +47,7 @@ async def train_model(): async with CraftExperiment.start("test_hook_experiment") as exp: # Create run with sync_metadata hook run = exp.run(train_model, post_run_hooks=[PostRunHookFn.sync_metadata]) - await exp.wait() + await exp.join() # Verify run completed assert run.result is not None @@ -77,7 +77,7 @@ async def task_with_string_result(): run = exp.run( task_with_string_result, post_run_hooks=[PostRunHookFn.sync_metadata] ) - await exp.wait() + await exp.join() # Verify metadata was not updated metadb = global_runtime().metadb @@ -103,7 +103,7 @@ async def task_without_metadata_key(): run = exp.run( task_without_metadata_key, post_run_hooks=[PostRunHookFn.sync_metadata] ) - await exp.wait() + await exp.join() # Verify metadata was not updated metadb = global_runtime().metadb @@ -132,7 +132,7 @@ async def task2(): async with CraftExperiment.start("test_exp_hooks", config=config) as exp: run1 = exp.run(task1) run2 = exp.run(task2) - await exp.wait() + await exp.join() # Verify both runs have metadata synced metadb = global_runtime().metadb @@ -174,7 +174,7 @@ async def train_model(): run = exp.run( train_model, post_run_hooks=[PostRunHookFn.sync_metadata, add_custom_info] ) - await exp.wait() + await exp.join() # Verify both hooks ran metadb = global_runtime().metadb @@ -209,7 +209,7 @@ async def train_model(): run_id=run.id, meta={"experiment_version": "v2", "notes": "test run"} ) - await exp.wait() + await exp.join() # Verify metadata was merged, not replaced run_obj = metadb.get_run(run_id=run.id) @@ -240,7 +240,7 @@ async def train_model(): run = exp.run( train_model, post_run_hooks=[buggy_hook, PostRunHookFn.sync_metadata] ) - await exp.wait() + await exp.join() # Run should still complete successfully assert run.result is not None @@ -269,7 +269,7 @@ async def train_model(): train_model, post_run_hooks=[PostRunHookFn.sync_metadata, PostRunHookFn.sync_status], ) - await exp.wait() + await exp.join() # Verify both hooks ran metadb = global_runtime().metadb @@ -297,7 +297,7 @@ async def task_with_none_result(): task_with_none_result, post_run_hooks=[PostRunHookFn.sync_metadata, PostRunHookFn.sync_status], ) - await exp.wait() + await exp.join() # Verify metadata was not updated metadb = global_runtime().metadb diff --git a/tests/unit/experiment/test_experiment.py b/tests/unit/experiment/test_experiment.py index 667ca6a0..ad80961d 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(), @@ -245,8 +245,41 @@ async def fake_work(): exp.run(fake_work) assert datetime.now() - start_time <= timedelta(seconds=1) + await exp.join() + assert datetime.now() - start_time >= timedelta(seconds=3) + + 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(): + """wait() 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() + # The run finishes after ~1s, but wait() 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 @@ -275,7 +308,7 @@ async def fake_work(exp_id: uuid.UUID): run2 = exp.run(lambda: fake_work(exp.id)) assert len(exp._runs) == 2 - await exp.wait() + await exp.join() assert datetime.now() - start_time >= timedelta(seconds=3) assert len(exp._runs) == 0 @@ -309,7 +342,7 @@ async def fake_work(timeout: int): run_3 = exp.run(lambda: fake_work(6)) # At this point, 4 runs are started. assert len(exp._runs) == 4 - await exp.wait() + await exp.join() assert len(exp._runs) == 0 run_0_obj = run_0._get_obj() @@ -338,7 +371,7 @@ async def test_create_experiment_with_max_execution_seconds(): name="first-experiment", config=experiment.ExperimentConfig(max_execution_seconds=2), ) as exp: - await exp.wait() + await exp.join() assert exp.is_done() exp_obj = exp._get_obj() @@ -364,7 +397,7 @@ async def fake_work(exp: CraftExperiment): ) as exp: exp.run(lambda: asyncio.sleep(5)) exp.run(partial(fake_work, exp)) - await exp.wait() + await exp.join() exp_obj = exp._get_obj() assert exp_obj.status == Status.INTERRUPTED @@ -391,7 +424,7 @@ async def fake_work(exp: CraftExperiment): ) as exp: exp.run(lambda: asyncio.sleep(5)) exp.run(partial(fake_work, exp)) - await exp.wait() + await exp.join() exp_obj = exp._get_obj() assert exp_obj.status == Status.CANCELLED From 49f661e9d6dd4cd1510569bcca4e67af4b660292 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Tue, 21 Jul 2026 19:08:46 +0100 Subject: [PATCH 2/5] add more logic to the joining Signed-off-by: kerthcet --- alphatrion/experiment/base.py | 3 ++- tests/unit/experiment/test_experiment.py | 23 +++++++++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/alphatrion/experiment/base.py b/alphatrion/experiment/base.py index 69c357c9..17be8720 100644 --- a/alphatrion/experiment/base.py +++ b/alphatrion/experiment/base.py @@ -190,7 +190,6 @@ def __init__(self, config: ExperimentConfig | None = None): # if join() is called, the experiment will auto stop when the runs # are all finished. self._joining = False - self._end_status = None self._stopped = asyncio.Event() self._received_signal: int | None = None self._signal_task: asyncio.Task | None = None @@ -390,6 +389,8 @@ def _timeout(self) -> int | None: # wait for them to complete. async def join(self): self._joining = True + if len(self._runs) == 0: + self.done() await self._context.wait() # wait blocks until the experiment is terminated, either by timeout or by diff --git a/tests/unit/experiment/test_experiment.py b/tests/unit/experiment/test_experiment.py index ad80961d..959b808f 100644 --- a/tests/unit/experiment/test_experiment.py +++ b/tests/unit/experiment/test_experiment.py @@ -252,6 +252,29 @@ 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.join(), 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(): """wait() must NOT auto-complete when all runs finish; it blocks until the From 1f41388efdedbe2da2d5b85738b227c5ecef7ed1 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Tue, 21 Jul 2026 19:36:22 +0100 Subject: [PATCH 3/5] rename functions Signed-off-by: kerthcet --- README.md | 4 +-- alphatrion/experiment/base.py | 22 +++++++-------- .../integration/server/test_graphql_query.py | 4 +-- tests/integration/test_craft_experiment.py | 2 +- tests/integration/test_log.py | 8 +++--- tests/integration/test_run_hooks.py | 18 ++++++------ tests/unit/experiment/test_experiment.py | 28 +++++++++++-------- 7 files changed, 45 insertions(+), 41 deletions(-) diff --git a/README.md b/README.md index eab4477b..f349e74b 100644 --- a/README.md +++ b/README.md @@ -94,7 +94,7 @@ async def my_task(): async with CraftExperiment.start(name="my_experiment") as exp: run = exp.run(my_task) - await exp.join() + await exp.wait() ``` ### 4. Launch Dashboard @@ -146,7 +146,7 @@ async with CraftExperiment.start("training") as exp: train_model, post_run_hooks=[PostRunHookFn.sync_metadata, PostRunHookFn.sync_status] ) - await exp.join() + await exp.wait() ``` ### 7. Cleanup diff --git a/alphatrion/experiment/base.py b/alphatrion/experiment/base.py index 17be8720..d864444d 100644 --- a/alphatrion/experiment/base.py +++ b/alphatrion/experiment/base.py @@ -168,7 +168,7 @@ class Experiment(ABC): "_total_runs_counter", # The end status, None, Err or Cancelled. "_end_status", - # True once join() is called; the experiment auto-completes when all + # True once wait() is called; the experiment auto-completes when all # runs are finished. "_joining", "_stopped", @@ -187,7 +187,7 @@ def __init__(self, config: ExperimentConfig | None = None): self._early_stopping_counter = 0 self._total_runs_counter = 0 self._end_status = None - # if join() is called, the experiment will auto stop when the runs + # if wait() is called, the experiment will auto stop when the runs # are all finished. self._joining = False self._stopped = asyncio.Event() @@ -384,21 +384,21 @@ def _timeout(self) -> int | None: return timeout - # join blocks until all the runs are finished, then the experiment is + # 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 join(self): + async def wait(self): self._joining = True if len(self._runs) == 0: self.done() await self._context.wait() - # wait blocks until the experiment is terminated, either by timeout or by - # calling done()/cancel() (e.g. from a signal handler). Unlike join(), 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(self): + # 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: @@ -499,7 +499,7 @@ def _post_run(self, run: Run): ): self.done() - # If the experiment is joining and all runs are finished, + # If the experiment is waiting and all runs are finished, # we can stop the experiment. if self._joining and len(self._runs) == 0: self.done() diff --git a/tests/integration/server/test_graphql_query.py b/tests/integration/server/test_graphql_query.py index 56c8a050..cf1f7a9f 100644 --- a/tests/integration/server/test_graphql_query.py +++ b/tests/integration/server/test_graphql_query.py @@ -340,7 +340,7 @@ async def test_query_single_run( run = exp.run(create_joke) run_id = run.id exp_id = exp.id - await exp.join() + await exp.wait() query = f""" query {{ @@ -510,7 +510,7 @@ async def test_query_experiment_with_usage( exp.run(create_joke) exp._on_signal(signal.SIGTERM) # Simulate sending a signal to trigger resume - await exp.join() + await exp.wait() query = f""" query {{ diff --git a/tests/integration/test_craft_experiment.py b/tests/integration/test_craft_experiment.py index 5b16a5f7..93c63fb1 100644 --- a/tests/integration/test_craft_experiment.py +++ b/tests/integration/test_craft_experiment.py @@ -29,7 +29,7 @@ async def fake_work(duration: int): exp.run(lambda: fake_work(5)) exp.run(lambda: fake_work(6)) - await exp.join() + await exp.wait() runtime = global_runtime() diff --git a/tests/integration/test_log.py b/tests/integration/test_log.py index 44264c9e..d4e1c748 100644 --- a/tests/integration/test_log.py +++ b/tests/integration/test_log.py @@ -504,7 +504,7 @@ async def fake_sleep(value: float): exp.run(lambda: fake_work(0.1)) exp.run(lambda: fake_work(0.2)) # trigger early stopping - await exp.join() + await exp.wait() assert ( len( @@ -542,7 +542,7 @@ async def fake_sleep(value: float): exp.run(lambda: fake_work(1)) exp.run(lambda: fake_sleep(5)) # running in parallel. - await exp.join() + await exp.wait() assert ( len(exp._runtime.metadb.list_metrics_by_experiment_id(experiment_id=exp.id)) @@ -607,7 +607,7 @@ async def fake_sleep(value: float): exp.run(lambda: fake_work(0.3)) exp.run(lambda: fake_sleep(0.4)) exp.run(lambda: fake_work(0.9)) - await exp.join() + await exp.wait() assert ( len( @@ -644,7 +644,7 @@ async def fake_sleep(value: float): exp.run(lambda: fake_work(0.3)) exp.run(lambda: fake_sleep(0.4)) exp.run(lambda: fake_work(0.2)) - await exp.join() + await exp.wait() assert ( len( diff --git a/tests/integration/test_run_hooks.py b/tests/integration/test_run_hooks.py index 826eac60..b0d5b9aa 100644 --- a/tests/integration/test_run_hooks.py +++ b/tests/integration/test_run_hooks.py @@ -47,7 +47,7 @@ async def train_model(): async with CraftExperiment.start("test_hook_experiment") as exp: # Create run with sync_metadata hook run = exp.run(train_model, post_run_hooks=[PostRunHookFn.sync_metadata]) - await exp.join() + await exp.wait() # Verify run completed assert run.result is not None @@ -77,7 +77,7 @@ async def task_with_string_result(): run = exp.run( task_with_string_result, post_run_hooks=[PostRunHookFn.sync_metadata] ) - await exp.join() + await exp.wait() # Verify metadata was not updated metadb = global_runtime().metadb @@ -103,7 +103,7 @@ async def task_without_metadata_key(): run = exp.run( task_without_metadata_key, post_run_hooks=[PostRunHookFn.sync_metadata] ) - await exp.join() + await exp.wait() # Verify metadata was not updated metadb = global_runtime().metadb @@ -132,7 +132,7 @@ async def task2(): async with CraftExperiment.start("test_exp_hooks", config=config) as exp: run1 = exp.run(task1) run2 = exp.run(task2) - await exp.join() + await exp.wait() # Verify both runs have metadata synced metadb = global_runtime().metadb @@ -174,7 +174,7 @@ async def train_model(): run = exp.run( train_model, post_run_hooks=[PostRunHookFn.sync_metadata, add_custom_info] ) - await exp.join() + await exp.wait() # Verify both hooks ran metadb = global_runtime().metadb @@ -209,7 +209,7 @@ async def train_model(): run_id=run.id, meta={"experiment_version": "v2", "notes": "test run"} ) - await exp.join() + await exp.wait() # Verify metadata was merged, not replaced run_obj = metadb.get_run(run_id=run.id) @@ -240,7 +240,7 @@ async def train_model(): run = exp.run( train_model, post_run_hooks=[buggy_hook, PostRunHookFn.sync_metadata] ) - await exp.join() + await exp.wait() # Run should still complete successfully assert run.result is not None @@ -269,7 +269,7 @@ async def train_model(): train_model, post_run_hooks=[PostRunHookFn.sync_metadata, PostRunHookFn.sync_status], ) - await exp.join() + await exp.wait() # Verify both hooks ran metadb = global_runtime().metadb @@ -297,7 +297,7 @@ async def task_with_none_result(): task_with_none_result, post_run_hooks=[PostRunHookFn.sync_metadata, PostRunHookFn.sync_status], ) - await exp.join() + await exp.wait() # Verify metadata was not updated metadb = global_runtime().metadb diff --git a/tests/unit/experiment/test_experiment.py b/tests/unit/experiment/test_experiment.py index 959b808f..a544737a 100644 --- a/tests/unit/experiment/test_experiment.py +++ b/tests/unit/experiment/test_experiment.py @@ -245,7 +245,7 @@ async def fake_work(): exp.run(fake_work) assert datetime.now() - start_time <= timedelta(seconds=1) - await exp.join() + await exp.wait() assert datetime.now() - start_time >= timedelta(seconds=3) exp_obj = exp._runtime.metadb.get_experiment(experiment_id=exp_id) @@ -268,7 +268,7 @@ async def test_experiment_join_with_no_runs(): exp_id = current_exp_id.get() # No runs launched; join() must return promptly instead of hanging. - await asyncio.wait_for(exp.join(), timeout=3) + await asyncio.wait_for(exp.wait(), timeout=3) assert exp.is_done() exp_obj = exp._runtime.metadb.get_experiment(experiment_id=exp_id) @@ -276,8 +276,8 @@ async def test_experiment_join_with_no_runs(): @pytest.mark.asyncio -async def test_experiment_with_wait(): - """wait() must NOT auto-complete when all runs finish; it blocks until the +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(), @@ -298,9 +298,10 @@ async def fake_work(): exp.run(fake_work) - await exp.wait() - # The run finishes after ~1s, but wait() keeps blocking until the - # timeout at ~3s instead of auto-completing when the run drains. + 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 @@ -331,7 +332,7 @@ async def fake_work(exp_id: uuid.UUID): run2 = exp.run(lambda: fake_work(exp.id)) assert len(exp._runs) == 2 - await exp.join() + await exp.wait() assert datetime.now() - start_time >= timedelta(seconds=3) assert len(exp._runs) == 0 @@ -365,7 +366,7 @@ async def fake_work(timeout: int): run_3 = exp.run(lambda: fake_work(6)) # At this point, 4 runs are started. assert len(exp._runs) == 4 - await exp.join() + await exp.wait() assert len(exp._runs) == 0 run_0_obj = run_0._get_obj() @@ -394,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.join() + # 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() @@ -420,7 +424,7 @@ async def fake_work(exp: CraftExperiment): ) as exp: exp.run(lambda: asyncio.sleep(5)) exp.run(partial(fake_work, exp)) - await exp.join() + await exp.wait() exp_obj = exp._get_obj() assert exp_obj.status == Status.INTERRUPTED @@ -447,7 +451,7 @@ async def fake_work(exp: CraftExperiment): ) as exp: exp.run(lambda: asyncio.sleep(5)) exp.run(partial(fake_work, exp)) - await exp.join() + await exp.wait() exp_obj = exp._get_obj() assert exp_obj.status == Status.CANCELLED From 72859ba93e2001a3baf8e1873988a9f1e0bd46db Mon Sep 17 00:00:00 2001 From: kerthcet Date: Tue, 21 Jul 2026 19:41:04 +0100 Subject: [PATCH 4/5] polish comments Signed-off-by: kerthcet --- alphatrion/experiment/base.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/alphatrion/experiment/base.py b/alphatrion/experiment/base.py index d864444d..aa7576ec 100644 --- a/alphatrion/experiment/base.py +++ b/alphatrion/experiment/base.py @@ -168,9 +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 + # True once wait_until_done() is called; the experiment auto-completes when all # runs are finished. - "_joining", + "_waiting", "_stopped", "_received_signal", "_signal_task", @@ -187,9 +187,9 @@ def __init__(self, config: ExperimentConfig | None = None): self._early_stopping_counter = 0 self._total_runs_counter = 0 self._end_status = None - # if wait() is called, the experiment will auto stop when the runs + # if wait_until_done() is called, the experiment will auto stop when the runs # are all finished. - self._joining = False + self._waiting = False self._stopped = asyncio.Event() self._received_signal: int | None = None self._signal_task: asyncio.Task | None = None @@ -388,7 +388,7 @@ def _timeout(self) -> int | None: # auto stopped. Use this when you have launched all the runs and want to # wait for them to complete. async def wait(self): - self._joining = True + self._waiting = True if len(self._runs) == 0: self.done() await self._context.wait() @@ -501,7 +501,7 @@ def _post_run(self, run: Run): # If the experiment is waiting and all runs are finished, # we can stop the experiment. - if self._joining and len(self._runs) == 0: + if self._waiting and len(self._runs) == 0: self.done() @classmethod From 1aca7002e4ae5a94e33d2ed207656264b8804e0c Mon Sep 17 00:00:00 2001 From: kerthcet Date: Tue, 21 Jul 2026 19:43:10 +0100 Subject: [PATCH 5/5] fix comments Signed-off-by: kerthcet --- alphatrion/experiment/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/alphatrion/experiment/base.py b/alphatrion/experiment/base.py index aa7576ec..60e9d3fb 100644 --- a/alphatrion/experiment/base.py +++ b/alphatrion/experiment/base.py @@ -168,7 +168,7 @@ class Experiment(ABC): "_total_runs_counter", # The end status, None, Err or Cancelled. "_end_status", - # True once wait_until_done() is called; the experiment auto-completes when all + # True once wait() is called; the experiment auto-completes when all # runs are finished. "_waiting", "_stopped", @@ -187,7 +187,7 @@ def __init__(self, config: ExperimentConfig | None = None): self._early_stopping_counter = 0 self._total_runs_counter = 0 self._end_status = None - # if wait_until_done() is called, the experiment will auto stop when the runs + # if wait() is called, the experiment will auto stop when the runs # are all finished. self._waiting = False self._stopped = asyncio.Event()