From b3e811bb8c160780c60369ac3667fda9920ae80e Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Wed, 27 Nov 2024 12:28:15 -0600 Subject: [PATCH] Handle platform dependent start method (#179) In Python multiprocessing, start method name is platform dependent [1]. For macOS, it is "spawn" for reasons listed in [2], vs. "fork" for Linux. [1] https://docs.python.org/3/library/multiprocessing.html#multiprocessing.get_start_method [2] https://github.com/python/cpython/issues/77906 --- python/test/unit/runtime/test_subproc.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/python/test/unit/runtime/test_subproc.py b/python/test/unit/runtime/test_subproc.py index 02777923303a..8f9c592a411f 100644 --- a/python/test/unit/runtime/test_subproc.py +++ b/python/test/unit/runtime/test_subproc.py @@ -49,7 +49,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' + # This can be either fork or spawn, depending on the platform. + assert multiprocessing.get_start_method() in ["fork", "spawn"] proc = multiprocessing.Process(target=compile_fn_dot, args=(config, )) proc.start() proc.join() @@ -92,7 +93,9 @@ 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' + # This can be either fork or spawn, depending on the platform. + assert multiprocessing.get_start_method() in ["fork", "spawn"] + proc = multiprocessing.Process(target=compile_empty_kernel_with_gc, args=(config, )) # stage 3.c