num_jobs.py

import math
import multiprocessing
from typing import Optional

import imagesize
import psutil
from psutil._common import bytes2human


def bytes_of_uint8_img(img_path: str, colored: bool) -> int:
    """Compute the number of bytes allocated by an image file
    once loaded in memory"""
    width, height = imagesize.get(img_path)
    n_bytes = width * height * 8
    return n_bytes*3 if colored else n_bytes


def get_num_jobs(bytes_per_process: Optional[int] = None) -> int:
    """Estimate the number of processes to run in parallel"""
    num_procs = math.floor(0.9*multiprocessing.cpu_count())

    if bytes_per_process is not None:
        available_bytes = getattr(psutil.virtual_memory(), 'available')
        print(f"Available memory: {bytes2human(available_bytes)}")
        print(f"Estimated memory per process: {bytes2human(bytes_per_process)}")

        num_procs = min(num_procs,
                        math.floor((0.9*available_bytes)/bytes_per_process))

    print(f"Use {num_procs} cpus out of {multiprocessing.cpu_count()}")
    return num_procs