Method __call__

tests/test_utils.py:577–609

(self, obj)

                warnings.warn(f"While closing process group: {e}.")

    def __call__(self, obj):
        if not torch.distributed.is_available():
            return unittest.skipIf(True, "Skipping distributed tests because not torch.distributed.is_available()")(obj)
        if torch.cuda.is_available() and torch.cuda.device_count() < self.nproc_per_node:
            return unittest.skipIf(
                True,
                f"Skipping distributed tests because it requires {self.nnodes} devices "
                f"but got {torch.cuda.device_count()}",
            )(obj)

        _cache_original_func(obj)

        @functools.wraps(obj)
        def _wrapper(*args, **kwargs):
            tmp = torch.multiprocessing.get_context(self.method)
            processes = []
            results = tmp.Queue()
            func = _call_original_func
            args = [obj.__name__, obj.__module__] + list(args)

            for proc_rank in range(self.nproc_per_node):
                run_args = (func, proc_rank, args, kwargs, results)
                p = tmp.Process(target=self.run_process, args=run_args, daemon=self.daemon)
                p.start()
                processes.append(p)

            for p in processes:
                p.join()
                pr = results.get(block=False)
                assert pr is True, f"Distributed call failed: {pr}"
            _del_original_func(obj)

        return _wrapper


class TimedCall: