feat: fix for issues related with hitting limit of sqlalchemy pool + TCP connections for

lightly_studio/src/lightly_studio/api/server.py

@@ -30,8 +30,19 @@ def __init__(self, host: str, port: int) -> None:
 
     def start(self) -> None:
         """Start the API server using Uvicorn."""
-        # start the app
-        uvicorn.run(app, host=self.host, port=self.port, http="h11")
+        # start the app with connection limits and timeouts
+        uvicorn.run(
+            app,
+            host=self.host,
+            port=self.port,
+            http="h11",
+            # https://uvicorn.dev/settings/#resource-limits
+            limit_concurrency=100,  # Max concurrent connections
+            limit_max_requests=10000,  # Max requests before worker restart
+            # https://uvicorn.dev/settings/#timeouts
+            timeout_keep_alive=5,  # Keep-alive timeout in seconds
+            timeout_graceful_shutdown=30,  # Graceful shutdown timeout
+        )
 
 
 def _get_available_port(host: str, preferred_port: int, max_tries: int = 50) -> int:

lightly_studio/src/lightly_studio/db_manager.py

@@ -8,8 +8,8 @@
 from typing import Generator
 
 from fastapi import Depends
+from sqlalchemy import StaticPool
 from sqlalchemy.engine import Engine
-from sqlalchemy.pool import Pool
 from sqlmodel import Session, SQLModel, create_engine
 from typing_extensions import Annotated
 
@@ -27,21 +27,33 @@ def __init__(
         self,
         engine_url: str | None = None,
         cleanup_existing: bool = False,
-        poolclass: type[Pool] | None = None,
+        single_threaded: bool = False,
     ) -> None:
         """Create a new instance of the DatabaseEngine.
 
         Args:
             engine_url: The database engine URL. If None, defaults to a local DuckDB file.
             cleanup_existing: If True, deletes the existing database file if it exists.
-            poolclass: The SQLAlchemy pool class to use. Use StaticPool for
-                in-memory databases for testing, otherwise different DB connections
-                connect to different in-memory databases.
+            single_threaded: If True, creates a single-threaded engine suitable for testing.
         """
         self._engine_url = engine_url if engine_url else "duckdb:///lightly_studio.db"
         if cleanup_existing:
             _cleanup_database_file(engine_url=self._engine_url)
-        self._engine = create_engine(url=self._engine_url, poolclass=poolclass)
+
+        self._engine = None
+
+        if single_threaded:
+            self._engine = create_engine(
+                url=self._engine_url,
+                poolclass=StaticPool,
+            )
+        else:
+            self._engine = create_engine(
+                url=self._engine_url,
+                pool_size=10,
+                max_overflow=40,
+            )
+
         SQLModel.metadata.create_all(self._engine)
 
     @contextmanager
@@ -113,7 +125,10 @@ def connect(db_file: str | None = None, cleanup_existing: bool = False) -> None:
             is used.
     """
     engine_url = f"duckdb:///{db_file}" if db_file is not None else None
-    engine = DatabaseEngine(engine_url=engine_url, cleanup_existing=cleanup_existing)
+    engine = DatabaseEngine(
+        engine_url=engine_url,
+        cleanup_existing=cleanup_existing,
+    )
     set_engine(engine=engine)
 
 

lightly_studio/tests/conftest.py

@@ -9,7 +9,6 @@
 import pytest
 from fastapi.testclient import TestClient
 from pydantic import BaseModel
-from sqlalchemy import StaticPool
 from sqlmodel import Session
 
 from lightly_studio import db_manager
@@ -50,7 +49,7 @@
 @pytest.fixture
 def db_session() -> Generator[Session, None, None]:
     """Create a test database manager session."""
-    test_manager = DatabaseEngine("duckdb:///:memory:", poolclass=StaticPool)
+    test_manager = DatabaseEngine("duckdb:///:memory:", single_threaded=True)
     with test_manager.session() as session:
         yield session
 

lightly_studio/tests/core/conftest.py

@@ -23,7 +23,10 @@ def patch_dataset(
     mocker.patch.object(
         db_manager,
         "get_engine",
-        return_value=db_manager.DatabaseEngine("duckdb:///:memory:"),
+        return_value=db_manager.DatabaseEngine(
+            engine_url="duckdb:///:memory:",
+            single_threaded=True,
+        ),
     )
 
     # Create a test-specific EmbeddingManager singleton.

lightly_studio/tests/test_db_manager.py

@@ -55,7 +55,7 @@ def test_set_engine__file_db(
 ) -> None:
     """Test set_engine function."""
     engine_url = f"duckdb:///{tmp_path / 'test.db'}"
-    engine = DatabaseEngine(engine_url=engine_url)
+    engine = DatabaseEngine(engine_url=engine_url, single_threaded=True)
     db_manager.set_engine(engine=engine)
 
     fetched_engine = db_manager.get_engine()
@@ -68,7 +68,7 @@ def test_set_engine__memory_db(
 ) -> None:
     """Test set_engine function."""
     engine_url = "duckdb:///:memory:"
-    engine = DatabaseEngine(engine_url=engine_url)
+    engine = DatabaseEngine(engine_url=engine_url, single_threaded=True)
     db_manager.set_engine(engine=engine)
 
     fetched_engine = db_manager.get_engine()
@@ -81,9 +81,11 @@ def test_set_engine__raises_if_already_set(
 ) -> None:
     """Test set_engine raises if the engine is already set."""
     engine_url = "duckdb:///:memory:"
-    db_manager.set_engine(engine=DatabaseEngine(engine_url=engine_url))
+    db_manager.set_engine(engine=DatabaseEngine(engine_url=engine_url, single_threaded=True))
     with pytest.raises(RuntimeError, match="Database engine is already set and cannot be changed."):
-        db_manager.set_engine(engine=DatabaseEngine("duckdb:///:memory:"))
+        db_manager.set_engine(
+            engine=DatabaseEngine(engine_url="duckdb:///:memory:", single_threaded=True)
+        )
 
 
 def test_connect(
@@ -116,7 +118,10 @@ def test_connect__db_file_none(
 
     engine = db_manager.get_engine()
     assert engine is mock_engine
-    mock_engine_class.assert_called_once_with(engine_url=None, cleanup_existing=True)
+    mock_engine_class.assert_called_once_with(
+        engine_url=None,
+        cleanup_existing=True,
+    )
 
 
 def test_session_data_consistency(mocker: MockerFixture, tmp_path: Path) -> None:

lightly_studio_view/e2e/fixtures/cocoDataset.ts

@@ -19,7 +19,7 @@ export const cocoDataset = {
     totalLabels: 71,
 
     /** Default page size when loading samples */
-    defaultPageSize: 100,
+    defaultPageSize: 50,
 
     /** Expected filename when exporting annotations */
     exportFilename: 'coco_export.json',

lightly_studio_view/src/lib/components/LazyTrigger/LazyTrigger.svelte

@@ -36,4 +36,4 @@
     });
 </script>
 
-<div bind:this={triggerRef}></div>
+<div bind:this={triggerRef} style="height: 1px; width: 1px;"></div>

lightly_studio_view/src/lib/hooks/useImagesInfinite/useImagesInfinite.ts

@@ -93,7 +93,7 @@ const buildRequestBody = (params: ImagesInfiniteParams, pageParam: number): Read
     const baseBody: ReadImagesRequest = {
         pagination: {
             offset: pageParam,
-            limit: 100
+            limit: 50
         },
         text_embedding: params.text_embedding,
         filters: {