feat: delete Dataset (InftyAI#201)

* Add delete dataset feature

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

* Add delete dataset feature

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

---------

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

Author: kerthcet

alphatrion/artifact/artifact.py

@@ -9,8 +9,7 @@
 
 
 class Artifact:
-    def __init__(self, team_id: str, insecure: bool = False):
-        self._team_id = team_id
+    def __init__(self, insecure: bool = False):
         self._url = get_registry_url()
         self._client = oras.client.OrasClient(
             hostname=self._url.strip("/"), auth_backend="token", insecure=insecure
@@ -50,7 +49,7 @@ def push(
         if version is None:
             version = utiltime.now_2_hash()
 
-        path = f"{self._team_id}/{repo_name}:{version}"
+        path = f"{repo_name}:{version}"
         target = f"{self._url}/{path}"
 
         try:
@@ -61,7 +60,7 @@ def push(
         return path
 
     def list_versions(self, repo_name: str) -> list[str]:
-        target = f"{self._url}/{self._team_id}/{repo_name}"
+        target = f"{self._url}/{repo_name}"
         try:
             tags = self._client.get_tags(target)
             return tags
@@ -91,7 +90,7 @@ def pull(
                            (defaults to ORAS temp directory)
         :return: list of absolute file paths that were downloaded
         """
-        path = f"{self._team_id}/{repo_name}:{version}"
+        path = f"{repo_name}:{version}"
         target = f"{self._url}/{path}"
 
         if output_dir:
@@ -115,7 +114,7 @@ def pull(
                 os.chdir(original_dir)
 
     def delete(self, repo_name: str, versions: str | list[str]):
-        target = f"{self._url}/{self._team_id}/{repo_name}"
+        target = f"{self._url}/{repo_name}"
 
         try:
             self._client.delete_tags(target, tags=versions)

alphatrion/log/log.py

@@ -10,9 +10,9 @@
 from alphatrion.snapshot.snapshot import (
     checkpoint_path,
 )
+from alphatrion.storage import runtime as storage_runtime
 
 BEST_RESULT_PATH = "best_result_path"
-EXECUTION_RESULT = "execution_result"
 
 
 async def log_artifact(
@@ -45,7 +45,7 @@ async def log_artifact(
     if runtime is None:
         raise RuntimeError("Runtime is not initialized. Please call init() first.")
 
-    if not runtime.artifact_storage_enabled():
+    if not storage_runtime.artifact_storage_enabled():
         raise RuntimeError(
             "Artifact storage is not enabled in the runtime."
             "Set ENABLE_ARTIFACT_STORAGE=true in the environment variables."
@@ -59,7 +59,7 @@ async def log_artifact(
 
     loop = asyncio.get_running_loop()
     return await loop.run_in_executor(
-        None, runtime._artifact.push, repo_name, paths, version
+        None, runtime._artifact.push, f"{runtime.team_id}/{repo_name}", paths, version
     )
 
 

alphatrion/runtime/runtime.py

@@ -3,7 +3,6 @@
 import uuid
 
 from alphatrion import envs
-from alphatrion.artifact.artifact import Artifact
 from alphatrion.storage import runtime as storage_runtime
 from alphatrion.storage.sqlstore import SQLStore
 
@@ -58,6 +57,7 @@ def __init__(
         storage_runtime.init()
         self._metadb = storage_runtime.storage_runtime().metadb
         self._tracestore = storage_runtime.storage_runtime().tracestore
+        self._artifact = storage_runtime.storage_runtime().artifact
 
         self._user_id = user_id
         self._team_id = team_id
@@ -74,18 +74,9 @@ def __init__(
             self._team_id = teams[0].uuid
 
         self._root_path = os.getenv(envs.ROOT_PATH, os.path.expanduser("~/.alphatrion"))
-
-        artifact_insecure = os.getenv(envs.ARTIFACT_INSECURE, "false").lower() == "true"
-
-        if self.artifact_storage_enabled():
-            self._artifact = Artifact(team_id=self._team_id, insecure=artifact_insecure)
-
         if not os.path.exists(self._root_path):
             os.makedirs(self._root_path, exist_ok=True)
 
-    def artifact_storage_enabled(self) -> bool:
-        return os.getenv(envs.ENABLE_ARTIFACT_STORAGE, "true").lower() == "true"
-
     @property
     def metadb(self) -> SQLStore:
         return self._metadb
@@ -94,6 +85,10 @@ def metadb(self) -> SQLStore:
     def tracestore(self):
         return self._tracestore
 
+    @property
+    def artifact(self):
+        return self._artifact
+
     @property
     def user_id(self) -> uuid.UUID:
         return self._user_id

alphatrion/server/graphql/resolvers.py

@@ -352,8 +352,10 @@ async def list_artifact_tags(
     ) -> list[ArtifactTag]:
         """List tags for a repository."""
 
-        arf = artifact.Artifact(team_id=team_id, insecure=True)
-        return [ArtifactTag(name=tag) for tag in arf.list_versions(repo_name)]
+        arf = runtime.storage_runtime().artifact
+        return [
+            ArtifactTag(name=tag) for tag in arf.list_versions(f"{team_id}/{repo_name}")
+        ]
 
     @staticmethod
     async def list_artifact_files(
@@ -362,8 +364,8 @@ async def list_artifact_files(
         """List files in an artifact without loading content."""
 
         try:
-            arf = artifact.Artifact(team_id=team_id, insecure=True)
-            file_paths = arf.pull(repo_name=repo_name, version=tag)
+            arf = runtime.storage_runtime().artifact
+            file_paths = arf.pull(repo_name=f"{team_id}/{repo_name}", version=tag)
 
             if not file_paths:
                 return []
@@ -405,11 +407,11 @@ async def get_artifact_content(
         """Get artifact content from registry."""
         try:
             # Initialize artifact client
-            arf = artifact.Artifact(team_id=team_id, insecure=True)
+            arf = runtime.storage_runtime().artifact
 
             # Pull the artifact - ORAS will manage temp directory
             # Returns absolute paths to files in ORAS temp directory
-            file_paths = arf.pull(repo_name=repo_name, version=tag)
+            file_paths = arf.pull(repo_name=f"{team_id}/{repo_name}", version=tag)
 
             if not file_paths:
                 raise RuntimeError("No files found in artifact")
@@ -875,4 +877,21 @@ def delete_experiments(experiment_ids: list[strawberry.ID]) -> int:
     @staticmethod
     def delete_dataset(dataset_id: strawberry.ID) -> bool:
         metadb = runtime.storage_runtime().metadb
+        artifact = runtime.storage_runtime().artifact
+        dataset = metadb.get_dataset(dataset_id=dataset_id)
+
+        # delete the artifact file as well
+        if dataset:
+            try:
+                repo_name, version = dataset.path.split(":", 1)
+                artifact.delete(repo_name=repo_name, versions=version)
+            except Exception as e:
+                print(f"Failed to delete artifact for dataset {dataset_id}: {e}")
+
         return metadb.delete_dataset(dataset_id=dataset_id)
+
+    @staticmethod
+    def delete_datasets(dataset_ids: list[strawberry.ID]) -> bool:
+        for id in dataset_ids:
+            GraphQLMutations.delete_dataset(dataset_id=id)
+        return True

alphatrion/server/graphql/schema.py

@@ -141,6 +141,7 @@ def datasets(
 
     dataset: Dataset | None = strawberry.field(resolver=GraphQLResolvers.get_dataset)
 
+
 @strawberry.type
 class Mutation:
     @strawberry.mutation
@@ -175,5 +176,9 @@ def delete_experiments(self, experiment_ids: list[strawberry.ID]) -> int:
     def delete_dataset(self, dataset_id: strawberry.ID) -> bool:
         return GraphQLMutations.delete_dataset(dataset_id=dataset_id)
 
+    @strawberry.mutation
+    def delete_datasets(self, dataset_ids: list[strawberry.ID]) -> bool:
+        return GraphQLMutations.delete_datasets(dataset_ids=dataset_ids)
+
 
 schema = strawberry.Schema(query=Query, mutation=Mutation)

alphatrion/server/graphql/types.py

@@ -253,6 +253,7 @@ class RemoveUserFromTeamInput:
     user_id: strawberry.ID
     team_id: strawberry.ID
 
+
 # Artifact types
 @strawberry.type
 class ArtifactRepository:

alphatrion/storage/runtime.py

@@ -6,6 +6,7 @@
 from traceloop.sdk import Traceloop
 
 from alphatrion import envs
+from alphatrion.artifact.artifact import Artifact
 from alphatrion.storage.sqlstore import SQLStore
 from alphatrion.storage.tracestore import TraceStore
 from alphatrion.tracing.clickhouse_exporter import ClickHouseSpanExporter
@@ -54,6 +55,10 @@ def __init__(self):
             tracer_provider = trace.get_tracer_provider()
             tracer_provider.add_span_processor(ContextAttributesSpanProcessor())
 
+        artifact_insecure = os.getenv(envs.ARTIFACT_INSECURE, "false").lower() == "true"
+        if artifact_storage_enabled():
+            self._artifact = Artifact(insecure=artifact_insecure)
+
         self._inited = True
 
     @property
@@ -70,6 +75,10 @@ def flush(self):
             if isinstance(tracer_provider, TracerProvider):
                 tracer_provider.force_flush(timeout_millis=5000)
 
+    @property
+    def artifact(self):
+        return self._artifact
+
 
 def init():
     """
@@ -85,3 +94,7 @@ def storage_runtime() -> StorageRuntime:
     if __STORAGE_RUNTIME__ is None:
         raise RuntimeError("StorageRuntime is not initialized. Call init() first.")
     return __STORAGE_RUNTIME__
+
+
+def artifact_storage_enabled() -> bool:
+    return os.getenv(envs.ENABLE_ARTIFACT_STORAGE, "true").lower() == "true"

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

@@ -0,0 +1,54 @@
+import { useMutation, useQueryClient } from '@tanstack/react-query';
+import { graphqlMutation, mutations } from '../lib/graphql-client';
+
+interface DeleteDatasetResponse {
+  deleteDataset: boolean;
+}
+
+interface DeleteDatasetsResponse {
+  deleteDatasets: number;
+}
+
+/**
+ * Hook to delete a single dataset
+ */
+export function useDeleteDataset() {
+  const queryClient = useQueryClient();
+
+  return useMutation({
+    mutationFn: async (datasetId: string) => {
+      const data = await graphqlMutation<DeleteDatasetResponse>(
+        mutations.deleteDataset,
+        { datasetId }
+      );
+      return data.deleteDataset;
+    },
+    onSuccess: () => {
+      // Invalidate datasets queries to refetch the list
+      queryClient.invalidateQueries({ queryKey: ['datasets'] });
+      queryClient.invalidateQueries({ queryKey: ['dataset'] });
+    },
+  });
+}
+
+/**
+ * Hook to delete multiple datasets in batch
+ */
+export function useDeleteDatasets() {
+  const queryClient = useQueryClient();
+
+  return useMutation({
+    mutationFn: async (datasetIds: string[]) => {
+      const data = await graphqlMutation<DeleteDatasetsResponse>(
+        mutations.deleteDatasets,
+        { datasetIds }
+      );
+      return data.deleteDatasets;
+    },
+    onSuccess: () => {
+      // Invalidate datasets queries to refetch the list
+      queryClient.invalidateQueries({ queryKey: ['datasets'] });
+      queryClient.invalidateQueries({ queryKey: ['dataset'] });
+    },
+  });
+}

dashboard/src/lib/graphql-client.ts

@@ -430,4 +430,16 @@ export const mutations = {
       deleteExperiments(experimentIds: $experimentIds)
     }
   `,
+
+  deleteDataset: `
+    mutation DeleteDataset($datasetId: ID!) {
+      deleteDataset(datasetId: $datasetId)
+    }
+  `,
+
+  deleteDatasets: `
+    mutation DeleteDatasets($datasetIds: [ID!]!) {
+      deleteDatasets(datasetIds: $datasetIds)
+    }
+  `,
 };