Function _worker_loop

python/oneflow/utils/data/_utils/worker.py:235–377

(
    dataset_kind,
    dataset,
    index_queue,
    data_queue,
    done_event,
    auto_collation,
    collate_fn,
    drop_last,
    base_seed,
    init_fn,
    worker_id,
    num_workers,
    persistent_workers,
)

def _worker_loop(
    dataset_kind,
    dataset,
    index_queue,
    data_queue,
    done_event,
    auto_collation,
    collate_fn,
    drop_last,
    base_seed,
    init_fn,
    worker_id,
    num_workers,
    persistent_workers,
):
    # See NOTE [ Data Loader Multiprocessing Shutdown Logic ] for details on the
    # logic of this function.
    try:

        def cleanup_shm_at_exit(num, frame):
            unlink_all_shared_memory()
            # Use os._exit() to handle the exit of the subprocess to avoid share memory leaks
            # caused by the subprocess continuing for a period of time after the parent process ends.
            os._exit(0)

        _prctl_pr_set_pdeathsig(signal.SIGINT)

        # Initialize C side signal handlers for SIGBUS and SIGSEGV. Python signal
        # module's handlers are executed after Python returns from C low-level
        # handlers, likely when the same fatal signal had already happened
        # again.
        # https://docs.python.org/3/library/signal.html#execution-of-python-signal-handlers
        signal_handling._set_worker_signal_handlers()
        signal.signal(signal.SIGTERM, cleanup_shm_at_exit)
        signal.signal(signal.SIGINT, cleanup_shm_at_exit)
        flow.set_num_threads(1)
        seed = base_seed + worker_id
        random.seed(seed)
        flow.manual_seed(seed)
        if HAS_NUMPY:
            np_seed = _generate_state(base_seed, worker_id)
            import numpy as np

            np.random.seed(np_seed)

        global _worker_info
        _worker_info = WorkerInfo(
            id=worker_id, num_workers=num_workers, seed=seed, dataset=dataset
        )

        from oneflow.utils.data import _DatasetKind

        init_exception = None

        try:
            if init_fn is not None:
                init_fn(worker_id)