Skip to content

Commit 7f4d1a7

Browse files
committed
fix: use multiprocessing whenever fork context is available
1 parent 928330f commit 7f4d1a7

File tree

2 files changed

+13
-5
lines changed

2 files changed

+13
-5
lines changed

src/taskgraph/generator.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77
import logging
88
import multiprocessing
99
import os
10-
import platform
1110
from concurrent.futures import (
1211
FIRST_COMPLETED,
1312
ProcessPoolExecutor,
@@ -441,7 +440,9 @@ def _run(self):
441440
# redone in the new processes. Ideally this would be fixed, or we
442441
# would take another approach to parallel kind generation. In the
443442
# meantime, it's not supported outside of Linux.
444-
if platform.system() != "Linux" or os.environ.get("TASKGRAPH_SERIAL"):
443+
if "fork" not in multiprocessing.get_all_start_methods() or os.environ.get(
444+
"TASKGRAPH_SERIAL"
445+
):
445446
all_tasks = self._load_tasks_serial(kinds, kind_graph, parameters)
446447
else:
447448
all_tasks = self._load_tasks_parallel(kinds, kind_graph, parameters)

test/test_generator.py

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
44

55

6+
import multiprocessing
67
from concurrent.futures import ProcessPoolExecutor
78

89
import pytest
@@ -12,6 +13,11 @@
1213
from taskgraph.generator import Kind, load_tasks_for_kind, load_tasks_for_kinds
1314
from taskgraph.loader.default import loader as default_loader
1415

16+
forkonly = pytest.mark.skipif(
17+
"fork" not in multiprocessing.get_all_start_methods(),
18+
reason="requires 'fork' multiprocessing support",
19+
)
20+
1521

1622
class FakePPE(ProcessPoolExecutor):
1723
loaded_kinds = []
@@ -21,6 +27,7 @@ def submit(self, kind_load_tasks, *args):
2127
return super().submit(kind_load_tasks, *args)
2228

2329

30+
@forkonly
2431
def test_kind_ordering(mocker, maketgg):
2532
"When task kinds depend on each other, they are loaded in postorder"
2633
mocked_ppe = mocker.patch.object(generator, "ProcessPoolExecutor", new=FakePPE)
@@ -272,9 +279,9 @@ def _get_loader(self):
272279
)
273280
def test_default_loader(config, expected_transforms):
274281
loader = Kind("", "", config, {})._get_loader()
275-
assert loader is default_loader, (
276-
"Default Kind loader should be taskgraph.loader.default.loader"
277-
)
282+
assert (
283+
loader is default_loader
284+
), "Default Kind loader should be taskgraph.loader.default.loader"
278285
loader("", "", config, {}, [], False)
279286

280287
assert config["transforms"] == expected_transforms

0 commit comments

Comments
 (0)