To run a test in a subprocess. In particular, this can avoid (GPU) memory issue. Args: test_case (`unittest.TestCase`): The test that will run `target_func`. target_func (`Callable`): The function implementing the actual testing logic. inputs
(test_case, target_func, inputs=None, timeout=None)
| 818 | |
| 819 | # Taken from: https://github.com/huggingface/transformers/blob/3658488ff77ff8d45101293e749263acf437f4d5/src/transformers/testing_utils.py#L1787 |
| 820 | def run_test_in_subprocess(test_case, target_func, inputs=None, timeout=None): |
| 821 | """ |
| 822 | To run a test in a subprocess. In particular, this can avoid (GPU) memory issue. |
| 823 | |
| 824 | Args: |
| 825 | test_case (`unittest.TestCase`): |
| 826 | The test that will run `target_func`. |
| 827 | target_func (`Callable`): |
| 828 | The function implementing the actual testing logic. |
| 829 | inputs (`dict`, *optional*, defaults to `None`): |
| 830 | The inputs that will be passed to `target_func` through an (input) queue. |
| 831 | timeout (`int`, *optional*, defaults to `None`): |
| 832 | The timeout (in seconds) that will be passed to the input and output queues. If not specified, the env. |
| 833 | variable `PYTEST_TIMEOUT` will be checked. If still `None`, its value will be set to `600`. |
| 834 | """ |
| 835 | if timeout is None: |
| 836 | timeout = int(os.environ.get("PYTEST_TIMEOUT", 600)) |
| 837 | |
| 838 | start_methohd = "spawn" |
| 839 | ctx = multiprocessing.get_context(start_methohd) |
| 840 | |
| 841 | input_queue = ctx.Queue(1) |
| 842 | output_queue = ctx.JoinableQueue(1) |
| 843 | |
| 844 | # We can't send `unittest.TestCase` to the child, otherwise we get issues regarding pickle. |
| 845 | input_queue.put(inputs, timeout=timeout) |
| 846 | |
| 847 | process = ctx.Process(target=target_func, args=(input_queue, output_queue, timeout)) |
| 848 | process.start() |
| 849 | # Kill the child process if we can't get outputs from it in time: otherwise, the hanging subprocess prevents |
| 850 | # the test to exit properly. |
| 851 | try: |
| 852 | results = output_queue.get(timeout=timeout) |
| 853 | output_queue.task_done() |
| 854 | except Exception as e: |
| 855 | process.terminate() |
| 856 | test_case.fail(e) |
| 857 | process.join(timeout=timeout) |
| 858 | |
| 859 | if results["error"] is not None: |
| 860 | test_case.fail(f'{results["error"]}') |
| 861 | |
| 862 | |
| 863 | class CaptureLogger: |
no test coverage detected