fix: use manual db session in background tasks (#33)

* feat: db manual session

* fix: preview dataset query

* fix: cache dataset

* fix: cache dataset

* feat: separate shard statistics table

* chore: build ui

* fix: delete task status if done

* fix: unified background task status cache

* chore: build ui

* fix: python openapi client

* fix: shard sync delete

* fix: test

Author: hanchchch

lavender_data/server/background_worker/__init__.py

@@ -4,7 +4,6 @@
 
 from .background_worker import (
     TaskStatus,
-    TaskMetadata,
     BackgroundWorker,
     get_background_worker,
     setup_background_worker,
@@ -19,7 +18,6 @@
 
 __all__ = [
     "TaskStatus",
-    "TaskMetadata",
     "BackgroundWorker",
     "get_background_worker",
     "setup_background_worker",

lavender_data/server/background_worker/background_worker.py

@@ -21,14 +21,6 @@ class TaskStatus(BaseModel):
     total: int
 
 
-class TaskMetadata(BaseModel):
-    uid: str
-    name: str
-    start_time: datetime
-    kwargs: dict
-    status: Optional[TaskStatus] = None
-
-
 def set_task_status(
     task_id: str,
     status: Optional[str] = None,
@@ -40,21 +32,21 @@ def set_task_status(
     if _status is None:
         _status = TaskStatus(status="", current=0, total=0)
 
-    next(get_cache()).set(
-        f"task-{task_id}",
+    next(get_cache()).hset(
+        f"background-worker:tasks",
+        task_id,
         json.dumps(
             {
-                "status": (status if status is not None else _status.status),
-                "current": (current if current is not None else _status.current),
+                "status": status if status is not None else _status.status,
+                "current": current if current is not None else _status.current,
                 "total": total if total is not None else _status.total,
             }
         ),
-        ex=ex,
     )
 
 
 def get_task_status(task_uid: str) -> Optional[TaskStatus]:
-    status = next(get_cache()).get(f"task-{task_uid}")
+    status = next(get_cache()).hget(f"background-worker:tasks", task_uid)
     if status is None:
         return None
 
@@ -67,7 +59,15 @@ def get_task_status(task_uid: str) -> Optional[TaskStatus]:
 
 
 def delete_task_status(task_uid: str):
-    next(get_cache()).delete(f"task-{task_uid}")
+    next(get_cache()).hdel(f"background-worker:tasks", task_uid)
+
+
+def all_task_statuses() -> dict[str, TaskStatus]:
+    tasks = next(get_cache()).hgetall(f"background-worker:tasks")
+    return {
+        task_id: TaskStatus.model_validate(json.loads(t))
+        for task_id, t in tasks.items()
+    }
 
 
 class Aborted(Exception):
@@ -129,7 +129,6 @@ def _run_task_no_status(
 
 
 class TaskItem(NamedTuple):
-    metadata: TaskMetadata
     future: Future
     abort_event: threading.Event
 
@@ -141,40 +140,35 @@ def __init__(self, num_workers: int):
 
         self._process_pool = ProcessPool(self._num_workers)
 
-        self._tasks: list[TaskItem] = []
-        self._tasks_lock = threading.Lock()
-
-        self._task_status: dict[str, TaskStatus] = {}
-
         self._executor = ThreadPoolExecutor(self._num_workers)
+        self._abort_events: dict[str, threading.Event] = {}
+        self._futures: dict[str, Future] = {}
 
         self._start_cleanup_thread()
 
     def process_pool(self) -> ProcessPool:
         return self._process_pool
 
     def _cleanup_tasks(self):
-        with self._tasks_lock:
-            for t in self._tasks:
-                if get_task_status(t.metadata.uid) is None:
-                    self._tasks.remove(t)
+        for task_id, status in all_task_statuses().items():
+            if status.status == "completed":
+                delete_task_status(task_id)
+
+            if status.status in ["completed", "aborted", "failed"]:
+                self._abort_events.pop(task_id, None)
+                self._futures.pop(task_id, None)
 
     def _start_cleanup_thread(self):
         def _cleanup_tasks():
             while True:
-                time.sleep(1)
+                time.sleep(10)
                 self._cleanup_tasks()
 
         threading.Thread(target=_cleanup_tasks, daemon=True).start()
 
-    def running_tasks(self) -> list[TaskMetadata]:
+    def list_tasks(self) -> dict[str, TaskStatus]:
         self._cleanup_tasks()
-        with self._tasks_lock:
-            tasks = [t.metadata for t in self._tasks]
-            tasks.sort(key=lambda t: t.start_time)
-            for task in tasks:
-                task.status = get_task_status(task.uid)
-            return tasks
+        return {task_id: status for task_id, status in all_task_statuses().items()}
 
     def get_task_status(self, task_id: str) -> Optional[TaskStatus]:
         return get_task_status(task_id)
@@ -204,19 +198,8 @@ def thread_pool_submit(
             **kwargs,
         )
 
-        with self._tasks_lock:
-            self._tasks.append(
-                TaskItem(
-                    metadata=TaskMetadata(
-                        uid=task_id,
-                        name=task_name or func.__name__,
-                        start_time=datetime.now(UTC),
-                        kwargs=kwargs,
-                    ),
-                    future=future,
-                    abort_event=abort_event,
-                )
-            )
+        self._abort_events[task_id] = abort_event
+        self._futures[task_id] = future
 
         return task_id
 
@@ -228,20 +211,20 @@ def process_pool_submit(
         return self.process_pool().submit(func, **kwargs)
 
     def abort(self, task_id: str):
-        with self._tasks_lock:
-            status = get_task_status(task_id)
-            if status is not None:
-                delete_task_status(task_id)
+        status = get_task_status(task_id)
+        if status is not None:
+            delete_task_status(task_id)
+
+        if task_id in self._abort_events:
+            self._abort_events[task_id].set()
 
-            task = next((t for t in self._tasks if t.metadata.uid == task_id), None)
-            if task is not None:
-                task.abort_event.set()
-                task.future.cancel()
-                self._tasks.remove(task)
+        if task_id in self._futures:
+            self._futures[task_id].cancel()
 
     def abort_all(self):
-        for t in self._tasks:
-            self.abort(t.metadata.uid)
+        for task_id, status in all_task_statuses().items():
+            if status.status == "running":
+                self.abort(task_id)
 
     def shutdown(self):
         self._logger.debug("Shutting down background worker")

lavender_data/server/cli/create_api_key.py

@@ -3,7 +3,7 @@
 
 from sqlmodel import select
 
-from lavender_data.server.db import get_session, setup_db
+from lavender_data.server.db import db_manual_session, setup_db
 from lavender_data.server.db.models import ApiKey
 from lavender_data.server.settings import get_settings
 
@@ -13,18 +13,21 @@ def create_api_key(
     expires_at: Optional[datetime] = None,
 ):
     setup_db(get_settings().lavender_data_db_url)
-    session = next(get_session())
 
-    api_key = None
-    if note:
-        api_key = session.exec(select(ApiKey).where(ApiKey.note == note)).one_or_none()
+    with db_manual_session() as session:
+        api_key = None
+        if note:
+            api_key = session.exec(
+                select(ApiKey).where(ApiKey.note == note)
+            ).one_or_none()
 
-    if api_key is None:
-        api_key = ApiKey(note=note, expires_at=expires_at)
-        session.add(api_key)
-    else:
-        api_key.expires_at = expires_at
+        if api_key is None:
+            api_key = ApiKey(note=note, expires_at=expires_at)
+            session.add(api_key)
+        else:
+            api_key.expires_at = expires_at
+
+        session.commit()
+        session.refresh(api_key)
 
-    session.commit()
-    session.refresh(api_key)
     return api_key

lavender_data/server/dataset/preview.py

@@ -2,15 +2,25 @@
 import time
 from typing import Any, Union
 
-from sqlalchemy.exc import NoResultFound
+from sqlmodel import select
+from sqlalchemy.orm import selectinload
+
 import filetype
 import hashlib
 import numpy as np
 import json
 
 from lavender_data.server.settings import files_dir
-from lavender_data.server.db import get_session
-from lavender_data.server.db.models import Dataset, Shard
+from lavender_data.server.db import db_manual_session
+from lavender_data.server.db.models import (
+    Dataset,
+    Shard,
+    Shardset,
+    DatasetPublic,
+    ShardPublic,
+    ShardsetPublic,
+    DatasetColumnPublic,
+)
 from lavender_data.server.cache import CacheClient, get_cache
 from lavender_data.server.reader import (
     get_reader_instance,
@@ -30,8 +40,17 @@
     torch = None
 
 
+class _Shardset(ShardsetPublic):
+    shards: list[ShardPublic]
+    columns: list[DatasetColumnPublic]
+
+
+class _Dataset(DatasetPublic):
+    shardsets: list[_Shardset]
+
+
 def _read_dataset(
-    dataset: Dataset,
+    dataset: _Dataset,
     index: int,
     reader: ReaderInstance,
     cache: CacheClient,
@@ -165,14 +184,36 @@ def preview_dataset(
     offset: int,
     limit: int,
 ) -> list[dict[str, Any]]:
-    session = next(get_session())
     cache = next(get_cache())
     reader = get_reader_instance()
 
-    try:
-        dataset = session.get_one(Dataset, dataset_id)
-    except NoResultFound:
-        raise ValueError(f"Dataset {dataset_id} not found")
+    cached_dataset = cache.hget(f"preview:{dataset_id}", "dataset")
+    if cached_dataset is None:
+        with db_manual_session() as session:
+            dataset = session.exec(
+                select(Dataset)
+                .where(Dataset.id == dataset_id)
+                .options(
+                    selectinload(Dataset.shardsets).options(
+                        selectinload(Shardset.columns),
+                        selectinload(Shardset.shards),
+                    )
+                )
+            ).one()
+
+        if dataset is None:
+            raise ValueError(f"Dataset {dataset_id} not found")
+
+        dataset = _Dataset.model_validate(dataset)
+        cache.hset(f"preview:{dataset_id}", "dataset", dataset.model_dump_json())
+        for shardset in dataset.shardsets:
+            cache.hset(
+                f"preview:{dataset_id}",
+                f"dataset.shardsets.{shardset.id}",
+                shardset.model_dump_json(),
+            )
+    else:
+        dataset = _Dataset.model_validate_json(cached_dataset)
 
     samples = []
     for index in range(offset, offset + limit):