Skip to content

Commit 5b1bbc2

Browse files
authored
fix: set NUM_WORKERS=0 for non-fork multiprocessing start methods (#4784)
fix #4686 <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Bug Fixes** - Improved compatibility by automatically adjusting worker settings when the multiprocessing start method is not supported, preventing potential issues with data loading. - Added a warning message to inform users when worker settings are changed due to multiprocessing constraints. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
1 parent 2731c5d commit 5b1bbc2

File tree

1 file changed

+11
-0
lines changed

1 file changed

+11
-0
lines changed

deepmd/pt/utils/env.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
# SPDX-License-Identifier: LGPL-3.0-or-later
2+
import logging
3+
import multiprocessing
24
import os
35

46
import numpy as np
@@ -22,6 +24,15 @@
2224
except AttributeError:
2325
ncpus = os.cpu_count()
2426
NUM_WORKERS = int(os.environ.get("NUM_WORKERS", min(4, ncpus)))
27+
if multiprocessing.get_start_method() != "fork":
28+
# spawn or forkserver does not support NUM_WORKERS > 0 for DataLoader
29+
log = logging.getLogger(__name__)
30+
log.warning(
31+
"NUM_WORKERS > 0 is not supported with spawn or forkserver start method. "
32+
"Setting NUM_WORKERS to 0."
33+
)
34+
NUM_WORKERS = 0
35+
2536
# Make sure DDP uses correct device if applicable
2637
LOCAL_RANK = os.environ.get("LOCAL_RANK")
2738
LOCAL_RANK = int(0 if LOCAL_RANK is None else LOCAL_RANK)

0 commit comments

Comments
 (0)