From 5d3b6d60729e473f73887f8ae2cf96cdfc5fb65f Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Wed, 17 Dec 2025 17:31:08 -0500 Subject: [PATCH] Updated zarr version to 3.0 --- bbconf/config_parser/bedbaseconfig.py | 31 +++++++++-- bbconf/modules/bedfiles.py | 63 ++++++++++++++++------ manual_testing.py | 75 ++++++++++++++++++++------- requirements/requirements-all.txt | 4 +- 4 files changed, 130 insertions(+), 43 deletions(-) diff --git a/bbconf/config_parser/bedbaseconfig.py b/bbconf/config_parser/bedbaseconfig.py index c96edca5..3b7feb48 100644 --- a/bbconf/config_parser/bedbaseconfig.py +++ b/bbconf/config_parser/bedbaseconfig.py @@ -24,6 +24,7 @@ from sentence_transformers import SparseEncoder from umap import UMAP from zarr import Group as Z_GROUP +from zarr.storage import FsspecStore from bbconf.config_parser.const import ( S3_BEDSET_PATH_FOLDER, @@ -190,20 +191,30 @@ def zarr_root(self) -> Union[Z_GROUP, None]: endpoint_url=self._config.s3.endpoint_url, key=self._config.s3.aws_access_key_id, secret=self._config.s3.aws_secret_access_key, + asynchronous=True, ) except BotoCoreError as e: _LOGGER.error(f"Error in creating s3fs object: {e}") warnings.warn(f"Error in creating s3fs object: {e}", UserWarning) return None - s3_path = f"s3://{self._config.s3.bucket}/{ZARR_TOKENIZED_FOLDER}" + s3_path = f"{self._config.s3.bucket}/{ZARR_TOKENIZED_FOLDER}" - zarr_store = s3fs.S3Map( - root=s3_path, s3=s3fc_obj, check=False, create=self._config.s3.modify_access + store = FsspecStore( + fs=s3fc_obj, + path=s3_path, ) - cache = zarr.LRUStoreCache(zarr_store, max_size=2**28) - return zarr.group(store=cache, overwrite=False) + try: + root = zarr.open_group( + store=store, + mode="a" if self._config.s3.modify_access else "r", + ) + return root + except Exception as e: + _LOGGER.error(f"Error opening zarr group: {e}") + warnings.warn(f"Error opening zarr group: {e}", UserWarning) + return None def _init_db_engine(self) -> BaseEngine: """ @@ -496,6 +507,11 @@ def _init_umap_model(self) -> Union[UMAP, None]: except requests.RequestException as e: _LOGGER.error(f"Error downloading UMAP model from URL: {e}") return None + except TypeError as e: + _LOGGER.error( + f"Error loading UMAP model from URL. Unable open pickle file. Error: {e}" + ) + return None else: try: with open(model_path, "rb") as file: @@ -504,6 +520,11 @@ def _init_umap_model(self) -> Union[UMAP, None]: except FileNotFoundError as e: _LOGGER.error(f"Error loading UMAP model from local path: {e}") return None + except TypeError as e: + _LOGGER.error( + f"Error loading UMAP model from URL. Unable open pickle file. Error: {e}" + ) + return None if not isinstance(umap_model, UMAP): _LOGGER.error(f"Loaded object is not a UMAP instance: {type(umap_model)}") diff --git a/bbconf/modules/bedfiles.py b/bbconf/modules/bedfiles.py index 4cfbf588..7505f01f 100644 --- a/bbconf/modules/bedfiles.py +++ b/bbconf/modules/bedfiles.py @@ -1626,14 +1626,15 @@ def add_tokenized( ) if self.exist_tokenized(bed_id, universe_id): + _LOGGER.info("Tokenized file already exists in the database.") if not overwrite: - if not overwrite: - raise TokenizeFileExistsError( - "Tokenized file already exists in the database. " - "Set overwrite to True to overwrite it." - ) - else: - self.delete_tokenized(bed_id, universe_id) + raise TokenizeFileExistsError( + "Tokenized file already exists in the database. " + "Set overwrite to True to overwrite it." + ) + else: + _LOGGER.info("Overwriting existing tokenized file in the database.") + self.delete_tokenized(bed_id, universe_id) path = self._add_zarr_s3( bed_id=bed_id, @@ -1664,16 +1665,32 @@ def _add_zarr_s3( :return: zarr path """ - univers_group = self.config.zarr_root.require_group(universe_id) + root = self.config.zarr_root - if not univers_group.get(bed_id): + # Handle group creation (require_group is deprecated in zarr 3.x) + try: + univers_group = root[universe_id] + except KeyError: + univers_group = root.create_group(universe_id) + + # Check existence + bed_exists = bed_id in univers_group + + if not bed_exists: _LOGGER.info("Saving tokenized vector to s3") - path = univers_group.create_dataset(bed_id, data=tokenized_vector).path + array = univers_group.create_array( + name=bed_id, + data=np.array(tokenized_vector, dtype="int64"), + ) + path = array.name elif overwrite: _LOGGER.info("Overwriting tokenized vector in s3") - path = univers_group.create_dataset( - bed_id, data=tokenized_vector, overwrite=True - ).path + del univers_group[bed_id] + array = univers_group.create_array( + name=bed_id, + data=np.array(tokenized_vector, dtype="int64"), + ) + path = array.name else: raise TokenizeFileExistsError( "Tokenized file already exists in the database. " @@ -1694,12 +1711,19 @@ def get_tokenized(self, bed_id: str, universe_id: str) -> TokenizedBedResponse: if not self.exist_tokenized(bed_id, universe_id): raise TokenizeFileNotExistError("Tokenized file not found in the database.") - univers_group = self.config.zarr_root.require_group(universe_id) + + root = self.config.zarr_root + + try: + univers_group = root[universe_id] + data = list(univers_group[bed_id][:]) # Explicit slice for full read + except KeyError: + raise TokenizeFileNotExistError("Tokenized file not found in the database.") return TokenizedBedResponse( universe_id=universe_id, bed_id=bed_id, - tokenized_bed=list(univers_group[bed_id]), + tokenized_bed=data, ) def delete_tokenized(self, bed_id: str, universe_id: str) -> None: @@ -1713,9 +1737,14 @@ def delete_tokenized(self, bed_id: str, universe_id: str) -> None: """ if not self.exist_tokenized(bed_id, universe_id): raise TokenizeFileNotExistError("Tokenized file not found in the database.") - univers_group = self.config.zarr_root.require_group(universe_id) - del univers_group[bed_id] + root = self.config.zarr_root + + try: + univers_group = root[universe_id] + del univers_group[bed_id] # Delete syntax unchanged + except KeyError: + raise TokenizeFileNotExistError("Tokenized file not found in the database.") with Session(self._sa_engine) as session: statement = delete(TokenizedBed).where( diff --git a/manual_testing.py b/manual_testing.py index ee1fd705..ee2ecdb2 100644 --- a/manual_testing.py +++ b/manual_testing.py @@ -5,6 +5,7 @@ import numpy as np import s3fs import zarr +from zarr.storage import FsspecStore # from dotenv import load_dotenv from geniml.io import RegionSet @@ -38,23 +39,33 @@ def zarr_local(): tokenized_name = "0dcdf8986a72a3d85805bbc9493a13026l" overwrite = True - root = zarr.group( - store="/home/bnt4me/virginia/repos/bbconf/zarr_test", overwrite=False + # Use zarr.open_group instead of zarr.group + root = zarr.open_group( + store="/home/bnt4me/virginia/repos/bbconf/zarr_test", mode="a" ) - univers_group = root.require_group("7126993b14054a32de2da4a0b9173be5") - if not univers_group.get(tokenized_name): + # Handle group creation (require_group is deprecated) + universe_id = "7126993b14054a32de2da4a0b9173be5" + try: + univers_group = root[universe_id] + except KeyError: + univers_group = root.create_group(universe_id) + + # Check existence and handle overwrite + if tokenized_name not in univers_group: print("not overwriting") - ua = univers_group.create_dataset(tokenized_name, data=tok_regions) + ua = univers_group.create_array( + name=tokenized_name, data=np.array(tok_regions, dtype="int64") + ) elif overwrite: print("overwriting") - ua = univers_group.create_dataset( - tokenized_name, data=tok_regions, overwrite=True + del univers_group[tokenized_name] + ua = univers_group.create_array( + name=tokenized_name, data=np.array(tok_regions, dtype="int64") ) else: raise ValueError("fff") ua = univers_group - univers_group._delitem_nosync() def zarr_s3(): @@ -87,16 +98,34 @@ def zarr_s3(): endpoint_url=os.getenv("AWS_ENDPOINT_URL"), key=os.getenv("AWS_ACCESS_KEY_ID"), secret=os.getenv("AWS_SECRET_ACCESS_KEY"), + asynchronous=False, + skip_instance_cache=True, ) print(os.getenv("AWS_SECRET_ACCESS_KEY")) - s3_path = "s3://bedbase/new/" + s3_path = "bedbase/new/" # Remove s3:// prefix for FsspecStore + + # Use FsspecStore instead of S3Map + LRUStoreCache + store = FsspecStore(fs=s3fc_obj, path=s3_path) - zarr_store = s3fs.S3Map(root=s3_path, s3=s3fc_obj, check=False, create=True) - cache = zarr.LRUStoreCache(zarr_store, max_size=2**28) + # Use zarr.open_group instead of zarr.group + root = zarr.open_group(store=store, mode="a") - root = zarr.group(store=cache, overwrite=False) - univers_group = root.require_group("7126993b14054a32de2da4a0b9173be5") - univers_group.create_dataset(tokenized_name, data=tok_regions, overwrite=True) + # Handle group creation (require_group is deprecated) + universe_id = "7126993b14054a32de2da4a0b9173be5" + try: + univers_group = root[universe_id] + except KeyError: + univers_group = root.create_group(universe_id) + + # Handle overwrite + if tokenized_name in univers_group: + if overwrite: + del univers_group[tokenized_name] + + # Use create_array instead of create_dataset + univers_group.create_array( + name=tokenized_name, data=np.array(tok_regions, dtype="int64") + ) f = univers_group[tokenized_name] @@ -109,16 +138,24 @@ def get_from_s3(): # endpoint_url="https://s3.us-west-002.backblazeb2.com/", # key=os.getenv("AWS_ACCESS_KEY_ID"), # secret=os.getenv("AWS_SECRET_ACCESS_KEY"), + asynchronous=False, + skip_instance_cache=True, ) import s3fs - s3fc_obj = s3fs.S3FileSystem(endpoint_url="https://s3.us-west-002.backblazeb2.com/") - s3_path = "s3://bedbase/tokenized.zarr/" - zarr_store = s3fs.S3Map(root=s3_path, s3=s3fc_obj, check=False, create=True) - cache = zarr.LRUStoreCache(zarr_store, max_size=2**28) + s3fc_obj = s3fs.S3FileSystem( + endpoint_url="https://s3.us-west-002.backblazeb2.com/", + asynchronous=False, + skip_instance_cache=True, + ) + s3_path = "bedbase/tokenized.zarr/" # Remove s3:// prefix for FsspecStore + + # Use FsspecStore instead of S3Map + LRUStoreCache + store = FsspecStore(fs=s3fc_obj, path=s3_path) - root = zarr.group(store=cache, overwrite=False) + # Use zarr.open_group instead of zarr.group + root = zarr.open_group(store=store, mode="r") # print(str(root.tree)) diff --git a/requirements/requirements-all.txt b/requirements/requirements-all.txt index cbf01e83..ee55ed00 100644 --- a/requirements/requirements-all.txt +++ b/requirements/requirements-all.txt @@ -1,7 +1,7 @@ yacman >= 0.9.1 sqlalchemy >= 2.0.0 gtars >= 0.4.0 -geniml[ml] >= 0.8.2 +geniml[ml] >= 0.8.3 psycopg >= 3.1.15 coloredlogs pydantic >= 2.9.0 @@ -9,7 +9,7 @@ botocore >= 1.34.0, < 1.36.0 boto3 >= 1.34.54, < 1.36.0 pephubclient >= 0.4.5 sqlalchemy_schemadisplay -zarr < 3.0.0 +zarr >= 3.0.0 pyyaml >= 6.0.1 # for s3fs because of the errors s3fs >= 2024.3.1 pandas >= 2.0.0