Actually make dask.distributed optional #22

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- fix bug by making dependency `distributed` optional ![\#27](https://github.com/mllam/mllam-data-prep/pull/27)
 - change config example to call validation split `val` instead of `validation` [\#28](https://github.com/mllam/mllam-data-prep/pull/28)
 - fix typo in install dependency `distributed` ![\#20](https://github.com/mllam/mllam-data-prep/pull/20)
 - add missing `psutil` requirement. [\#21](https://github.com/mllam/mllam-data-prep/pull/21).

mllam_data_prep/__main__.py

@@ -1,13 +1,19 @@
 import os
 from pathlib import Path
 
-import psutil
-from dask.diagnostics import ProgressBar
-from dask.distributed import LocalCluster
 from loguru import logger
 
 from .create_dataset import create_dataset_zarr
 
+# Attempt to import psutil and dask.distributed modules
+DASK_DISTRIBUTED_AVAILABLE = True
+try:
+    import psutil
+    from dask.diagnostics import ProgressBar
+    from dask.distributed import LocalCluster
+except ImportError or ModuleNotFoundError:
+    DASK_DISTRIBUTED_AVAILABLE = False
+
 if __name__ == "__main__":
     import argparse
 
@@ -36,6 +42,13 @@
         ProgressBar().register()
 
     if args.dask_distributed_local_core_fraction > 0.0:
+        # Only run this block if dask.distributed is available
+        if not DASK_DISTRIBUTED_AVAILABLE:
+            raise ModuleNotFoundError(
+                "Currently dask.distributed isn't installed and therefore can't "
+                "be used in mllam-data-prep. Please install the optional dependency "
+                'with `python -m pip install "mllam-data-prep[dask-distributed]"`'
+            )
         # get the number of system cores
         n_system_cores = os.cpu_count()
         # compute the number of cores to use