avoid executing external tools by using Python's built-in os module to interact with the filesystem directly

This commit is contained in:
Sergey Krashevich 2023-04-28 01:33:37 +03:00
parent 83006eeb65
commit 2e549514a5
No known key found for this signature in database
GPG Key ID: 625171324E7D3856

View File

@ -9,6 +9,7 @@ import signal
import traceback
import urllib.parse
import yaml
import os
from abc import ABC, abstractmethod
from collections import Counter
@ -740,98 +741,77 @@ def escape_special_characters(path: str) -> str:
def get_cgroups_version() -> str:
"""Determine what version of cgroups is enabled"""
"""Determine what version of cgroups is enabled."""
stat_command = ["stat", "-fc", "%T", "/sys/fs/cgroup"]
cgroup_path = "/sys/fs/cgroup"
p = sp.run(
stat_command,
encoding="ascii",
capture_output=True,
)
try:
stat_info = os.stat(cgroup_path)
value = os.statvfs(cgroup_path).f_type
if p.returncode == 0:
value: str = p.stdout.strip().lower()
if value == "cgroup2fs":
if value == os.fsencode('cgroup2fs'):
return "cgroup2"
elif value == "tmpfs":
elif value == os.fsencode('tmpfs'):
return "cgroup"
else:
logger.debug(
f"Could not determine cgroups version: unhandled filesystem {value}"
)
else:
logger.debug(f"Could not determine cgroups version: {p.stderr}")
logger.debug(f"Could not determine cgroups version: unhandled filesystem {value}")
except Exception as e:
logger.debug(f"Could not determine cgroups version: {e}")
return "unknown"
def get_docker_memlimit_bytes() -> int:
"""Get mem limit in bytes set in docker if present. Returns -1 if no limit detected"""
"""Get mem limit in bytes set in docker if present. Returns -1 if no limit detected."""
# check running a supported cgroups version
if get_cgroups_version() == "cgroup2":
memlimit_command = ["cat", "/sys/fs/cgroup/memory.max"]
memlimit_path = "/sys/fs/cgroup/memory.max"
p = sp.run(
memlimit_command,
encoding="ascii",
capture_output=True,
)
if p.returncode == 0:
value: str = p.stdout.strip()
try:
with open(memlimit_path, 'r') as f:
value = f.read().strip()
if value.isnumeric():
return int(value)
elif value.lower() == "max":
return -1
else:
logger.debug(f"Unable to get docker memlimit: {p.stderr}")
except Exception as e:
logger.debug(f"Unable to get docker memlimit: {e}")
return -1
def get_cpu_stats() -> dict[str, dict]:
"""Get cpu usages for each process id"""
"""Get cpu usages for each process id."""
usages = {}
# -n=2 runs to ensure extraneous values are not included
top_command = ["top", "-b", "-n", "2"]
docker_memlimit = get_docker_memlimit_bytes() / 1024
p = sp.run(
top_command,
encoding="ascii",
capture_output=True,
)
if p.returncode != 0:
logger.error(p.stderr)
return usages
else:
lines = p.stdout.split("\n")
for line in lines:
stats = list(filter(lambda a: a != "", line.strip().split(" ")))
for pid in os.listdir('/proc'):
if pid.isdigit():
try:
if docker_memlimit > 0:
mem_res = int(stats[5])
mem_pct = str(
round((float(mem_res) / float(docker_memlimit)) * 100, 1)
)
else:
mem_pct = stats[9]
with open(f'/proc/{pid}/stat', 'r') as f:
stat_info = f.read().split()
usages[stats[0]] = {
"cpu": stats[8],
with open('/proc/meminfo', 'r') as f:
mem_info = f.readlines()
total_mem = int(mem_info[0].split()[1])
mem_res = int(stat_info[23]) / 1024
if docker_memlimit > 0:
mem_pct = round((float(mem_res) / float(docker_memlimit)) * 100, 1)
else:
mem_pct = round((float(mem_res) / float(total_mem)) * 100, 1)
cpu_total_time = sum(int(x) for x in stat_info[13:17])
cpu_usage = round(cpu_total_time / os.sysconf("SC_CLK_TCK"), 1)
usages[pid] = {
"cpu": cpu_usage,
"mem": mem_pct,
}
except:
continue
return usages
return usages
def get_amd_gpu_stats() -> dict[str, str]: