Repository: NVIDIA-Merlin/NVTabular (language: python)

Issue:
[BUG] Missing Dependency `pynvml`
**Describe the bug**
The package `pynvml` is missing from the dependencies.

**Steps/Code to reproduce bug**
- Installing  nvtabular=0.8  and cudatoolkit=11.0
- Running https://github.com/NVIDIA-Merlin/NVTabular/blob/main/examples/getting-started-movielens/02-ETL-with-NVTabular.ipynb
- Upon Dataset instantiation fails with `No module named 'pynvml'` while calling `_pynvml_mem_size()` in `utils.py`

**Expected behavior**
`pynvml` should be installed as a dependency of nvtabular

**Environment details (please complete the following information):**
 - Environment location: [Docker]
 - Method of NVTabular install: [conda, Docker]




Discussion / hints:
Thanks for the bug report Simon! Glad to see you giving NVTabular a try.

How did you install nvtabular? Was it through conda or pip? 

Also, Can you run the following and tell me what it reports?
```python
import numba.cuda
print(numba.__version__)
print(numba.cuda.current_context().get_memory_info())
```

I don't think that we should be directly depending on pynvml, its currently only used as a backup method if we can't get the GPU memory from ```numba.cuda```:

https://github.com/NVIDIA-Merlin/NVTabular/blob/22df264c3b36514755492879aa1a1d8e7732e235/nvtabular/utils.py#L70-L79

I'm thinking we should probably remove this backup method entirely and not depend on pynvml at all, but to do that we need to make sure that we can rely on the implementation from numba. 

@rjzamora what do you think?


Hey @benfred, 

yeah, I like it so far and hope that NVTabular facilitates the switch to NN-based Recommender Systems (away from LightFM) but let's see ... until now I'm constantly fighting Cuda Memory error and don't want to take on the complexity of distributed GPU ... 😄.

**To the topic:**
I installed via Conda in a Docker and I am running on a Kubeflow Cluster.

I did some further investigation, the problem seems to be that `import nvtabular` messes up `numba.cuda.current_context()`

Running this:
```python
import numba.cuda
import nvtabular as nvt

print(numba.__version__)
print(numba.cuda.current_context().get_memory_info())
```
Results in the NotImplementedError which leads to the fallback to pynvml:
```python
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
/tmp/ipykernel_767/2083043560.py in <module>
      2 
      3 print(numba.__version__)
----> 4 print(numba.cuda.current_context().get_memory_info())

/opt/conda/envs/jovyan_env/lib/python3.8/site-packages/numba/cuda/cudadrv/driver.py in get_memory_info(self)
   1246         """Returns (free, total) memory in bytes in the context.
   1247         """
-> 1248         return self.memory_manager.get_memory_info()
   1249 
   1250     def get_active_blocks_per_multiprocessor(self, func, blocksize, memsize,

/opt/conda/envs/jovyan_env/lib/python3.8/site-packages/rmm/rmm.py in get_memory_info(self)
    147 
    148     def get_memory_info(self):
--> 149         raise NotImplementedError()
    150 
    151     @property

NotImplementedError: 
```

But running (after Kernel restart of course):
```python
import numba.cuda
print(numba.__version__)
print(numba.cuda.current_context().get_memory_info())
import nvtabular as nvt
```
Gives (no `pynvml` needed):
```python
0.55.0rc1
MemoryInfo(free=11523588096, total=16945512448)
```

The problem is in the `memory_manager` attribute.

In the failing case this:
```python
ctx = numba.cuda.current_context()
ctx.memory_manager
```
is of type `rmm.rmm.RMMNumbaManager` whereas without importing `nvtabular` first, the manager is of type `numba.cuda.cudadrv.driver.NumbaCUDAMemoryManager` ... nasty 😀. I do not know enough about nvtabular, numba, cuda and co to further debug but I hope this helps.
@rjzamora , please take a look

Fix patch:
diff --git a/conda/environments/nvtabular_dev_cuda11.0.yml b/conda/environments/nvtabular_dev_cuda11.0.yml
index 8c2e19c4b04..038d77ebd63 100644
--- a/conda/environments/nvtabular_dev_cuda11.0.yml
+++ b/conda/environments/nvtabular_dev_cuda11.0.yml
@@ -8,14 +8,15 @@ dependencies:
   - nvtabular
   - python>=3.7
   - cudatoolkit=11.0
-  - cudf>=21.08.*
-  - dask-cuda>=21.08.*
-  - dask-cudf>=21.08.*
-  - rmm>=21.08.*
-  - dask==2021.7.1
-  - distributed>=2021.7.1
+  - cudf>=21.10.*
+  - dask-cuda>=21.10.*
+  - dask-cudf>=21.10.*
+  - rmm>=21.10.*
+  - dask==2021.9.1
+  - distributed>=2021.9.1
   - nvtx>=0.2.1
   - numba>=0.53.0
   - dlpack
   - scikit-learn
   - asvdb
+  - pynvml
diff --git a/conda/environments/nvtabular_dev_cuda11.2.yml b/conda/environments/nvtabular_dev_cuda11.2.yml
index 04569f23916..89ee56b09ca 100644
--- a/conda/environments/nvtabular_dev_cuda11.2.yml
+++ b/conda/environments/nvtabular_dev_cuda11.2.yml
@@ -8,15 +8,16 @@ dependencies:
   - nvtabular
   - python>=3.7
   - cudatoolkit=11.2
-  - cudf>=21.08.*
-  - dask-cuda>=21.08.*
-  - dask-cudf>=21.08.*
-  - rmm>=21.08.*
-  - dask==2021.7.1
-  - distributed>=2021.7.1
+  - cudf>=21.10.*
+  - dask-cuda>=21.10.*
+  - dask-cudf>=21.10.*
+  - rmm>=21.10.*
+  - dask==2021.9.1
+  - distributed>=2021.9.1
   - fsspec>=0.6.0
   - nvtx>=0.2.1
   - numba>=0.53.0
   - dlpack
   - scikit-learn
   - asvdb
+  - pynvml
diff --git a/conda/recipes/meta.yaml b/conda/recipes/meta.yaml
index 1c355f010f3..f41494e1e2d 100644
--- a/conda/recipes/meta.yaml
+++ b/conda/recipes/meta.yaml
@@ -32,10 +32,11 @@ requirements:
     - {{ req }}
   {% endfor %}
     - python
-    - cudf>=21.08.*
-    - dask-cudf>=21.08.*
+    - cudf>=21.10.*
+    - dask-cudf>=21.10.*
     - cupy>=7.2.0,<10.0.0a
     - nvtx>=0.2.1
+    - pynvml
 
 about:
   home: https://github.com/NVIDIA/NVTabular
diff --git a/requirements.txt b/requirements.txt
index 87923196a34..1b88304a231 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,5 @@
-dask==2021.07.1
-distributed==2021.07.1
+dask==2021.09.1
+distributed==2021.09.1
 pandas>=1.2.0
 numba>=0.54.0
 tqdm>=4.0


Test patch:
diff --git a/ci/build_and_test.sh b/ci/build_and_test.sh
index dbfaac44fdc..f0e63d18305 100755
--- a/ci/build_and_test.sh
+++ b/ci/build_and_test.sh
@@ -9,7 +9,7 @@ cd $nvt_directory
 echo "Installing NVTabular"
 python -m pip install --user --upgrade pip setuptools wheel pybind11 numpy==1.20.3 setuptools==59.4.0
 python -m pip uninstall nvtabular -y
-python setup.py develop --user
+python setup.py develop --user --no-deps
 
 # following checks requirement requirements-dev.txt to be installed
 echo "Running black --check"