Allow pyspy to attach in docker without SYS_PTRACE

Author: gjoseph92

distributed_pyspy/distributed_pyspy.py

@@ -1,14 +1,16 @@
 import asyncio
-from contextlib import contextmanager
-import tempfile
+import logging
 import os
 import signal
-from typing import List, Optional, Iterable
-import logging
+import tempfile
+from contextlib import contextmanager
+from typing import Iterable, List, Optional, Union
 
 import distributed
 from distributed.diagnostics import SchedulerPlugin
 
+from .prctl import allow_ptrace
+
 logger = logging.getLogger(__name__)
 
 
@@ -46,6 +48,7 @@ def __init__(
         self.pyspy_args.extend(extra_pyspy_args)
         self.proc = None
         self._tempfile = None
+        self._run_failed_msg = None
 
     def __repr__(self) -> str:
         return f"<{type(self).__name__} {self.pyspy_args}>"
@@ -69,6 +72,17 @@ async def start(self, scheduler):
             scheduler.handlers[self._HANDLER_NAME] = self._get_py_spy_profile
 
         pid = os.getpid()
+
+        try:
+            # Allow subprocesses of this process to ptrace it.
+            # Since we'll start py-spy as a subprocess, it will be below the current PID
+            # in the process tree, and therefore allowed to trace its parent.
+            allow_ptrace(pid)
+        except OSError as e:
+            self._run_failed_msg = str(e)
+        else:
+            self._run_failed_msg = None
+
         self.proc = await asyncio.create_subprocess_exec(
             "py-spy",
             "record",
@@ -88,9 +102,10 @@ async def _stop(self) -> Optional[int]:
         try:
             self.proc.send_signal(signal.SIGINT)
         except ProcessLookupError:
-            logger.warning(
-                f"py-spy subprocess {self.proc.pid} already terminated (it probably never ran?)."
-            )
+            msg = f"py-spy subprocess {self.proc.pid} already terminated (it probably never ran?)."
+            if self._run_failed_msg:
+                msg += "\nNOTE: " + self._run_failed_msg
+            logger.warning(msg)
 
         stdout, stderr = await self.proc.communicate()  # TODO timeout
         retcode = self.proc.returncode
@@ -147,7 +162,7 @@ def start_pyspy_on_scheduler(
     """
     client = client or distributed.worker.get_client()
 
-    async def _inject(dask_scheduler: distributed.Scheduler):
+    async def _inject_pyspy(dask_scheduler: distributed.Scheduler):
         plugin = PySpyScheduler(
             output=output,
             format=format,
@@ -164,11 +179,11 @@ async def _inject(dask_scheduler: distributed.Scheduler):
         await plugin.start(dask_scheduler)
         dask_scheduler.add_plugin(plugin)
 
-    client.run_on_scheduler(_inject)
+    client.run_on_scheduler(_inject_pyspy)
 
 
 def get_profile_from_scheduler(
-    path: str, client: Optional[distributed.Client] = None
+    path: Union[str, os.PathLike], client: Optional[distributed.Client] = None
 ) -> None:
     """
     Stop the current `PySpyScheduler` plugin, send back its profile data, and write it to ``path``.
@@ -188,7 +203,7 @@ async def _get_profile():
 
 @contextmanager
 def pyspy_on_scheduler(
-    output: str,
+    output: Union[str, os.PathLike],
     format: str = "speedscope",
     rate: int = 100,
     subprocesses: bool = True,

distributed_pyspy/prctl.py

@@ -0,0 +1,48 @@
+"""
+Enable other processes to ptrace this one via the ``prctl`` system call.
+
+Only needed on Ubuntu. See
+
+* https://www.kernel.org/doc/Documentation/admin-guide/LSM/Yama.rst
+* https://wiki.ubuntu.com/SecurityTeam/Roadmap/KernelHardening#ptrace_Protection
+* https://man7.org/linux/man-pages/man2/prctl.2.html
+
+Note that we could have used https://github.com/seveas/python-prctl, but since that package doesn't
+have prebuilt wheels, it would be more of a pain to install than making the one syscall ourselves via libc.
+"""
+import ctypes
+import ctypes.util
+from typing import Optional
+
+# https://github.com/torvalds/linux/blob/master/include/uapi/linux/prctl.h#L155-L156
+PR_SET_PTRACER = 0x59616D61
+PR_SET_PTRACER_ANY = -1
+
+
+def allow_ptrace(pid: Optional[int] = None):
+    """
+    Allow the PID to ptrace this process.
+
+    If ``pid`` is None, any process is allowed.
+    """
+    libc = ctypes.CDLL(ctypes.util.find_library("c"))
+
+    try:
+        prctl = libc.prctl
+    except AttributeError:
+        raise OSError(
+            "Your OS does not support the `prctl` syscall (only available on Linux), "
+            "so we can't automatically allow py-spy to run without elevated permissions.\n"
+            "You'll need to run the scheduler with sudo access (or if in Docker, just `--cap-add SYS_PTRACE`).\n"
+            "https://github.com/benfred/py-spy#when-do-you-need-to-run-as-sudo"
+        )
+
+    prctl.argtypes = [
+        ctypes.c_int,
+        ctypes.c_ulong,
+        ctypes.c_ulong,
+        ctypes.c_ulong,
+        ctypes.c_ulong,
+    ]
+
+    prctl(PR_SET_PTRACER, pid if pid is not None else PR_SET_PTRACER_ANY, 0, 0, 0)