Tools for profiling the number of processes/threads created during our tests.

monitor_test_threads.py

#!/usr/bin/env python
# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
# SPDX-License-Identifier: Apache-2.0
r"""
Run a command (typically pytest) while sampling thread counts from ``/proc``.

No root or ptrace is required: the monitor walks the process tree rooted at
the launched command and records ``Threads`` (nlwp) for each live process.

Example::

    python python/cudf_polars/devtools/monitor_test_threads.py \\
        -o /tmp/threads.csv --interval 1.0 -- \\
        pytest python/cudf_polars/tests/streaming/ -n 4 --dist=loadgroup

CSV columns
-----------
timestamp_iso
    Wall-clock sample time (UTC, ISO-8601).
elapsed_s
    Seconds since the monitored command started.
sample
    Monotonic sample index (0-based).
pid, ppid
    Process / parent process IDs.
nlwp
    Number of threads in the process (``/proc/<pid>/status`` ``Threads``).
comm
    Short process name (``/proc/<pid>/comm``).
cmdline
    NUL-joined argv, truncated (``/proc/<pid>/cmdline``).
"""

from __future__ import annotations

import argparse
import csv
import os
import signal
import subprocess
import sys
import time
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path


@dataclass(frozen=True)
class ProcInfo:
    pid: int
    ppid: int
    nlwp: int
    comm: str
    cmdline: str


def _read_text(path: Path) -> str | None:
    try:
        return path.read_text(encoding="utf-8", errors="replace")
    except (FileNotFoundError, ProcessLookupError, PermissionError, OSError):
        return None


def _parse_status(pid: int) -> tuple[int, int] | None:
    """Return ``(ppid, nlwp)`` from ``/proc/<pid>/status``, or ``None``."""
    text = _read_text(Path(f"/proc/{pid}/status"))
    if text is None:
        return None
    ppid: int | None = None
    nlwp: int | None = None
    for line in text.splitlines():
        if line.startswith("PPid:"):
            ppid = int(line.split()[1])
        elif line.startswith("Threads:"):
            nlwp = int(line.split()[1])
        if ppid is not None and nlwp is not None:
            return ppid, nlwp
    return None


def _read_comm(pid: int) -> str:
    text = _read_text(Path(f"/proc/{pid}/comm"))
    return text.strip() if text is not None else ""


def _read_cmdline(pid: int, *, max_len: int = 200) -> str:
    raw = _read_text(Path(f"/proc/{pid}/cmdline"))
    if raw is None:
        return ""
    cmd = raw.replace("\x00", " ").strip()
    if len(cmd) > max_len:
        return cmd[: max_len - 3] + "..."
    return cmd


def snapshot_process_tree(root_pid: int) -> list[ProcInfo]:
    """Return live descendants of ``root_pid`` (including ``root_pid``)."""
    # Build parent -> children from a single /proc scan, then walk from root.
    children: dict[int, list[int]] = {}
    info: dict[int, tuple[int, int]] = {}  # pid -> (ppid, nlwp)

    try:
        entries = list(Path("/proc").iterdir())
    except OSError:
        return []

    for entry in entries:
        name = entry.name
        if not name.isdigit():
            continue
        pid = int(name)
        parsed = _parse_status(pid)
        if parsed is None:
            continue
        ppid, nlwp = parsed
        info[pid] = (ppid, nlwp)
        children.setdefault(ppid, []).append(pid)

    if root_pid not in info and root_pid > 0:
        # Root may have exited between spawn and first sample.
        return []

    out: list[ProcInfo] = []
    stack = [root_pid]
    seen: set[int] = set()
    while stack:
        pid = stack.pop()
        if pid in seen:
            continue
        seen.add(pid)
        meta = info.get(pid)
        if meta is None:
            # Re-read in case it appeared after the scan, or skip if gone.
            parsed = _parse_status(pid)
            if parsed is None:
                continue
            ppid, nlwp = parsed
        else:
            ppid, nlwp = meta
        out.append(
            ProcInfo(
                pid=pid,
                ppid=ppid,
                nlwp=nlwp,
                comm=_read_comm(pid),
                cmdline=_read_cmdline(pid),
            )
        )
        stack.extend(children.get(pid, ()))
    out.sort(key=lambda p: (-p.nlwp, p.pid))
    return out


def _utc_iso() -> str:
    return datetime.now(UTC).isoformat(timespec="milliseconds")


def run_monitored(
    command: list[str],
    *,
    output: Path,
    interval_s: float,
    cwd: Path | None,
) -> int:
    output.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = [
        "timestamp_iso",
        "elapsed_s",
        "sample",
        "pid",
        "ppid",
        "nlwp",
        "comm",
        "cmdline",
    ]

    # Start in its own process group so we can signal the whole tree on Ctrl-C.
    proc = subprocess.Popen(
        command,
        cwd=cwd,
        start_new_session=True,
    )
    root_pid = proc.pid
    t0 = time.monotonic()
    sample_i = 0

    with output.open("w", newline="", encoding="utf-8") as fh:
        writer = csv.DictWriter(fh, fieldnames=fieldnames)
        writer.writeheader()
        fh.flush()

        try:
            while True:
                rc = proc.poll()
                procs = snapshot_process_tree(root_pid)
                elapsed = time.monotonic() - t0
                ts = _utc_iso()
                for p in procs:
                    writer.writerow(
                        {
                            "timestamp_iso": ts,
                            "elapsed_s": f"{elapsed:.3f}",
                            "sample": sample_i,
                            "pid": p.pid,
                            "ppid": p.ppid,
                            "nlwp": p.nlwp,
                            "comm": p.comm,
                            "cmdline": p.cmdline,
                        }
                    )
                # Always write a summary row for the tree even if empty,
                # so gaps in sampling are visible as sample indices.
                if not procs and rc is None:
                    writer.writerow(
                        {
                            "timestamp_iso": ts,
                            "elapsed_s": f"{elapsed:.3f}",
                            "sample": sample_i,
                            "pid": root_pid,
                            "ppid": "",
                            "nlwp": "",
                            "comm": "",
                            "cmdline": "<process tree empty / not yet visible>",
                        }
                    )
                fh.flush()
                sample_i += 1

                if rc is not None:
                    return rc

                time.sleep(interval_s)
        except KeyboardInterrupt:
            try:
                os.killpg(root_pid, signal.SIGINT)
            except ProcessLookupError:
                pass
            try:
                return proc.wait(timeout=10)
            except subprocess.TimeoutExpired:
                os.killpg(root_pid, signal.SIGKILL)
                return proc.wait()


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Run a command while sampling per-process thread counts from /proc "
            "(no root required)."
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument(
        "-o",
        "--output",
        type=Path,
        default=Path("thread_samples.csv"),
        help="CSV output path (default: thread_samples.csv)",
    )
    parser.add_argument(
        "-i",
        "--interval",
        type=float,
        default=1.0,
        help="Sampling interval in seconds (default: 1.0)",
    )
    parser.add_argument(
        "--cwd",
        type=Path,
        default=None,
        help="Working directory for the monitored command",
    )
    parser.add_argument(
        "command",
        nargs=argparse.REMAINDER,
        help="Command to run after -- (e.g. -- pytest -n 4 ...)",
    )
    args = parser.parse_args(argv)

    command = list(args.command)
    if command and command[0] == "--":
        command = command[1:]
    if not command:
        parser.error("provide a command after -- , e.g. -- pytest -n 4 tests/")

    if args.interval <= 0:
        parser.error("--interval must be positive")

    print(
        f"monitoring threads every {args.interval}s -> {args.output}\n"
        f"command: {' '.join(command)}",
        file=sys.stderr,
    )
    rc = run_monitored(
        command,
        output=args.output,
        interval_s=args.interval,
        cwd=args.cwd,
    )
    print(f"command exited {rc}; wrote {args.output}", file=sys.stderr)
    return rc


if __name__ == "__main__":
    raise SystemExit(main())

profile_ray_engine_threads.py

#!/usr/bin/env python
# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
# SPDX-License-Identifier: Apache-2.0
r"""
Capture process/thread snapshots around ``RayEngine`` construction and a tiny query.

Phases (written under ``--outdir``):

1. ``01_before_imports.txt`` — before importing Ray / cudf-polars / Polars
2. ``02_after_ray_engine.txt`` — after constructing :class:`~cudf_polars.engine.ray.RayEngine`
3. ``03_after_collect.txt`` — after ``pl.LazyFrame(...).collect(engine=...)``

Each snapshot walks the process tree rooted at this script's PID via ``/proc``
(no root required) and records per-process ``nlwp`` plus thread ``comm`` counts.

Example::

    CUDA_VISIBLE_DEVICES=1 \\
    RAPIDSMPF_NUM_STREAMING_THREADS=1 \\
    CUDF_POLARS__EXECUTOR__NUM_PY_EXECUTORS=1 \\
    LIBCUDF_NUM_HOST_WORKERS=2 KVIKIO_NTHREADS=2 \\
    POLARS_MAX_THREADS=1 OMP_NUM_THREADS=1 \\
    python python/cudf_polars/devtools/profile_ray_engine_threads.py \\
        --outdir /tmp/ray-thread-profile
"""

from __future__ import annotations

import argparse
import os
import subprocess
import sys
import time
from collections import Counter
from datetime import UTC, datetime
from pathlib import Path


def _read_text(path: Path) -> str | None:
    try:
        return path.read_text(encoding="utf-8", errors="replace")
    except OSError:
        return None


def _thread_names(pid: int) -> Counter[str]:
    counts: Counter[str] = Counter()
    task_dir = Path(f"/proc/{pid}/task")
    try:
        tids = list(task_dir.iterdir())
    except OSError:
        return counts
    for tid_path in tids:
        text = _read_text(tid_path / "comm")
        if text is not None:
            counts[text.strip()] += 1
    return counts


def _cmdline(pid: int, *, max_len: int = 160) -> str:
    raw = Path(f"/proc/{pid}/cmdline").read_bytes()
    cmd = raw.replace(b"\x00", b" ").decode(errors="replace").strip()
    if len(cmd) > max_len:
        return cmd[: max_len - 3] + "..."
    return cmd


def _status_fields(pid: int) -> dict[str, str]:
    out: dict[str, str] = {}
    text = _read_text(Path(f"/proc/{pid}/status"))
    if text is None:
        return out
    for line in text.splitlines():
        for key in (
            "Name:",
            "Pid:",
            "PPid:",
            "Threads:",
            "Cpus_allowed_list:",
            "VmRSS:",
        ):
            if line.startswith(key):
                out[key[:-1]] = line.split(":", 1)[1].strip()
    return out


def _selected_env(pid: int) -> list[str]:
    prefixes = (
        "CUDA_VISIBLE",
        "RAY_",
        "POLARS_",
        "RAYON_",
        "OMP_",
        "OPENBLAS_",
        "MKL_",
        "NUMEXPR_",
        "LIBCUDF_",
        "KVIKIO_",
        "RAPIDSMPF_",
        "CUDF_POLARS__",
    )
    try:
        raw = Path(f"/proc/{pid}/environ").read_bytes().split(b"\x00")
    except OSError:
        return []
    rows: list[str] = []
    for item in raw:
        if not item:
            continue
        text = item.decode(errors="replace")
        if any(text.startswith(p) for p in prefixes):
            rows.append(text)
    return sorted(rows)


def snapshot_tree(root_pid: int, *, label: str, path: Path) -> dict[str, int]:
    """Write a human-readable tree snapshot; return summary stats."""
    ps = subprocess.check_output(
        ["ps", "-eo", "pid,ppid,nlwp,etime,cmd", "--no-headers"],
        text=True,
    )
    kids: dict[int, list[int]] = {}
    info: dict[int, tuple[int, int, str, str]] = {}
    for line in ps.splitlines():
        parts = line.split(None, 4)
        if len(parts) < 5:
            continue
        pid, ppid, nlwp = int(parts[0]), int(parts[1]), int(parts[2])
        kids.setdefault(ppid, []).append(pid)
        info[pid] = (ppid, nlwp, parts[3], parts[4])

    stack = [root_pid]
    seen: set[int] = set()
    rows: list[tuple[int, int, int, str, str]] = []
    while stack:
        pid = stack.pop()
        if pid in seen:
            continue
        seen.add(pid)
        if pid in info:
            rows.append((pid, *info[pid]))
        stack.extend(kids.get(pid, ()))
    rows.sort(key=lambda r: (-r[2], r[0]))

    ts = datetime.now(UTC).isoformat(timespec="milliseconds")
    agg: Counter[str] = Counter()
    for pid, *_ in rows:
        agg.update(_thread_names(pid))

    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as fh:
        print(f"snapshot={label}", file=fh)
        print(f"timestamp_utc={ts}", file=fh)
        print(f"root_pid={root_pid}", file=fh)
        print(f"processes={len(rows)} sum_nlwp={sum(r[2] for r in rows)}", file=fh)
        print(file=fh)

        print("=== root process ===", file=fh)
        for key, value in _status_fields(root_pid).items():
            print(f"{key}: {value}", file=fh)
        print(f"cmdline: {_cmdline(root_pid)}", file=fh)
        print(file=fh)

        print("=== selected env ===", file=fh)
        env_rows = _selected_env(root_pid)
        if env_rows:
            for row in env_rows:
                print(row, file=fh)
        else:
            print("<none of the selected prefixes set>", file=fh)
        print(file=fh)

        print("=== process tree (by nlwp desc) ===", file=fh)
        for pid, ppid, nlwp, etime, cmd in rows:
            print(f"pid={pid} ppid={ppid} nlwp={nlwp} etime={etime}", file=fh)
            print(f"  cmd={cmd[:160]}", file=fh)
            names = _thread_names(pid)
            if names:
                print(f"  thread_names={dict(names.most_common())}", file=fh)
        print(file=fh)

        print("=== aggregate thread names across tree ===", file=fh)
        print(f"total_threads={sum(agg.values())}", file=fh)
        for name, count in agg.most_common():
            print(f"  {count:4d}  {name}", file=fh)

    summary = {
        "processes": len(rows),
        "sum_nlwp": sum(r[2] for r in rows),
        "root_nlwp": int(_status_fields(root_pid).get("Threads", "0") or 0),
    }
    print(
        f"[{label}] wrote {path} "
        f"(processes={summary['processes']} sum_nlwp={summary['sum_nlwp']} "
        f"root_nlwp={summary['root_nlwp']})",
        file=sys.stderr,
    )
    return summary


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--outdir",
        type=Path,
        default=Path("/tmp/ray-thread-profile"),
        help="Directory for snapshot files (default: /tmp/ray-thread-profile)",
    )
    parser.add_argument(
        "--settle-seconds",
        type=float,
        default=2.0,
        help="Seconds to wait after engine create / collect before snapshot",
    )
    parser.add_argument(
        "--num-ranks",
        type=int,
        default=None,
        help=(
            "Forwarded to RayEngine(num_ranks=...). Default None uses one rank "
            "per available GPU. Use with --allow-gpu-sharing when > visible GPUs."
        ),
    )
    parser.add_argument(
        "--allow-gpu-sharing",
        action="store_true",
        help="Pass engine_options allow_gpu_sharing=True (needed with --num-ranks).",
    )
    parser.add_argument(
        "--no-shutdown",
        action="store_true",
        help="Leave RayEngine alive after the final snapshot (default: shutdown).",
    )
    args = parser.parse_args(argv)

    root_pid = os.getpid()
    outdir: Path = args.outdir
    outdir.mkdir(parents=True, exist_ok=True)

    # ------------------------------------------------------------------
    # Phase 1: before importing Ray / cudf-polars / Polars
    # ------------------------------------------------------------------
    snapshot_tree(
        root_pid,
        label="01_before_imports",
        path=outdir / "01_before_imports.txt",
    )

    print("importing polars / cudf_polars / RayEngine...", file=sys.stderr)
    import polars as pl
    from cudf_polars.engine.ray import RayEngine

    engine_options: dict = {}
    if args.allow_gpu_sharing:
        engine_options["allow_gpu_sharing"] = True

    print(
        f"constructing RayEngine(num_ranks={args.num_ranks}, "
        f"engine_options={engine_options or None})...",
        file=sys.stderr,
    )
    engine = RayEngine(
        num_ranks=args.num_ranks,
        engine_options=engine_options or None,
    )
    print(f"RayEngine.creation_stats={RayEngine.creation_stats()}", file=sys.stderr)
    time.sleep(args.settle_seconds)

    # ------------------------------------------------------------------
    # Phase 2: after RayEngine construction (ray.init + actors)
    # ------------------------------------------------------------------
    snapshot_tree(
        root_pid,
        label="02_after_ray_engine",
        path=outdir / "02_after_ray_engine.txt",
    )
    (outdir / "ray_engine_stats_after_create.txt").write_text(
        f"{RayEngine.creation_stats()}\n", encoding="utf-8"
    )

    print("running tiny LazyFrame.collect(engine=RayEngine)...", file=sys.stderr)
    pl.LazyFrame({"a": [1, 2]}).collect(engine=engine)
    print(f"RayEngine.creation_stats={RayEngine.creation_stats()}", file=sys.stderr)
    time.sleep(args.settle_seconds)

    # ------------------------------------------------------------------
    # Phase 3: after collect
    # ------------------------------------------------------------------
    snapshot_tree(
        root_pid,
        label="03_after_collect",
        path=outdir / "03_after_collect.txt",
    )
    (outdir / "ray_engine_stats_after_collect.txt").write_text(
        f"{RayEngine.creation_stats()}\n", encoding="utf-8"
    )

    if not args.no_shutdown:
        print("shutting down RayEngine...", file=sys.stderr)
        engine.shutdown()
        print(f"RayEngine.creation_stats={RayEngine.creation_stats()}", file=sys.stderr)
        (outdir / "ray_engine_stats_after_shutdown.txt").write_text(
            f"{RayEngine.creation_stats()}\n", encoding="utf-8"
        )
    else:
        print("leaving RayEngine running (--no-shutdown)", file=sys.stderr)

    print(f"done; snapshots in {outdir}", file=sys.stderr)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())