Implementation:Datajuicer Data juicer Process Utils

Knowledge Sources	Datajuicer_Data_juicer
Domains	Resource Management, Parallelism, GPU Computing
Last Updated	2026-02-14 16:00 GMT

Overview

Multiprocessing and resource allocation utilities for configuring worker processes, thread limits, and parallelism levels based on available CPU, GPU, and memory resources.

Description

The process_utils module provides critical resource management for Data-Juicer, ensuring efficient use of system resources without overcommitting memory. Key components include:

Thread Management:

• setup_worker_threads -- Configures per-worker thread limits (PyTorch threads) to prevent thread over-subscription when running with multiple worker processes. Only configured once per process via a global flag.

Multiprocessing Setup:

• setup_mp -- Configures the multiprocessing start method (fork, forkserver, spawn), respecting the MP_START_METHOD environment variable and attempting methods in preference order.

Resource Queries:

• get_min_cuda_memory -- Queries minimum available CUDA memory across all GPUs via nvidia-smi.

Parallelism Calculation (Default Mode):

• calculate_np -- Determines the optimal number of processes for an operator by considering:
  • CPU count and per-operator CPU requirements.
  • Available system memory and per-operator memory requirements.
  • Available GPU memory and per-operator GPU requirements (for CUDA operators).
  • Falls back to the number of available nodes if resources are insufficient.

Parallelism Calculation (Ray Mode):

• calculate_ray_np -- Advanced concurrency calculator for Ray Data operators that:
  • Handles both task-based (CPU) and actor-based (GPU) operators.
  • Respects user-specified concurrency settings.
  • Calculates resource ratios (CPU, GPU, memory fractions) for each operator.
  • Separates fixed allocations (user-specified) from dynamic allocations (auto-scaling).
  • Uses _find_optimal_concurrency to find balanced concurrency allocations across all operators simultaneously.
  • Returns tuple (min_concurrency, max_concurrency) for actor operators.
  • Validates resource availability and raises errors for insufficient resources.
  • Applies a safety fraction (_OPS_MEMORY_LIMIT_FRACTION) to avoid OOM conditions.

Optimal Concurrency Search:

• _find_optimal_concurrency -- Mathematical optimization function that searches for the optimal concurrency allocation to maximize total resource utilization while maintaining balanced processing capacity across operators. Uses Lagrangian-style optimization with candidate enumeration.

Usage

Use this module to determine how many parallel workers to use for each operator in the pipeline. It is called automatically by the framework during pipeline execution to configure operator parallelism.

Code Reference

Source Location

• Repository: Datajuicer_Data_juicer
• File: data_juicer/utils/process_utils.py

Signature

def setup_worker_threads(num_threads: int = 1) -> None: ...
def setup_mp(method=None) -> None: ...
def get_min_cuda_memory() -> float: ...
def calculate_np(name: str, memory: float, num_cpus: int,
                 use_cuda: bool = False, num_gpus: int = 0) -> int: ...
def calculate_ray_np(operators: list) -> list: ...

Import

from data_juicer.utils.process_utils import (
    setup_worker_threads, setup_mp, calculate_np, calculate_ray_np
)

I/O Contract

Inputs

Name	Type	Required	Description
name	str	Yes	Operator name for logging purposes.
memory	float	No	Required memory per process in GB.
num_cpus	int	No	Required CPU cores per process.
use_cuda	bool	No	Whether the operator uses CUDA acceleration.
num_gpus	int	No	Required number of GPUs per process.
operators	list	Yes	List of operator instances for Ray mode concurrency calculation.
num_threads	int	No	Number of threads per worker process. Default 1.

Outputs

Name	Type	Description
auto_num_proc	int	Calculated optimal number of parallel processes for the operator.
operators	list	Modified operator list with assigned num_proc, num_cpus, and num_gpus (Ray mode).

Usage Examples

from data_juicer.utils.process_utils import (
    calculate_np, setup_mp, setup_worker_threads
)

# Configure multiprocessing start method
setup_mp(method="spawn")

# Calculate parallelism for a CPU operator
num_proc = calculate_np(
    name="text_length_filter",
    memory=0.5,   # 0.5 GB per process
    num_cpus=1,
    use_cuda=False
)
print(f"Optimal num_proc: {num_proc}")

# Calculate parallelism for a GPU operator
num_proc = calculate_np(
    name="image_aesthetics_filter",
    memory=2.0,   # 2 GB GPU memory per process
    num_cpus=1,
    use_cuda=True,
    num_gpus=1
)
print(f"Optimal num_proc (GPU): {num_proc}")

# In worker processes, limit threads
setup_worker_threads(num_threads=1)

Related Pages

• Datajuicer_Data_juicer_Model_Utils -- Uses setup_worker_threads in get_model for worker processes