From a75f85331dafab94350c888ec09b0054185876a3 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Fri, 3 Jul 2026 00:25:22 +0100 Subject: [PATCH] Add interrupted status Signed-off-by: kerthcet --- alphatrion/experiment/base.py | 69 ++++--- alphatrion/server/graphql/resolvers.py | 26 ++- alphatrion/server/graphql/schema.py | 4 +- alphatrion/server/graphql/types.py | 18 +- alphatrion/storage/sql_models.py | 21 +- docs/architecture/exp_state_machine.md | 98 +++++++++ .../server/test_graphql_mutation.py | 43 ++-- .../integration/server/test_graphql_query.py | 7 +- tests/unit/experiment/test_experiment.py | 192 +++++++++++++++++- 9 files changed, 405 insertions(+), 73 deletions(-) create mode 100644 docs/architecture/exp_state_machine.md diff --git a/alphatrion/experiment/base.py b/alphatrion/experiment/base.py index 8449bdaf..0089b5d8 100644 --- a/alphatrion/experiment/base.py +++ b/alphatrion/experiment/base.py @@ -13,6 +13,7 @@ from alphatrion.runtime.contextvars import current_exp_id from alphatrion.runtime.runtime import global_runtime from alphatrion.storage import runtime as storage_runtime +from alphatrion.storage import sql_models from alphatrion.storage.sql_models import Status from alphatrion.types import CallableEntry, PostRunHookFn from alphatrion.utils import context @@ -168,6 +169,7 @@ class Experiment(ABC): # The end status, None, Err or Cancelled. "_end_status", "_stopped", + "_received_signal", "_signal_task", ) @@ -187,6 +189,7 @@ def __init__(self, config: ExperimentConfig | None = None): self._start_waiting = False self._end_status = None self._stopped = asyncio.Event() + self._received_signal: int | None = None self._signal_task: asyncio.Task | None = None async def __aenter__(self): @@ -199,7 +202,6 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): self.done_with_err() else: self.done() - self._end_status = None if self._signal_task: # Already done, will not update the status again. @@ -225,22 +227,7 @@ def _start( name=name, team_id=self._runtime.team_id ) - # Just in case of kubernetes pod restarts, we want to make sure the experiment - # can be resumed if it is not completed, instead of creating a new experiment - # with the same name. If the experiment is already completed, we raise an error - # to avoid confusion. - if exp_obj and exp_obj.status != Status.COMPLETED: - self._id = exp_obj.uuid - self._runtime._metadb.update_experiment( - experiment_id=self._id, - status=Status.RUNNING, - ) - elif exp_obj and exp_obj.status == Status.COMPLETED: - raise RuntimeError( - f"Experiment with name '{name}' already exists and is completed. \ - Please choose a different name or delete the existing experiment." - ) - else: + if not exp_obj: self._id = self._runtime._metadb.create_experiment( name=name, org_id=self._runtime._org_id, @@ -253,6 +240,21 @@ def _start( params=params, status=Status.RUNNING, ) + # Just in case of kubernetes pod restarts, we want to make sure the experiment + # can be resumed if it is not terminated, instead of creating a new experiment + # with the same name. If the experiment is already terminated, we raise an error + # to avoid confusion. + elif exp_obj.status not in sql_models.TERMINAL_STATUS: + self._id = exp_obj.uuid + self._runtime._metadb.update_experiment( + experiment_id=self._id, + status=Status.RUNNING, + ) + else: + raise RuntimeError( + f"Experiment with name '{name}' already exists and is terminated with status '{exp_obj.status}'. \ + Please choose a different name or delete the existing experiment." + ) self._context = context.Context( cancel_func=self._stop, @@ -397,11 +399,19 @@ def done(self): self._cancel() def done_with_err(self): - self._end_status = "Err" + self._end_status = Status.FAILED self.done() def done_with_cancel(self): - self._end_status = "Cancelled" + self._end_status = Status.CANCELLED + self.done() + + def done_with_interrupt(self): + self._end_status = Status.INTERRUPTED + self.done() + + def done_with_abort(self): + self._end_status = Status.ABORTED self.done() def _cancel(self): @@ -422,10 +432,8 @@ def _stop(self): ).total_seconds() status = Status.COMPLETED - if self._end_status == "Err": - status = Status.FAILED - elif self._end_status == "Cancelled": - status = Status.CANCELLED + if self._end_status is not None: + status = self._end_status self._runtime.metadb.update_experiment( experiment_id=self._id, status=status, duration=duration @@ -500,15 +508,22 @@ def _start_signal_handlers(self): loop = asyncio.get_running_loop() # Handle SIGINT and SIGTERM to allow graceful shutdown. - # Make sure to call done() on receiving the signal. + # SIGINT (Ctrl+C) → CANCELLED (user intent to stop) + # SIGTERM (kill/K8s) → INTERRUPTED (resumable) for sig in (signal.SIGINT, signal.SIGTERM): - loop.add_signal_handler(sig, self._on_signal) + loop.add_signal_handler(sig, lambda s=sig: self._on_signal(s)) return asyncio.create_task(self._wait_for_stop()) - def _on_signal(self): + def _on_signal(self, sig: int): + self._received_signal = sig self._stopped.set() async def _wait_for_stop(self): await self._stopped.wait() - self.done_with_cancel() + + # Distinguish between SIGINT (user Ctrl+C) and SIGTERM (system/K8s) + if self._received_signal == signal.SIGINT: + self.done_with_cancel() # User explicitly stopped with Ctrl+C + else: + self.done_with_interrupt() # SIGTERM or other, treat as resumable diff --git a/alphatrion/server/graphql/resolvers.py b/alphatrion/server/graphql/resolvers.py index 5078d93a..7d0386b4 100644 --- a/alphatrion/server/graphql/resolvers.py +++ b/alphatrion/server/graphql/resolvers.py @@ -2027,11 +2027,13 @@ def delete_datasets( return True @staticmethod - def abort_experiment( + def stop_experiment( info: Info[GraphQLContext, None], experiment_id: strawberry.ID ) -> Experiment: - """Abort an experiment by changing its status to ABORTED. - Only works if the experiment is in PENDING status.""" + """Stop an experiment. + - PENDING → ABORTED (validation failure before start) + - RUNNING → CANCELLED (user-initiated stop) + """ user_id = uuid.UUID(info.context.user_id) experiment_id_uuid = uuid.UUID(experiment_id) @@ -2048,23 +2050,27 @@ def abort_experiment( "Not allowed to update experiment in team that user does not belong to" ) - # Only abort if experiment is in PENDING status - if exp.status != Status.PENDING: + # Determine target status based on current status + if exp.status == Status.PENDING: + new_status = Status.ABORTED + elif exp.status == Status.RUNNING: + new_status = Status.CANCELLED + else: raise RuntimeError( - f"Cannot abort experiment with status '{StatusMap[Status(exp.status)]}'. " - "Only experiments in PENDING status can be aborted." + f"Cannot stop experiment with status '{StatusMap[Status(exp.status)]}'. " + "Only experiments in PENDING or RUNNING status can be stopped." ) - # Update status to ABORTED + # Update status metadb.update_experiment( experiment_id=experiment_id_uuid, - status=Status.ABORTED, + status=new_status, ) # Get the updated experiment updated_exp = metadb.get_experiment(experiment_id=experiment_id_uuid) if not updated_exp: - raise RuntimeError("Failed to retrieve aborted experiment") + raise RuntimeError("Failed to retrieve stopped experiment") return Experiment( id=updated_exp.uuid, diff --git a/alphatrion/server/graphql/schema.py b/alphatrion/server/graphql/schema.py index 542514da..92cf2bc6 100644 --- a/alphatrion/server/graphql/schema.py +++ b/alphatrion/server/graphql/schema.py @@ -302,10 +302,10 @@ def update_experiment( return GraphQLMutations.update_experiment(info=info, input=input) @strawberry.mutation - def abort_experiment( + def stop_experiment( self, experiment_id: strawberry.ID, info: Info[GraphQLContext, None] ) -> Experiment: - return GraphQLMutations.abort_experiment(info=info, experiment_id=experiment_id) + return GraphQLMutations.stop_experiment(info=info, experiment_id=experiment_id) schema = strawberry.Schema(query=Query, mutation=Mutation) diff --git a/alphatrion/server/graphql/types.py b/alphatrion/server/graphql/types.py index d347699c..687a0c2f 100644 --- a/alphatrion/server/graphql/types.py +++ b/alphatrion/server/graphql/types.py @@ -148,13 +148,25 @@ def teams(self, info: Info) -> list[Team] | None: class GraphQLStatus(Enum): + """Status enum shared between Experiments and Runs. + + Run status is a subset: RUNNING, COMPLETED, CANCELLED, FAILED. + Experiment status includes all values. + + Terminal states (cannot resume): COMPLETED, CANCELLED, ABORTED + Recoverable states (can resume): FAILED, INTERRUPTED + """ + UNKNOWN = "UNKNOWN" PENDING = "PENDING" RUNNING = "RUNNING" - CANCELLED = "CANCELLED" COMPLETED = "COMPLETED" - FAILED = "FAILED" - ABORTED = "ABORTED" + CANCELLED = "CANCELLED" # User intentionally stopped, cannot resume + FAILED = "FAILED" # Recoverable: can retry after fixing error + ABORTED = "ABORTED" # Programmatically aborted, cannot resume + INTERRUPTED = ( + "INTERRUPTED" # Recoverable: system stopped (e.g. preemption), can resume + ) GraphQLStatusEnum = strawberry.enum(GraphQLStatus) diff --git a/alphatrion/storage/sql_models.py b/alphatrion/storage/sql_models.py index 6761b742..ca17d7af 100644 --- a/alphatrion/storage/sql_models.py +++ b/alphatrion/storage/sql_models.py @@ -19,14 +19,16 @@ Base = declarative_base() +# Status includes both run and experiment status, run status is a subset of experiment status. class Status(enum.IntEnum): UNKNOWN = 0 PENDING = 1 RUNNING = 2 COMPLETED = 9 - CANCELLED = 10 - FAILED = 11 - ABORTED = 12 + CANCELLED = 10 # User intentionally stopped, cannot resume + FAILED = 11 # Recoverable: can retry after fixing error + ABORTED = 12 # Programmatically aborted, cannot resume + INTERRUPTED = 13 # Recoverable: system stopped (e.g. preemption), can resume StatusMap = { @@ -37,9 +39,20 @@ class Status(enum.IntEnum): Status.COMPLETED: "COMPLETED", Status.FAILED: "FAILED", Status.ABORTED: "ABORTED", + Status.INTERRUPTED: "INTERRUPTED", } -FINISHED_STATUS = [Status.COMPLETED, Status.FAILED, Status.CANCELLED, Status.ABORTED] +# Terminal states: cannot transition to any other state +TERMINAL_STATUS = [Status.COMPLETED, Status.CANCELLED, Status.ABORTED] + +# Recoverable states: can transition back to RUNNING +RECOVERABLE_STATUS = [Status.FAILED, Status.INTERRUPTED] + +# All finished states (terminal + recoverable) +FINISHED_STATUS = TERMINAL_STATUS + RECOVERABLE_STATUS + +# Run status is a subset of experiment status to make it simpler. +RUN_STATUS = [Status.RUNNING, Status.COMPLETED, Status.CANCELLED, Status.FAILED] class Organization(Base): diff --git a/docs/architecture/exp_state_machine.md b/docs/architecture/exp_state_machine.md new file mode 100644 index 00000000..fbbc2a90 --- /dev/null +++ b/docs/architecture/exp_state_machine.md @@ -0,0 +1,98 @@ +# Experiment Status State Machine + +## Overview + +The `Status` enum is shared between **Experiments** and **Runs**: + +```python +class Status(enum.IntEnum): + UNKNOWN = 0 + PENDING = 1 + RUNNING = 2 + COMPLETED = 9 + CANCELLED = 10 # User intentionally stopped, cannot resume + FAILED = 11 # Recoverable: can retry after fixing error + ABORTED = 12 # Programmatically aborted, cannot resume + INTERRUPTED = 13 # Recoverable: system stopped (e.g. preemption), can resume +``` + +**Experiment statuses**: All values above +**Run statuses**: Subset only (RUNNING, COMPLETED, CANCELLED, FAILED) + +## State Diagram + +``` + ┌─────────────┐ + │ Experiment │ + └──────┬──────┘ + │ + ▼ + ┌─────────────┐ + ┌──────────│ PENDING │ + │ └──────┬──────┘ + │ │ + │ abort() │ start() + │ │ + ▼ ▼ + ┌────────────┐ ┌─────────────┐ + │ ABORTED │ │ RUNNING │ ◀───────────────────────────────────────────┐ + └────────────┘ └──────┬──────┘ │ + (terminal) │ │ + │ │ + ┌───────────────────────────┼───────────────────────────────────────┐ │ + │ │ │ │ │ + │ │ │ │ │ +cancel()/signal(SIGINT) done() exception signal(SIGTERM) │ + │ │ (unintentional) │ │ + ▼ ▼ ▼ ▼ │ + ┌────────────┐ ┌────────────┐ ┌────────────┐ ┌──────────────┐ │ + │ CANCELLED │ │ COMPLETED │ │ FAILED │ │ INTERRUPTED │ │ + └────────────┘ └────────────┘ └─────┬──────┘ └──────┬───────┘ │ + (terminal) (terminal) │ │ │ + │ │ │ + └────────────────────┘ │ + │ resume │ + └────────────────────────┘ + + +──────────────────────────────────────────────────────────────────────────────────────── +Terminal States (No Resume): + • COMPLETED - Success, finished normally + • CANCELLED - User stopped permanently (SIGINT/Ctrl+C or stopExperiment API) + • ABORTED - Validation failed before starting (only from PENDING state) + +Recoverable States (Can Resume): + • FAILED - Exception/error during execution, can retry after fix + • INTERRUPTED - System stopped (SIGTERM/K8s pod termination), can resume after restart +──────────────────────────────────────────────────────────────────────────────────────── +``` + +## State Categories + +| Category | States | Can Resume? | Description | +|----------|--------|-------------|-------------| +| **Active** | PENDING, RUNNING | N/A | Normal execution flow | +| **Terminal** | COMPLETED, CANCELLED, ABORTED | ❌ No | Final states, immutable | +| **Recoverable** | FAILED, INTERRUPTED | ✅ Yes | Can resume to RUNNING | + +### Semantics + +**Resumability is the key distinction:** + +- **COMPLETED**: Success, nothing more to do +- **CANCELLED**: User decided to stop permanently (SIGINT in local dev, or stopExperiment API) +- **ABORTED**: Validation failed before starting (only from PENDING via stopExperiment API) +- **FAILED**: Unintentional error, retry after fixing +- **INTERRUPTED**: System event (SIGTERM from K8s pod termination, preemption), resume when restarted + +## Implementation + +### Constants + +```python +# In sql_models.py +TERMINAL_STATUS = [Status.COMPLETED, Status.CANCELLED, Status.ABORTED] +RECOVERABLE_STATUS = [Status.FAILED, Status.INTERRUPTED] +FINISHED_STATUS = TERMINAL_STATUS + RECOVERABLE_STATUS +RUN_STATUS = [Status.RUNNING, Status.COMPLETED, Status.CANCELLED, Status.FAILED] +``` diff --git a/tests/integration/server/test_graphql_mutation.py b/tests/integration/server/test_graphql_mutation.py index d581f6b5..e654b38f 100644 --- a/tests/integration/server/test_graphql_mutation.py +++ b/tests/integration/server/test_graphql_mutation.py @@ -1412,10 +1412,10 @@ def test_update_experiment_not_found( assert "not found" in str(response.errors[0]).lower() -def test_abort_experiment_pending( +def test_stop_experiment_pending( execute_graphql, test_org_id, test_user_id, test_team_id ): - """Test aborting a pending experiment""" + """Test stopping a pending experiment""" runtime.init() metadb = runtime.storage_runtime().metadb @@ -1428,10 +1428,10 @@ def test_abort_experiment_pending( status=Status.PENDING, ) - # Abort the experiment + # Stop the experiment mutation = f""" mutation {{ - abortExperiment(experimentId: "{exp_id}") {{ + stopExperiment(experimentId: "{exp_id}") {{ id name status @@ -1444,17 +1444,17 @@ def test_abort_experiment_pending( user_id=test_user_id, ) assert response.errors is None - assert response.data["abortExperiment"]["status"] == "ABORTED" + assert response.data["stopExperiment"]["status"] == "ABORTED" # Verify in database exp = metadb.get_experiment(experiment_id=exp_id) assert exp.status == Status.ABORTED -def test_abort_experiment_running_fails( +def test_stop_experiment_running( execute_graphql, test_org_id, test_user_id, test_team_id ): - """Test that aborting a running experiment fails""" + """Test stopping a running experiment changes it to CANCELLED""" runtime.init() metadb = runtime.storage_runtime().metadb @@ -1470,10 +1470,10 @@ def test_abort_experiment_running_fails( status=Status.RUNNING, ) - # Try to abort the running experiment + # Stop the running experiment mutation = f""" mutation {{ - abortExperiment(experimentId: "{exp_id}") {{ + stopExperiment(experimentId: "{exp_id}") {{ id status }} @@ -1484,19 +1484,18 @@ def test_abort_experiment_running_fails( org_id=test_org_id, user_id=test_user_id, ) - # Should return an error - assert response.errors is not None - assert "Cannot abort" in str(response.errors[0]) + assert response.errors is None + assert response.data["stopExperiment"]["status"] == "CANCELLED" - # Verify status is still RUNNING + # Verify status is CANCELLED exp = metadb.get_experiment(experiment_id=exp_id) - assert exp.status == Status.RUNNING + assert exp.status == Status.CANCELLED -def test_abort_experiment_completed_fails( +def test_stop_experiment_completed_fails( execute_graphql, test_org_id, test_user_id, test_team_id ): - """Test that aborting a completed experiment fails""" + """Test that stopping a completed experiment fails""" runtime.init() metadb = runtime.storage_runtime().metadb @@ -1512,10 +1511,10 @@ def test_abort_experiment_completed_fails( status=Status.COMPLETED, ) - # Try to abort the completed experiment + # Try to stop the completed experiment mutation = f""" mutation {{ - abortExperiment(experimentId: "{exp_id}") {{ + stopExperiment(experimentId: "{exp_id}") {{ id status }} @@ -1528,17 +1527,17 @@ def test_abort_experiment_completed_fails( ) # Should return an error assert response.errors is not None - assert "Cannot abort" in str(response.errors[0]) + assert "Cannot stop" in str(response.errors[0]) # Verify status is still COMPLETED exp = metadb.get_experiment(experiment_id=exp_id) assert exp.status == Status.COMPLETED -def test_abort_experiment_not_found( +def test_stop_experiment_not_found( execute_graphql, test_org_id, test_user_id, test_team_id ): - """Test aborting a non-existent experiment when user has team access""" + """Test stopping a non-existent experiment when user has team access""" runtime.init() metadb = runtime.storage_runtime().metadb @@ -1549,7 +1548,7 @@ def test_abort_experiment_not_found( fake_exp_id = uuid.uuid4() mutation = f""" mutation {{ - abortExperiment(experimentId: "{fake_exp_id}") {{ + stopExperiment(experimentId: "{fake_exp_id}") {{ id status }} diff --git a/tests/integration/server/test_graphql_query.py b/tests/integration/server/test_graphql_query.py index 56eb77a6..cf1f7a9f 100644 --- a/tests/integration/server/test_graphql_query.py +++ b/tests/integration/server/test_graphql_query.py @@ -4,6 +4,7 @@ import asyncio import os +import signal import uuid from datetime import datetime, timedelta @@ -508,7 +509,7 @@ async def test_query_experiment_with_usage( # before get the response from the ollama, the spans will be empty. exp.run(create_joke) - exp._on_signal() # Simulate sending a signal to trigger resume + exp._on_signal(signal.SIGTERM) # Simulate sending a signal to trigger resume await exp.wait() query = f""" @@ -534,7 +535,7 @@ async def test_query_experiment_with_usage( ) assert response.errors is None - assert response.data["experiment"]["status"] == "CANCELLED" + assert response.data["experiment"]["status"] == "INTERRUPTED" assert response.data["experiment"]["aggregatedUsage"] is not None assert response.data["experiment"]["aggregatedUsage"]["totalTokens"] is not None assert response.data["experiment"]["aggregatedUsage"]["inputTokens"] is not None @@ -550,7 +551,7 @@ async def test_query_experiment_with_usage( assert response.data["experiment"]["aggregatedUsage"]["totalCost"] is not None exp_obj = runtime.storage_runtime().metadb.get_experiment(experiment_id=exp.id) - assert exp_obj.status == Status.CANCELLED + assert exp_obj.status == Status.INTERRUPTED # resume the experiment async with CraftExperiment.start(name="integration_test_experiment_resume") as exp: diff --git a/tests/unit/experiment/test_experiment.py b/tests/unit/experiment/test_experiment.py index 83275752..667ca6a0 100644 --- a/tests/unit/experiment/test_experiment.py +++ b/tests/unit/experiment/test_experiment.py @@ -1,6 +1,7 @@ # ruff: noqa: E501 import asyncio +import signal import unittest import uuid from datetime import datetime, timedelta @@ -354,7 +355,8 @@ async def test_experiment_with_signal(): async def fake_work(exp: CraftExperiment): await asyncio.sleep(2) - exp._on_signal() + # Simulate SIGTERM (system/K8s termination) + exp._on_signal(signal.SIGTERM) start_time = datetime.now() async with CraftExperiment.start( @@ -365,11 +367,36 @@ async def fake_work(exp: CraftExperiment): await exp.wait() exp_obj = exp._get_obj() - assert exp_obj.status == Status.CANCELLED + assert exp_obj.status == Status.INTERRUPTED assert (datetime.now() - start_time).total_seconds() >= 2 assert (datetime.now() - start_time).total_seconds() < 5 +@pytest.mark.asyncio +async def test_experiment_with_sigint_cancelled(): + """Test that SIGINT (Ctrl+C) results in CANCELLED status.""" + init( + team_id=uuid.uuid4(), + user_id=uuid.uuid4(), + org_id=uuid.uuid4(), + ) + + async def fake_work(exp: CraftExperiment): + await asyncio.sleep(2) + # Simulate SIGINT (Ctrl+C) + exp._on_signal(signal.SIGINT) + + async with CraftExperiment.start( + name="experiment-with-sigint", + ) as exp: + exp.run(lambda: asyncio.sleep(5)) + exp.run(partial(fake_work, exp)) + await exp.wait() + + exp_obj = exp._get_obj() + assert exp_obj.status == Status.CANCELLED + + @pytest.mark.asyncio async def test_experiment_with_result_return(): init( @@ -443,3 +470,164 @@ async def test_experiment_with_tags(): experiment_id=exp.id, ) assert len(all_tags) == 2 + + +@pytest.mark.asyncio +async def test_experiment_done_with_interrupt(test_team_id, test_user_id, test_org_id): + """Test that done_with_interrupt() marks experiment as INTERRUPTED.""" + init(team_id=test_team_id, user_id=test_user_id, org_id=test_org_id) + + exp_id = None + async with CraftExperiment.start(name="interrupt-test") as exp: + exp_id = exp.id + exp.done_with_interrupt() + + # Verify experiment status is INTERRUPTED + exp_obj = global_runtime()._metadb.get_experiment(experiment_id=exp_id) + assert exp_obj.status == Status.INTERRUPTED + assert exp_obj.duration is not None + + +@pytest.mark.asyncio +async def test_experiment_done_with_abort(test_team_id, test_user_id, test_org_id): + """Test that done_with_abort() marks experiment as ABORTED.""" + init(team_id=test_team_id, user_id=test_user_id, org_id=test_org_id) + + exp_id = None + async with CraftExperiment.start(name="abort-test") as exp: + exp_id = exp.id + exp.done_with_abort() + + # Verify experiment status is ABORTED + exp_obj = global_runtime()._metadb.get_experiment(experiment_id=exp_id) + assert exp_obj.status == Status.ABORTED + assert exp_obj.duration is not None + + +@pytest.mark.asyncio +async def test_experiment_resume_from_failed(test_team_id, test_user_id, test_org_id): + """Test that an experiment in FAILED state can be resumed.""" + init(team_id=test_team_id, user_id=test_user_id, org_id=test_org_id) + + exp_name = "failed-resume-test" + exp_id = None + + # Create an experiment and mark it as FAILED + async with CraftExperiment.start(name=exp_name) as exp: + exp_id = exp.id + exp.done_with_err() + + # Verify it's FAILED + exp_obj = global_runtime()._metadb.get_experiment(experiment_id=exp_id) + assert exp_obj.status == Status.FAILED + + # Resume the experiment - should work and transition to RUNNING + async with CraftExperiment.start(name=exp_name) as exp: + exp_obj = global_runtime()._metadb.get_experiment(experiment_id=exp.id) + assert exp_obj.status == Status.RUNNING + + # After exit, should be COMPLETED + exp_obj = global_runtime()._metadb.get_experiment(experiment_id=exp.id) + assert exp_obj.status == Status.COMPLETED + + +@pytest.mark.asyncio +async def test_experiment_resume_from_interrupted( + test_team_id, test_user_id, test_org_id +): + """Test that an experiment in INTERRUPTED state can be resumed.""" + init(team_id=test_team_id, user_id=test_user_id, org_id=test_org_id) + + exp_name = "interrupted-resume-test" + exp_id = None + + # Create an experiment and mark it as INTERRUPTED + async with CraftExperiment.start(name=exp_name) as exp: + exp_id = exp.id + exp.done_with_interrupt() + + # Verify it's INTERRUPTED + exp_obj = global_runtime()._metadb.get_experiment(experiment_id=exp_id) + assert exp_obj.status == Status.INTERRUPTED + + # Resume the experiment - should work and transition to RUNNING + async with CraftExperiment.start(name=exp_name) as exp: + exp_obj = global_runtime()._metadb.get_experiment(experiment_id=exp.id) + assert exp_obj.status == Status.RUNNING + + # After exit, should be COMPLETED + exp_obj = global_runtime()._metadb.get_experiment(experiment_id=exp.id) + assert exp_obj.status == Status.COMPLETED + + +@pytest.mark.asyncio +async def test_experiment_cannot_resume_completed( + test_team_id, test_user_id, test_org_id +): + """Test that a COMPLETED experiment cannot be resumed.""" + init(team_id=test_team_id, user_id=test_user_id, org_id=test_org_id) + + exp_name = "completed-no-resume" + + # Create and complete an experiment + async with CraftExperiment.start(name=exp_name) as exp: + exp_id = exp.id + + # Verify it's COMPLETED + exp_obj = global_runtime()._metadb.get_experiment(experiment_id=exp_id) + assert exp_obj.status == Status.COMPLETED + + # Try to resume - should raise error + with pytest.raises(RuntimeError, match="already exists and is terminated"): + async with CraftExperiment.start(name=exp_name) as exp: + pass + + +@pytest.mark.asyncio +async def test_experiment_cannot_resume_cancelled( + test_team_id, test_user_id, test_org_id +): + """Test that a CANCELLED experiment cannot be resumed.""" + init(team_id=test_team_id, user_id=test_user_id, org_id=test_org_id) + + exp_name = "cancelled-no-resume" + exp_id = None + + # Create an experiment and cancel it + async with CraftExperiment.start(name=exp_name) as exp: + exp_id = exp.id + exp.done_with_cancel() + + # Verify it's CANCELLED + exp_obj = global_runtime()._metadb.get_experiment(experiment_id=exp_id) + assert exp_obj.status == Status.CANCELLED + + # Try to resume - should raise error + with pytest.raises(RuntimeError, match="already exists and is terminated"): + async with CraftExperiment.start(name=exp_name) as exp: + pass + + +@pytest.mark.asyncio +async def test_experiment_cannot_resume_aborted( + test_team_id, test_user_id, test_org_id +): + """Test that an ABORTED experiment cannot be resumed.""" + init(team_id=test_team_id, user_id=test_user_id, org_id=test_org_id) + + exp_name = "aborted-no-resume" + exp_id = None + + # Create an experiment and abort it + async with CraftExperiment.start(name=exp_name) as exp: + exp_id = exp.id + exp.done_with_abort() + + # Verify it's ABORTED + exp_obj = global_runtime()._metadb.get_experiment(experiment_id=exp_id) + assert exp_obj.status == Status.ABORTED + + # Try to resume - should raise error + with pytest.raises(RuntimeError, match="already exists and is terminated"): + async with CraftExperiment.start(name=exp_name) as exp: + pass