Remove nvenc/dec for xenna 0.1.6

nemo_curator/backends/experimental/ray_data/adapter.py

@@ -79,10 +79,6 @@ 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.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 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 loguru import logger
 
 from nemo_curator.backends.base import BaseStageAdapter, NodeInfo, WorkerMetadata
@@ -44,9 +44,6 @@ 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:  # noqa: C901
+    def describe(self) -> str:
         """Get a detailed description of the pipeline stages and their requirements."""
         lines = [
             f"Pipeline: {self.name}",
@@ -148,10 +148,6 @@ def describe(self) -> str:  # noqa: C901
                 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,17 +36,13 @@ class Resources:
     Attributes:
         cpus: Number of CPU cores required
         gpu_memory_gb: GPU memory required in GB (Only for single-GPU stages)
-        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
 
@@ -70,6 +66,9 @@ 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,13 +19,12 @@
 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, _get_gpu_memory_gb
+from nemo_curator.stages.resources import Resources
 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
@@ -79,12 +78,7 @@ 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
-                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
-                )
+                self.resources = Resources(gpus=1.0 / 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.2",
+    "cosmos-xenna>=0.1.8",
     "fsspec",
     "jieba==0.42.1",
     "loguru",

tests/stages/video/clipping/test_clip_transcoding_stage.py

@@ -34,13 +34,6 @@ 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."""
 
@@ -280,6 +273,28 @@ 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 = []