diff --git a/python/dolma/cli/__main__.py b/python/dolma/cli/__main__.py index 1300e13d..1e6a62dd 100644 --- a/python/dolma/cli/__main__.py +++ b/python/dolma/cli/__main__.py @@ -1,3 +1,4 @@ +import multiprocessing from argparse import ArgumentParser from pathlib import Path from typing import List, Optional @@ -27,6 +28,14 @@ def main(argv: Optional[List[str]] = None): + try: + # attempting to set start method to spawn in case it is not set + multiprocessing.set_start_method("spawn") + except RuntimeError: + # method already set, check if it is set to spawn + if multiprocessing.get_start_method() != "spawn": + raise RuntimeError("Multiprocessing start method must be set to spawn") + parser = ArgumentParser( prog="dolma", usage="dolma [command] [options]", diff --git a/scripts/make_wikipedia.py b/scripts/make_wikipedia.py index f4422e2d..80b614a9 100644 --- a/scripts/make_wikipedia.py +++ b/scripts/make_wikipedia.py @@ -286,4 +286,7 @@ def main(): if __name__ == "__main__": + # setting multiprocessing start method to spawn + multiprocessing.set_start_method("spawn") + main()