Updated zarr requirementes to v3.0 [FUTURE RELEASE]

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)}")

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(

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))
 
 

requirements/requirements-all.txt

@@ -1,15 +1,15 @@
 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
 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