您最多选择25个主题
主题必须以中文或者字母或数字开头，可以包含连字符 (-)，并且长度不得超过35个字符
41 行
1.4 KiB
41 行
1.4 KiB
from typing import Optional
|
|
|
|
import os
|
|
|
|
|
|
def get_num_threads_to_use() -> Optional[int]:
|
|
"""
|
|
Gets the number of threads to use. For most problems, 4 is all you
|
|
need, but for smaller machines, we'd like to scale to less than that.
|
|
By default, PyTorch uses 1/2 of the available cores.
|
|
"""
|
|
num_cpus = _get_num_available_cpus()
|
|
return max(min(num_cpus // 2, 4), 1) if num_cpus is not None else None
|
|
|
|
|
|
def _get_num_available_cpus() -> Optional[int]:
|
|
"""
|
|
Returns number of CPUs using cgroups if possible. This accounts
|
|
for Docker containers that are limited in cores.
|
|
"""
|
|
period = _read_in_integer_file("/sys/fs/cgroup/cpu/cpu.cfs_period_us")
|
|
quota = _read_in_integer_file("/sys/fs/cgroup/cpu/cpu.cfs_quota_us")
|
|
share = _read_in_integer_file("/sys/fs/cgroup/cpu/cpu.shares")
|
|
is_kubernetes = os.getenv("KUBERNETES_SERVICE_HOST") is not None
|
|
|
|
if period > 0 and quota > 0:
|
|
return int(quota // period)
|
|
elif period > 0 and share > 0 and is_kubernetes:
|
|
# In kubernetes, each requested CPU is 1024 CPU shares
|
|
# https://kubernetes.io/docs/concepts/configuration/manage-resources-containers/#how-pods-with-resource-limits-are-run
|
|
return int(share // 1024)
|
|
else:
|
|
return os.cpu_count()
|
|
|
|
|
|
def _read_in_integer_file(filename: str) -> int:
|
|
try:
|
|
with open(filename) as f:
|
|
return int(f.read().rstrip())
|
|
except FileNotFoundError:
|
|
return -1
|