From 5fda9254b1e7f82ba72f363640e7a82c3c589611 Mon Sep 17 00:00:00 2001 From: ky-hiverge Date: Mon, 15 Jun 2026 16:26:59 +0800 Subject: [PATCH 1/2] Support to delete running/pending Experiment Signed-off-by: ky-hiverge --- alphatrion/storage/sqlstore.py | 25 +++- .../server/test_graphql_mutation.py | 28 ++-- tests/unit/storage/test_sql.py | 132 ++++++++++++++++++ 3 files changed, 168 insertions(+), 17 deletions(-) diff --git a/alphatrion/storage/sqlstore.py b/alphatrion/storage/sqlstore.py index 4c1e890..c3492bd 100644 --- a/alphatrion/storage/sqlstore.py +++ b/alphatrion/storage/sqlstore.py @@ -743,6 +743,17 @@ def list_exps_by_timeframe( return exps def delete_experiment(self, experiment_id: uuid.UUID) -> bool: + """ + Delete a single experiment by setting is_del flag. + Also deletes all associated runs. + + Status transitions on deletion: + - PENDING experiments -> ABORTED + - RUNNING experiments -> CANCELLED + - Other statuses remain unchanged + """ + from sqlalchemy import case + session = self._session() # Try to delete the experiment @@ -752,17 +763,19 @@ def delete_experiment(self, experiment_id: uuid.UUID) -> bool: .first() ) - if exp and exp.status == Status.RUNNING: - raise ValueError( - "Cannot delete a running experiment. Please stop it first." - ) - # Delete all runs associated with this experiment - # (regardless of experiment status) session.query(Run).filter(Run.experiment_id == experiment_id).update( {Run.is_del: 1}, synchronize_session=False ) + if exp: + # Update status based on current state + if exp.status == Status.PENDING: + exp.status = Status.ABORTED + elif exp.status == Status.RUNNING: + exp.status = Status.CANCELLED + # Other statuses remain unchanged + exp.is_del = 1 session.commit() session.close() diff --git a/tests/integration/server/test_graphql_mutation.py b/tests/integration/server/test_graphql_mutation.py index 34cdfc2..d581f6b 100644 --- a/tests/integration/server/test_graphql_mutation.py +++ b/tests/integration/server/test_graphql_mutation.py @@ -897,10 +897,10 @@ def test_delete_experiments_batch_deletes_runs( assert metadb.get_run(run_id=run_2_2) is None -def test_delete_running_experiment_fails( +def test_delete_running_experiment_succeeds( execute_graphql, test_org_id, test_user_id, test_team_id ): - """Test that deleting a running experiment raises an error""" + """Test that deleting a running experiment marks it as CANCELLED""" runtime.init() metadb = runtime.storage_runtime().metadb @@ -913,8 +913,6 @@ def test_delete_running_experiment_fails( # Set experiment status to RUNNING metadb.update_experiment( - org_id=test_org_id, - team_id=test_team_id, experiment_id=experiment_id, status=Status.RUNNING, ) @@ -924,7 +922,7 @@ def test_delete_running_experiment_fails( assert exp is not None assert exp.status == Status.RUNNING - # Try to delete running experiment via mutation + # Delete running experiment via mutation mutation = f""" mutation {{ deleteExperiment(experimentId: "{experiment_id}") @@ -936,14 +934,22 @@ def test_delete_running_experiment_fails( user_id=test_user_id, ) - # Should return an error - assert response.errors is not None - assert "Cannot delete a running experiment" in str(response.errors[0]) + # Should succeed without errors + assert response.errors is None + assert response.data["deleteExperiment"] is True - # Verify experiment still exists + # Verify experiment is deleted (get_experiment filters out deleted ones) exp = metadb.get_experiment(experiment_id=experiment_id) - assert exp is not None - assert exp.status == Status.RUNNING + assert exp is None + + # Verify status is CANCELLED by querying directly + from alphatrion.storage.sql_models import Experiment + + with metadb._session() as session: + exp = session.query(Experiment).filter(Experiment.uuid == experiment_id).first() + assert exp is not None + assert exp.is_del == 1 + assert exp.status == Status.CANCELLED def test_delete_experiments_with_running_and_pending( diff --git a/tests/unit/storage/test_sql.py b/tests/unit/storage/test_sql.py index a40fcc2..0e938c6 100644 --- a/tests/unit/storage/test_sql.py +++ b/tests/unit/storage/test_sql.py @@ -510,3 +510,135 @@ def test_delete_experiments_already_deleted(db): # Try to delete again deleted_count = db.delete_experiments([exp_id]) assert deleted_count == 0 + + +def test_delete_single_experiment_running(db): + """Test that delete_experiment marks RUNNING experiment as CANCELLED""" + from alphatrion.storage.sql_models import Experiment + + org_id = uuid.uuid4() + team_id = uuid.uuid4() + user_id = uuid.uuid4() + + # Create running experiment + exp_id = db.create_experiment( + org_id=org_id, + team_id=team_id, + user_id=user_id, + name="running_exp", + status=Status.RUNNING, + ) + + # Delete the running experiment + result = db.delete_experiment(exp_id) + assert result is True + + # Verify experiment is deleted + exp = db.get_experiment(exp_id) + assert exp is None + + # Verify status is CANCELLED by querying directly + with db._session() as session: + exp = session.query(Experiment).filter(Experiment.uuid == exp_id).first() + assert exp is not None + assert exp.is_del == 1 + assert exp.status == Status.CANCELLED + + +def test_delete_single_experiment_pending(db): + """Test that delete_experiment marks PENDING experiment as ABORTED""" + from alphatrion.storage.sql_models import Experiment + + org_id = uuid.uuid4() + team_id = uuid.uuid4() + user_id = uuid.uuid4() + + # Create pending experiment + exp_id = db.create_experiment( + org_id=org_id, + team_id=team_id, + user_id=user_id, + name="pending_exp", + status=Status.PENDING, + ) + + # Delete the pending experiment + result = db.delete_experiment(exp_id) + assert result is True + + # Verify experiment is deleted + exp = db.get_experiment(exp_id) + assert exp is None + + # Verify status is ABORTED by querying directly + with db._session() as session: + exp = session.query(Experiment).filter(Experiment.uuid == exp_id).first() + assert exp is not None + assert exp.is_del == 1 + assert exp.status == Status.ABORTED + + +def test_delete_single_experiment_completed(db): + """Test that delete_experiment keeps COMPLETED status unchanged""" + from alphatrion.storage.sql_models import Experiment + + org_id = uuid.uuid4() + team_id = uuid.uuid4() + user_id = uuid.uuid4() + + # Create completed experiment + exp_id = db.create_experiment( + org_id=org_id, + team_id=team_id, + user_id=user_id, + name="completed_exp", + status=Status.COMPLETED, + ) + + # Delete the completed experiment + result = db.delete_experiment(exp_id) + assert result is True + + # Verify experiment is deleted + exp = db.get_experiment(exp_id) + assert exp is None + + # Verify status remains COMPLETED + with db._session() as session: + exp = session.query(Experiment).filter(Experiment.uuid == exp_id).first() + assert exp is not None + assert exp.is_del == 1 + assert exp.status == Status.COMPLETED + + +def test_delete_single_experiment_with_runs(db): + """Test that delete_experiment also deletes associated runs""" + org_id = uuid.uuid4() + team_id = uuid.uuid4() + user_id = uuid.uuid4() + + # Create experiment with runs + exp_id = db.create_experiment( + org_id=org_id, + team_id=team_id, + user_id=user_id, + name="exp_with_runs", + status=Status.RUNNING, + ) + + run_id_1 = db.create_run( + org_id=org_id, team_id=team_id, user_id=user_id, experiment_id=exp_id + ) + run_id_2 = db.create_run( + org_id=org_id, team_id=team_id, user_id=user_id, experiment_id=exp_id + ) + + # Delete the experiment + result = db.delete_experiment(exp_id) + assert result is True + + # Verify runs are also deleted + run1 = db.get_run(run_id_1) + run2 = db.get_run(run_id_2) + assert run1 is None + assert run2 is None From 06a770d42405f2b64374e9ae12726793cf343fc9 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Mon, 15 Jun 2026 16:28:51 +0800 Subject: [PATCH 2/2] fix lint Signed-off-by: kerthcet --- alphatrion/storage/sqlstore.py | 1 - alphatrion/tracing/tracing.py | 1 - 2 files changed, 2 deletions(-) diff --git a/alphatrion/storage/sqlstore.py b/alphatrion/storage/sqlstore.py index c3492bd..20d73de 100644 --- a/alphatrion/storage/sqlstore.py +++ b/alphatrion/storage/sqlstore.py @@ -752,7 +752,6 @@ def delete_experiment(self, experiment_id: uuid.UUID) -> bool: - RUNNING experiments -> CANCELLED - Other statuses remain unchanged """ - from sqlalchemy import case session = self._session() diff --git a/alphatrion/tracing/tracing.py b/alphatrion/tracing/tracing.py index 6cada59..55a11ec 100644 --- a/alphatrion/tracing/tracing.py +++ b/alphatrion/tracing/tracing.py @@ -1,4 +1,3 @@ - from opentelemetry.semconv_ai import TraceloopSpanKindValues from traceloop.sdk.decorators import agent as _agent from traceloop.sdk.decorators import task as _task