diff --git a/dashboard/modules/reporter/profile_manager.py b/dashboard/modules/reporter/profile_manager.py index 9dcdb67d91cb0..2cc7377c16c00 100644 --- a/dashboard/modules/reporter/profile_manager.py +++ b/dashboard/modules/reporter/profile_manager.py @@ -1,4 +1,5 @@ import asyncio +import shutil import subprocess import sys from pathlib import Path @@ -49,11 +50,12 @@ def _format_failed_pyspy_command(cmd, stdout, stderr) -> str: # 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: - process = await asyncio.create_subprocess_shell( - "sudo -n true", + process = await asyncio.create_subprocess_exec( + "sudo", + "-n", + "true", stdout=subprocess.PIPE, stderr=subprocess.PIPE, - shell=True, ) _, _ = await process.communicate() return process.returncode == 0 @@ -65,17 +67,20 @@ def __init__(self, profile_dir_path: str): self.profile_dir_path.mkdir(exist_ok=True) async def trace_dump(self, pid: int, native: bool = False) -> (bool, str): - cmd = f"py-spy dump -p {pid}" + pyspy = shutil.which("py-spy") + if pyspy is None: + return False, "py-spy is not installed" + + cmd = [pyspy, "dump", "-p", str(pid)] # We if sys.platform == "linux" and native: - cmd += " --native" + cmd.append("--native") if await _can_passwordless_sudo(): - cmd = "sudo -n " + cmd.replace("py-spy", "$(which py-spy)") - process = await asyncio.create_subprocess_shell( - cmd, + cmd = ["sudo", "-n"] + cmd + process = await asyncio.create_subprocess_exec( + *cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, - shell=True, ) stdout, stderr = await process.communicate() if process.returncode != 0: @@ -86,6 +91,16 @@ async def trace_dump(self, pid: int, native: bool = False) -> (bool, str): async def cpu_profile( self, pid: int, format="flamegraph", duration: float = 5, native: bool = False ) -> (bool, str): + pyspy = shutil.which("py-spy") + if pyspy is None: + return False, "py-spy is not installed" + + if format not in ("flamegraph", "raw", "speedscope"): + return ( + False, + f"Invalid format {format}, " + "must be [flamegraph, raw, speedscope]", + ) + if format == "flamegraph": extension = "svg" else: @@ -93,19 +108,26 @@ async def cpu_profile( profile_file_path = ( self.profile_dir_path / f"{format}_{pid}_cpu_profiling.{extension}" ) - cmd = ( - f"py-spy record " - f"-o {profile_file_path} -p {pid} -d {duration} -f {format}" - ) + cmd = [ + pyspy, + "record", + "-o", + profile_file_path, + "-p", + str(pid), + "-d", + str(duration), + "-f", + format, + ] if sys.platform == "linux" and native: - cmd += " --native" + cmd.append("--native") if await _can_passwordless_sudo(): - cmd = "sudo -n " + cmd.replace("py-spy", "$(which py-spy)") - process = await asyncio.create_subprocess_shell( - cmd, + cmd = ["sudo", "-n"] + cmd + process = await asyncio.create_subprocess_exec( + *cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, - shell=True, ) stdout, stderr = await process.communicate() if process.returncode != 0: