Method _wrapper

tests/test_utils.py:590–607

(*args, **kwargs)

        @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