Fix multiprocessing testcase

[k-artem]

Need to force context re-creating inside testcase otherwise runtime exception raises:

    def set_start_method(self, method, force=False):
        if self._actual_context is not None and not force:
>           raise RuntimeError('context has already been set')
E           RuntimeError: context has already been set

Command for verification:
pytest -v tests/unit/runtime/test_ds_initialize.py

[k-artem]

@sfc-gh-truwase please have a look

[sfc-gh-truwase]

@k-artem can you share a bit about the error that this fixes? In my experience set_start_method is very flaky and environment-dependent.

@tohtana, @loadams, any thoughts?

[k-artem]

@k-artem can you share a bit about the error that this fixes? In my experience set_start_method is very flaky and environment-dependent.

@tohtana, @loadams, any thoughts?

sure, on my side it’s consistently reproducible. torch.multiprocessing is a wrapper around Python’s multiprocessing, and its start method can only be set once per process via set_start_method and after that the multiprocessing context is locked and cannot be changed. However, the test calls set_start_method three times with different methods via @pytest.mark.parametrize('method', ['spawn', 'fork', 'forkserver']) , which triggers an exception after first call. In my fix, I force the multiprocessing context to be recreated between runs.

tests/unit/runtime/test_ds_initialize.py::test_start_method_safety[spawn] PASSED                                                                                                                          [ 33%]
tests/unit/runtime/test_ds_initialize.py::test_start_method_safety[fork] FAILED                                                                                                                           [ 66%]
tests/unit/runtime/test_ds_initialize.py::test_start_method_safety[forkserver] FAILED                                                                                                                     [100%]

=================================================================================================== FAILURES ====================================================================================================
________________________________________________________________________________________ test_start_method_safety[fork] _________________________________________________________________________________________

method = 'fork'

    @pytest.mark.parametrize('method', ['spawn', 'fork', 'forkserver'])
    def test_start_method_safety(method):
        import torch.multiprocessing as mp
        print(os.getpid())
>       mp.set_start_method(method)

tests/unit/runtime/test_ds_initialize.py:33:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <multiprocessing.context.DefaultContext object at 0x7620fe759280>, method = 'fork', force = False

    def set_start_method(self, method, force=False):
        if self._actual_context is not None and not force:
>           raise RuntimeError('context has already been set')
E           RuntimeError: context has already been set

/usr/lib/python3.12/multiprocessing/context.py:247: RuntimeError
--------------------------------------------------------------------------------------------- Captured stdout call ----------------------------------------------------------------------------------------------

[k-artem]

hi @sfc-gh-truwase please approve if no objections

[sfc-gh-truwase]

@k-artem, the approval from @loadams is actually sufficient.