Enabled nightly build using RAPIDS nightly (NVIDIA#237)

* Sign and add Signed-off-by: Praateek Mahajan <[email protected]> * implement strtobool Signed-off-by: Praateek Mahajan <[email protected]> * pre-commit Signed-off-by: Praateek Mahajan <[email protected]> * Update readme Signed-off-by: Praateek Mahajan <[email protected]> --------- Signed-off-by: Praateek Mahajan <[email protected]>
sarahyurick · Sep 18, 2024 · c89c115 · c89c115
1 parent 5c6ccb4
commit c89c115
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -83,7 +83,12 @@ Before installing NeMo Curator, ensure that the following requirements are met:
   - Volta™ or higher ([compute capability 7.0+](https://developer.nvidia.com/cuda-gpus))
   - CUDA 12 (or above)
 
-You can install NeMo-Curator from PyPi, from source or get it through the NeMo Framework container.
+You can install NeMo-Curator
+1. from PyPi
+2. from source
+3. get it through the [NeMo Framework container](https://github.com/NVIDIA/NeMo?tab=readme-ov-file#docker-containers).
+
+
 
 #### From PyPi
 
@@ -126,6 +131,21 @@ pip install --extra-index-url https://pypi.nvidia.com nemo-curator[cuda12x]
     pip install --extra-index-url https://pypi.nvidia.com ".[cuda12x]"
     ```
 
+#### Using Nightly Dependencies for Rapids
+
+You can also install NeMo Curator using the Rapids nightly, to do so you can set the environment variable `RAPIDS_NIGHTLY=1`.
+
+
+```bash
+# installing from pypi
+RAPIDS_NIGHTLY=1 pip install --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple "nemo-curator[cuda12x]"
+
+# installing from source
+RAPIDS_NIGHTLY=1 pip install --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple ".[cuda12x]"
+```
+
+When the environment variable set to 0 or not set (default behavior) it'll use the stable version of Rapids.
+
 #### From the NeMo Framework Container
 
 The latest release of NeMo Curator comes preinstalled in the [NeMo Framework Container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo/tags). If you want the latest commit inside the container, you can reinstall NeMo Curator using:

diff --git a/requirements/requirements_rapids_nightly.txt b/requirements/requirements_rapids_nightly.txt
@@ -0,0 +1,6 @@
+cudf-cu12>=24.10.0a0,<=24.10
+cugraph-cu12>=24.10.0a0,<=24.10
+cuml-cu12>=24.10.0a0,<=24.10
+dask-cuda>=24.10.0a0,<=24.10
+dask-cudf-cu12>=24.10.0a0,<=24.10
+spacy[cuda12x]>=3.6.0, <4.0.0
diff --git a/setup.py b/setup.py
@@ -21,6 +21,13 @@
 long_description = (here / "README.md").read_text(encoding="utf-8")
 
 
+def strtobool(value: str) -> bool:
+    value = value.lower()
+    if value in ("y", "yes", "1", "true"):
+        return True
+    return False
+
+
 def req_file(filename, folder="requirements"):
     with open(os.path.join(folder, filename), encoding="utf-8") as f:
         content = f.readlines()
@@ -29,8 +36,14 @@ def req_file(filename, folder="requirements"):
 
 install_requires = req_file("requirements.txt")
 
-extras_require = {
-    "cuda12x": req_file("requirements_cuda12x.txt"),
+cuda12_requirements_filename = (
+    "requirements_rapids_nightly.txt"
+    if strtobool(os.getenv("RAPIDS_NIGHTLY", "false"))
+    else "requirements_cuda12x.txt"
+)
+
+extras_require: dict = {
+    "cuda12x": req_file(cuda12_requirements_filename),
     "image": req_file("requirements_image.txt"),
 }