Rebasing Kvikio checks

Signed-off-by: Julio Faracco <[email protected]>

dasf/datasets/base.py

@@ -29,13 +29,15 @@
 
     from kvikio.nvcomp_codec import NvCompBatchCodec
     from kvikio.zarr import GDSStore
-    USE_KVIKIO=True
 except ImportError: # pragma: no cover
-    USE_KVIKIO=False
+    pass
 
 from pathlib import Path
 
 from dasf.utils.funcs import human_readable_size
+from dasf.utils.funcs import is_kvikio_supported
+from dasf.utils.funcs import is_gds_supported
+from dasf.utils.funcs import is_nvcomp_codec_supported
 from dasf.utils.decorators import task_handler
 from dasf.utils.types import is_array
 from dasf.utils.types import is_dask_array
@@ -493,7 +495,8 @@ def _lazy_load(self, xp, **kwargs):
             The data (or a Future load object, for `_lazy` operations).
 
         """
-        if self._backend == "kvikio" and USE_KVIKIO:
+        if (self._backend == "kvikio" and is_kvikio_supported() and
+            is_gds_supported() and is_nvcomp_codec_supported()):
             store = GDSStore(self._root_file)
             meta = json.loads(store[".zarray"])
             meta["compressor"] = NvCompBatchCodec("lz4").get_config()

dasf/utils/funcs.py

@@ -39,6 +39,18 @@
 except ImportError: # pragma: no cover
     JAX_SUPPORTED = False
 
+try:
+    import kvikio
+    KVIKIO_SUPPORTED = True
+except ImportError: # pragma: no cover
+    KVIKIO_SUPPORTED = False
+
+try:
+    from kvikio.nvcomp_codec import NvCompBatchCodec
+    NV_COMP_BATCH_CODEC_SUPPORTED = True
+except ImportError: # pragma: no cover
+    NV_COMP_BATCH_CODEC_SUPPORTED = False
+
 
 def human_readable_size(size, decimal=3) -> str:
     """
@@ -298,6 +310,31 @@ def is_gpu_supported() -> bool:
     return GPU_SUPPORTED and get_gpu_count() >= 1
 
 
+def is_kvikio_supported() -> bool:
+    """
+    Return if kvikio is supported (installed).
+    """
+    return KVIKIO_SUPPORTED
+
+
+def is_gds_supported() -> bool:
+    """
+    Return if GPU Direct Store is supported.
+    """
+    if is_kvikio_supported():
+        props = kvikio.DriverProperties()
+        return props.is_gds_available
+
+    return False
+
+
+def is_nvcomp_codec_supported() -> bool:
+    """
+    Return if NVidia Compressor Codecs are supported.
+    """
+    return NV_COMP_BATCH_CODEC_SUPPORTED
+
+
 def is_jax_supported() -> bool:
     """
     Return if JAX is supported.