python-wml 3.0.0__py3-none-any.whl

wml/wtorch/data/multi_processing_data_loader_iter.py

@@ -0,0 +1,866 @@
+import os
+import threading
+import itertools
+import warnings
+from typing import Any, Callable, Iterable, TypeVar, Generic, Sequence, List, Optional
+import multiprocessing as python_multiprocessing
+import torch
+import torch.multiprocessing as multiprocessing
+from torch._utils import ExceptionWrapper
+if torch.__version__ < "1.9.0":
+    from torch._six import queue, container_abcs
+else:
+    import queue
+    import collections as container_abcs
+import wml.wtorch.utils as wtu
+import time
+from . import IterableDataset, Sampler, SequentialSampler, RandomSampler, BatchSampler, Dataset
+from . import _utils, _BaseDataLoaderIter
+import wml.wml_utils as wmlu
+import traceback
+import sys
+import math
+
+class _MultiProcessingDataLoaderIter(_BaseDataLoaderIter):
+    r"""Iterates once over the DataLoader's dataset, as specified by the sampler"""
+
+    # NOTE [ Data Loader Multiprocessing Shutdown Logic ]
+    #
+    # Preliminary:
+    #
+    # Our data model looks like this (queues are indicated with curly brackets):
+    #
+    #                main process                              ||
+    #                     |                                    ||
+    #               {index_queue}                              ||
+    #                     |                                    ||
+    #              worker processes                            ||     DATA
+    #                     |                                    ||
+    #            {worker_result_queue}                         ||     FLOW
+    #                     |                                    ||
+    #      pin_memory_thread of main process                   ||   DIRECTION
+    #                     |                                    ||
+    #               {data_queue}                               ||
+    #                     |                                    ||
+    #                data output                               \/
+    #
+    # P.S. `worker_result_queue` and `pin_memory_thread` part may be omitted if
+    #      `pin_memory=False`.
+    #
+    #
+    # Terminating multiprocessing logic requires very careful design. In
+    # particular, we need to make sure that
+    #
+    #   1. The iterator gracefully exits the workers when its last reference is
+    #      gone or it is depleted.
+    #
+    #      In this case, the workers should be gracefully exited because the
+    #      main process may still need to continue to run, and we want cleaning
+    #      up code in the workers to be executed (e.g., releasing GPU memory).
+    #      Naturally, we implement the shutdown logic in `__del__` of
+    #      DataLoaderIterator.
+    #
+    #      We delay the discussion on the logic in this case until later.
+    #
+    #   2. The iterator exits the workers when the loader process and/or worker
+    #      processes exits normally or with error.
+    #
+    #      We set all workers and `pin_memory_thread` to have `daemon=True`.
+    #
+    #      You may ask, why can't we make the workers non-daemonic, and
+    #      gracefully exit using the same logic as we have in `__del__` when the
+    #      iterator gets deleted (see 1 above)?
+    #
+    #      First of all, `__del__` is **not** guaranteed to be called when
+    #      interpreter exits. Even if it is called, by the time it executes,
+    #      many Python core library resources may alreay be freed, and even
+    #      simple things like acquiring an internal lock of a queue may hang.
+    #      Therefore, in this case, we actually need to prevent `__del__` from
+    #      being executed, and rely on the automatic termination of daemonic
+    #      children.
+    #
+    #      Thus, we register an `atexit` hook that sets a global flag
+    #      `_utils.python_exit_status`. Since `atexit` hooks are executed in the
+    #      reverse order of registration, we are guaranteed that this flag is
+    #      set before library resources we use are freed (which, at least in
+    #      CPython, is done via an `atexit` handler defined in
+    #      `multiprocessing/util.py`
+    #      https://github.com/python/cpython/blob/c606624af8d4cb3b4a052fb263bb983b3f87585b/Lib/multiprocessing/util.py#L320-L362
+    #      registered when an object requiring this mechanism is first
+    #      created, e.g., `mp.Queue`
+    #      https://github.com/python/cpython/blob/c606624af8d4cb3b4a052fb263bb983b3f87585b/Lib/multiprocessing/context.py#L100-L103
+    #      https://github.com/python/cpython/blob/c606624af8d4cb3b4a052fb263bb983b3f87585b/Lib/multiprocessing/queues.py#L29
+    #      )
+    #
+    #      So in `__del__`, we check if `_utils.python_exit_status` is set or
+    #      `None` (freed), and perform no-op if so.
+    #
+    #      However, simply letting library clean-up codes run can also be bad,
+    #      because such codes (i.e., `multiprocessing.util._exit_function()`)
+    #      include join putting threads for `mp.Queue`, which can be blocking.
+    #      Hence, the main process putting threads are called with
+    #      `cancel_join_thread` at creation.  See later section
+    #      [ 3b. A process won't hang when putting into a queue; ]
+    #      for more details.
+    #
+    #      Here are two example cases where library clean-up codes can run
+    #      before `__del__` is called:
+    #
+    #        1. If we hold onto a reference to the iterator, it more often
+    #           than not tries to do `multiprocessing` library cleaning before
+    #           clearing the alive referenced objects (https://github.com/pytorch/pytorch/issues/48666)
+    #           and thus prevents our cleaning-up code to run first.
+    #
+    #        2. A similar issue araises when a `DataLoader` is used in a subprocess.
+    #           When a process ends, it shuts the all its daemonic children
+    #           down with a SIGTERM (instead of joining them without a timeout).
+    #           Simiarly for threads, but by a different mechanism. This fact,
+    #           together with a few implementation details of multiprocessing, forces
+    #           us to make workers daemonic. All of our problems arise when a
+    #           DataLoader is used in a subprocess, and are caused by multiprocessing
+    #           code which looks more or less like this:
+    #
+    #               try:
+    #                   your_function_using_a_dataloader()
+    #               finally:
+    #                   multiprocessing.util._exit_function()
+    #
+    #           The joining/termination mentioned above happens inside
+    #           `_exit_function()`. Now, if `your_function_using_a_dataloader()`
+    #           throws, the stack trace stored in the exception will prevent the
+    #           frame which uses `DataLoaderIter` to be freed. If the frame has any
+    #           reference to the `DataLoaderIter` (e.g., in a method of the iter),
+    #           its  `__del__`, which starts the shutdown procedure, will not be
+    #           called. That, in turn, means that workers aren't notified. Attempting
+    #           to join in `_exit_function` will then result in a hang.
+    #
+    #           For context, `_exit_function` is also registered as an `atexit` call.
+    #           So it is unclear to me (@ssnl) why this is needed in a finally block.
+    #           The code dates back to 2008 and there is no comment on the original
+    #           PEP 371 or patch https://bugs.python.org/issue3050 (containing both
+    #           the finally block and the `atexit` registration) that explains this.
+    #
+    #
+    #      Finally, another choice is to just shutdown workers with logic in 1
+    #      above whenever we see an error in `next`. This isn't ideal because
+    #        a. It prevents users from using try-catch to resume data loading.
+    #        b. It doesn't prevent hanging if users have references to the
+    #           iterator.
+    #
+    #   3. All processes exit if any of them die unexpectedly by fatal signals.
+    #
+    #      As shown above, the workers are set as daemonic children of the main
+    #      process. However, automatic cleaning-up of such child processes only
+    #      happens if the parent process exits gracefully (e.g., not via fatal
+    #      signals like SIGKILL). So we must ensure that each process will exit
+    #      even the process that should send/receive data to/from it were
+    #      killed, i.e.,
+    #
+    #        a. A process won't hang when getting from a queue.
+    #
+    #           Even with carefully designed data dependencies (i.e., a `put()`
+    #           always corresponding to a `get()`), hanging on `get()` can still
+    #           happen when data in queue is corrupted (e.g., due to
+    #           `cancel_join_thread` or unexpected exit).
+    #
+    #           For child exit, we set a timeout whenever we try to get data
+    #           from `data_queue`, and check the workers' status on each timeout
+    #           and error.
+    #           See `_DataLoaderiter._get_batch()` and
+    #           `_DataLoaderiter._try_get_data()` for details.
+    #
+    #           Additionally, for child exit on non-Windows platforms, we also
+    #           register a SIGCHLD handler (which is supported on Windows) on
+    #           the main process, which checks if any of the workers fail in the
+    #           (Python) handler. This is more efficient and faster in detecting
+    #           worker failures, compared to only using the above mechanism.
+    #           See `DataLoader.cpp` and `_utils/signal_handling.py` for details.
+    #
+    #           For `.get()` calls where the sender(s) is not the workers, we
+    #           guard them with timeouts, and check the status of the sender
+    #           when timeout happens:
+    #             + in the workers, the `_utils.worker.ManagerWatchdog` class
+    #               checks the status of the main process.
+    #             + if `pin_memory=True`, when getting from `pin_memory_thread`,
+    #               check `pin_memory_thread` status periodically until `.get()`
+    #               returns or see that `pin_memory_thread` died.
+    #
+    #        b. A process won't hang when putting into a queue;
+    #
+    #           We use `mp.Queue` which has a separate background thread to put
+    #           objects from an unbounded buffer array. The background thread is
+    #           daemonic and usually automatically joined when the process
+    #           *exits*.
+    #
+    #           In case that the receiver has ended abruptly while
+    #           reading from the pipe, the join will hang forever.  The usual
+    #           solution for this in Python is calling  `q.cancel_join_thread`,
+    #           which prevents automatically joining it when finalizing
+    #           (exiting).
+    #
+    #           Nonetheless, `cancel_join_thread` must only be called when the
+    #           queue is **not** going to be read from or write into by another
+    #           process, because it may hold onto a lock or leave corrupted data
+    #           in the queue, leading other readers/writers to hang.
+    #
+    #           Hence,
+    #             + For worker processes, we only do so (for their output
+    #               queues, i.e., `worker_result_queue`) before exiting.
+    #             + For `pin_memory_thread`, its output queue `data_queue` is a
+    #               `queue.Queue` that does blocking `put` if the queue is full.
+    #               So there is no above problem, but as a result, in
+    #               `_pin_memory_loop`, we do need to  wrap the `put` in a loop
+    #               that breaks not only upon success, but also when the main
+    #               process stops reading, i.e., is shutting down.
+    #             + For loader process, we `cancel_join_thread()` for all
+    #               `_index_queues` because the whole purpose of workers and
+    #               `pin_memory_thread` is to serve the loader process.  If
+    #               loader process is already exiting, we don't really care if
+    #               the queues are corrupted.
+    #
+    #
+    # Now let's get back to 1:
+    #   how we gracefully exit the workers when the last reference to the
+    #   iterator is gone.
+    #
+    # To achieve this, we implement the following logic along with the design
+    # choices mentioned above:
+    #
+    # `workers_done_event`:
+    #   A `multiprocessing.Event` shared among the main process and all worker
+    #   processes. This is used to signal the workers that the iterator is
+    #   shutting down. After it is set, they will not send processed data to
+    #   queues anymore, and only wait for the final `None` before exiting.
+    #   `done_event` isn't strictly needed. I.e., we can just check for `None`
+    #   from the input queue, but it allows us to skip wasting resources
+    #   processing data if we are already shutting down.
+    #
+    # `pin_memory_thread_done_event`:
+    #   A `threading.Event` for a similar purpose to that of
+    #   `workers_done_event`, but is for the `pin_memory_thread`. The reason
+    #   that separate events are needed is that `pin_memory_thread` reads from
+    #   the output queue of the workers. But the workers, upon seeing that
+    #   `workers_done_event` is set, only wants to see the final `None`, and is
+    #   not required to flush all data in the output queue (e.g., it may call
+    #   `cancel_join_thread` on that queue if its `IterableDataset` iterator
+    #   happens to exhaust coincidentally, which is out of the control of the
+    #   main process). Thus, since we will exit `pin_memory_thread` before the
+    #   workers (see below), two separete events are used.
+    #
+    # NOTE: In short, the protocol is that the main process will set these
+    #       `done_event`s and then the corresponding processes/threads a `None`,
+    #       and that they may exit at any time after receiving the `None`.
+    #
+    # NOTE: Using `None` as the final signal is valid, since normal data will
+    #       always be a 2-tuple with the 1st element being the index of the data
+    #       transferred (different from dataset index/key), and the 2nd being
+    #       either the dataset key or the data sample (depending on which part
+    #       of the data model the queue is at).
+    #
+    # [ worker processes ]
+    #   While loader process is alive:
+    #     Get from `index_queue`.
+    #       If get anything else,
+    #          Check `workers_done_event`.
+    #            If set, continue to next iteration
+    #                    i.e., keep getting until see the `None`, then exit.
+    #            Otherwise, process data:
+    #                If is fetching from an `IterableDataset` and the iterator
+    #                    is exhausted, send an `_IterableDatasetStopIteration`
+    #                    object to signal iteration end. The main process, upon
+    #                    receiving such an object, will send `None` to this
+    #                    worker and not use the corresponding `index_queue`
+    #                    anymore.
+    #       If timed out,
+    #          No matter `workers_done_event` is set (still need to see `None`)
+    #          or not, must continue to next iteration.
+    #   (outside loop)
+    #   If `workers_done_event` is set,  (this can be False with `IterableDataset`)
+    #     `data_queue.cancel_join_thread()`.  (Everything is ending here:
+    #                                          main process won't read from it;
+    #                                          other workers will also call
+    #                                          `cancel_join_thread`.)
+    #
+    # [ pin_memory_thread ]
+    #   # No need to check main thread. If this thread is alive, the main loader
+    #   # thread must be alive, because this thread is set as daemonic.
+    #   While `pin_memory_thread_done_event` is not set:
+    #     Get from `index_queue`.
+    #       If timed out, continue to get in the next iteration.
+    #       Otherwise, process data.
+    #       While `pin_memory_thread_done_event` is not set:
+    #         Put processed data to `data_queue` (a `queue.Queue` with blocking put)
+    #         If timed out, continue to put in the next iteration.
+    #         Otherwise, break, i.e., continuing to the out loop.
+    #
+    #   NOTE: we don't check the status of the main thread because
+    #           1. if the process is killed by fatal signal, `pin_memory_thread`
+    #              ends.
+    #           2. in other cases, either the cleaning-up in __del__ or the
+    #              automatic exit of daemonic thread will take care of it.
+    #              This won't busy-wait either because `.get(timeout)` does not
+    #              busy-wait.
+    #
+    # [ main process ]
+    #   In the DataLoader Iter's `__del__`
+    #     b. Exit `pin_memory_thread`
+    #          i.   Set `pin_memory_thread_done_event`.
+    #          ii   Put `None` in `worker_result_queue`.
+    #          iii. Join the `pin_memory_thread`.
+    #          iv.  `worker_result_queue.cancel_join_thread()`.
+    #
+    #     c. Exit the workers.
+    #          i.   Set `workers_done_event`.
+    #          ii.  Put `None` in each worker's `index_queue`.
+    #          iii. Join the workers.
+    #          iv.  Call `.cancel_join_thread()` on each worker's `index_queue`.
+    #
+    #        NOTE: (c) is better placed after (b) because it may leave corrupted
+    #              data in `worker_result_queue`, which `pin_memory_thread`
+    #              reads from, in which case the `pin_memory_thread` can only
+    #              happen at timeing out, which is slow. Nonetheless, same thing
+    #              happens if a worker is killed by signal at unfortunate times,
+    #              but in other cases, we are better off having a non-corrupted
+    #              `worker_result_queue` for `pin_memory_thread`.
+    #
+    #   NOTE: If `pin_memory=False`, there is no `pin_memory_thread` and (b)
+    #         can be omitted
+    #
+    # NB: `done_event`s isn't strictly needed. E.g., we can just check for
+    #     `None` from `index_queue`, but it allows us to skip wasting resources
+    #     processing indices already in `index_queue` if we are already shutting
+    #     down.
+
+    def __init__(self, loader,batch_split_nr=1):
+        super(_MultiProcessingDataLoaderIter, self).__init__(loader)
+        print(f"Use _WMultiProcessingDataLoaderIter, batch split nr = {batch_split_nr}, pin memory= {loader.pin_memory}, num_works= {self._num_workers}")
+
+        assert self._num_workers > 0
+        assert self._prefetch_factor > 0
+
+
+        if loader.multiprocessing_context is None:
+            multiprocessing_context = multiprocessing
+        else:
+            multiprocessing_context = loader.multiprocessing_context
+
+        self._worker_init_fn = loader.worker_init_fn
+        self._worker_queue_idx_cycle = itertools.cycle(range(self._num_workers))
+        # No certainty which module multiprocessing_context is
+        self._worker_result_queue = multiprocessing_context.Queue(max(self._prefetch_factor*batch_split_nr*2,min(self._num_workers,6)*self._prefetch_factor*batch_split_nr))  # type: ignore
+        self._worker_pids_set = False
+        self._shutdown = False
+        self._workers_done_event = multiprocessing_context.Event()
+        self.stop_iteration = False
+
+        self._index_queues = []
+        self._workers = []
+        self.batch_split_nr = max(1,batch_split_nr)
+        for i in range(self._num_workers):
+            # No certainty which module multiprocessing_context is
+            index_queue = multiprocessing_context.Queue()  # type: ignore
+            # Need to `cancel_join_thread` here!
+            # See sections (2) and (3b) above.
+            index_queue.cancel_join_thread()
+            w = multiprocessing_context.Process(
+                target=_utils.worker._worker_loop,
+                args=(self._dataset_kind, self._dataset, index_queue,
+                      self._worker_result_queue, self._workers_done_event,
+                      self._auto_collation, self._collate_fn, self._drop_last,
+                      self._base_seed + i, self._worker_init_fn, i, self._num_workers,
+                      self._persistent_workers))
+            w.daemon = True
+            # NB: Process.start() actually take some time as it needs to
+            #     start a process and pass the arguments over via a pipe.
+            #     Therefore, we only add a worker to self._workers list after
+            #     it started, so that we do not call .join() if program dies
+            #     before it starts, and __del__ tries to join but will get:
+            #     AssertionError: can only join a started process.
+            w.start()
+            self._index_queues.append(index_queue)
+            self._workers.append(w)
+
+        if self._pin_memory:
+            self._pin_memory_thread_done_event = threading.Event()
+
+            # Queue is not type-annotated
+            self._data_queue = queue.Queue(4*self.batch_split_nr)  # type: ignore
+            self.pin_memory_stream = torch.cuda.Stream()
+            pin_memory_thread = threading.Thread(
+                target=_utils.pin_memory._pin_memory_loop_stream,
+                args=(self._worker_result_queue, self._data_queue,
+                      torch.cuda.current_device(),
+                      self._pin_memory_thread_done_event,
+                      self.pin_memory_stream))
+            pin_memory_thread.daemon = True
+            pin_memory_thread.start()
+            # Similar to workers (see comment above), we only register
+            # pin_memory_thread once it is started.
+            self._pin_memory_thread = pin_memory_thread
+            self._try_get_data_imp = self._try_get_data_imp_stream
+        else:
+            self._data_queue = self._worker_result_queue
+            self._try_get_data_imp = self._try_get_data_imp_no_stream
+
+        # .pid can be None only before process is spawned (not the case, so ignore)
+        _utils.signal_handling._set_worker_pids(id(self), tuple(w.pid for w in self._workers))  # type: ignore
+        _utils.signal_handling._set_SIGCHLD_handler()
+        self._worker_pids_set = True
+        self._reset(loader, first_iter=True)
+        self.datas_cache = []
+
+        all_pids = f"{os.getpid()}"
+        for x in self._workers:
+            all_pids += f" {x.ident}"
+        print(f"All data workers process PID: {all_pids}")
+        print(all_pids.replace(" ",","))
+
+    def _reset(self, loader, first_iter=False):
+        super()._reset(loader, first_iter)
+        self.stop_iteration = False
+        self._send_idx = 0  # idx of the next task to be sent to workers
+        self._rcvd_idx = 0  # idx of the next task to be returned in __next__
+        # information about data not yet yielded, i.e., tasks w/ indices in range [rcvd_idx, send_idx).
+        # map: task idx => - (worker_id,)        if data isn't fetched (outstanding)
+        #                  \ (worker_id, data)   if data is already fetched (out-of-order)
+        self._task_info = {}
+        self._tasks_outstanding = 0  # always equal to count(v for v in task_info.values() if len(v) == 1)
+        # A list of booleans representing whether each worker still has work to
+        # do, i.e., not having exhausted its iterable dataset object. It always
+        # contains all `True`s if not using an iterable-style dataset
+        # (i.e., if kind != Iterable).
+        # Not that this indicates that a worker still has work to do *for this epoch*.
+        # It does not mean that a worker is dead. In case of `_persistent_workers`,
+        # the worker will be reset to available in the next epoch.
+        self._workers_status = [True for i in range(self._num_workers)]
+        # prime the prefetch loop
+        self._try_put_index_first()
+
+    def _try_get_data(self, timeout=_utils.MP_STATUS_CHECK_INTERVAL):
+        if self.batch_split_nr<=1:
+            res,data = self._try_get_data_imp(timeout=timeout)
+            if res is False:
+                if self.stop_iteration:
+                    for q in self._index_queues:
+                        if q.qsize() > 0:
+                            return False, None
+                    raise StopIteration()
+                else:
+                    return res,data
+            else:
+                return res,data
+        else:
+            try_nr = self.batch_split_nr+2
+            while try_nr>0 and len(self.datas_cache)<self.batch_split_nr:
+                res,data = self._try_get_data_imp(timeout=max(1,timeout/self.batch_split_nr))
+                if res:
+                    if isinstance(data[1], ExceptionWrapper):
+                        e = data[1]
+                        msg = "ERROR: Caught {} {}.\nOriginal {}".format(
+                            e.exc_type.__name__, e.where, e.exc_msg)
+                        print(msg)
+                    elif data[1] is not None:
+                        self.datas_cache.append(data[1])
+                    elif self.stop_iteration:
+                        break
+                try_nr -= 1
+            if len(self.datas_cache)>=self.batch_split_nr:
+                try:
+                    data = wtu.concat_datas(self.datas_cache[:self.batch_split_nr],dim=0)
+                    self.datas_cache = self.datas_cache[self.batch_split_nr:]
+                except Exception as e:
+                    print(f"ERROR: Concat datas faild, {e}.")
+                    self.datas_cache = []
+                    return False,None
+                return True,(0,data)
+            if self.stop_iteration:
+                for q in self._index_queues:
+                    if q.qsize()>0:
+                        return False,None
+                raise StopIteration()
+            else:
+                return False,None
+
+    @staticmethod
+    def record_stream(data,stream):
+        if torch.is_tensor(data):
+            if data.dtype != torch.int16: #hack: 不处理int16
+                data.record_stream(stream)
+        elif isinstance(data,container_abcs.Mapping):
+            for k,v in data.items():
+                _MultiProcessingDataLoaderIter.record_stream(v,stream)
+        elif isinstance(data,Iterable):
+            for x in data:
+                _MultiProcessingDataLoaderIter.record_stream(x,stream)
+
+    def _try_get_data_imp_stream(self, timeout=_utils.MP_STATUS_CHECK_INTERVAL):
+        torch.cuda.current_stream().wait_stream(self.pin_memory_stream)
+        res = self._try_get_data_imp_no_stream(timeout=timeout)
+        if res[0]:
+            _MultiProcessingDataLoaderIter.record_stream(res[1][1],torch.cuda.current_stream())
+        return res
+
+    def _try_get_data_imp_no_stream(self, timeout=_utils.MP_STATUS_CHECK_INTERVAL):
+        # Tries to fetch data from `self._data_queue` once for a given timeout.
+        # This can also be used as inner loop of fetching without timeout, with
+        # the sender status as the loop condition.
+        #
+        # This raises a `RuntimeError` if any worker died expectedly. This error
+        # can come from either the SIGCHLD handler in `_utils/signal_handling.py`
+        # (only for non-Windows platforms), or the manual check below on errors
+        # and timeouts.
+        #
+        # Returns a 2-tuple:
+        #   (bool: whether successfully get data, any: data if successful else None)
+        try:
+            data = self._data_queue.get(timeout=timeout)
+            return (True, data)
+        except Exception as e:
+            # At timeout and error, we manually check whether any worker has
+            # failed. Note that this is the only mechanism for Windows to detect
+            # worker failures.
+            failed_workers = []
+            for worker_id, w in enumerate(self._workers):
+                if self._workers_status[worker_id] and not w.is_alive():
+                    failed_workers.append(w)
+                    self._mark_worker_as_unavailable(worker_id)
+            if len(failed_workers) > 0:
+                pids_str = ', '.join(str(w.pid) for w in failed_workers)
+                raise RuntimeError('DataLoader worker (pid(s) {}) exited unexpectedly'.format(pids_str)) from e
+            if isinstance(e, queue.Empty):
+                return (False, None)
+            import tempfile
+            import errno
+            try:
+                # Raise an exception if we are this close to the FDs limit.
+                # Apparently, trying to open only one file is not a sufficient
+                # test.
+                # See NOTE [ DataLoader on Linux and open files limit ]
+                fds_limit_margin = 10
+                fs = [tempfile.NamedTemporaryFile() for i in range(fds_limit_margin)]
+            except OSError as e:
+                if e.errno == errno.EMFILE:
+                    print(f"Open files nr")
+                    cmd = f"ls /proc/{os.getpid()}/fd | wc -l"
+                    os.system(cmd)
+                    print(f"ulimit -n")
+                    cmd = f"ulimit -n"
+                    os.system(cmd)
+                    raise RuntimeError(
+                        "Too many open files. Communication with the"
+                        " workers is no longer possible. Please increase the"
+                        " limit using `ulimit -n` in the shell or change the"
+                        " sharing strategy by calling"
+                        " `torch.multiprocessing.set_sharing_strategy('file_system')`"
+                        " at the beginning of your code") from None
+            raise
+
+# NOTE [ DataLoader on Linux and open files limit ]
+#
+# On Linux when DataLoader is used with multiprocessing we pass the data between
+# the root process and the workers through SHM files. We remove those files from
+# the filesystem as soon as they are created and keep them alive by
+# passing around their file descriptors through AF_UNIX sockets. (See
+# docs/source/multiprocessing.rst and 'Multiprocessing Technical Notes` in
+# the wiki (https://github.com/pytorch/pytorch/wiki).)
+#
+# This sometimes leads us to exceeding the open files limit. When that happens,
+# and the offending file descriptor is coming over a socket, the `socket` Python
+# package silently strips the file descriptor from the message, setting only the
+# `MSG_CTRUNC` flag (which might be a bit misleading since the manpage says that
+# it _indicates that some control data were discarded due to lack of space in
+# the buffer for ancillary data_). This might reflect the C implementation of
+# AF_UNIX sockets.
+#
+# This behaviour can be reproduced with the script and instructions at the
+# bottom of this note.
+#
+# When that happens, the standard Python `multiprocessing` (and not
+# `torch.multiprocessing`) raises a `RuntimeError: received 0 items of ancdata`
+#
+# Sometimes, instead of the FD being stripped, you may get an `OSError:
+# Too many open files`, both in the script below and in DataLoader. However,
+# this is rare and seems to be nondeterministic.
+#
+#
+#   #!/usr/bin/env python3
+#   import sys
+#   import socket
+#   import os
+#   import array
+#   import shutil
+#   import socket
+#
+#
+#   if len(sys.argv) != 4:
+#       print("Usage: ", sys.argv[0], " tmp_dirname iteration (send|recv)")
+#       sys.exit(1)
+#
+#   if __name__ == '__main__':
+#       dirname = sys.argv[1]
+#       sock_path = dirname + "/sock"
+#       iterations = int(sys.argv[2])
+#       def dummy_path(i):
+#           return dirname + "/" + str(i) + ".dummy"
+#
+#
+#       if sys.argv[3] == 'send':
+#           while not os.path.exists(sock_path):
+#               pass
+#           client = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
+#           client.connect(sock_path)
+#           for i in range(iterations):
+#               fd = os.open(dummy_path(i), os.O_WRONLY | os.O_CREAT)
+#               ancdata = array.array('i', [fd])
+#               msg = bytes([i % 256])
+#               print("Sending fd ", fd, " (iteration #", i, ")")
+#               client.sendmsg([msg], [(socket.SOL_SOCKET, socket.SCM_RIGHTS, ancdata)])
+#
+#
+#       else:
+#           assert sys.argv[3] == 'recv'
+#
+#           if os.path.exists(dirname):
+#               raise Exception("Directory exists")
+#
+#           os.mkdir(dirname)
+#
+#           print("Opening socket...")
+#           server = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
+#           server.bind(sock_path)
+#
+#           print("Listening...")
+#           for i in range(iterations):
+#               a = array.array('i')
+#               msg, ancdata, flags, addr = server.recvmsg(1, socket.CMSG_SPACE(a.itemsize))
+#               assert(len(ancdata) == 1)
+#               cmsg_level, cmsg_type, cmsg_data = ancdata[0]
+#               a.frombytes(cmsg_data)
+#               print("Received fd ", a[0], " (iteration #", i, ")")
+#
+#           shutil.rmtree(dirname)
+#
+# Steps to reproduce:
+#
+# 1. Run two shells and set lower file descriptor limit in the receiving one:
+# (shell1) ulimit -n 1020
+# (shell2) ulimit -n 1022
+#
+# 2. Run the script above with the `recv` option in the first shell
+# (shell1) ./test_socket.py sock_tmp 1017 recv
+#
+# 3. Run the script with the `send` option in the second shell:
+# (shell2) ./test_socket.py sock_tmp 1017 send
+
+    def _get_data(self):
+        # Fetches data from `self._data_queue`.
+        #
+        # We check workers' status every `MP_STATUS_CHECK_INTERVAL` seconds,
+        # which we achieve by running `self._try_get_data(timeout=MP_STATUS_CHECK_INTERVAL)`
+        # in a loop. This is the only mechanism to detect worker failures for
+        # Windows. For other platforms, a SIGCHLD handler is also used for
+        # worker failure detection.
+        #
+        # If `pin_memory=True`, we also need check if `pin_memory_thread` had
+        # died at timeouts.
+        if self._timeout > 0:
+            success, data = self._try_get_data(self._timeout)
+            if success:
+                return data
+            else:
+                raise RuntimeError('DataLoader timed out after {} seconds'.format(self._timeout))
+        elif self._pin_memory:
+            while self._pin_memory_thread.is_alive():
+                success, data = self._try_get_data()
+                if success:
+                    return data
+            else:
+                # while condition is false, i.e., pin_memory_thread died.
+                raise RuntimeError('Pin memory thread exited unexpectedly')
+            # In this case, `self._data_queue` is a `queue.Queue`,. But we don't
+            # need to call `.task_done()` because we don't use `.join()`.
+        else:
+            while True:
+                success, data = self._try_get_data()
+                if success:
+                    return data
+
+    def _next_data(self):
+        while True:
+            idx, data = self._get_data()
+            return self._process_data(data)
+
+    def _try_put_index(self):
+        try:
+            dnr = self._prefetch_factor
+            hdnr = max(1,dnr//2)
+            for i,index_queue in enumerate(self._index_queues):
+                if self._workers_status[i] and index_queue.qsize()<hdnr:
+                    nr = dnr-index_queue.qsize()
+                    for _ in range(nr):
+                        index = self._next_index()
+                        if self.batch_split_nr>1:
+                            if len(index)%self.batch_split_nr!=0:
+                                print(f"ERROR: batch_split_nr = {self.batch_split_nr}, batch size = {len(index)}")
+                            indexs = wmlu.list_to_2dlistv2(index,self.batch_split_nr)
+                            for index in indexs:
+                                index_queue.put((0,index))
+                        else:
+                            index_queue.put((0,index))
+        except StopIteration:
+            self.stop_iteration = True
+            return
+
+    def _try_put_index_first(self):
+        try:
+            dnr = self._prefetch_factor
+            nr = max(1,math.ceil(dnr/len(self._index_queues)))
+            for i,index_queue in enumerate(self._index_queues):
+                for _ in range(nr):
+                    index = self._next_index()
+                    if self.batch_split_nr>1:
+                        if len(index)%self.batch_split_nr!=0:
+                            print(f"ERROR: batch_split_nr = {self.batch_split_nr}, batch size = {len(index)}")
+                        indexs = wmlu.list_to_2dlistv2(index,self.batch_split_nr)
+                        for index in indexs:
+                            index_queue.put((0,index))
+                    else:
+                        index_queue.put((0,index))
+        except StopIteration:
+            self.stop_iteration = True
+            return
+
+    def _process_data(self, data):
+        self._try_put_index()
+        if isinstance(data, ExceptionWrapper):
+            data.reraise()
+        return data
+
+    def _mark_worker_as_unavailable(self, worker_id, shutdown=False):
+        # Mark a worker as having finished its work e.g., due to
+        # exhausting an `IterableDataset`. This should be used only when this
+        # `_MultiProcessingDataLoaderIter` is going to continue running.
+
+        assert self._workers_status[worker_id] or (self._persistent_workers and shutdown)
+
+        # Signal termination to that specific worker.
+        q = self._index_queues[worker_id]
+        # Indicate that no more data will be put on this queue by the current
+        # process.
+        q.put(None)
+
+        # Note that we don't actually join the worker here, nor do we remove the
+        # worker's pid from C side struct because (1) joining may be slow, and
+        # (2) since we don't join, the worker may still raise error, and we
+        # prefer capturing those, rather than ignoring them, even though they
+        # are raised after the worker has finished its job.
+        # Joinning is deferred to `_shutdown_workers`, which it is called when
+        # all workers finish their jobs (e.g., `IterableDataset` replicas) or
+        # when this iterator is garbage collected.
+
+        self._workers_status[worker_id] = False
+        print(f"WARNING: make process {worker_id} to unavaiable.")
+
+        assert self._workers_done_event.is_set() == shutdown
+
+    def _clear_index_queue(self):
+        for index_queue in self._index_queues:
+            try:
+                while True:
+                    index_queue.get_nowait()
+            except:
+                pass
+    
+    def _clear_worker_results_queue(self):
+        try:
+            while True:
+                self._worker_result_queue.get_nowait()
+        except Exception as e:
+            pass
+
+    def _shutdown_workers(self):
+        # Called when shutting down this `_MultiProcessingDataLoaderIter`.
+        # See NOTE [ Data Loader Multiprocessing Shutdown Logic ] for details on
+        # the logic of this function.
+        print(f"Shutdown data loader workers.")
+        traceback.print_exc(file=sys.stdout)
+        sys.stdout.flush()
+        python_exit_status = _utils.python_exit_status
+        if python_exit_status is True or python_exit_status is None:
+            # See (2) of the note. If Python is shutting down, do no-op.
+            return
+        # Normal exit when last reference is gone / iterator is depleted.
+        # See (1) and the second half of the note.
+        if not self._shutdown:
+            self._shutdown = True
+            try:
+                # Normal exit when last reference is gone / iterator is depleted.
+                # See (1) and the second half of the note.
+
+                # Exit `pin_memory_thread` first because exiting workers may leave
+                # corrupted data in `worker_result_queue` which `pin_memory_thread`
+                # reads from.
+
+                # Exit workers now.
+                self._clear_index_queue()
+                self._workers_done_event.set()
+                self._clear_worker_results_queue()
+                try:
+                    self._worker_result_queue.put_nowait((None, None))
+                except:
+                    pass
+
+                for worker_id in range(len(self._workers)):
+                    # Get number of workers from `len(self._workers)` instead of
+                    # `self._num_workers` in case we error before starting all
+                    # workers.
+                    # If we are using workers_status with persistent_workers
+                    # we have to shut it down because the worker is paused
+                    if self._persistent_workers or self._workers_status[worker_id]:
+                        self._mark_worker_as_unavailable(worker_id, shutdown=True)
+
+
+                if hasattr(self, '_pin_memory_thread'):
+                    # Use hasattr in case error happens before we set the attribute.
+                    self._pin_memory_thread_done_event.set()
+                    # Send something to pin_memory_thread in case it is waiting
+                    # so that it can wake up and check `pin_memory_thread_done_event`
+                    self._pin_memory_thread.join()
+                    self._worker_result_queue.cancel_join_thread()
+                    self._worker_result_queue.close()
+
+                for w in self._workers:
+                    # We should be able to join here, but in case anything went
+                    # wrong, we set a timeout and if the workers fail to join,
+                    # they are killed in the `finally` block.
+                    w.join(timeout=_utils.MP_STATUS_CHECK_INTERVAL)
+                for q in self._index_queues:
+                    q.cancel_join_thread()
+                    q.close()
+            finally:
+                # Even though all this function does is putting into queues that
+                # we have called `cancel_join_thread` on, weird things can
+                # happen when a worker is killed by a signal, e.g., hanging in
+                # `Event.set()`. So we need to guard this with SIGCHLD handler,
+                # and remove pids from the C side data structure only at the
+                # end.
+                #
+                # FIXME: Unfortunately, for Windows, we are missing a worker
+                #        error detection mechanism here in this function, as it
+                #        doesn't provide a SIGCHLD handler.
+                if self._worker_pids_set:
+                    _utils.signal_handling._remove_worker_pids(id(self))
+                    self._worker_pids_set = False
+                for w in self._workers:
+                    if w.is_alive():
+                        # Existing mechanisms try to make the workers exit
+                        # peacefully, but in case that we unfortunately reach
+                        # here, which we shouldn't, (e.g., pytorch/pytorch#39570),
+                        # we kill the worker.
+                        w.terminate()
+
+    def __del__(self):
+        try:
+            self._shutdown_workers()
+        except:
+            pass