wehub-resource-sync
379 lines
13 KiB
Python
379 lines
13 KiB
Python
import asyncio
|
|
import logging
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Tuple, Union
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
DARWIN_SET_CHOWN_CMD = "sudo chown root: `which {profiler}`"
|
|
LINUX_SET_CHOWN_CMD = "sudo chown root:root `which {profiler}`"
|
|
|
|
PROFILER_PERMISSIONS_ERROR_MESSAGE = """
|
|
Note that this command requires `{profiler}` to be installed with root permissions. You
|
|
can install `{profiler}` and give it root permissions as follows:
|
|
$ pip install {profiler}
|
|
$ {set_chown_command}
|
|
$ sudo chmod u+s `which {profiler}`
|
|
|
|
Alternatively, you can start Ray with passwordless sudo / root permissions.
|
|
|
|
"""
|
|
|
|
|
|
def decode(string: Union[str, bytes]):
|
|
if isinstance(string, bytes):
|
|
return string.decode("utf-8")
|
|
return string
|
|
|
|
|
|
def _format_failed_profiler_command(cmd, profiler, stdout, stderr) -> str:
|
|
stderr_str = decode(stderr)
|
|
extra_message = ""
|
|
|
|
# If some sort of permission error returned, show a message about how
|
|
# to set up permissions correctly.
|
|
if "permission" in stderr_str.lower():
|
|
set_chown_command = (
|
|
DARWIN_SET_CHOWN_CMD.format(profiler=profiler)
|
|
if sys.platform == "darwin"
|
|
else LINUX_SET_CHOWN_CMD.format(profiler=profiler)
|
|
)
|
|
extra_message = PROFILER_PERMISSIONS_ERROR_MESSAGE.format(
|
|
profiler=profiler, set_chown_command=set_chown_command
|
|
)
|
|
|
|
return f"""Failed to execute `{cmd}`.
|
|
{extra_message}
|
|
=== stderr ===
|
|
{decode(stderr)}
|
|
|
|
=== stdout ===
|
|
{decode(stdout)}
|
|
"""
|
|
|
|
|
|
# If we can sudo, always try that. Otherwise, py-spy will only work if the user has
|
|
# root privileges or has configured setuid on the py-spy script.
|
|
async def _can_passwordless_sudo() -> bool:
|
|
try:
|
|
process = await asyncio.create_subprocess_exec(
|
|
"sudo",
|
|
"-n",
|
|
"true",
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
)
|
|
except FileNotFoundError:
|
|
return False
|
|
else:
|
|
_, _ = await process.communicate()
|
|
return process.returncode == 0
|
|
|
|
|
|
class CpuProfilingManager:
|
|
def __init__(self, profile_dir_path: str):
|
|
self.profile_dir_path = Path(profile_dir_path)
|
|
self.profile_dir_path.mkdir(exist_ok=True)
|
|
self.profiler_name = "py-spy"
|
|
|
|
async def trace_dump(
|
|
self, pid: int, native: bool = False, subprocesses: bool = False
|
|
) -> Tuple[bool, str]:
|
|
"""Capture and dump a trace for a specified process.
|
|
|
|
Args:
|
|
pid: The process ID (PID) of the target process for trace capture.
|
|
native: If True, includes native (C/C++) stack frames.
|
|
Default is False.
|
|
subprocesses: If True, also dumps stack traces for
|
|
child processes of the target process. Default is False.
|
|
|
|
Returns:
|
|
A tuple containing a boolean indicating the success of the trace
|
|
capture operation and a string with the trace data or an error message.
|
|
"""
|
|
pyspy = shutil.which(self.profiler_name)
|
|
if pyspy is None:
|
|
return False, "Failed to execute: py-spy is not installed"
|
|
|
|
cmd = [pyspy, "dump", "-p", str(pid)]
|
|
# We
|
|
if sys.platform == "linux" and native:
|
|
cmd.append("--native")
|
|
if subprocesses:
|
|
cmd.append("--subprocesses")
|
|
if await _can_passwordless_sudo():
|
|
cmd = ["sudo", "-n"] + cmd
|
|
process = await asyncio.create_subprocess_exec(
|
|
*cmd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
)
|
|
stdout, stderr = await process.communicate()
|
|
if process.returncode != 0:
|
|
return False, _format_failed_profiler_command(
|
|
cmd, self.profiler_name, stdout, stderr
|
|
)
|
|
else:
|
|
return True, decode(stdout)
|
|
|
|
async def cpu_profile(
|
|
self,
|
|
pid: int,
|
|
format: str = "flamegraph",
|
|
duration: float = 5,
|
|
native: bool = False,
|
|
idle: bool = False,
|
|
subprocesses: bool = False,
|
|
) -> Tuple[bool, str]:
|
|
"""Perform CPU profiling on a specified process.
|
|
|
|
Args:
|
|
pid: The process ID (PID) of the target process to be profiled.
|
|
format: The format of the CPU profile output. Default is "flamegraph".
|
|
duration: The duration of the profiling session in seconds.
|
|
Default is 5 seconds.
|
|
native: If True, includes native (C/C++) stack frames. Default is False.
|
|
idle: If True, includes off-CPU / sleeping threads (e.g. threads
|
|
blocked on locks, I/O, or CUDA syncs). Default is False.
|
|
subprocesses: If True, also profiles child
|
|
processes of the target process (e.g. data loader or
|
|
multiprocess inference workers). Default is False.
|
|
|
|
Returns:
|
|
A tuple containing a boolean indicating the success of the profiling
|
|
operation and a string with the profile data or an error message.
|
|
"""
|
|
pyspy = shutil.which(self.profiler_name)
|
|
if pyspy is None:
|
|
return False, "Failed to execute: py-spy is not installed"
|
|
|
|
if format not in ("flamegraph", "raw", "speedscope"):
|
|
return (
|
|
False,
|
|
f"Failed to execute: Invalid format {format}, "
|
|
+ "must be [flamegraph, raw, speedscope]",
|
|
)
|
|
|
|
if format == "flamegraph":
|
|
extension = "svg"
|
|
else:
|
|
extension = "txt"
|
|
profile_file_path = (
|
|
self.profile_dir_path / f"{format}_{pid}_cpu_profiling.{extension}"
|
|
)
|
|
cmd = [
|
|
pyspy,
|
|
"record",
|
|
"-o",
|
|
profile_file_path,
|
|
"-p",
|
|
str(pid),
|
|
"-d",
|
|
str(duration),
|
|
"-f",
|
|
format,
|
|
]
|
|
if sys.platform == "linux" and native:
|
|
cmd.append("--native")
|
|
if idle:
|
|
cmd.append("--idle")
|
|
if subprocesses:
|
|
cmd.append("--subprocesses")
|
|
if await _can_passwordless_sudo():
|
|
cmd = ["sudo", "-n"] + cmd
|
|
process = await asyncio.create_subprocess_exec(
|
|
*cmd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
)
|
|
stdout, stderr = await process.communicate()
|
|
if process.returncode != 0:
|
|
return False, _format_failed_profiler_command(
|
|
cmd, self.profiler_name, stdout, stderr
|
|
)
|
|
else:
|
|
return True, open(profile_file_path, "rb").read()
|
|
|
|
|
|
class MemoryProfilingManager:
|
|
def __init__(self, profile_dir_path: str):
|
|
self.profile_dir_path = Path(profile_dir_path) / "memray"
|
|
self.profile_dir_path.mkdir(exist_ok=True)
|
|
self.profiler_name = "memray"
|
|
|
|
async def get_profile_result(
|
|
self,
|
|
pid: int,
|
|
profiler_filename: str,
|
|
format: str = "flamegraph",
|
|
leaks: bool = False,
|
|
) -> Tuple[bool, str]:
|
|
"""Convert the Memray profile result to specified format.
|
|
|
|
Args:
|
|
pid: The process ID (PID) associated with the profiling operation.
|
|
profiler_filename: The filename of the profiler output to be processed.
|
|
format: The format of the profile result. Default is "flamegraph".
|
|
leaks: If True, include memory leak information in the profile result.
|
|
|
|
Returns:
|
|
A tuple containing a boolean indicating the success of the operation
|
|
and a string with the processed profile result or an error message.
|
|
"""
|
|
memray = shutil.which(self.profiler_name)
|
|
if memray is None:
|
|
return False, "Failed to execute: memray is not installed"
|
|
|
|
profile_file_path = self.profile_dir_path / profiler_filename
|
|
if not Path(profile_file_path).is_file():
|
|
return False, f"Failed to execute: process {pid} has not been profiled"
|
|
|
|
profiler_name, _ = os.path.splitext(profiler_filename)
|
|
profile_visualize_path = self.profile_dir_path / f"{profiler_name}.html"
|
|
if format == "flamegraph":
|
|
visualize_cmd = [
|
|
memray,
|
|
"flamegraph",
|
|
"-o",
|
|
profile_visualize_path,
|
|
"-f",
|
|
]
|
|
elif format == "table":
|
|
visualize_cmd = [
|
|
memray,
|
|
"table",
|
|
"-o",
|
|
profile_visualize_path,
|
|
"-f",
|
|
]
|
|
else:
|
|
return (
|
|
False,
|
|
f"Failed to execute: Report with format: {format} is not supported",
|
|
)
|
|
|
|
if leaks:
|
|
visualize_cmd.append("--leaks")
|
|
visualize_cmd.append(profile_file_path)
|
|
|
|
process = await asyncio.create_subprocess_exec(
|
|
*visualize_cmd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
)
|
|
stdout, stderr = await process.communicate()
|
|
if process.returncode != 0:
|
|
return False, _format_failed_profiler_command(
|
|
visualize_cmd, self.profiler_name, stdout, stderr
|
|
)
|
|
|
|
return True, open(profile_visualize_path, "rb").read()
|
|
|
|
async def attach_profiler(
|
|
self,
|
|
pid: int,
|
|
native: bool = False,
|
|
trace_python_allocators: bool = False,
|
|
verbose: bool = False,
|
|
) -> Tuple[bool, str]:
|
|
"""Attach a Memray profiler to a specified process.
|
|
|
|
Args:
|
|
pid: The process ID (PID) of the target process which
|
|
the profiler attached to.
|
|
native: If True, includes native (C/C++) stack frames. Default is False.
|
|
trace_python_allocators: If True, includes Python stack frames.
|
|
Default is False.
|
|
verbose: If True, enables verbose output. Default is False.
|
|
|
|
Returns:
|
|
A tuple containing a boolean indicating the success of the operation
|
|
and a string of a success message or an error message.
|
|
"""
|
|
memray = shutil.which(self.profiler_name)
|
|
if memray is None:
|
|
return False, None, "Failed to execute: memray is not installed"
|
|
|
|
timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
|
|
profiler_filename = f"{pid}_memory_profiling_{timestamp}.bin"
|
|
profile_file_path = self.profile_dir_path / profiler_filename
|
|
cmd = [memray, "attach", str(pid), "-o", profile_file_path]
|
|
|
|
if native:
|
|
cmd.append("--native")
|
|
if trace_python_allocators:
|
|
cmd.append("--trace-python-allocators")
|
|
if verbose:
|
|
cmd.append("--verbose")
|
|
if await _can_passwordless_sudo():
|
|
cmd = ["sudo", "-n"] + cmd
|
|
|
|
process = await asyncio.create_subprocess_exec(
|
|
*cmd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
)
|
|
|
|
stdout, stderr = await process.communicate()
|
|
if process.returncode != 0:
|
|
return (
|
|
False,
|
|
None,
|
|
_format_failed_profiler_command(
|
|
cmd, self.profiler_name, stdout, stderr
|
|
),
|
|
)
|
|
else:
|
|
return (
|
|
True,
|
|
profiler_filename,
|
|
f"Success attaching memray to process {pid}",
|
|
)
|
|
|
|
async def detach_profiler(
|
|
self,
|
|
pid: int,
|
|
verbose: bool = False,
|
|
) -> Tuple[bool, str]:
|
|
"""Detach a profiler from a specified process.
|
|
|
|
Args:
|
|
pid: The process ID (PID) of the target process the
|
|
profiler detached from.
|
|
verbose: If True, enables verbose output. Default is False.
|
|
|
|
Returns:
|
|
A tuple containing a boolean indicating the success of the operation
|
|
and a string of a success message or an error message.
|
|
"""
|
|
memray = shutil.which(self.profiler_name)
|
|
if memray is None:
|
|
return False, "Failed to execute: memray is not installed"
|
|
|
|
cmd = [memray, "detach"]
|
|
if verbose:
|
|
cmd.append("--verbose")
|
|
cmd.append(str(pid))
|
|
|
|
if await _can_passwordless_sudo():
|
|
cmd = ["sudo", "-n"] + cmd
|
|
|
|
process = await asyncio.create_subprocess_exec(
|
|
*cmd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
)
|
|
stdout, stderr = await process.communicate()
|
|
if process.returncode != 0:
|
|
return False, _format_failed_profiler_command(
|
|
cmd, self.profiler_name, stdout, stderr
|
|
)
|
|
else:
|
|
return True, f"Success detaching memray from process {pid}"
|