Don't use fork method if it's not available on the platform (#5051)

Signed-off-by: Anatoly Myachev <[email protected]>

Author: anmyachev

python/test/unit/runtime/test_subproc.py

@@ -7,6 +7,7 @@
 from triton.compiler import ASTSource
 
 target = triton.runtime.driver.active.get_current_target()
+start_method = 'fork' if 'fork' in multiprocessing.get_all_start_methods() else 'spawn'
 
 
 def compile_fn(attrs):
@@ -27,8 +28,8 @@ def kernel_sub(a, b, o, N: tl.constexpr):
 
 def test_compile_in_subproc() -> None:
     config = AttrsDescriptor.from_hints({i: 16 for i in range(4)})
-    multiprocessing.set_start_method('fork')
-    proc = multiprocessing.Process(target=compile_fn, args=(config, ))
+    mp_ctx = multiprocessing.get_context(start_method)
+    proc = mp_ctx.Process(target=compile_fn, args=(config, ))
     proc.start()
     proc.join()
     assert proc.exitcode == 0
@@ -49,8 +50,8 @@ def kernel_dot(Z):
 
 def test_compile_in_forked_subproc(fresh_triton_cache) -> None:
     config = AttrsDescriptor.from_hints({0: 16})
-    assert multiprocessing.get_start_method() == 'fork'
-    proc = multiprocessing.Process(target=compile_fn_dot, args=(config, ))
+    mp_ctx = multiprocessing.get_context(start_method)
+    proc = mp_ctx.Process(target=compile_fn_dot, args=(config, ))
     proc.start()
     proc.join()
     assert proc.exitcode == 0
@@ -92,8 +93,8 @@ def test_compile_in_forked_subproc_with_forced_gc(fresh_triton_cache) -> None:
 
     # stage 2.p
     shutil.rmtree(fresh_triton_cache)
-    assert multiprocessing.get_start_method() == 'fork'
-    proc = multiprocessing.Process(target=compile_empty_kernel_with_gc, args=(config, ))
+    mp_ctx = multiprocessing.get_context(start_method)
+    proc = mp_ctx.Process(target=compile_empty_kernel_with_gc, args=(config, ))
 
     # stage 3.c
     proc.start()