Source code for data_juicer.utils.process_utils

import math
import os
import subprocess

import multiprocess as mp
import psutil
from loguru import logger

from data_juicer import cuda_device_count




def setup_mp(method=None):
    if mp.current_process().name != 'MainProcess':
        return

    if method is None:
        method = ['fork', 'forkserver', 'spawn']
    if not isinstance(method, (list, tuple)):
        method = [method]
    method = [m.lower() for m in method]

    env_method = os.getenv('MP_START_METHOD', '').lower()
    if env_method in method:
        method = [env_method]

    available_methods = mp.get_all_start_methods()
    for m in method:
        if m in available_methods:
            try:
                logger.debug(f"Setting multiprocess start method to '{m}'")
                mp.set_start_method(m, force=True)
            except RuntimeError as e:
                logger.warning(f'Error setting multiprocess start method: {e}')
            break





def get_min_cuda_memory():
    # get cuda memory info using "nvidia-smi" command
    import torch
    min_cuda_memory = torch.cuda.get_device_properties(
        0).total_memory / 1024**2
    nvidia_smi_output = subprocess.check_output([
        'nvidia-smi', '--query-gpu=memory.free',
        '--format=csv,noheader,nounits'
    ]).decode('utf-8')
    for line in nvidia_smi_output.strip().split('\n'):
        free_memory = int(line)
        min_cuda_memory = min(min_cuda_memory, free_memory)
    return min_cuda_memory





def calculate_np(name,
                 mem_required,
                 cpu_required,
                 num_proc=None,
                 use_cuda=False):
    """Calculate the optimum number of processes for the given OP"""
    eps = 1e-9  # about 1 byte

    if use_cuda:
        auto_num_proc = None
        cuda_mem_available = get_min_cuda_memory() / 1024
        if mem_required == 0:
            logger.warning(f'The required cuda memory of Op[{name}] '
                           f'has not been specified. '
                           f'Please specify the mem_required field in the '
                           f'config file, or you might encounter CUDA '
                           f'out of memory error. You can reference '
                           f'the mem_required field in the '
                           f'config_all.yaml file.')
        else:
            auto_num_proc = math.floor(
                cuda_mem_available / mem_required) * cuda_device_count()
            if cuda_mem_available / mem_required < 1.0:
                logger.warning(
                    f'The required cuda memory:{mem_required}GB might '
                    f'be more than the available cuda memory:'
                    f'{cuda_mem_available}GB.'
                    f'This Op[{name}] might '
                    f'require more resource to run.')

        if auto_num_proc and num_proc:
            op_proc = min(auto_num_proc, num_proc)
            if num_proc > auto_num_proc:
                logger.warning(
                    f'The given num_proc: {num_proc} is greater than '
                    f'the value {auto_num_proc} auto calculated based '
                    f'on the mem_required of Op[{name}]. '
                    f'Set the `num_proc` to {auto_num_proc}.')
        elif not auto_num_proc and not num_proc:
            op_proc = cuda_device_count()
            logger.warning(
                f'Both mem_required and num_proc of Op[{name}] are not set.'
                f'Set the `num_proc` to number of GPUs {op_proc}.')
        else:
            op_proc = auto_num_proc if auto_num_proc else num_proc

        op_proc = max(op_proc, 1)
        return op_proc
    else:
        if num_proc is None:
            num_proc = psutil.cpu_count()

        op_proc = num_proc
        cpu_available = psutil.cpu_count()
        mem_available = psutil.virtual_memory().available
        mem_available = mem_available / 1024**3
        op_proc = min(op_proc, math.floor(cpu_available / cpu_required + eps))
        op_proc = min(op_proc,
                      math.floor(mem_available / (mem_required + eps)))
        if op_proc < 1.0:
            logger.warning(f'The required CPU number:{cpu_required} '
                           f'and memory:{mem_required}GB might '
                           f'be more than the available CPU:{cpu_available} '
                           f'and memory :{mem_available}GB.'
                           f'This Op [{name}] might '
                           f'require more resource to run.')
        op_proc = max(op_proc, 1)
        return op_proc