feat(brick): geNomad score and annotation rings

New feature: geNomad score and annotations rings [#35]

app/src/lib/stores/RingStore.ts

@@ -44,9 +44,7 @@ function toggleRingVisibility(id: string) {
     rings.update(currentRings => {
         return currentRings.map(ring => {
             if (ring.id === id) {
-                console.log(ring.visible)
                 ring.visible = !ring.visible;
-                console.log(ring.visible)
             }
             return ring;
         });

brick/api/schemas.py

@@ -9,7 +9,7 @@
 from .core.config import settings
 
 from ..rings import BlastRing, AnnotationRing, LabelRing, ReferenceRing, GenomadRing
-from ..rings import LabelSegment, RingReference, RingType, Ring
+from ..rings import RingSegment, RingReference, RingType, Ring
 from ..rings import GenomadPredictionClass
 
 SessionID = Annotated[
@@ -351,7 +351,7 @@ class AnnotationRingResponse(BaseModel):
 
 class LabelRingSchema(RingSchema):
     tsv_id: SessionFileID | None = None
-    labels: List[LabelSegment] = []
+    labels: List[RingSegment] = []
 
     @field_validator("tsv_id")
     @classmethod

brick/api/tasks.py

@@ -1,18 +1,15 @@
-import pandas
 import json
+import uuid
 import shutil
-import subprocess
+import pandas
 import tempfile
+import subprocess
 import contextlib
-import uuid
-
-from typing import List
 
 from Bio import SeqIO
-from typing import Tuple, Annotated, Optional
 from pathlib import Path
 from datetime import datetime
-from celery.exceptions import TimeLimitExceeded, SoftTimeLimitExceeded
+from typing import List, Tuple, Annotated, Optional
 
 from .core.config import settings
 from .core.celery import celery_app
@@ -35,7 +32,7 @@
     GenomadRingSchema,
 )
 from ..rings import BlastRing, AnnotationRing, LabelRing, ReferenceRing, GenomadRing
-from ..rings import Ring, RingSegment, RingType, RingReference
+from ..rings import Ring, RingSegment, RingType, RingReference, LabelSegment
 
 from ..utils import slice_fasta_sequences
 
@@ -308,7 +305,18 @@ def process_label_ring(
             )
 
         if ring_schema.labels:
-            ring.add_custom_labels(labels=ring_schema.labels, sanitize=True)
+            ring.add_custom_labels(
+                labels=[
+                    LabelSegment(
+                        start=ring_segment.start,
+                        end=ring_segment.end,
+                        text=ring_segment.text,
+                        labelIdentifier=str(uuid.uuid4()),
+                    )
+                    for ring_segment in ring_schema.labels
+                ],
+                sanitize=True,
+            )
 
         result = check_or_update_label_ring(
             reference=ring_schema.reference, new_segments=ring.data
@@ -319,7 +327,7 @@ def process_label_ring(
                 session_id=ring_schema.reference.session_id, ring=ring
             )
         else:
-            ring = result
+            ring = result.model_copy()
 
         return {"success": True, "result": ring.model_dump()}
     except Exception as e:
@@ -773,7 +781,7 @@ def check_or_update_label_ring(
             "id": reference.session_id,
             "rings": {
                 "$elemMatch": {
-                    "type": "label",  # Replace with the actual value used for RingType.LABEL
+                    "type": RingType.LABEL,
                     "reference.reference_id": reference.reference_id,
                     "reference.sequence.id": reference.sequence.id,
                 }

brick/rings.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
-from pydantic import BaseModel, Field, validator, ValidationError
+from pydantic import BaseModel, Field, validator
 from statistics import mean
 from strenum import StrEnum
 from pathlib import Path
-from typing import List
+from typing import List, Generator, Tuple
 from Bio import SeqIO
 
 import pandas
@@ -279,7 +279,7 @@ def parse_blastn_output(
 
 
 # Generator function
-def parse_aggregated_genomad_output(file: Path) -> GenomadEntry:
+def parse_aggregated_genomad_output(file: Path) -> Generator[GenomadEntry]:
     """
     Parses the aggregated_classification output from a sliced genome file
     """
@@ -306,7 +306,7 @@ def extract_genomad_contiguous_segments(
     min_segment_length: int,  # should be a multiple of slice length
     prediction_classes: List[GenomadPredictionClass],
     segment_type: RingSegmentType,
-) -> List[RingSegment] | List[LabelSegment]:
+) -> Generator[RingSegment] | Generator[LabelSegment]:
     """Extracts segments of high probabilty for each prediction class with a minimum total length for label or annotation rings"""
 
     segments = []
@@ -385,7 +385,7 @@ def process_column(column: str):
 
 def get_start_end_from_seq_name(
     seq_name: str, name_split: str = "__", range_split: str = ".."
-) -> tuple(int, int, str):
+) -> Tuple[int, int, str]:
 
     try:
         seq_id: str = seq_name.split(name_split)[0]

docker/docker-compose.localhost.yml

@@ -13,7 +13,6 @@ services:
     build:
       context: .
       dockerfile: docker/Dockerfile.app
-    restart: unless-stopped
     expose:
       - "5173"
     security_opt:
@@ -42,7 +41,6 @@ services:
     build:
       context: .
       dockerfile: docker/Dockerfile.app.dev
-    restart: unless-stopped
     volumes:
       - ./app:/usr/src/app
       - /usr/src/app/node_modules
@@ -76,11 +74,11 @@ services:
     build: 
       context: .
       dockerfile: docker/Dockerfile.server
-    restart: unless-stopped
     command: uvicorn brick.api.main:app --host 0.0.0.0 --port 8080
     volumes:
       - api:/data
       - work:/tmp
+      - databases:/data
     env_file: docker/brick.env
     environment:
       CORS_ORIGINS: http://app:5173
@@ -94,6 +92,7 @@ services:
     depends_on:
       - redis
       - mongodb
+      - databases
     secrets:
       - brick_db_user
       - brick_db_pwd
@@ -106,11 +105,11 @@ services:
     build: 
       context: .
       dockerfile: docker/Dockerfile.server
-    restart: unless-stopped
     command: uvicorn brick.api.main:app --host 0.0.0.0 --port 8080
     volumes:
       - api_dev:/data
       - work:/tmp
+      - databases:/data
     env_file: docker/brick.env
     environment:
       CORS_ORIGINS: http://app-dev:5174
@@ -124,6 +123,7 @@ services:
     depends_on:
       - redis
       - mongodb
+      - databases
     secrets:
       - brick_db_user
       - brick_db_pwd
@@ -145,10 +145,10 @@ services:
     build: 
       context: .
       dockerfile: docker/Dockerfile.server
-    restart: unless-stopped
     command: celery -A brick.api.core.celery.celery_app worker --loglevel=INFO
     volumes:
       - work:/tmp
+      - databases:/data
     env_file: docker/brick.env
     environment:
       CORS_ORIGINS: http://app:5173
@@ -161,6 +161,7 @@ services:
       - api
       - redis
       - mongodb
+      - databases
     secrets:
       - brick_db_user
       - brick_db_pwd
@@ -173,10 +174,10 @@ services:
     build: 
       context: .
       dockerfile: docker/Dockerfile.server
-    restart: unless-stopped
     command: celery -A brick.api.core.celery.celery_app worker --loglevel=INFO
     volumes:
       - work:/tmp
+      - databases:/data
     env_file: docker/brick.env
     environment:
       CORS_ORIGINS: http://app-dev:5174
@@ -189,12 +190,20 @@ services:
       - api-dev
       - redis
       - mongodb
+      - databases
     secrets:
       - brick_db_user
       - brick_db_pwd     
     networks:
       - internal 
-
+    deploy:
+      resources:
+        limits:            # process is killed
+          cpus: 8          # either number of cores or fraction
+          memory: 42G
+        reservations:      # resources reserved for container
+          cpus: 8
+          memory: 20G      # determined by the largest database 
 
   # ========= 
   # DATABASES 
@@ -210,7 +219,6 @@ services:
       - dev
       - test
     image: redis:7.0.12-alpine
-    restart: unless-stopped
     expose:
       - "6379"
     volumes:
@@ -249,13 +257,35 @@ services:
     networks:
       - internal
 
+  # ================
+  # DATABASE STORAGE
+  # ================
+
+  # Specification of the image tag will cause the 
+  # image to not rebuild when using --build. This
+  # is intended as we want to update the storage
+  # volume with the database manually, rather than
+  # downloading on each --build
+
+  databases:
+    image: brick-database-storage:latest
+    profiles:
+      - prod
+      - dev
+    build: 
+      context: .
+      dockerfile: docker/Dockerfile.dbs
+    tty: true
+    volumes:
+      - databases:/data  # /data/genomad_db
+    networks:
+      - internal
 
   # ================================================
   # DELETES SESSIONS AND UPLOADED FILES PERIODICALLY
   # ================================================
 
   data-cleaner:
-    restart: unless-stopped
     profiles:
       - server
     build: 
@@ -280,7 +310,6 @@ services:
       - internal
 
   data-cleaner-dev:
-    restart: unless-stopped
     profiles:
       - server-dev
     build: 
@@ -344,6 +373,8 @@ volumes:
     driver: local
   mongo_db_cfg:
     driver: local
+  databases:
+    driver: local
 
 secrets:
   mongo_root_user: