Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,883 Bytes
d1ed09d |
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 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
from __future__ import annotations
import logging
import sys
import psutil
__all__ = ("memory_limit", "MEMORY_LIMIT")
logger = logging.getLogger(__name__)
def memory_limit() -> int:
"""Get the memory limit (in bytes) for this system.
Takes the minimum value from the following locations:
- Total system host memory
- Cgroups limit (if set)
- RSS rlimit (if set)
"""
limit = psutil.virtual_memory().total
# Check cgroups if available
# Note: can't use LINUX and WINDOWS constants as they upset mypy
if sys.platform == "linux":
path_used = None
for path in [
"/sys/fs/cgroup/memory/memory.limit_in_bytes", # cgroups v1 hard limit
"/sys/fs/cgroup/memory/memory.soft_limit_in_bytes", # cgroups v1 soft limit
"/sys/fs/cgroup/memory.max", # cgroups v2 hard limit
"/sys/fs/cgroup/memory.high", # cgroups v2 soft limit
]:
try:
with open(path) as f:
cgroups_limit = int(f.read())
if cgroups_limit > 0:
path_used = path
limit = min(limit, cgroups_limit)
except Exception:
pass
if path_used:
logger.debug(
"Setting system memory limit based on cgroup value defined in %s",
path_used,
)
# Check rlimit if available
if sys.platform != "win32":
try:
import resource
hard_limit = resource.getrlimit(resource.RLIMIT_RSS)[1]
if 0 < hard_limit < limit:
logger.debug(
"Limiting system memory based on RLIMIT_RSS to %s", hard_limit
)
limit = hard_limit
except (ImportError, OSError):
pass
return limit
MEMORY_LIMIT = memory_limit()
|