fix(profiling): make TraceLens and dynolog survive real GPU traces

pyproject.toml

@@ -61,6 +61,11 @@ Homepage = "https://github.com/ROCm/madengine"
 Issues = "https://github.com/ROCm/madengine/issues"
 
 
+[tool.hatch.metadata]
+# The tracelens extra installs TraceLens from git; hatchling rejects such direct
+# references during metadata generation unless they are allowed here.
+allow-direct-references = true
+
 [tool.hatch.build.targets.wheel]
 # scripts/ is listed in .gitignore (to exclude external MAD project scripts/ dirs during dev).
 # artifacts bypasses VCS exclusion without risk of duplicate-file errors from force-include.

src/madengine/scripts/common/post_scripts/dynolog_stop.sh

@@ -84,10 +84,15 @@ if [ "$trace_count" -gt 0 ]; then
     ls -la "$OUTPUT_DIR" || true
 else
     echo "⚠️  No torch.profiler traces were captured in ${OUTPUT_DIR}"
-    if [ -f "$RESULT_FILE" ] && [ "$(cat "$RESULT_FILE")" = "no_process" ]; then
+    result=""
+    [ -f "$RESULT_FILE" ] && result=$(cat "$RESULT_FILE")
+    if [ "$result" = "no_process" ]; then
         echo "⚠️  The trigger never matched a PyTorch process. Most likely causes:"
         echo "⚠️    - the workload is not PyTorch, or predates PyTorch 1.13"
         echo "⚠️    - the model run finished before TORCH_PROFILE_WARMUP_S elapsed"
+    elif [ "$result" = "request_rejected" ]; then
+        echo "⚠️  dyno rejected the trace request itself, so no trace was ever"
+        echo "⚠️  requested. See the trigger log below for the dyno output."
     fi
 fi
 

src/madengine/scripts/common/tools/dynolog_trigger.sh

@@ -8,7 +8,11 @@
 # Runs in the background for the lifetime of the model run. `dyno gputrace` can
 # only configure PyTorch processes that have already registered with the daemon,
 # and there is no way to know when the workload reaches steady state, so this
-# polls with --fail-on-no-process until a request is accepted.
+# polls until a request is accepted.
+#
+# `dyno gputrace` exits 0 whether or not it matched anything, so the outcome has
+# to be read from its output: the response carries the matched pids, and an empty
+# `processesMatched` list means the workload has not registered yet.
 
 set -u
 
@@ -54,16 +58,30 @@ attempt=0
 while [ "$attempt" -lt "$MAX_ATTEMPTS" ]; do
     attempt=$((attempt + 1))
     echo "[dynolog-trigger] attempt ${attempt}/${MAX_ATTEMPTS}: requesting trace -> ${LOG_FILE}"
-    if dyno --port "$PORT" gputrace \
+    response=$(dyno --port "$PORT" gputrace \
         --job-id "$JOB_ID" \
         --log-file "$LOG_FILE" \
         --process-limit "$PROCESS_LIMIT" \
-        --fail-on-no-process \
-        "${OPTS[@]}"; then
+        "${OPTS[@]}" 2>&1)
+    echo "$response"
+
+    if echo "$response" | grep -q '"processesMatched":\[[0-9]'; then
         echo "[dynolog-trigger] trace request accepted on attempt ${attempt}"
         echo "accepted" > "$RESULT_FILE"
         exit 0
     fi
+
+    # A response that reports no matches is the expected case while the workload
+    # is still starting up. Anything else means dyno rejected the request itself
+    # (an unsupported flag, an unreachable daemon), which retrying cannot fix.
+    if ! echo "$response" | grep -q 'processesMatched'; then
+        echo "[dynolog-trigger] dyno rejected the request; not retrying."
+        echo "[dynolog-trigger] Check the dyno output above against the installed"
+        echo "[dynolog-trigger] dynolog version ('dyno gputrace --help')."
+        echo "request_rejected" > "$RESULT_FILE"
+        exit 1
+    fi
+
     echo "[dynolog-trigger] no PyTorch process matched yet; retrying in ${RETRY_INTERVAL_S}s"
     sleep "$RETRY_INTERVAL_S"
 done

src/madengine/scripts/common/tools/tracelens_analyze.py

@@ -16,6 +16,7 @@
 """
 
 import argparse
+import codecs
 import csv
 import glob
 import gzip
@@ -25,6 +26,7 @@
 import shutil
 import subprocess
 import sys
+import tempfile
 from typing import Dict, List, Optional, Sequence, Tuple
 
 # Trace kinds, in discovery precedence order. The first pattern set that claims a
@@ -92,6 +94,9 @@
     "TraceLens_compare_perf_reports_pytorch": "TraceLens.Reporting.compare_perf_reports_pytorch",
 }
 
+# Read size used when checking a trace for undecodable bytes and rewriting it.
+_SANITIZE_CHUNK_BYTES = 1 << 20
+
 SUMMARY_CSV_FIELDS = (
     "trace_file",
     "kind",
@@ -205,18 +210,20 @@ def _resolve_python(python: Optional[str]) -> str:
 def _build_command(python: str, script_name: str, args: Sequence[str]) -> List[str]:
     """Return argv invoking a TraceLens entry point with ``args``.
 
-    Prefers the installed console script (clearer logs, honours the package's
-    own entry-point wiring) and falls back to importing the module's ``main``.
-    The fallback exits with ``main()``'s return value, the same way the console
-    scripts pip generates do, so a report failure is not silently swallowed.
+    Prefers the console script installed alongside ``python`` (clearer logs,
+    honours the package's own entry-point wiring) and falls back to importing the
+    module's ``main`` with that same interpreter. The fallback exits with
+    ``main()``'s return value, the same way the console scripts pip generates do,
+    so a report failure is not silently swallowed.
+
+    Both forms stay inside the environment the caller asked for. Searching PATH
+    instead would defeat the isolation ``--python`` exists to provide: TraceLens
+    pins protobuf and xprof, and is installed in a venv of its own.
     """
     bindir = os.path.dirname(os.path.abspath(python))
     candidate = os.path.join(bindir, script_name)
     if os.path.isfile(candidate) and os.access(candidate, os.X_OK):
         return [candidate, *args]
-    on_path = shutil.which(script_name)
-    if on_path:
-        return [on_path, *args]
     module = _ENTRY_POINTS[script_name]
     return [
         python,
@@ -247,6 +254,75 @@ def _run(command: Sequence[str], cwd: Optional[str] = None) -> Tuple[int, str]:
     return completed.returncode, output
 
 
+def _has_invalid_utf8(path: str) -> bool:
+    """Return True when ``path`` is not decodable as UTF-8."""
+    decoder = codecs.getincrementaldecoder("utf-8")()
+    try:
+        with open(path, "rb") as handle:
+            while True:
+                chunk = handle.read(_SANITIZE_CHUNK_BYTES)
+                if not chunk:
+                    decoder.decode(b"", final=True)
+                    return False
+                decoder.decode(chunk)
+    except UnicodeDecodeError:
+        return True
+    except OSError:
+        return False
+
+
+def _write_sanitized_copy(path: str, destination: str) -> None:
+    """Copy ``path`` to ``destination`` with undecodable bytes replaced."""
+    decoder = codecs.getincrementaldecoder("utf-8")("replace")
+    with open(path, "rb") as source:
+        with open(destination, "w", encoding="utf-8") as target:
+            while True:
+                chunk = source.read(_SANITIZE_CHUNK_BYTES)
+                if not chunk:
+                    target.write(decoder.decode(b"", final=True))
+                    return
+                target.write(decoder.decode(chunk))
+
+
+def _sanitized_trace(trace: str, kind: str, workspace: List[Optional[str]]) -> str:
+    """Return a trace path TraceLens can load, sanitizing bytes if it must.
+
+    rocprofv3 copies HIP API ``const char *`` arguments into its JSON verbatim, so
+    arguments that do not point at a string (``fname``, ``kname``) leave raw bytes
+    in the trace. TraceLens loads traces with orjson, which rejects the whole file
+    when any byte is not valid UTF-8, so a single stray pointer costs the entire
+    report. Analyzing a sanitized copy keeps the original trace untouched.
+
+    Args:
+        trace: Path to the discovered trace.
+        kind: Discovered trace kind.
+        workspace: Single-element list caching the scratch directory, so it is
+            created only once and only when a trace actually needs sanitizing.
+
+    Returns:
+        The path to analyze: ``trace`` itself, or a sanitized copy of it.
+    """
+    if kind not in (KIND_PYTORCH, KIND_ROCPROF_JSON) or not trace.endswith(".json"):
+        return trace
+    if not _has_invalid_utf8(trace):
+        return trace
+
+    if workspace[0] is None:
+        workspace[0] = tempfile.mkdtemp(prefix="madengine-tracelens-")
+    destination = os.path.join(workspace[0], os.path.basename(trace))
+    print(
+        f"[tracelens] {trace} is not valid UTF-8 (rocprofv3 writes raw pointer "
+        "bytes for some HIP API string arguments); analyzing a sanitized copy",
+        flush=True,
+    )
+    try:
+        _write_sanitized_copy(trace, destination)
+    except OSError as exc:
+        print(f"[tracelens] could not sanitize {trace}: {exc}", flush=True)
+        return trace
+    return destination
+
+
 def _failure_detail(returncode: int, output: str) -> str:
     """Return a one-line explanation for a failed TraceLens invocation."""
     lines = [line for line in output.strip().splitlines() if line.strip()]
@@ -340,21 +416,30 @@ def _pftrace_jobs(
 
 
 def _rank_regex() -> str:
-    """Return the rank-extraction regex covering madengine trace filenames.
+    """Return the rank-extraction regex covering rank-labelled trace filenames.
 
-    Matches dynolog output (``libkineto_trace_<pid>.json``, where madengine
-    renames per rank), torch.profiler defaults (``..._rank0_...``), and the
-    ``rank[N]`` form used by ``tensorboard_trace_handler``.
+    Matches the torch.profiler default (``..._rank0_...``) and the ``rank[N]``
+    form written by ``tensorboard_trace_handler``. Traces captured on demand
+    through dynolog are named after the process id instead, and carry no rank.
     """
     return r"rank[\[\-_/]?(?P<rank>\d+)"
 
 
+def _is_rank_labelled(trace: str) -> bool:
+    """Return True when the rank of ``trace`` can be read from its filename."""
+    return re.search(_rank_regex(), os.path.basename(trace)) is not None
+
+
 def _collective_args(
-    root: str, out_base: str, world_size: int, extra: Sequence[str]
+    traces: Sequence[str], out_base: str, world_size: int, extra: Sequence[str]
 ) -> List[str]:
+    # TraceLens takes a glob rather than a list of traces, so scope it to the tree
+    # the per-rank traces were found in; a wider one sweeps up unrelated JSON, and
+    # rocprofv3 results are hundreds of megabytes each.
+    directory = os.path.commonpath([os.path.dirname(t) for t in traces])
     return [
         "--trace_glob",
-        os.path.join(root, "**", "*.json*"),
+        os.path.join(directory, "**", "*.json*"),
         "--rank_regex",
         _rank_regex(),
         "--world_size",
@@ -415,20 +500,24 @@ def analyze(
         "auto": {KIND_PYTORCH, KIND_ROCPROF_JSON, KIND_PFTRACE},
     }[mode]
 
+    # Scratch directory for sanitized trace copies, created on first need.
+    sanitize_workspace: List[Optional[str]] = [None]
+
     jobs: List[Tuple[str, str, str, List[str]]] = []
     for kind, paths in sorted(traces.items()):
         if kind not in wanted:
             continue
         for trace in paths:
             stem = _report_stem(trace, root)
             out_base = os.path.join(output_dir, stem)
+            readable = _sanitized_trace(trace, kind, sanitize_workspace)
             if kind == KIND_PYTORCH and mode != "collective":
                 jobs.append(
                     (
                         trace,
                         kind,
                         "TraceLens_generate_perf_report_pytorch",
-                        _pytorch_args(trace, out_base, gpu_arch, extra_args),
+                        _pytorch_args(readable, out_base, gpu_arch, extra_args),
                     )
                 )
             elif kind == KIND_ROCPROF_JSON:
@@ -437,30 +526,42 @@ def analyze(
                         trace,
                         kind,
                         "TraceLens_generate_perf_report_rocprof",
-                        _rocprof_args(trace, out_base, extra_args),
+                        _rocprof_args(readable, out_base, extra_args),
                     )
                 )
             elif kind == KIND_PFTRACE:
-                for tool, args in _pftrace_jobs(trace, out_base, extra_args):
+                for tool, args in _pftrace_jobs(readable, out_base, extra_args):
                     jobs.append((trace, kind, tool, args))
 
-    # A multi-rank collective report needs at least two per-rank PyTorch traces.
+    # A multi-rank collective report needs at least two per-rank PyTorch traces,
+    # and TraceLens reads each trace's rank from its filename.
     pytorch_traces = traces.get(KIND_PYTORCH, [])
-    ranks = world_size or len(pytorch_traces)
-    if mode in ("auto", "collective") and len(pytorch_traces) > 1 and ranks > 1:
+    ranked = [trace for trace in pytorch_traces if _is_rank_labelled(trace)]
+    unrankable: List[Tuple[str, str]] = []
+    ranks = world_size or len(ranked)
+    if mode in ("auto", "collective") and len(ranked) > 1 and ranks > 1:
         jobs.append(
             (
-                f"{len(pytorch_traces)} per-rank traces",
+                f"{len(ranked)} per-rank traces",
                 KIND_PYTORCH,
                 "TraceLens_generate_multi_rank_collective_report_pytorch",
                 _collective_args(
-                    root,
+                    ranked,
                     os.path.join(output_dir, "multi_rank_collective"),
                     ranks,
                     extra_args,
                 ),
             )
         )
+    elif mode in ("auto", "collective") and len(pytorch_traces) > 1:
+        unrankable.append(
+            (
+                f"{len(pytorch_traces)} PyTorch traces",
+                "the collective report needs the rank in each trace's filename, "
+                "and none of these carry one. Traces captured on demand through "
+                "dynolog are named after the process id.",
+            )
+        )
 
     results: List[Dict[str, str]] = []
     for trace, kind, tool, args in jobs:
@@ -479,6 +580,9 @@ def analyze(
             }
         )
 
+    if sanitize_workspace[0] is not None:
+        shutil.rmtree(sanitize_workspace[0], ignore_errors=True)
+
     for path, reason in unsupported:
         print(f"[tracelens] skipping {path}: {reason}", flush=True)
         results.append(
@@ -492,6 +596,21 @@ def analyze(
             }
         )
 
+    for label, reason in unrankable:
+        print(f"[tracelens] skipping the collective report: {reason}", flush=True)
+        results.append(
+            {
+                "trace_file": label,
+                "kind": KIND_PYTORCH,
+                "tracelens_tool": (
+                    "TraceLens_generate_multi_rank_collective_report_pytorch"
+                ),
+                "status": "SKIPPED",
+                "output": "",
+                "detail": reason,
+            }
+        )
+
     summary_csv = os.path.join(output_dir, "tracelens_summary.csv")
     with open(summary_csv, "w", newline="", encoding="utf-8") as handle:
         writer = csv.DictWriter(handle, fieldnames=list(SUMMARY_CSV_FIELDS))