Delete experiments should delete runs as well (#191)

* Delete experiments should delete runs as well

Signed-off-by: kerthcet <kerthcet@gmail.com>

* fix lint

Signed-off-by: kerthcet <kerthcet@gmail.com>

* fix lint

Signed-off-by: kerthcet <kerthcet@gmail.com>

* Add 4 tags per line

Signed-off-by: kerthcet <kerthcet@gmail.com>

* update layout

Signed-off-by: kerthcet <kerthcet@gmail.com>

* fix test

Signed-off-by: kerthcet <kerthcet@gmail.com>

---------

Signed-off-by: kerthcet <kerthcet@gmail.com>

Author: kerthcet

alphatrion/server/graphql/resolvers.py

@@ -692,3 +692,21 @@ def remove_user_from_team(input: RemoveUserFromTeamInput) -> bool:
 
         # Remove user from team (deletes TeamMember entry)
         return metadb.remove_user_from_team(user_id=user_id, team_id=team_id)
+
+    @staticmethod
+    # TODO: We should have the team_id in the header for authz, and verify the
+    # team_id matches the experiment's team_id before allowing deletion.
+    def delete_experiment(experiment_id: strawberry.ID) -> bool:
+        metadb = runtime.storage_runtime().metadb
+        # Soft delete experiment by setting is_del flag
+        return metadb.delete_experiment(experiment_id=experiment_id)
+
+    @staticmethod
+    # TODO: We should have the team_id in the header for authz, and verify the
+    # team_id matches the experiment's team_id before allowing deletion.
+    def delete_experiments(experiment_ids: list[strawberry.ID]) -> int:
+        metadb = runtime.storage_runtime().metadb
+        # Convert strawberry IDs to UUIDs
+        uuids = [uuid.UUID(exp_id) for exp_id in experiment_ids]
+        # Soft delete experiments by setting is_del flag
+        return metadb.delete_experiments(experiment_ids=uuids)

alphatrion/server/graphql/schema.py

@@ -126,5 +126,13 @@ def add_user_to_team(self, input: AddUserToTeamInput) -> bool:
     def remove_user_from_team(self, input: RemoveUserFromTeamInput) -> bool:
         return GraphQLMutations.remove_user_from_team(input=input)
 
+    @strawberry.mutation
+    def delete_experiment(self, experiment_id: strawberry.ID) -> bool:
+        return GraphQLMutations.delete_experiment(experiment_id=experiment_id)
+
+    @strawberry.mutation
+    def delete_experiments(self, experiment_ids: list[strawberry.ID]) -> int:
+        return GraphQLMutations.delete_experiments(experiment_ids=experiment_ids)
+
 
 schema = strawberry.Schema(query=Query, mutation=Mutation)

alphatrion/storage/sql_models.py

@@ -114,7 +114,9 @@ class Experiment(Base):
 
     uuid = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
     team_id = Column(UUID(as_uuid=True), nullable=False)
-    user_id = Column(UUID(as_uuid=True), nullable=True)
+    user_id = Column(
+        UUID(as_uuid=True), nullable=True, comment="User who created the experiment"
+    )
     name = Column(String, nullable=False)
     description = Column(String, nullable=True)
     meta = Column(
@@ -171,7 +173,9 @@ class Run(Base):
     uuid = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
     team_id = Column(UUID(as_uuid=True), nullable=False)
     experiment_id = Column(UUID(as_uuid=True), nullable=False)
-    user_id = Column(UUID(as_uuid=True), nullable=True)
+    user_id = Column(
+        UUID(as_uuid=True), nullable=True, comment="User who created the run"
+    )
     meta = Column(
         MutableDict.as_mutable(JSON),
         nullable=True,

alphatrion/storage/sqlstore.py

@@ -346,6 +346,20 @@ def create_experiment(
         uid = uuid.uuid4()
 
         session = self._session()
+        # TODO: add back the validation.
+        # # verify user is in the team
+        # membership = (
+        #     session.query(TeamMember)
+        #     .filter(
+        #         TeamMember.user_id == user_id,
+        #         TeamMember.team_id == team_id,
+        #     )
+        #     .first()
+        # )
+        # if membership is None:
+        #     session.close()
+        #     raise ValueError("User must be a member of the team to create experiment")
+
         new_exp = Experiment(
             uuid=uid,
             team_id=team_id,
@@ -396,22 +410,22 @@ def get_experiment(self, experiment_id: uuid.UUID) -> Experiment | None:
         return exp
 
     # Different team may have the same experiment name.
-    def get_exp_by_name(self, name: str, team_id: uuid.UUID) -> Experiment | None:
+    def get_exp_by_name(
+        self, name: str, team_id: uuid.UUID, include_deleted: bool = False
+    ) -> Experiment | None:
         # make sure the team exists
         team = self.get_team(team_id)
         if team is None:
             return None
 
         session = self._session()
-        trial = (
-            session.query(Experiment)
-            .filter(
-                Experiment.name == name,
-                Experiment.team_id == team_id,
-                Experiment.is_del == 0,
-            )
-            .first()
+        query = session.query(Experiment).filter(
+            Experiment.name == name,
+            Experiment.team_id == team_id,
         )
+        if not include_deleted:
+            query = query.filter(Experiment.is_del == 0)
+        trial = query.first()
         session.close()
         return trial
 
@@ -532,6 +546,70 @@ def list_exps_by_timeframe(
         session.close()
         return exps
 
+    def delete_experiment(self, experiment_id: uuid.UUID) -> bool:
+        session = self._session()
+
+        # Try to delete the experiment
+        exp = (
+            session.query(Experiment)
+            .filter(Experiment.uuid == experiment_id, Experiment.is_del == 0)
+            .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:
+            exp.is_del = 1
+            session.commit()
+            session.close()
+            return True
+
+        # Even if experiment doesn't exist, commit the run deletions
+        session.commit()
+        session.close()
+        return False
+
+    def delete_experiments(self, experiment_ids: list[uuid.UUID]) -> int:
+        """
+        Batch delete experiments by setting is_del flag.
+        Also deletes all associated runs.
+        Returns the number of experiments successfully deleted.
+        """
+        session = self._session()
+        # Delete the experiments
+        # if experiment is running, skip deletion for that experiment
+        filtered_exps = (
+            session.query(Experiment.uuid)
+            .filter(
+                Experiment.uuid.in_(experiment_ids),
+                Experiment.is_del == 0,
+                Experiment.status != Status.RUNNING,
+            )
+            .all()
+        )
+        filtered_exp_ids = [exp_id for (exp_id,) in filtered_exps]  # unpack tuples
+
+        deleted_count = (
+            session.query(Experiment)
+            .filter(Experiment.uuid.in_(filtered_exp_ids))
+            .update({Experiment.is_del: 1}, synchronize_session=False)
+        )
+        # Delete all runs associated with these experiments
+        session.query(Run).filter(Run.experiment_id.in_(filtered_exp_ids)).update(
+            {Run.is_del: 1}, synchronize_session=False
+        )
+        session.commit()
+        session.close()
+        return deleted_count
+
     # ---------- Run APIs ----------
 
     def create_run(

dashboard/src/components/ui/checkbox.tsx

@@ -0,0 +1,27 @@
+import * as React from 'react';
+import { cn } from '../../lib/utils';
+
+export interface CheckboxProps
+  extends React.InputHTMLAttributes<HTMLInputElement> {}
+
+const Checkbox = React.forwardRef<HTMLInputElement, CheckboxProps>(
+  ({ className, ...props }, ref) => {
+    return (
+      <input
+        type="checkbox"
+        className={cn(
+          'h-4 w-4 rounded border-gray-300 text-primary cursor-pointer',
+          'focus:ring-0 focus:ring-offset-0 focus:outline-none',
+          'checked:border-primary checked:bg-primary',
+          className
+        )}
+        ref={ref}
+        {...props}
+      />
+    );
+  }
+);
+
+Checkbox.displayName = 'Checkbox';
+
+export { Checkbox };

dashboard/src/hooks/use-experiment-mutations.ts

@@ -0,0 +1,54 @@
+import { useMutation, useQueryClient } from '@tanstack/react-query';
+import { graphqlMutation, mutations } from '../lib/graphql-client';
+
+interface DeleteExperimentResponse {
+  deleteExperiment: boolean;
+}
+
+interface DeleteExperimentsResponse {
+  deleteExperiments: number;
+}
+
+/**
+ * Hook to delete a single experiment
+ */
+export function useDeleteExperiment() {
+  const queryClient = useQueryClient();
+
+  return useMutation({
+    mutationFn: async (experimentId: string) => {
+      const data = await graphqlMutation<DeleteExperimentResponse>(
+        mutations.deleteExperiment,
+        { experimentId }
+      );
+      return data.deleteExperiment;
+    },
+    onSuccess: () => {
+      // Invalidate experiments queries to refetch the list
+      queryClient.invalidateQueries({ queryKey: ['experiments'] });
+      queryClient.invalidateQueries({ queryKey: ['experiment'] });
+    },
+  });
+}
+
+/**
+ * Hook to delete multiple experiments in batch
+ */
+export function useDeleteExperiments() {
+  const queryClient = useQueryClient();
+
+  return useMutation({
+    mutationFn: async (experimentIds: string[]) => {
+      const data = await graphqlMutation<DeleteExperimentsResponse>(
+        mutations.deleteExperiments,
+        { experimentIds }
+      );
+      return data.deleteExperiments;
+    },
+    onSuccess: () => {
+      // Invalidate experiments queries to refetch the list
+      queryClient.invalidateQueries({ queryKey: ['experiments'] });
+      queryClient.invalidateQueries({ queryKey: ['experiment'] });
+    },
+  });
+}

dashboard/src/lib/graphql-client.ts

@@ -3,10 +3,8 @@ import axios from 'axios';
 /**
  * GraphQL client for AlphaTrion backend
  *
- * The backend provides a read-only GraphQL API at /graphql
- * with queries for teams, experiments, runs, and metrics.
- *
- * No subscriptions or mutations are currently supported.
+ * The backend provides a GraphQL API at /graphql
+ * with queries and mutations for teams, experiments, runs, and metrics.
  */
 
 // Use relative URL to work with proxy in development
@@ -63,6 +61,17 @@ export async function graphqlQuery<T>(
   }
 }
 
+/**
+ * Execute a GraphQL mutation
+ */
+export async function graphqlMutation<T>(
+  mutation: string,
+  variables?: Record<string, unknown>
+): Promise<T> {
+  // Mutations use the same endpoint and logic as queries
+  return graphqlQuery<T>(mutation, variables);
+}
+
 // GraphQL query templates
 export const queries = {
   listTeams: `
@@ -343,3 +352,18 @@ export const queries = {
   `,
 
 };
+
+// GraphQL mutation templates
+export const mutations = {
+  deleteExperiment: `
+    mutation DeleteExperiment($experimentId: ID!) {
+      deleteExperiment(experimentId: $experimentId)
+    }
+  `,
+
+  deleteExperiments: `
+    mutation DeleteExperiments($experimentIds: [ID!]!) {
+      deleteExperiments(experimentIds: $experimentIds)
+    }
+  `,
+};

dashboard/src/pages/experiments/[id].tsx

@@ -366,7 +366,7 @@ export function ExperimentDetailPage() {
                       <TableRow className="hover:bg-transparent border-b">
                         <TableHead className="h-11 text-xs font-semibold uppercase tracking-wider text-muted-foreground bg-muted/50">UUID</TableHead>
                         <TableHead className="h-11 text-xs font-semibold uppercase tracking-wider text-muted-foreground bg-muted/50">Status</TableHead>
-                        <TableHead className="h-11 text-xs font-semibold uppercase tracking-wider text-muted-foreground bg-muted/50 text-right">Created</TableHead>
+                        <TableHead className="h-11 text-xs font-semibold uppercase tracking-wider text-muted-foreground bg-muted/50">Created</TableHead>
                       </TableRow>
                     </TableHeader>
                     <TableBody>
@@ -385,7 +385,7 @@ export function ExperimentDetailPage() {
                               {run.status}
                             </Badge>
                           </TableCell>
-                          <TableCell className="py-3 text-sm text-muted-foreground text-right">
+                          <TableCell className="py-3 text-sm text-muted-foreground">
                             {formatDistanceToNow(new Date(run.createdAt), {
                               addSuffix: true,
                             })}