test_multi_processes.py

from __future__ import annotations

import multiprocessing
import time
from concurrent.futures import ProcessPoolExecutor

from file_progress import (
    MultiWorkerFileProgress,
    get_sub_progress_process,
)


def _run_worker(args: tuple[int, int]):
    worker_id, steps = args
    sub_progress = get_sub_progress_process(worker_id)
    for step in range(1, steps + 1):
        sub_progress.update(
            f"process worker {worker_id}",
            step,
            steps,
            extra=f"task={step} phase=sleep",
        )
        time.sleep(1)
        sub_progress.increment_total(1)
    return worker_id


if __name__ == "__main__":
    worker_steps = [3, 4, 5]
    ctx = multiprocessing.get_context("spawn")
    with MultiWorkerFileProgress(
        desc="process pool demo",
        total=sum(worker_steps),
        interval_seconds=0.0,
        cleanup_on_success=False,
        verbose=2,
    ) as progress:
        initializer, initargs = progress.process_initializer(ctx)
        with ProcessPoolExecutor(
            max_workers=len(worker_steps),
            mp_context=ctx,
            initializer=initializer,
            initargs=initargs,
        ) as executor:
            list(executor.map(_run_worker, enumerate(worker_steps, start=1)))