52 lines
1.5 KiB
Python
52 lines
1.5 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding:utf-8 -*-
|
|
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
|
|
|
|
import cv2
|
|
|
|
import os
|
|
import subprocess
|
|
|
|
__all__ = ["configure_nccl", "configure_module"]
|
|
|
|
|
|
def configure_nccl():
|
|
"""Configure multi-machine environment variables of NCCL."""
|
|
os.environ["NCCL_LAUNCH_MODE"] = "PARALLEL"
|
|
os.environ["NCCL_IB_HCA"] = subprocess.getoutput(
|
|
"pushd /sys/class/infiniband/ > /dev/null; for i in mlx5_*; "
|
|
"do cat $i/ports/1/gid_attrs/types/* 2>/dev/null "
|
|
"| grep v >/dev/null && echo $i ; done; popd > /dev/null"
|
|
)
|
|
os.environ["NCCL_IB_GID_INDEX"] = "3"
|
|
os.environ["NCCL_IB_TC"] = "106"
|
|
|
|
|
|
def configure_module(ulimit_value=8192):
|
|
"""
|
|
Configure pytorch module environment. setting of ulimit and cv2 will be set.
|
|
|
|
Args:
|
|
ulimit_value(int): default open file number on linux. Default value: 8192.
|
|
"""
|
|
# system setting
|
|
try:
|
|
import resource
|
|
|
|
rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
|
|
resource.setrlimit(resource.RLIMIT_NOFILE, (ulimit_value, rlimit[1]))
|
|
except Exception:
|
|
# Exception might be raised in Windows OS or rlimit reaches max limit number.
|
|
# However, set rlimit value might not be necessary.
|
|
pass
|
|
|
|
# cv2
|
|
# multiprocess might be harmful on performance of torch dataloader
|
|
os.environ["OPENCV_OPENCL_RUNTIME"] = "disabled"
|
|
try:
|
|
cv2.setNumThreads(0)
|
|
cv2.ocl.setUseOpenCL(False)
|
|
except Exception:
|
|
# cv2 version mismatch might rasie exceptions.
|
|
pass
|