ci: Re-organize curator dep (#1004)

* Re-organize curator dep

Signed-off-by: Dong Hyuk Chang <[email protected]>

* Fix syntax error

Signed-off-by: Dong Hyuk Chang <[email protected]>

* Update uv lock file

Signed-off-by: Dong Hyuk Chang <[email protected]>

* Update grouping based on feedback

Signed-off-by: Dong Hyuk Chang <[email protected]>

* Update based on feedback

Signed-off-by: Dong Hyuk Chang <[email protected]>

* Update pyproject and uv lock

Signed-off-by: Dong Hyuk Chang <[email protected]>

* Update dockerfile

Signed-off-by: Dong Hyuk Chang <[email protected]>

* Fix syantx error and add back install build first

Signed-off-by: Dong Hyuk Chang <[email protected]>

---------

Signed-off-by: Dong Hyuk Chang <[email protected]>

Author: thomasdhc

.github/workflows/cicd-main.yml

@@ -151,7 +151,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - name: Run tests (CPU)
         run: |
-          uv sync --link-mode copy --locked --extra text --extra audio --extra video --extra internvideo2 --group test
+          uv sync --link-mode copy --locked --extra audio_cpu --extra text_cpu --extra video_cpu --group test
           uv run coverage run --branch --source=nemo_curator -m pytest -v tests -m "not gpu"
 
       - name: Generate report

docker/Dockerfile

@@ -67,11 +67,7 @@ FROM nemo_curator_dep AS nemo_curator
 COPY . /opt/Curator
 RUN cd /opt/Curator && \
   uv sync --locked --only-group build && \
-  if [ $CURATOR_ENVIRONMENT = "dev" ]; then \
-      uv sync --link-mode copy --locked --extra all --extra video_nvenc --all-groups; \
-  else \
-      uv sync --link-mode copy --locked --extra all --all-groups; \
-  fi;
+  uv sync --link-mode copy --locked --extra all --all-groups
 
 COPY <<EOF /opt/venv/env.sh
 export UV_PROJECT_ENVIRONMENT=/opt/venv

pyproject.toml

@@ -56,7 +56,38 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-text = [
+cuda12 = ["gpustat", "pynvml"]
+
+# Installs CPU + GPU text curation modules
+deduplication_cuda12 = [
+    "cudf-cu12==25.6.*",
+    "cugraph-cu12==25.6.*",
+    "cuml-cu12==25.6.*",
+    "nx-cugraph-cu12==25.6.*",
+    "pylibraft-cu12==25.6.*",
+    "raft-dask-cu12==25.6.*",
+    "rapidsmpf-cu12==25.6.*",
+]
+
+audio_cpu = [
+    "nemo_toolkit[asr]",
+]
+audio_cuda12 = [
+    "nemo_curator[audio_cpu]",
+    "nemo_curator[cuda12]",
+]
+
+image_cpu = []
+
+# NVIDIA DALI (simplified; update the package to match your CUDA version if needed)
+image_cuda12 = [
+    "nemo_curator[cuda12]",
+    "nemo_curator[deduplication_cuda12]",
+    "nvidia-dali-cuda120",
+]
+
+# Text Curation Dependencies
+text_cpu = [
     # Download / Extract
     "beautifulsoup4",
     "justext",
@@ -75,40 +106,13 @@ text = [
     "ftfy==6.1.1",
 ]
 
-# Installs CPU + GPU text curation modules
-deduplication_cuda12x = [
-    "cudf-cu12==25.6.*",
-    "cugraph-cu12==25.6.*",
-    "cuml-cu12==25.6.*",
-    "nx-cugraph-cu12==25.6.*",
-    "pylibraft-cu12==25.6.*",
-    "raft-dask-cu12==25.6.*",
-    "rapidsmpf-cu12==25.6.*",
-]
-
-audio = [
-    "nemo_toolkit[asr]",
-]
-
-video = [
-    "av==13.1.0",
-    "opencv-python",
-    "torchvision",
-    "einops",
-    "easydict",
-]
-
-video_cuda = [
-    "cvcuda_cu12",
-    "pycuda",
-    "vllm==0.9.2; (platform_machine == 'x86_64' and platform_system != 'Darwin')",
-    "flash-attn<=2.8.3; (platform_machine == 'x86_64' and platform_system != 'Darwin')",
-    "torch<=2.8.0",
-    "torchaudio",
+text_cuda12 = [
+    "nemo_curator[cuda12]",
+    "nemo_curator[deduplication_cuda12]",
+    "nemo_curator[text_cpu]",
 ]
 
-video_nvenc = ["PyNvVideoCodec==2.0.2; (platform_machine == 'x86_64' and platform_system != 'Darwin')"]
-
+# Video Curation Dependencies
 internvideo2 = [
     "av>=13.0.0",
     "deepspeed>=0.15.1",
@@ -136,19 +140,33 @@ internvideo2 = [
     "wandb>=0.18.3",
 ]
 
-# NVIDIA DALI (simplified; update the package to match your CUDA version if needed)
-dali = [
-    "nvidia-dali-cuda120",
+video_cpu = [
+    "nemo_curator[internvideo2]",
+    "av==13.1.0",
+    "opencv-python",
+    "torchvision",
+    "einops",
+    "easydict",
+]
+
+video_cuda12 = [
+    "nemo_curator[video_cpu]",
+    "nemo_curator[cuda12]",
+    "cvcuda_cu12",
+    "flash-attn<=2.8.3; (platform_machine == 'x86_64' and platform_system != 'Darwin')",
+    "pycuda",
+    "PyNvVideoCodec==2.0.2; (platform_machine == 'x86_64' and platform_system != 'Darwin')",
+    "torch<=2.8.0",
+    "torchaudio",
+    "vllm==0.9.2; (platform_machine == 'x86_64' and platform_system != 'Darwin')",
 ]
 
+# All dependencies
 all = [
-    "nemo_curator[dali]",
-    "nemo_curator[deduplication_cuda12x]",
-    "nemo_curator[text]",
-    "nemo_curator[audio]",
-    "nemo_curator[internvideo2]",
-    "nemo_curator[video]",
-    "nemo_curator[video_cuda]",
+    "nemo_curator[audio_cuda12]",
+    "nemo_curator[image_cuda12]",
+    "nemo_curator[text_cuda12]",
+    "nemo_curator[video_cuda12]",
 ]
 
 [dependency-groups]