From 96d8c1213d339aa53bb152842454546faa68ac17 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Thu, 28 May 2026 16:34:48 +0100 Subject: [PATCH 1/2] fix experiment exit with wrong status Signed-off-by: kerthcet --- alphatrion/experiment/base.py | 5 ++++ tests/unit/experiment/test_experiment.py | 30 ++++++++++++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/alphatrion/experiment/base.py b/alphatrion/experiment/base.py index 376bce9..ee17eeb 100644 --- a/alphatrion/experiment/base.py +++ b/alphatrion/experiment/base.py @@ -194,6 +194,11 @@ async def __aenter__(self): return self async def __aexit__(self, exc_type, exc_val, exc_tb): + # If we are here because of an exception, we want to mark the experiment as done with error. + if exc_type is not None: + self.done_with_err() + else: + self.done() self.done() self._end_status = None diff --git a/tests/unit/experiment/test_experiment.py b/tests/unit/experiment/test_experiment.py index 7544600..8327575 100644 --- a/tests/unit/experiment/test_experiment.py +++ b/tests/unit/experiment/test_experiment.py @@ -160,6 +160,36 @@ async def test_experiment_with_done_with_err(): assert global_runtime().metadb.get_run(run_id=run_id).status == Status.CANCELLED +@pytest.mark.asyncio +async def test_experiment_exception_handling(): + """Test that exceptions in experiment context automatically mark it as failed.""" + init( + team_id=uuid.uuid4(), + user_id=uuid.uuid4(), + org_id=uuid.uuid4(), + ) + + exp_id = None + run_id = None + + with pytest.raises(ValueError, match="Simulated error"): + async with CraftExperiment.start(name="failing-experiment") as exp: + exp_id = exp.id + run = exp.run(lambda: asyncio.sleep(2)) + run_id = run.id + raise ValueError("Simulated error") + + # Verify experiment was marked as FAILED due to exception + exp_obj = global_runtime().metadb.get_experiment(experiment_id=exp_id) + assert exp_obj is not None + assert exp_obj.status == Status.FAILED + assert exp_obj.duration is not None + + # Verify run was cancelled + run_obj = global_runtime().metadb.get_run(run_id=run_id) + assert run_obj.status == Status.CANCELLED + + @pytest.mark.asyncio async def test_experiment_with_resume(): init( From 80b026fa9233ef770fe8f140ca75ee875128ace2 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Thu, 28 May 2026 16:37:16 +0100 Subject: [PATCH 2/2] fix Signed-off-by: kerthcet --- alphatrion/experiment/base.py | 1 - 1 file changed, 1 deletion(-) diff --git a/alphatrion/experiment/base.py b/alphatrion/experiment/base.py index ee17eeb..77d3952 100644 --- a/alphatrion/experiment/base.py +++ b/alphatrion/experiment/base.py @@ -199,7 +199,6 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): self.done_with_err() else: self.done() - self.done() self._end_status = None if self._signal_task: