Skip to content

Commit

Permalink
Handle platform dependent start method (#179)
Browse files Browse the repository at this point in the history
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] python/cpython#77906
  • Loading branch information
digantdesai authored Nov 27, 2024
1 parent 5fa72f8 commit b3e811b
Showing 1 changed file with 5 additions and 2 deletions.
7 changes: 5 additions & 2 deletions python/test/unit/runtime/test_subproc.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit b3e811b

Please sign in to comment.