move cupyx import to not happen on driver (#588)

Signed-off-by: Erik Ordentlich <[email protected]>
NVIDIA · Mar 8, 2024 · f36bbc2 · f36bbc2
1 parent 9e03ac7
commit f36bbc2
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/python/src/spark_rapids_ml/classification.py b/python/src/spark_rapids_ml/classification.py
@@ -39,7 +39,6 @@
     import cupy as cp
     from pyspark.ml._typing import ParamMap
 
-import cupyx
 import numpy as np
 import pandas as pd
 import scipy
@@ -960,6 +959,7 @@ def _logistic_regression_fit(
             dfs: FitInputType,
             params: Dict[str, Any],
         ) -> Dict[str, Any]:
+            import cupyx
             from cuml.linear_model.logistic_regression_mg import LogisticRegressionMG
 
             X_list = [x for (x, _, _) in dfs]

diff --git a/python/src/spark_rapids_ml/utils.py b/python/src/spark_rapids_ml/utils.py
@@ -21,8 +21,8 @@
 if TYPE_CHECKING:
     import cudf
     import cupy as cp
+    import cupyx
 
-import cupyx
 import numpy as np
 import pandas as pd
 import scipy
@@ -218,6 +218,8 @@ def _concat_and_free(
 
     if the type of input arrays is scipy or cupyx csr_matrix, 'order' parameter will not be used.
     """
+    import cupyx
+
     if isinstance(array_list[0], scipy.sparse.csr_matrix):
         concated = scipy.sparse.vstack(array_list)
     elif isinstance(array_list[0], cupyx.scipy.sparse.csr_matrix):