chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,221 @@
|
||||
# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
|
||||
import paddle
|
||||
from paddle.distributed.utils.log_utils import get_logger
|
||||
|
||||
logger = get_logger("INFO", "root")
|
||||
|
||||
SUCCESS_CODE = 0
|
||||
FAIL_CODE = 1
|
||||
|
||||
|
||||
def _get_cpu_info(numa_id):
|
||||
"""
|
||||
get cpu info from lscpu
|
||||
"""
|
||||
|
||||
def _process_raw_cpu_info(i):
|
||||
processed_cpu_info = []
|
||||
cpu_ranges = i.split(',')
|
||||
for cpu_range in cpu_ranges:
|
||||
start, end = (
|
||||
int(cpu_range.split("-")[0]),
|
||||
int(cpu_range.split("-")[1]),
|
||||
)
|
||||
processed_cpu_info.extend(list(range(start, end + 1)))
|
||||
return processed_cpu_info
|
||||
|
||||
try:
|
||||
cpus = None
|
||||
cmd = ["lscpu"]
|
||||
output = subprocess.check_output(cmd).decode("utf-8").split(os.linesep)
|
||||
numa_key = f"node{numa_id}"
|
||||
for line in output:
|
||||
if line.find(numa_key) >= 0:
|
||||
raw_cpu_info = line.strip().split()[3]
|
||||
cpus = _process_raw_cpu_info(raw_cpu_info)
|
||||
break
|
||||
return cpus
|
||||
except Exception as e:
|
||||
logger.warning(f"_get_cpu_info failed, reason:{e}")
|
||||
return None
|
||||
|
||||
|
||||
def _has_nvidia_smi():
|
||||
"""
|
||||
check if nvidia-smi is available
|
||||
"""
|
||||
return shutil.which("nvidia-smi")
|
||||
|
||||
|
||||
def _has_xpu_smi():
|
||||
"""
|
||||
check if xpu-smi is available
|
||||
"""
|
||||
return shutil.which("xpu-smi")
|
||||
|
||||
|
||||
def _get_xpu_device_from_env(str_device_list, local_rank):
|
||||
if len(str_device_list.strip()) == 0:
|
||||
return None
|
||||
visible_devices = str_device_list.split(',')
|
||||
if len(visible_devices) <= local_rank:
|
||||
return None
|
||||
return visible_devices[local_rank]
|
||||
|
||||
|
||||
def _get_xpu_device(local_rank):
|
||||
"""
|
||||
get currently used xpu physical device id
|
||||
"""
|
||||
# NOTE(lijin23): priority XPULINK_VISIBLE_DEVICES > XPU_VISIBLE_DEVICES >
|
||||
# CUDA_VISIBLE_DEVICES
|
||||
xpulink_visible_devices = os.getenv("XPULINK_VISIBLE_DEVICES")
|
||||
if xpulink_visible_devices is not None:
|
||||
return _get_xpu_device_from_env(xpulink_visible_devices, local_rank)
|
||||
|
||||
xpu_visible_devices = os.getenv("XPU_VISIBLE_DEVICES")
|
||||
if xpu_visible_devices is not None:
|
||||
return _get_xpu_device_from_env(xpu_visible_devices, local_rank)
|
||||
|
||||
cuda_visible_devices = os.getenv("CUDA_VISIBLE_DEVICES")
|
||||
if cuda_visible_devices is not None:
|
||||
return _get_xpu_device_from_env(cuda_visible_devices, local_rank)
|
||||
|
||||
return str(local_rank)
|
||||
|
||||
|
||||
def _get_gpu_device(local_rank):
|
||||
"""
|
||||
get currently used gpu physical device id
|
||||
"""
|
||||
cuda_visible_devices = os.getenv("CUDA_VISIBLE_DEVICES")
|
||||
if cuda_visible_devices is None or cuda_visible_devices == "":
|
||||
return str(local_rank)
|
||||
cuda_visible_devices = cuda_visible_devices.split(',')
|
||||
if len(cuda_visible_devices) <= local_rank:
|
||||
return None
|
||||
return cuda_visible_devices[local_rank]
|
||||
|
||||
|
||||
def _get_gpu_numa_info(gpu_id):
|
||||
"""
|
||||
get gpu numa info from nvidia-smi
|
||||
"""
|
||||
try:
|
||||
cmd = ["nvidia-smi", "topo", "-C", "-i", gpu_id]
|
||||
output = subprocess.check_output(cmd, timeout=3).decode("utf-8")
|
||||
numa_id = output.strip().split()[-1]
|
||||
return numa_id
|
||||
except Exception as e:
|
||||
logger.warning(f"_get_cpu_info failed, reason:{e}")
|
||||
return None
|
||||
|
||||
|
||||
def _get_xpu_affinity_mask(xpu_id):
|
||||
xpu_id = int(xpu_id)
|
||||
cmd = ["xpu-smi", "topo", "-m"]
|
||||
if os.getenv("CUDA_DEVICE_ORDER") == "OAM_ID":
|
||||
# NOTE(lijin23): if CUDA_DEVICE_ORDER is set to OAM_ID,
|
||||
# we need to get the cpu affinity using OAM_ID
|
||||
cmd = ["xpu-smi", "topo", "-mo"]
|
||||
output = subprocess.check_output(cmd, timeout=60).decode("utf-8")
|
||||
cpu_affinity = output.splitlines()[xpu_id + 1].split()[-2]
|
||||
affinity_mask = []
|
||||
for affinity_range in cpu_affinity.split(','):
|
||||
start, end = affinity_range.split('-')
|
||||
affinity_mask.extend(range(int(start), int(end) + 1))
|
||||
return affinity_mask
|
||||
|
||||
|
||||
def set_affinity_gpu():
|
||||
"""
|
||||
set affinity for gpu
|
||||
"""
|
||||
if not _has_nvidia_smi():
|
||||
logger.warning(
|
||||
"nvidia-smi is not available, set_affinity is aborted, plz check your environment."
|
||||
)
|
||||
return FAIL_CODE
|
||||
local_rank = max(int(os.getenv("PADDLE_LOCAL_RANK", "0")), 0)
|
||||
device_id = _get_gpu_device(local_rank)
|
||||
if device_id is None:
|
||||
logger.warning(
|
||||
"Failed to get device id from cuda_visible_devices, set_affinity is aborted, plz check your environment."
|
||||
)
|
||||
return FAIL_CODE
|
||||
numa_id = _get_gpu_numa_info(device_id)
|
||||
if numa_id is None:
|
||||
logger.warning(
|
||||
"Failed to get numa info, set_affinity is aborted, plz check your environment."
|
||||
)
|
||||
return FAIL_CODE
|
||||
if numa_id == "N/A":
|
||||
logger.warning(
|
||||
"nvidia-smi topo return numa id as N/A, set_affinity is aborted, plz check your environment. (Notice: This is expected behavior when executed on single numa node environment)"
|
||||
)
|
||||
return FAIL_CODE
|
||||
affinity_mask = _get_cpu_info(numa_id)
|
||||
if affinity_mask is None:
|
||||
logger.warning(
|
||||
"Failed to get cpu info, set_affinity is aborted, plz check your environment."
|
||||
)
|
||||
return FAIL_CODE
|
||||
affinity = os.sched_getaffinity(0)
|
||||
logger.info(f"Check affinity before setting: {affinity}")
|
||||
os.sched_setaffinity(0, affinity_mask)
|
||||
affinity = os.sched_getaffinity(0)
|
||||
logger.info(f"check affinity after setting: {affinity}")
|
||||
return SUCCESS_CODE
|
||||
|
||||
|
||||
def set_affinity_xpu():
|
||||
"""
|
||||
set affinity for xpu
|
||||
"""
|
||||
if not _has_xpu_smi():
|
||||
logger.warning(
|
||||
"xpu-smi is not available, set_affinity is aborted, plz check your environment."
|
||||
)
|
||||
return FAIL_CODE
|
||||
local_rank = max(int(os.getenv("PADDLE_LOCAL_RANK", "0")), 0)
|
||||
device_id = _get_xpu_device(local_rank)
|
||||
if device_id is None:
|
||||
logger.warning(
|
||||
"Failed to get device id, set_affinity is aborted, plz check your environment."
|
||||
)
|
||||
return FAIL_CODE
|
||||
affinity_mask = _get_xpu_affinity_mask(device_id)
|
||||
affinity = os.sched_getaffinity(0)
|
||||
logger.info(f"Check affinity before setting: {affinity}")
|
||||
os.sched_setaffinity(0, affinity_mask)
|
||||
affinity = os.sched_getaffinity(0)
|
||||
logger.info(f"Check affinity after setting: {affinity}")
|
||||
return SUCCESS_CODE
|
||||
|
||||
|
||||
def set_affinity():
|
||||
if paddle.device.is_compiled_with_cuda():
|
||||
return set_affinity_gpu()
|
||||
elif paddle.device.is_compiled_with_xpu():
|
||||
return set_affinity_xpu()
|
||||
else:
|
||||
# TODO(@gexiao): supports other devices if needed
|
||||
logger.warning("Currently set_affinity only supports gpu env.")
|
||||
return FAIL_CODE
|
||||
Reference in New Issue
Block a user