Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions alphatrion/server/graphql/resolvers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
8 changes: 8 additions & 0 deletions alphatrion/server/graphql/schema.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
8 changes: 6 additions & 2 deletions alphatrion/storage/sql_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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,
Expand Down
96 changes: 87 additions & 9 deletions alphatrion/storage/sqlstore.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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(
Expand Down
27 changes: 27 additions & 0 deletions dashboard/src/components/ui/checkbox.tsx
Original file line number Diff line number Diff line change
@@ -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 };
54 changes: 54 additions & 0 deletions dashboard/src/hooks/use-experiment-mutations.ts
Original file line number Diff line number Diff line change
@@ -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'] });
},
});
}
32 changes: 28 additions & 4 deletions dashboard/src/lib/graphql-client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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: `
Expand Down Expand Up @@ -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)
}
`,
};
4 changes: 2 additions & 2 deletions dashboard/src/pages/experiments/[id].tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand All @@ -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,
})}
Expand Down
Loading
Loading