Stricter check for query planning. (NVIDIA#107)

* Stricter query planning checks with newer versions of dask

Signed-off-by: Ayush Dattagupta <[email protected]>

* Add checks to tests/__init__

Signed-off-by: Ayush Dattagupta <[email protected]>

* Check sys.modules to ensure dask-expr is not enabled

Signed-off-by: Ayush Dattagupta <[email protected]>

* Search for "dask_expr" in sys modules

Co-authored-by: Richard (Rick) Zamora <[email protected]>
Signed-off-by: Ayush Dattagupta <[email protected]>

* use dask_expr instead of dask-expr

Signed-off-by: Ayush Dattagupta <[email protected]>

---------

Signed-off-by: Ayush Dattagupta <[email protected]>
Co-authored-by: Richard (Rick) Zamora <[email protected]>

Author: 2 people

examples/fuzzy_deduplication.py

@@ -16,7 +16,6 @@
 import time
 
 import dask
-from dask import dataframe as dd
 
 from nemo_curator import FuzzyDuplicates, FuzzyDuplicatesConfig
 from nemo_curator.datasets import DocumentDataset
@@ -49,6 +48,8 @@ def main(args):
 
         t0 = time.time()
         if filetype == "parquet":
+            from dask import dataframe as dd
+
             input_dataset = DocumentDataset(
                 dd.read_parquet(
                     dataset_dir,

examples/slurm/start-slurm.sh

@@ -67,6 +67,7 @@ export CUDF_SPILL="1"
 export RMM_SCHEDULER_POOL_SIZE="1GB"
 export RMM_WORKER_POOL_SIZE="72GiB"
 export LIBCUDF_CUFILE_POLICY=OFF
+export DASK_DATAFRAME__QUERY_PLANNING=False
 
 
 # =================================================================

nemo_curator/__init__.py

@@ -12,15 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import sys
+
 import dask
 
 # Disable query planning if possible
 # https://github.com/NVIDIA/NeMo-Curator/issues/73
-if dask.config.get("dataframe.query-planning") is True:
+if dask.config.get("dataframe.query-planning") is True or "dask_expr" in sys.modules:
     raise NotImplementedError(
-        "NeMo Curator does not support query planning yet. "
-        "Please disable query planning before importing "
-        "`nemo_curator`, `dask.dataframe` or `dask_cudf`."
+        """
+        NeMo Curator does not support query planning yet.
+        Please disable query planning before importing
+        `dask.dataframe` or `dask_cudf`. This can be done via:
+        `export DASK_DATAFRAME__QUERY_PLANNING=False`, or
+        importing `dask.dataframe/dask_cudf` after importing
+        `nemo_curator`.
+        """
     )
 else:
     dask.config.set({"dataframe.query-planning": False})

tests/__init__.py

@@ -12,15 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import sys
+
 import dask
 
-# Disable query planning before any tests are loaded
+# Disable query planning if possible
 # https://github.com/NVIDIA/NeMo-Curator/issues/73
-if dask.config.get("dataframe.query-planning") is True:
+if dask.config.get("dataframe.query-planning") is True or "dask_expr" in sys.modules:
     raise NotImplementedError(
-        "NeMo Curator does not support query planning yet. "
-        "Please disable query planning before importing "
-        "`nemo_curator`, `dask.dataframe` or `dask_cudf`."
+        """
+        NeMo Curator does not support query planning yet.
+        Please disable query planning before importing
+        `dask.dataframe` or `dask_cudf`. This can be done via:
+        `export DASK_DATAFRAME__QUERY_PLANNING=False`, or
+        importing `dask.dataframe/dask_cudf` after importing
+        `nemo_curator`.
+        """
     )
 else:
     dask.config.set({"dataframe.query-planning": False})