From ba35f8e68afa5983d1e5ee1a4dd78cd844c53481 Mon Sep 17 00:00:00 2001
From: Bradley Dice <bdice@bradleydice.com>
Date: Fri, 20 Dec 2024 12:42:35 -0600
Subject: [PATCH] Reduce dependencies on numba. (#1761)

This PR makes `numba` an optional dependency of RMM.

We are keeping `numba` as a hard dependency in tests, though I explored what it would look like as a soft dependency in e2ff7f1. It turns out that the current RMM test suite relies on `numba` for about 90% of the tests, as a way to copy data from host to device and back (to verify that the allocations are valid and usable).

Closes #1760.

Authors:
  - Bradley Dice (https://github.com/bdice)

Approvers:
  - Matthew Murray (https://github.com/Matt711)
  - Mark Harris (https://github.com/harrism)
  - Vyas Ramasubramani (https://github.com/vyasr)

URL: https://github.com/rapidsai/rmm/pull/1761
---
 dependencies.yaml               |  4 ++--
 python/rmm/pyproject.toml       |  2 +-
 python/rmm/rmm/_cuda/stream.pyx | 15 +++++++++------
 3 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/dependencies.yaml b/dependencies.yaml
index 671809a66..7f94f3f6e 100644
--- a/dependencies.yaml
+++ b/dependencies.yaml
@@ -274,7 +274,6 @@ dependencies:
     common:
       - output_types: [conda, requirements, pyproject]
         packages:
-          - numba>=0.57
           - numpy>=1.23,<3.0a0
     specific:
       - output_types: [conda, requirements, pyproject]
@@ -295,6 +294,7 @@ dependencies:
     common:
       - output_types: [conda, requirements, pyproject]
         packages:
+          - numba>=0.57
           - pytest
           - pytest-cov
     specific:
@@ -309,7 +309,7 @@ dependencies:
               - cuda-nvcc
           - matrix:
             packages:
-      - output_types: [conda, requirements]
+      - output_types: [conda, requirements, pyproject]
         # Define additional constraints for testing with oldest dependencies.
         matrices:
           - matrix:
diff --git a/python/rmm/pyproject.toml b/python/rmm/pyproject.toml
index 6ff45a337..87795974f 100644
--- a/python/rmm/pyproject.toml
+++ b/python/rmm/pyproject.toml
@@ -31,7 +31,6 @@ license = { text = "Apache 2.0" }
 requires-python = ">=3.10"
 dependencies = [
     "cuda-python>=11.8.5,<12.0a0",
-    "numba>=0.57",
     "numpy>=1.23,<3.0a0",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 classifiers = [
@@ -47,6 +46,7 @@ classifiers = [
 
 [project.optional-dependencies]
 test = [
+    "numba>=0.57",
     "pytest",
     "pytest-cov",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
diff --git a/python/rmm/rmm/_cuda/stream.pyx b/python/rmm/rmm/_cuda/stream.pyx
index 951f3f40c..e51f28721 100644
--- a/python/rmm/rmm/_cuda/stream.pyx
+++ b/python/rmm/rmm/_cuda/stream.pyx
@@ -90,12 +90,15 @@ cdef class Stream:
         return self.c_is_default()
 
     def _init_from_numba_stream(self, obj):
-        from numba import cuda
-        if isinstance(obj, cuda.cudadrv.driver.Stream):
-            self._cuda_stream = <cudaStream_t><uintptr_t>(int(obj))
-            self._owner = obj
-        else:
-            raise TypeError(f"Cannot create stream from {type(obj)}")
+        try:
+            from numba import cuda
+            if isinstance(obj, cuda.cudadrv.driver.Stream):
+                self._cuda_stream = <cudaStream_t><uintptr_t>(int(obj))
+                self._owner = obj
+                return
+        except ImportError:
+            pass
+        raise TypeError(f"Cannot create stream from {type(obj)}")
 
     def _init_from_cupy_stream(self, obj):
         try: