Revert "Remove nvenc/dec for xenna 0.1.6"

nemo_curator/backends/experimental/ray_data/adapter.py

@@ -79,6 +79,10 @@ def process_dataset(self, dataset: Dataset, ignore_head_node: bool = False) -> D
         Returns:
             Dataset: Processed Ray Data dataset
         """
+        # TODO: Support nvdecs / nvencs
+        if self.stage.resources.gpus <= 0 and (self.stage.resources.nvdecs > 0 or self.stage.resources.nvencs > 0):
+            msg = "Ray Data does not support nvdecs / nvencs. Please use gpus instead."
+            raise ValueError(msg)
 
         is_actor_stage_ = self.stage.ray_stage_spec().get(RayStageSpecKeys.IS_ACTOR_STAGE, is_actor_stage(self.stage))
 

nemo_curator/backends/xenna/adapter.py

@@ -13,9 +13,9 @@
 # limitations under the License.
 
 from cosmos_xenna.pipelines import v1 as pipelines_v1
-from cosmos_xenna.pipelines.private.resources import NodeInfo as XennaNodeInfo
-from cosmos_xenna.pipelines.private.resources import Resources as XennaResources
-from cosmos_xenna.pipelines.private.resources import WorkerMetadata as XennaWorkerMetadata
+from cosmos_xenna.ray_utils.resources import NodeInfo as XennaNodeInfo
+from cosmos_xenna.ray_utils.resources import Resources as XennaResources
+from cosmos_xenna.ray_utils.resources import WorkerMetadata as XennaWorkerMetadata
 from loguru import logger
 
 from nemo_curator.backends.base import BaseStageAdapter, NodeInfo, WorkerMetadata
@@ -44,6 +44,9 @@ def required_resources(self) -> XennaResources:
         return XennaResources(
             cpus=self.processing_stage.resources.cpus,
             gpus=self.processing_stage.resources.gpus,
+            nvdecs=self.processing_stage.resources.nvdecs,
+            nvencs=self.processing_stage.resources.nvencs,
+            entire_gpu=self.processing_stage.resources.entire_gpu,
         )
 
     @property

nemo_curator/pipeline/pipeline.py

@@ -129,7 +129,7 @@ def __repr__(self) -> str:
         stage_info = ", ".join([f"{s.name}({s.__class__.__name__})" for s in self.stages])
         return f"Pipeline(name='{self.name}', stages=[{stage_info}])"
 
-    def describe(self) -> str:
+    def describe(self) -> str:  # noqa: C901
         """Get a detailed description of the pipeline stages and their requirements."""
         lines = [
             f"Pipeline: {self.name}",
@@ -148,6 +148,10 @@ def describe(self) -> str:
                 lines.append(f"  Resources: {stage.resources.cpus} CPUs")
                 if stage.resources.requires_gpu:
                     lines.append(f"    GPU Memory: {stage.resources.gpu_memory_gb} GB ({stage.resources.gpus} GPUs)")
+                if stage.resources.nvdecs > 0:
+                    lines.append(f"    NVDEC: {stage.resources.nvdecs}")
+                if stage.resources.nvencs > 0:
+                    lines.append(f"    NVENC: {stage.resources.nvencs}")
 
                 lines.append(f"  Batch size: {stage.batch_size}")
 

nemo_curator/stages/resources.py

@@ -36,13 +36,17 @@ class Resources:
     Attributes:
         cpus: Number of CPU cores required
         gpu_memory_gb: GPU memory required in GB (Only for single-GPU stages)
-        entire_gpu: Whether to allocate entire GPU regardless of memory
+        nvdecs: Number of NVDEC units required
+        nvencs: Number of NVENC units required
+        entire_gpu: Whether to allocate entire GPU regardless of memory (This also gives you nvdecs and nvencs of that GPU)
         gpus: Number of GPUs required (Only for multi-GPU stages)
     """
 
     # TODO : Revisit this gpu_memory_gb, gpus, entire_gpu too many variables for gpu
     cpus: float = 1.0
     gpu_memory_gb: float = 0.0
+    nvdecs: int = 0
+    nvencs: int = 0
     entire_gpu: bool = False
     gpus: float = 0.0
 
@@ -66,9 +70,6 @@ def __post_init__(self):
                 error_message += "Please use gpus for multi-GPU stages."
                 raise ValueError(error_message)
 
-        if self.entire_gpu:
-            self.gpus = 1.0
-
     @property
     def requires_gpu(self) -> bool:
         """Check if this stage requires GPU resources."""

nemo_curator/stages/video/clipping/clip_extraction_stages.py

@@ -19,12 +19,13 @@
 from dataclasses import dataclass
 from typing import Any
 
+from cosmos_xenna.ray_utils.resources import _get_local_gpu_info, _make_gpu_resources_from_gpu_name
 from loguru import logger
 
 from nemo_curator.backends.base import WorkerMetadata
 from nemo_curator.backends.experimental.utils import RayStageSpecKeys
 from nemo_curator.stages.base import ProcessingStage
-from nemo_curator.stages.resources import Resources
+from nemo_curator.stages.resources import Resources, _get_gpu_memory_gb
 from nemo_curator.tasks.video import Clip, Video, VideoTask
 from nemo_curator.utils import grouping
 from nemo_curator.utils.operation_utils import make_pipeline_temporary_dir
@@ -78,7 +79,12 @@ def __post_init__(self) -> None:
         if self.encoder == "h264_nvenc" or self.use_hwaccel:
             if self.nb_streams_per_gpu > 0:
                 # Assume that we have same type of GPUs
-                self.resources = Resources(gpus=1.0 / self.nb_streams_per_gpu)
+                gpu_info = _get_local_gpu_info()[0]
+                nvencs = _make_gpu_resources_from_gpu_name(gpu_info.name).num_nvencs
+                gpu_memory_gb = _get_gpu_memory_gb()
+                self.resources = Resources(
+                    nvencs=nvencs // self.nb_streams_per_gpu, gpu_memory_gb=gpu_memory_gb // self.nb_streams_per_gpu
+                )
             else:
                 self.resources = Resources(gpus=1)
         else:

pyproject.toml

@@ -53,7 +53,7 @@ classifiers = [
 dependencies = [
     "absl-py>=2.0.0,<3.0.0",
     "comment_parser",
-    "cosmos-xenna==0.1.8",
+    "cosmos-xenna==0.1.2",
     "fsspec",
     "hydra-core",
     "jieba==0.42.1",

tests/stages/video/clipping/test_clip_transcoding_stage.py

@@ -34,6 +34,13 @@ def __init__(self, index: int, name: str):
         self.name = name
 
 
+# Mock GPU resources class to simulate GPU resources
+class MockGpuResources:
+    def __init__(self, num_nvencs: int = 3, num_nvdecs: int = 3):
+        self.num_nvencs = num_nvencs
+        self.num_nvdecs = num_nvdecs
+
+
 class TestClipTranscodingStage:
     """Test cases for ClipTranscodingStage."""
 
@@ -273,28 +280,6 @@ def test_add_hwaccel_options_disabled(self) -> None:
         # Should not add any hwaccel options
         assert "-hwaccel" not in command
 
-    def test_add_hwaccel_options_enabled(self) -> None:
-        """Test hardware acceleration options when enabled."""
-        command = []
-        stage = ClipTranscodingStage(use_hwaccel=True, encoder="h264_nvenc", nb_streams_per_gpu=1)
-
-        stage._add_hwaccel_options(command)
-
-        assert "-hwaccel" in command
-        assert "-hwaccel_output_format" in command
-        assert stage.resources.gpus == 1.0
-
-    def test_add_hwaccel_options_enabled_multiple_streams(self) -> None:
-        """Test hardware acceleration options when enabled."""
-        command = []
-        stage = ClipTranscodingStage(use_hwaccel=True, encoder="h264_nvenc", nb_streams_per_gpu=4)
-
-        stage._add_hwaccel_options(command)
-
-        assert "-hwaccel" in command
-        assert "-hwaccel_output_format" in command
-        assert stage.resources.gpus == 0.25
-
     def test_add_input_options(self) -> None:
         """Test adding input options to FFmpeg command."""
         command = []