chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,378 @@
|
||||
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}"
|
||||
Reference in New Issue
Block a user