Spaces:
Running
Running
File size: 1,464 Bytes
e11e4fe |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
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
|