Move ray actor pool out of experimental

benchmarking/scripts/utils.py

@@ -25,7 +25,7 @@
 import pyarrow.compute as pc
 import pyarrow.parquet as pq
 
-from nemo_curator.backends.experimental.ray_actor_pool.executor import RayActorPoolExecutor
+from nemo_curator.backends.ray_actor_pool.executor import RayActorPoolExecutor
 from nemo_curator.backends.ray_data import RayDataExecutor
 from nemo_curator.backends.xenna import XennaExecutor
 from nemo_curator.utils.file_utils import get_all_file_paths_and_size_under

benchmarking/tools/generate_ci_tests.py

@@ -158,9 +158,7 @@ def generate_pipeline(curator_dir: str, scope: str, test_paths: str) -> dict:
 
 
 def main() -> None:
-    parser = argparse.ArgumentParser(
-        description="Generate GitLab CI jobs for Curator benchmarks"
-    )
+    parser = argparse.ArgumentParser(description="Generate GitLab CI jobs for Curator benchmarks")
     parser.add_argument(
         "--curator-dir",
         type=str,

docs/_extensions/rich_metadata/__init__.py

@@ -27,6 +27,7 @@
 # Import YAML for frontmatter parsing
 try:
     import yaml
+
     YAML_AVAILABLE = True
 except ImportError:
     YAML_AVAILABLE = False
@@ -250,7 +251,7 @@ def build_meta_tags(metadata: dict[str, Any], context: dict[str, Any]) -> dict[s
         "basic": _add_basic_fields(metadata),
         "opengraph": _add_opengraph_fields(metadata, context),
         "twitter": _add_twitter_fields(metadata, context),
-        "custom": _add_custom_fields(metadata)
+        "custom": _add_custom_fields(metadata),
     }
 
 
@@ -474,4 +475,3 @@ def setup(app: Sphinx) -> dict[str, Any]:
         "parallel_read_safe": True,
         "parallel_write_safe": True,
     }
-

docs/_extensions/rich_metadata/verify_metadata.py

@@ -126,9 +126,9 @@ def verify_html_file(html_path: Path) -> bool:
         print(f"❌ File not found: {html_path}")
         return False
 
-    print(f"\n{'='*80}")
+    print(f"\n{'=' * 80}")
     print(f"Verifying: {html_path.name}")
-    print(f"{'='*80}\n")
+    print(f"{'=' * 80}\n")
 
     html_content = html_path.read_text(encoding="utf-8")
 
@@ -154,9 +154,7 @@ def verify_html_file(html_path: Path) -> bool:
 
 def main() -> None:
     """Main entry point for the verification script."""
-    parser = argparse.ArgumentParser(
-        description="Verify rich metadata injection in built HTML files"
-    )
+    parser = argparse.ArgumentParser(description="Verify rich metadata injection in built HTML files")
     parser.add_argument(
         "html_files",
         nargs="+",
@@ -177,16 +175,15 @@ def main() -> None:
         if not verify_html_file(html_file):
             all_passed = False
 
-    print(f"\n{'='*80}")
+    print(f"\n{'=' * 80}")
     if all_passed:
         print("✅ All files verified successfully!")
     else:
         print("⚠️  Some files are missing metadata")
-    print(f"{'='*80}\n")
+    print(f"{'=' * 80}\n")
 
     sys.exit(0 if all_passed else 1)
 
 
 if __name__ == "__main__":
     main()
-

fern/_fix_broken_links.py

@@ -125,7 +125,10 @@
     ("../../curate-audio/process-data/audio-analysis/index.md", "/curate-audio/process-data/audio-analysis"),
     ("../../curate-audio/process-data/text-integration/index.md", "/curate-audio/process-data/text-integration"),
     ("../../curate-text/index.md", "/curate-text"),
-    ("../../curate-text/process-data/quality-assessment/distributed-classifier.md", "/curate-text/process-data/quality-assessment/distributed-classifier"),
+    (
+        "../../curate-text/process-data/quality-assessment/distributed-classifier.md",
+        "/curate-text/process-data/quality-assessment/distributed-classifier",
+    ),
     ("../../curate-text/process-data/deduplication/semdedup.md", "/curate-text/process-data/deduplication/semdedup"),
     ("../../curate-images/index.md", "/curate-images"),
     ("../../curate-images/process-data/embeddings/index.md", "/curate-images/process-data/embeddings"),

fern/substitute_variables.py

@@ -74,7 +74,9 @@ def main() -> int:
         default="versions/v26.02",
         help="Path under fern/ containing MDX (e.g. versions/v25.09, versions/v26.02)",
     )
-    parser.add_argument("--version", help="Version string for version/container_version/current_release (e.g. 25.09, 26.02)")
+    parser.add_argument(
+        "--version", help="Version string for version/container_version/current_release (e.g. 25.09, 26.02)"
+    )
     parser.add_argument("--dry-run", action="store_true", help="Show what would be changed without modifying files")
     args = parser.parse_args()
 

fern/versions/v25.09/pages/api-reference/executors/experimental.mdx

@@ -46,7 +46,7 @@ Uses Ray Actor Pool for distributed execution.
 ### Import
 
 ```python
-from nemo_curator.backends.experimental import RayActorPoolExecutor
+from nemo_curator.backends.ray_actor_pool import RayActorPoolExecutor
 ```
 
 ### Usage

fern/versions/v25.09/pages/reference/infrastructure/execution-backends.mdx

@@ -124,7 +124,7 @@ results = pipeline.run(executor)
 Executor using Ray Actor pools for custom distributed processing patterns such as deduplication.
 
 ```python
-from nemo_curator.backends.experimental.ray_actor_pool import RayActorPoolExecutor
+from nemo_curator.backends.ray_actor_pool import RayActorPoolExecutor
 
 executor = RayActorPoolExecutor()
 results = pipeline.run(executor)

fern/versions/v26.02/pages/api-reference/executors/experimental.mdx

@@ -16,7 +16,7 @@ Uses Ray Data for distributed execution.
 ### Import
 
 ```python
-from nemo_curator.backends.experimental import RayDataExecutor
+from nemo_curator.backends.ray_data import RayDataExecutor
 ```
 
 ### Usage
@@ -46,7 +46,7 @@ Uses Ray Actor Pool for distributed execution.
 ### Import
 
 ```python
-from nemo_curator.backends.experimental import RayActorPoolExecutor
+from nemo_curator.backends.ray_actor_pool import RayActorPoolExecutor
 ```
 
 ### Usage
@@ -149,8 +149,8 @@ class MyCustomExecutor(BaseExecutor):
 | Executor | Best For | Considerations |
 |----------|----------|----------------|
 | `XennaExecutor` | Production workloads | Default choice, most stable |
-| `RayDataExecutor` | Ray-native environments | Experimental |
-| `RayActorPoolExecutor` | Fine-grained actor control | Experimental |
+| `RayDataExecutor` | Ray-native environments | Production |
+| `RayActorPoolExecutor` | Fine-grained actor control | Production |
 
 ## Source Code
 

nemo_curator/backends/experimental/ray_actor_pool/adapter.py → nemo_curator/backends/ray_actor_pool/adapter.py

@@ -15,7 +15,7 @@
 from loguru import logger
 
 from nemo_curator.backends.base import BaseStageAdapter
-from nemo_curator.backends.experimental.utils import get_worker_metadata_and_node_id
+from nemo_curator.backends.utils import get_worker_metadata_and_node_id
 from nemo_curator.stages.base import ProcessingStage
 
 

nemo_curator/backends/experimental/ray_actor_pool/executor.py → nemo_curator/backends/ray_actor_pool/executor.py

@@ -23,8 +23,7 @@
 from tqdm import tqdm
 
 from nemo_curator.backends.base import BaseExecutor
-from nemo_curator.backends.experimental.utils import RayStageSpecKeys, execute_setup_on_node
-from nemo_curator.backends.utils import register_loguru_serializer
+from nemo_curator.backends.utils import RayStageSpecKeys, execute_setup_on_node, register_loguru_serializer
 from nemo_curator.tasks import EmptyTask, Task
 
 from .adapter import RayActorPoolStageAdapter

nemo_curator/backends/experimental/ray_actor_pool/raft_adapter.py → nemo_curator/backends/ray_actor_pool/raft_adapter.py

@@ -18,7 +18,7 @@
 from loguru import logger
 
 from nemo_curator.backends.base import BaseStageAdapter
-from nemo_curator.backends.experimental.utils import get_worker_metadata_and_node_id
+from nemo_curator.backends.utils import get_worker_metadata_and_node_id
 from nemo_curator.stages.base import ProcessingStage
 
 if TYPE_CHECKING:

nemo_curator/backends/experimental/ray_actor_pool/shuffle_adapter.py → nemo_curator/backends/ray_actor_pool/shuffle_adapter.py

@@ -19,7 +19,7 @@
 from loguru import logger
 
 from nemo_curator.backends.base import BaseStageAdapter
-from nemo_curator.backends.experimental.utils import RayStageSpecKeys, get_worker_metadata_and_node_id
+from nemo_curator.backends.utils import RayStageSpecKeys, get_worker_metadata_and_node_id
 from nemo_curator.tasks import FileGroupTask
 
 if TYPE_CHECKING:

nemo_curator/backends/experimental/ray_actor_pool/utils.py → nemo_curator/backends/ray_actor_pool/utils.py

@@ -18,7 +18,7 @@
 import ray
 from loguru import logger
 
-from nemo_curator.backends.experimental.utils import get_available_cpu_gpu_resources
+from nemo_curator.backends.utils import get_available_cpu_gpu_resources
 
 if TYPE_CHECKING:
     from ray.actor import ActorClass

nemo_curator/backends/ray_data/adapter.py

@@ -20,7 +20,7 @@
 from ray.data import Dataset
 
 from nemo_curator.backends.base import BaseStageAdapter
-from nemo_curator.backends.experimental.utils import RayStageSpecKeys, get_worker_metadata_and_node_id
+from nemo_curator.backends.utils import RayStageSpecKeys, get_worker_metadata_and_node_id
 from nemo_curator.stages.base import ProcessingStage
 
 from .utils import calculate_concurrency_for_actors_for_stage, is_actor_stage

nemo_curator/backends/ray_data/executor.py

@@ -19,8 +19,7 @@
 from ray.data import DataContext, Dataset
 
 from nemo_curator.backends.base import BaseExecutor
-from nemo_curator.backends.experimental.utils import execute_setup_on_node
-from nemo_curator.backends.utils import register_loguru_serializer
+from nemo_curator.backends.utils import execute_setup_on_node, register_loguru_serializer
 from nemo_curator.tasks import EmptyTask, Task
 
 from .adapter import RayDataStageAdapter

nemo_curator/backends/ray_data/utils.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from nemo_curator.backends.experimental.utils import get_available_cpu_gpu_resources
+from nemo_curator.backends.utils import get_available_cpu_gpu_resources
 from nemo_curator.stages.base import ProcessingStage
 
 

nemo_curator/backends/utils.py

@@ -12,15 +12,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+import time
 from copy import deepcopy
-from typing import TYPE_CHECKING
+from enum import Enum
+from typing import TYPE_CHECKING, Any
 
 import ray
 from loguru import logger
+from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
+
+from nemo_curator.backends.base import NodeInfo, WorkerMetadata
+from nemo_curator.stages.base import ProcessingStage
 
 if TYPE_CHECKING:
     import loguru
 
+# Global variable to cache head node ID
+_HEAD_NODE_ID_CACHE = None
+
 
 def _logger_custom_serializer(
     _: "loguru.Logger",
@@ -112,3 +122,115 @@ def warn_on_env_var_override(existing_config: dict | None, merged_config: dict |
             "Update the executor configuration before running if this is unintended.",
             overridden_keys,
         )
+
+
+def is_head_node(node: dict[str, Any]) -> bool:
+    """Check if a node is the head node."""
+    return "node:__internal_head__" in node.get("Resources", {})
+
+
+def get_head_node_id() -> str | None:
+    """Get the head node ID from the Ray cluster, with lazy evaluation and caching.
+
+    Returns:
+        The head node ID if a head node exists, otherwise None.
+    """
+    global _HEAD_NODE_ID_CACHE  # noqa: PLW0603
+
+    if _HEAD_NODE_ID_CACHE is not None:
+        return _HEAD_NODE_ID_CACHE
+
+    # Compute head node ID
+    for node in ray.nodes():
+        if is_head_node(node):
+            _HEAD_NODE_ID_CACHE = node["NodeID"]
+            return _HEAD_NODE_ID_CACHE
+
+    return None
+
+
+class RayStageSpecKeys(str, Enum):
+    """String enum of different flags that define keys inside ray_stage_spec."""
+
+    IS_ACTOR_STAGE = "is_actor_stage"
+    IS_FANOUT_STAGE = "is_fanout_stage"
+    IS_RAFT_ACTOR = "is_raft_actor"
+    IS_LSH_STAGE = "is_lsh_stage"
+    IS_SHUFFLE_STAGE = "is_shuffle_stage"
+    MAX_CALLS_PER_WORKER = "max_calls_per_worker"
+
+
+def get_worker_metadata_and_node_id() -> tuple[NodeInfo, WorkerMetadata]:
+    """Get the worker metadata and node id from the runtime context."""
+    ray_context = ray.get_runtime_context()
+    return NodeInfo(node_id=ray_context.get_node_id()), WorkerMetadata(worker_id=ray_context.get_worker_id())
+
+
+def get_available_cpu_gpu_resources(
+    init_and_shutdown: bool = False, ignore_head_node: bool = False
+) -> tuple[int, int]:
+    """Get available CPU and GPU resources from Ray."""
+    if init_and_shutdown:
+        ray.init(ignore_reinit_error=True)
+    time.sleep(0.2)  # ray.available_resources() returns might have a lag
+    # available resources can be different from total resources, however curator assumes
+    # entire cluster is available for use and only one pipeline is being run at a time.
+    # therefore available resources should match total resources.
+    available_resources = ray.available_resources()
+    available_cpus = available_resources.get("CPU", 0)
+    available_gpus = available_resources.get("GPU", 0)
+    if ignore_head_node:
+        head_node_id = get_head_node_id()
+        if head_node_id is not None:
+            total_resources = ray.state.total_resources_per_node().get(head_node_id, {})
+            head_node_cpus = total_resources.get("CPU", 0)
+            head_node_gpus = total_resources.get("GPU", 0)
+            logger.info(
+                f"Ignoring head node {head_node_id} with {head_node_cpus} CPUs and {head_node_gpus} GPUs for resource calculation"
+            )
+            available_cpus = max(0, available_cpus - head_node_cpus)
+            available_gpus = max(0, available_gpus - head_node_gpus)
+        else:
+            logger.warning("ignore_head_node=True but no head node found in the cluster")
+    if init_and_shutdown:
+        ray.shutdown()
+    return (available_cpus, available_gpus)
+
+
+@ray.remote
+def _setup_stage_on_node(stage: ProcessingStage, node_info: NodeInfo, worker_metadata: WorkerMetadata) -> None:
+    """Ray remote function to execute setup_on_node for a stage.
+
+    This runs as a Ray remote task (not an actor).
+    vLLM's auto-detection only forces the spawn multiprocessing method inside Ray actors,
+    not in Ray tasks. Without this override, vLLM defaults to fork in tasks and hits
+    RuntimeError: Cannot re-initialize CUDA in forked subprocess.
+    We explicitly set the environment variable to spawn to prevent this.
+    """
+    os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+    stage.setup_on_node(node_info, worker_metadata)
+
+
+def execute_setup_on_node(stages: list[ProcessingStage], ignore_head_node: bool = False) -> None:
+    """Execute setup on node for a stage."""
+    head_node_id = get_head_node_id()
+    ray_tasks = []
+    for node in ray.nodes():
+        node_id = node["NodeID"]
+        node_info = NodeInfo(node_id=node_id)
+        worker_metadata = WorkerMetadata(worker_id="", allocation=None)
+        if ignore_head_node and node_id == head_node_id:
+            logger.info(f"Ignoring setup on head node {node_id}")
+            continue
+
+        logger.info(f"Executing setup on node {node_id} for {len(stages)} stages")
+
+        for stage in stages:
+            ray_tasks.append(
+                _setup_stage_on_node.options(
+                    num_cpus=stage.resources.cpus if stage.resources is not None else 1,
+                    num_gpus=stage.resources.gpus if stage.resources is not None else 0,
+                    scheduling_strategy=NodeAffinitySchedulingStrategy(node_id=node_id, soft=False),
+                ).remote(stage, node_info, worker_metadata)
+            )
+    ray.get(ray_tasks)

nemo_curator/models/transnetv2.py

@@ -39,6 +39,7 @@
 _TRANSNETV2_MODEL_WEIGHTS: Final = "transnetv2-pytorch-weights.pth"
 _TRANSNETV2_MODEL_REVISION: Final = "db6ceab"
 
+
 class _TransNetV2(nn.Module):
     def __init__(  # noqa: PLR0913
         self,