Coverage for bbconf/config_parser/bedbaseconfig.py: 69%
278 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 22:05 +0000
1import io
2import logging
3import os
4import warnings
5from pathlib import Path
6from typing import Literal
8import boto3
9import joblib
10import qdrant_client
11import requests
12import s3fs
13import yacman
14import zarr
15from botocore.client import BaseClient
16from botocore.exceptions import BotoCoreError, EndpointConnectionError
17from fastembed import TextEmbedding
18from geniml.region2vec.main import Region2VecExModel
19from geniml.search import BED2BEDSearchInterface
20from geniml.search.backends import BiVectorBackend, QdrantBackend
21from geniml.search.interfaces import BiVectorSearchInterface
22from geniml.search.query2vec import BED2Vec
23from qdrant_client import QdrantClient, models
24from sentence_transformers import SparseEncoder
25from umap import UMAP
26from zarr import Group as Z_GROUP
28from bbconf.config_parser.const import (
29 S3_BEDSET_PATH_FOLDER,
30 S3_FILE_PATH_FOLDER,
31 S3_PLOTS_PATH_FOLDER,
32)
33from bbconf.config_parser.models import ConfigFile
34from bbconf.const import PKG_NAME, ZARR_TOKENIZED_FOLDER
35from bbconf.db_utils import BaseEngine
36from bbconf.exceptions import (
37 BadAccessMethodError,
38 BedBaseConfError,
39 BedbaseS3ConnectionError,
40)
41from bbconf.helpers import get_absolute_path, get_bedbase_cfg
42from bbconf.models.base_models import FileModel
43from bbconf.models.bed_models import BedFiles, BedPlots
44from bbconf.models.bedset_models import BedSetPlots
45from bbconf.models.drs_models import AccessMethod, AccessURL
47_LOGGER = logging.getLogger(PKG_NAME)
50class BedBaseConfig:
51 """
52 Class to handle BEDbase configuration file and create objects for different modules.
53 """
55 def __init__(self, config: Path | str, init_ml: bool = True):
56 """
57 Initialize BedBaseConfig object.
59 Args:
60 config: Path to the configuration file.
61 init_ml: Initialize machine learning models used for search.
62 """
64 self.cfg_path = get_bedbase_cfg(config)
65 self._config = self._read_config_file(self.cfg_path)
66 self._db_engine = self._init_db_engine()
68 try:
69 self.qdrant_client: QdrantClient = self._init_qdrant_client()
70 except Exception as err:
71 _LOGGER.error(
72 f"Unable to create Qdrant client. Skipping ML model initialization. Error: {err}"
73 )
74 init_ml = False
76 if init_ml:
77 self.dense_encoder: TextEmbedding = self._init_dense_encoder()
78 self.sparse_encoder: SparseEncoder | None = self._init_sparce_model()
79 self.umap_encoder: UMAP | None = self._init_umap_model()
80 self.r2v_encoder: Region2VecExModel | None = self._init_r2v_encoder()
82 self._init_qdrant_hybrid(
83 qdrant_cl=self.qdrant_client,
84 dense_encoder=self.dense_encoder,
85 )
87 self.qdrant_file_backend: QdrantBackend | None = (
88 self._init_qdrant_file_backend(qdrant_cl=self.qdrant_client)
89 ) # used for bivec search
90 self._qdrant_text_backend: QdrantBackend | None = (
91 self._init_qdrant_text_backend(
92 qdrant_cl=self.qdrant_client,
93 dense_encoder=self.dense_encoder,
94 )
95 ) # used for bivec search
97 self.b2b_search_interface = self._init_b2b_search_interface(
98 qdrant_file_backend=self.qdrant_file_backend,
99 region_encoder=self.r2v_encoder,
100 )
102 self.bivec_search_interface = self._init_bivec_interface(
103 qdrant_file_backend=self.qdrant_file_backend,
104 qdrant_text_backend=self._qdrant_text_backend,
105 text_encoder=self.dense_encoder,
106 )
107 else:
108 _LOGGER.info(
109 "Skipping initialization of ML models, init_ml parameter set to False."
110 )
111 self.r2v_encoder = None
112 self.b2b_search_interface = None
113 self.bivec_search_interface = None
114 self.umap_encoder: UMAP | None = None
115 self.sparse_encoder = None
117 self._boto3_client = self._init_boto3_client()
119 @staticmethod
120 def _read_config_file(config_path: str) -> ConfigFile:
121 """
122 Read configuration file and insert default values if not set.
124 Args:
125 config_path: Configuration file path.
127 Returns:
128 Parsed ConfigFile object.
129 """
131 _config = yacman.YAMLConfigManager.from_yaml_file(filepath=config_path).exp
132 return ConfigFile(**_config)
134 @property
135 def config(self) -> ConfigFile:
136 """
137 Get configuration.
139 Returns:
140 Configuration object.
141 """
142 return self._config
144 @property
145 def db_engine(self) -> BaseEngine:
146 """
147 Get database engine.
149 Returns:
150 Database engine.
151 """
152 return self._db_engine
154 @property
155 def boto3_client(self) -> boto3.client:
156 """
157 Get boto3 client.
159 Returns:
160 Boto3 client.
161 """
162 return self._boto3_client
164 @property
165 def zarr_root(self) -> Z_GROUP | None:
166 """
167 Get zarr root object (Group).
169 Returns:
170 Zarr root group object.
171 """
173 try:
174 s3fc_obj = s3fs.S3FileSystem(
175 endpoint_url=self._config.s3.endpoint_url,
176 key=self._config.s3.aws_access_key_id,
177 secret=self._config.s3.aws_secret_access_key,
178 )
179 except BotoCoreError as e:
180 _LOGGER.error(f"Error in creating s3fs object: {e}")
181 warnings.warn(f"Error in creating s3fs object: {e}", UserWarning)
182 return None
184 s3_path = f"s3://{self._config.s3.bucket}/{ZARR_TOKENIZED_FOLDER}"
186 zarr_store = s3fs.S3Map(
187 root=s3_path, s3=s3fc_obj, check=False, create=self._config.s3.modify_access
188 )
189 cache = zarr.LRUStoreCache(zarr_store, max_size=2**28)
191 return zarr.group(store=cache, overwrite=False)
193 def _init_db_engine(self) -> BaseEngine:
194 """
195 Create database engine object using credentials provided in config file
196 """
198 _LOGGER.info("Initializing database engine...")
199 return BaseEngine(
200 host=self._config.database.host,
201 port=self._config.database.port,
202 database=self._config.database.database,
203 user=self._config.database.user,
204 password=self._config.database.password,
205 drivername=f"{self._config.database.dialect}+{self._config.database.driver}",
206 run_migrations=self._config.database.run_migrations,
207 )
209 def _init_qdrant_client(self) -> QdrantClient:
210 """
211 Create qdrant client object using credentials provided in config file
212 """
214 _LOGGER.info("Initializing qdrant client...")
216 try:
217 qdrant_cl = QdrantClient(
218 url=self.config.qdrant.host,
219 port=self.config.qdrant.port,
220 api_key=self.config.qdrant.api_key,
221 )
223 except qdrant_client.http.exceptions.ResponseHandlingException as err:
224 raise BedBaseConfError(
225 f"Error in Connection to qdrant! skipping... Error: {err}"
226 )
228 return qdrant_cl
230 def _init_qdrant_file_backend(
231 self, qdrant_cl: QdrantClient
232 ) -> QdrantBackend | None:
233 """
234 Create qdrant client object using credentials provided in config file.
236 Args:
237 qdrant_cl: QdrantClient object.
239 Returns:
240 QdrantClient or None
241 """
243 _LOGGER.info("Initializing qdrant bivec file backend...")
245 if not isinstance(qdrant_cl, QdrantClient):
246 _LOGGER.error(
247 "Unable to create Qdrant bivec file collection, qdrant client is None."
248 )
249 return None
251 try:
252 return QdrantBackend(
253 qdrant_client=qdrant_cl,
254 collection=self.config.qdrant.file_collection,
255 )
256 except Exception as e:
257 _LOGGER.error(f"Unable to create Qdrant collection: {e}")
258 return None
260 def _init_qdrant_text_backend(
261 self, qdrant_cl: QdrantClient, dense_encoder: TextEmbedding
262 ) -> QdrantBackend | None:
263 """
264 Create qdrant client text embedding object using credentials provided in config file.
266 Args:
267 qdrant_cl: QdrantClient object.
268 dense_encoder: TextEmbedding model for encoding text queries.
270 Returns:
271 QdrantClient or None
272 """
274 _LOGGER.info("Initializing qdrant bivec text backend...")
276 if not isinstance(qdrant_cl, QdrantClient):
277 _LOGGER.error(
278 "Unable to create Qdrant bivec text collection, qdrant client is None."
279 )
280 return None
281 if not isinstance(dense_encoder, TextEmbedding):
282 _LOGGER.error(
283 "Unable to create Qdrant bivec text collection, dense encoder is None."
284 )
285 return None
287 dimensions = int(dense_encoder.get_embedding_size(self._config.path.text2vec))
288 try:
289 return QdrantBackend(
290 qdrant_client=qdrant_cl,
291 dim=dimensions,
292 collection=self.config.qdrant.text_collection,
293 )
294 except Exception as e:
295 _LOGGER.error(f"Unable to create Qdrant collection: {e}")
296 return None
298 def _init_qdrant_hybrid(
299 self, qdrant_cl: QdrantClient, dense_encoder: TextEmbedding
300 ) -> None:
301 """
302 Create qdrant client with sparse and text embedding object using credentials provided in config file.
304 Args:
305 qdrant_cl: QdrantClient object.
306 dense_encoder: TextEmbedding model for encoding text queries.
308 Returns:
309 None, Initializes or creates the hybrid collection on the provided QdrantClient.
310 """
312 _LOGGER.info("Initializing qdrant sparse collection...")
314 if not isinstance(qdrant_cl, QdrantClient):
315 _LOGGER.error(
316 "Unable to create Qdrant hybrid collection, qdrant client is None."
317 )
318 return None
319 if not isinstance(dense_encoder, TextEmbedding):
320 _LOGGER.error(
321 "Unable to create Qdrant hybrid collection, dense encoder is None."
322 )
323 return None
325 dimensions = int(dense_encoder.get_embedding_size(self._config.path.text2vec))
326 collection_name = self.config.qdrant.hybrid_collection
328 try:
329 if not qdrant_cl.collection_exists(collection_name):
330 _LOGGER.info(
331 "Collection 'bedbase_query_search' does not exist, creating it."
332 )
334 qdrant_cl.create_collection(
335 collection_name=collection_name,
336 vectors_config={
337 "dense": models.VectorParams(
338 size=dimensions, distance=models.Distance.COSINE
339 ),
340 },
341 sparse_vectors_config={
342 "sparse": models.SparseVectorParams(
343 index=models.SparseIndexParams(
344 on_disk=False,
345 )
346 )
347 },
348 quantization_config=models.ScalarQuantization(
349 scalar=models.ScalarQuantizationConfig(
350 type=models.ScalarType.INT8,
351 quantile=0.99,
352 always_ram=True,
353 ),
354 ),
355 )
357 qdrant_cl.create_payload_index(
358 collection_name=collection_name,
359 field_name="assay",
360 field_type=models.PayloadSchemaType.KEYWORD,
361 )
362 qdrant_cl.create_payload_index(
363 collection_name=collection_name,
364 field_name="genome_alias",
365 field_type=models.PayloadSchemaType.KEYWORD,
366 )
368 except Exception as err:
369 _LOGGER.error(
370 f"Error in creating Qdrant hybrid collection! skipping... Error: {err}. Qdrant host: {self._config.qdrant.host}"
371 )
372 warnings.warn(
373 f"error in creating Qdrant hybrid collection! skipping... Error: {err}",
374 UserWarning,
375 )
376 return None
378 def _init_bivec_interface(
379 self,
380 qdrant_file_backend: QdrantBackend,
381 qdrant_text_backend: QdrantBackend,
382 text_encoder: TextEmbedding,
383 ) -> BiVectorSearchInterface | None:
384 """
385 Create BiVectorSearchInterface object using credentials provided in config file.
387 Args:
388 qdrant_file_backend: QdrantBackend for file vectors.
389 qdrant_text_backend: QdrantBackend for text vectors.
390 text_encoder: TextEmbedding model for encoding text queries.
392 Returns:
393 BiVectorSearchInterface.
394 """
396 _LOGGER.info("Initializing BiVectorBackend...")
397 search_backend = BiVectorBackend(
398 metadata_backend=qdrant_text_backend, bed_backend=qdrant_file_backend
399 )
400 _LOGGER.info("Initializing BiVectorSearchInterface...")
401 search_interface = BiVectorSearchInterface(
402 backend=search_backend,
403 query2vec=text_encoder,
404 )
405 return search_interface
407 def _init_b2b_search_interface(
408 self,
409 qdrant_file_backend: QdrantBackend,
410 region_encoder: Region2VecExModel | str,
411 ) -> BED2BEDSearchInterface | None:
412 """
413 Create Bed 2 BED search interface and return this object.
415 Returns:
416 Bed2BEDSearchInterface object.
417 """
418 try:
419 _LOGGER.info("Initializing search bed 2 bed search interfaces...")
420 return BED2BEDSearchInterface(
421 backend=qdrant_file_backend,
422 query2vec=BED2Vec(model=region_encoder),
423 )
424 except Exception as e:
425 _LOGGER.error("Error in creating BED2BEDSearchInterface object: " + str(e))
426 warnings.warn(
427 "Error in creating BED2BEDSearchInterface object: " + str(e),
428 UserWarning,
429 )
430 return None
432 def _init_r2v_encoder(self) -> Region2VecExModel | None:
433 """
434 Create Region2VecExModel object using credentials provided in config file
435 """
436 try:
437 _LOGGER.info(
438 f"Initializing region2vec encoder... Model used: {self.config.path.region2vec}"
439 )
440 return Region2VecExModel(self.config.path.region2vec)
441 except Exception as e:
442 _LOGGER.error(f"Error in creating Region2VecExModel object: {e}")
443 warnings.warn(
444 f"Error in creating Region2VecExModel object: {e}", UserWarning
445 )
446 return None
448 def _init_dense_encoder(self) -> TextEmbedding | None:
449 """
450 Initialize dense model from the specified path or huggingface model hub
451 """
453 _LOGGER.info(
454 f"Initializing dense encoder... Model used: {self.config.path.text2vec}"
455 )
456 dense_encoder = TextEmbedding(self.config.path.text2vec)
457 return dense_encoder
459 def _init_sparce_model(self) -> SparseEncoder | None:
460 """
461 Initialize SparseEncoder model from the specified path or huggingface model hub
462 """
463 try:
464 _LOGGER.info(
465 f"Initializing sparse encoder... Model used: {self.config.path.sparse_model}"
466 )
467 sparse_encoder = SparseEncoder(self.config.path.sparse_model)
468 except Exception as e:
469 _LOGGER.error(f"Error in creating SparseEncoder object: {e}")
470 warnings.warn(f"Error in creating SparseEncoder object: {e}", UserWarning)
471 return None
472 return sparse_encoder
474 def _init_umap_model(self) -> UMAP | None:
475 """
476 Load UMAP model from the specified path, or url
477 """
479 if not self.config.path.umap_model:
480 _LOGGER.warning(
481 "UMAP model path is not specified in the configuration, and won't be used."
482 )
483 return None
485 model_path = self.config.path.umap_model
486 umap_model = None
487 if model_path.startswith(("http://", "https://")):
488 try:
489 response = requests.get(model_path)
490 response.raise_for_status()
491 buffer = io.BytesIO(response.content)
492 umap_model = joblib.load(buffer)
493 _LOGGER.info(f"UMAP model loaded from URL: {model_path}")
494 except requests.RequestException as e:
495 _LOGGER.error(f"Error downloading UMAP model from URL: {e}")
496 return None
497 else:
498 try:
499 with open(model_path, "rb") as file:
500 umap_model = joblib.load(file)
501 _LOGGER.info(f"UMAP model loaded from local path: {model_path}")
502 except FileNotFoundError as e:
503 _LOGGER.error(f"Error loading UMAP model from local path: {e}")
504 return None
506 if not isinstance(umap_model, UMAP):
507 _LOGGER.error(f"Loaded object is not a UMAP instance: {type(umap_model)}")
508 return None
509 # np.random.seed(42)
510 umap_model.random_state = 42
511 return umap_model
513 def _init_boto3_client(
514 self,
515 ) -> BaseClient | None:
516 """
517 Create boto3 client object using credentials provided in config file.
519 Returns:
520 Boto3 client.
521 """
522 try:
523 return boto3.client(
524 "s3",
525 endpoint_url=self._config.s3.endpoint_url,
526 aws_access_key_id=self._config.s3.aws_access_key_id,
527 aws_secret_access_key=self._config.s3.aws_secret_access_key,
528 )
529 except Exception as e:
530 _LOGGER.error(f"Error in creating boto3 client object: {e}")
531 warnings.warn(f"Error in creating boto3 client object: {e}", UserWarning)
532 return None
534 def upload_s3(self, file_path: str, s3_path: Path | str) -> None:
535 """
536 Upload file to s3.
538 Args:
539 file_path: Local path to the file.
540 s3_path: Path to the file in s3 with file name.
542 Returns:
543 None.
544 """
545 if not self._boto3_client:
546 _LOGGER.warning(
547 "Could not upload file to s3. Connection to s3 not established. Skipping.."
548 )
549 raise BedbaseS3ConnectionError(
550 "Could not upload file to s3. Connection error."
551 )
552 if not os.path.exists(file_path):
553 raise BedBaseConfError(f"File {os.path.abspath(file_path)} does not exist.")
554 _LOGGER.info(f"Uploading file to s3: {s3_path}")
555 return self._boto3_client.upload_file(file_path, self.config.s3.bucket, s3_path)
557 def upload_files_s3(
558 self,
559 identifier: str,
560 files: BedFiles | BedPlots | BedSetPlots,
561 base_path: str,
562 type: Literal["files", "plots", "bedsets"] = "files",
563 ) -> BedFiles | BedPlots | BedSetPlots:
564 """
565 Upload files to s3.
567 Args:
568 identifier: Bed file identifier.
569 files: Dictionary with files to upload.
570 base_path: Local path to the output files.
571 type: Type of files to upload [files, plots, bedsets].
573 Returns:
574 None.
575 """
577 if type == "files":
578 s3_output_base_folder = S3_FILE_PATH_FOLDER
579 elif type == "plots":
580 s3_output_base_folder = S3_PLOTS_PATH_FOLDER
581 elif type == "bedsets":
582 s3_output_base_folder = S3_BEDSET_PATH_FOLDER
583 else:
584 raise BedBaseConfError(
585 f"Invalid type: {type}. Should be 'files', 'plots', or 'bedsets'"
586 )
588 for key, value in files:
589 if not value:
590 continue
591 file_base_name = os.path.basename(value.path)
592 file_path = get_absolute_path(value.path, base_path)
593 s3_path = os.path.join(
594 s3_output_base_folder,
595 identifier[0],
596 identifier[1],
597 file_base_name,
598 )
599 self.upload_s3(file_path, s3_path=s3_path)
601 setattr(value, "name", key)
602 setattr(value, "size", os.path.getsize(file_path))
603 setattr(value, "path", s3_path)
605 if value.path_thumbnail:
606 file_base_name_thumbnail = os.path.basename(value.path_thumbnail)
607 file_path_thumbnail = get_absolute_path(value.path_thumbnail, base_path)
608 s3_path_thumbnail = os.path.join(
609 s3_output_base_folder,
610 identifier[0],
611 identifier[1],
612 file_base_name_thumbnail,
613 )
614 self.upload_s3(file_path_thumbnail, s3_path=s3_path_thumbnail)
615 setattr(value, "path_thumbnail", s3_path_thumbnail)
617 return files
619 def delete_s3(self, s3_path: str) -> None:
620 """
621 Delete file from s3.
623 Args:
624 s3_path: Path to the file in s3.
626 Returns:
627 None.
628 """
629 if not self._boto3_client:
630 _LOGGER.warning(
631 "Could not delete file from s3. Connection to s3 not established. Skipping.."
632 )
633 raise BedbaseS3ConnectionError(
634 "Could not delete file from s3. Connection error."
635 )
636 try:
637 _LOGGER.info(f"Deleting file from s3: {s3_path}")
638 return self._boto3_client.delete_object(
639 Bucket=self.config.s3.bucket, Key=s3_path
640 )
641 except EndpointConnectionError:
642 raise BedbaseS3ConnectionError(
643 "Could not delete file from s3. Connection error."
644 )
646 def delete_files_s3(self, files: list[FileModel]) -> None:
647 """
648 Delete files from s3.
650 Args:
651 files: List of file objects.
653 Returns:
654 None.
655 """
656 for file in files:
657 self.delete_s3(file.path)
658 if file.path_thumbnail:
659 self.delete_s3(file.path_thumbnail)
660 return None
662 def get_prefixed_uri(self, postfix: str, access_id: str) -> str:
663 """
664 Return uri with correct prefix (schema).
666 Args:
667 postfix: Postfix of the uri (or everything after uri schema).
668 access_id: Access method name, e.g. http, s3, etc.
670 Returns:
671 Full uri path.
672 """
674 if self.config.access_methods is None:
675 _LOGGER.error("No access methods are defined in the config.")
676 raise BadAccessMethodError("No access methods are defined in the config.")
677 try:
678 prefix = getattr(self.config.access_methods, access_id).prefix
679 return os.path.join(prefix, postfix)
680 except (KeyError, AttributeError):
681 _LOGGER.error(f"Access method {access_id} is not defined.")
682 raise BadAccessMethodError(f"Access method {access_id} is not defined.")
684 def construct_access_method_list(self, rel_path: str) -> list[AccessMethod]:
685 """
686 Construct access method list for a given record.
688 Args:
689 rel_path: Relative path to the record.
691 Returns:
692 List of access methods.
693 """
694 if self.config.access_methods is None:
695 return []
696 access_methods = []
697 for access_id in self.config.access_methods.model_dump().keys():
698 access_dict = AccessMethod(
699 type=access_id,
700 access_id=access_id,
701 access_url=AccessURL(url=self.get_prefixed_uri(rel_path, access_id)),
702 region=self.config.access_methods.model_dump()[access_id].get(
703 "region", None
704 ),
705 )
706 access_methods.append(access_dict)
707 return access_methods