From db8bade9cfe02f4b4b83f7fb3d6bfb4a412a8acd Mon Sep 17 00:00:00 2001 From: Cemberk Date: Wed, 12 Aug 2026 20:19:46 -0500 Subject: [PATCH] fix(profiling): make TraceLens and dynolog survive real GPU traces Everything here came out of running the tools on an AMD GPU node against real rocprofv3 and dynolog captures. None of it is reachable without hardware, which is why the no-GPU test suite went green over broken code. - rocprofv3 copies HIP API `const char *` arguments into its JSON verbatim, so an argument that does not point at a string leaves raw bytes behind. TraceLens loads traces with orjson, which rejects the whole document, so twelve stray bytes cost the entire report on a 295 MB trace. The analyzer now detects undecodable bytes incrementally and analyzes a sanitized copy, leaving the original trace untouched. - `dyno gputrace` has no `--fail-on-no-process` flag in the pinned v0.5.0, and exits 0 whether or not it matched a process. Every attempt therefore looked like "no PyTorch process registered yet" and the trigger retried a request dyno had refused, twelve times. It now reads `processesMatched` out of the response, and gives up immediately when dyno rejected the request rather than when the workload is merely slow to register. `dynolog_stop.sh` reports that outcome instead of blaming the workload. - The multi-rank collective report globbed the whole run directory and assumed every PyTorch trace carried its rank in its filename. Traces captured on demand are named after the process id, so TraceLens matched three files and then rejected all of them. The report is now built only from rank-labelled traces, from a glob scoped to their own directory, and is skipped with an explanation when no trace carries a rank. - `_build_command` fell back to searching PATH for TraceLens console scripts, which defeats the isolation `--python` exists to provide: TraceLens pins protobuf and xprof, and is installed in a venv of its own. - `pip install -e .` failed outright on this branch: the `tracelens` extra installs from git, and hatchling rejects direct references unless they are allowed explicitly. Co-authored-by: Cursor --- pyproject.toml | 5 + .../common/post_scripts/dynolog_stop.sh | 7 +- .../scripts/common/tools/dynolog_trigger.sh | 26 ++- .../scripts/common/tools/tracelens_analyze.py | 161 +++++++++++++++--- tests/e2e/test_dynolog_dummy_pipeline.py | 113 ++++++++++-- tests/e2e/test_tracelens_dummy_pipeline.py | 41 +++++ .../TraceLens/Reporting/_dummy.py | 11 ++ tests/unit/test_tracelens_analyze.py | 41 ++++- 8 files changed, 366 insertions(+), 39 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index cc845c66..db40256d 100644 --- a/pyproject.toml +++ b/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. diff --git a/src/madengine/scripts/common/post_scripts/dynolog_stop.sh b/src/madengine/scripts/common/post_scripts/dynolog_stop.sh index 7f9ebb22..cb8546db 100644 --- a/src/madengine/scripts/common/post_scripts/dynolog_stop.sh +++ b/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 diff --git a/src/madengine/scripts/common/tools/dynolog_trigger.sh b/src/madengine/scripts/common/tools/dynolog_trigger.sh index 9999e36c..11d90258 100644 --- a/src/madengine/scripts/common/tools/dynolog_trigger.sh +++ b/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 diff --git a/src/madengine/scripts/common/tools/tracelens_analyze.py b/src/madengine/scripts/common/tools/tracelens_analyze.py index 63e30587..cd23dc9b 100644 --- a/src/madengine/scripts/common/tools/tracelens_analyze.py +++ b/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_.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\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,6 +500,9 @@ 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: @@ -422,13 +510,14 @@ def analyze( 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)) diff --git a/tests/e2e/test_dynolog_dummy_pipeline.py b/tests/e2e/test_dynolog_dummy_pipeline.py index 5084a087..79224e22 100644 --- a/tests/e2e/test_dynolog_dummy_pipeline.py +++ b/tests/e2e/test_dynolog_dummy_pipeline.py @@ -60,27 +60,68 @@ while : ; do sleep 0.2; done """ -# Records each trace request, rejects the first $DUMMY_DYNO_REJECTS of them the -# way `--fail-on-no-process` does before the workload has registered, then -# accepts and writes the trace Kineto would have written. +# Stands in for `dyno gputrace` from the pinned dynolog release, and holds to the +# same contract, because both halves of that contract are easy to get wrong: +# +# * it validates its options the way clap does, rejecting anything it does not +# know with exit 2 and a usage message +# * it exits 0 whether or not it matched a process, reporting the matched pids in +# its response instead +# +# $DUMMY_DYNO_REJECTS requests report no match, as they do before the workload has +# registered. $DUMMY_DYNO_UNSUPPORTED drops one option from the supported set, the +# way an older or newer dynolog than madengine expects would. DYNO_STUB = """#!/bin/sh printf '%s\\n' "$*" >> "$DUMMY_DYNO_LOG" -requests=$(wc -l < "$DUMMY_DYNO_LOG") -if [ "$requests" -le "${DUMMY_DYNO_REJECTS:-0}" ]; then - echo "No processes were matched, exiting" >&2 - exit 1 + +supported=" --duration-ms --iterations --job-id --log-file --pids --process-limit\ + --profile-memory --profile-start-iteration-roundup --profile-start-time\ + --record-shapes --with-flops --with-modules --with-stacks " +if [ -n "${DUMMY_DYNO_UNSUPPORTED:-}" ]; then + supported=$(printf '%s' "$supported" | sed "s| ${DUMMY_DYNO_UNSUPPORTED} | |") fi + +for arg in "$@"; do + case "$arg" in + --port|gputrace) continue;; + --*) + case "$supported" in + *" $arg "*) ;; + *) + echo "error: Found argument '$arg' which wasn't expected, or isn't valid in this context" + echo "" + echo "USAGE:" + echo " dyno gputrace --log-file --job-id --process-limit " + echo "" + echo "For more information try --help" + exit 2;; + esac;; + esac +done + log_file="" previous="" for arg in "$@"; do if [ "$previous" = "--log-file" ]; then log_file=$arg; fi previous=$arg done + +echo "Kineto config = " +echo "ACTIVITIES_LOG_FILE=$log_file" +requests=$(wc -l < "$DUMMY_DYNO_LOG") +if [ "$requests" -le "${DUMMY_DYNO_REJECTS:-0}" ]; then + echo 'response = {"activityProfilersBusy":0,"activityProfilersTriggered":[],"eventProfilersBusy":0,"eventProfilersTriggered":[],"processesMatched":[]}' + echo "No processes were matched, please check --job-id or --pids flags" + exit 0 +fi +echo 'response = {"activityProfilersBusy":0,"activityProfilersTriggered":[4242],"eventProfilersBusy":0,"eventProfilersTriggered":[],"processesMatched":[4242]}' +echo "Matched 1 processes" +echo "Trace output files will be written to:" +echo " ${log_file%.json}_4242.json" if [ "${DUMMY_DYNO_WRITE_TRACE:-0}" = "1" ] && [ -n "$log_file" ]; then # Kineto appends the process id to the requested filename. printf '{"traceEvents": [], "schemaVersion": 1}' > "${log_file%.json}_4242.json" fi -echo "response length: 1" exit 0 """ @@ -243,7 +284,11 @@ def run_trigger(work: Path, dummy: DummyDynolog, **env: str): return run_script(TRIGGER_SCRIPT, work, dummy.environ(**settings)) def test_trigger_retries_until_the_workload_registers(self, workdir, dummy_dynolog): - """A trace cannot be requested until PyTorch has registered, so we poll.""" + """A trace cannot be requested until PyTorch has registered, so we poll. + + dyno reports that in its response and still exits 0, so a request that + matched nothing has to be told apart from one that succeeded by output. + """ result = self.run_trigger(workdir, dummy_dynolog, DUMMY_DYNO_REJECTS="2") assert result.returncode == 0, result.stdout @@ -251,10 +296,37 @@ def test_trigger_retries_until_the_workload_registers(self, workdir, dummy_dynol assert "accepted on attempt 3" in result.stdout assert RESULT_FILE.read_text().strip() == "accepted" + def test_an_option_dyno_rejects_fails_fast_and_says_so( + self, workdir, dummy_dynolog + ): + """A request dyno refuses to parse can never succeed, so stop retrying it. + + Retrying it instead spends every attempt on a permanent error and then + blames the workload for never registering. + """ + result = self.run_trigger( + workdir, + dummy_dynolog, + DUMMY_DYNO_UNSUPPORTED="--with-modules", + TORCH_PROFILE_MAX_ATTEMPTS="5", + ) + + assert result.returncode != 0 + assert len(dummy_dynolog.requests()) == 1, dummy_dynolog.requests() + assert "not retrying" in result.stdout + assert "wasn't expected" in result.stdout, result.stdout + assert "no PyTorch process matched" not in result.stdout + assert RESULT_FILE.read_text().strip() == "request_rejected" + def test_request_carries_the_data_tracelens_needs(self, workdir, dummy_dynolog): - """Shapes, stacks, and modules are what make the TraceLens reports useful.""" - self.run_trigger(workdir, dummy_dynolog) + """Shapes, stacks, and modules are what make the TraceLens reports useful. + + Every option here also has to exist in the dynolog release the pre-script + installs; the dyno stand-in rejects anything else. + """ + result = self.run_trigger(workdir, dummy_dynolog) + assert result.returncode == 0, result.stdout request = dummy_dynolog.requests()[0] for flag in ( "--record-shapes", @@ -262,7 +334,6 @@ def test_request_carries_the_data_tracelens_needs(self, workdir, dummy_dynolog): "--with-modules", "--iterations 5", "--process-limit 64", - "--fail-on-no-process", ): assert flag in request, request # An absolute path, because the workload's working directory is its own. @@ -356,3 +427,21 @@ def test_missing_traces_are_explained(self, workdir, dummy_dynolog): assert result.returncode == 0, result.stdout assert "No torch.profiler traces were captured" in result.stdout assert "never matched a PyTorch process" in result.stdout + + def test_a_rejected_request_is_reported_as_such(self, workdir, dummy_dynolog): + """A request dyno never accepted is a different problem from a quiet workload.""" + run_script( + START_SCRIPT, workdir, dummy_dynolog.environ(TORCH_PROFILE_WARMUP_S="30") + ) + TestTraceRequest.run_trigger( + workdir, + dummy_dynolog, + DUMMY_DYNO_UNSUPPORTED="--record-shapes", + TORCH_PROFILE_MAX_ATTEMPTS="1", + ) + + result = run_script(STOP_SCRIPT, workdir, dummy_dynolog.environ()) + + assert result.returncode == 0, result.stdout + assert "dyno rejected the trace request" in result.stdout + assert "never matched a PyTorch process" not in result.stdout diff --git a/tests/e2e/test_tracelens_dummy_pipeline.py b/tests/e2e/test_tracelens_dummy_pipeline.py index 444a5f2a..b3da5679 100644 --- a/tests/e2e/test_tracelens_dummy_pipeline.py +++ b/tests/e2e/test_tracelens_dummy_pipeline.py @@ -350,6 +350,47 @@ def test_summary_records_each_report_and_the_unreadable_db( # The skip has to say which preset the user should have run instead. assert "rocprofv3_lightweight" in skipped[0]["detail"] + def test_trace_with_undecodable_bytes_is_still_analyzed( + self, tmp_path, dummy_tracelens + ): + """rocprofv3 traces are not always valid UTF-8, and TraceLens demands it. + + rocprofv3 copies HIP API ``const char *`` arguments into its JSON as they + are, so an argument that does not point at a string leaves raw bytes in an + otherwise perfectly good 300MB trace. TraceLens loads traces with orjson, + which rejects the whole document over those few bytes. + """ + work = tmp_path / "raw-bytes" + trace = write_rocprof_json(work / "rocprof_output" / "510_results.json") + original = trace.read_bytes().replace( + b'"gemm_kernel(float*, float*)"', b'"\x90{-\xad\xcb\x7f"' + ) + trace.write_bytes(original) + + result = run_analyzer(work, dummy_tracelens) + + assert result.returncode == 0, result.stdout + rows = summary_rows(work) + assert [r["status"] for r in rows] == ["SUCCESS"], rows + # Analyzed through a copy, so the trace the run collected is untouched. + assert trace.read_bytes() == original + analyzed = only(dummy_tracelens.invocations(), ROCPROF_REPORT)["args"][ + "profile_json_path" + ] + assert analyzed != str(trace) + assert "sanitized copy" in result.stdout + + def test_a_valid_trace_is_analyzed_where_it_lies(self, profiled_run, dummy_tracelens): + """Traces TraceLens can already read are not copied; they can be huge.""" + result = run_analyzer(profiled_run, dummy_tracelens) + + assert result.returncode == 0, result.stdout + analyzed = only(dummy_tracelens.invocations(), ROCPROF_REPORT)["args"][ + "profile_json_path" + ] + assert analyzed == str(profiled_run / "rocprof_output" / "1234_results.json") + assert "sanitized copy" not in result.stdout + def test_gzipped_kineto_trace_is_analyzed(self, tmp_path, dummy_tracelens): """tensorboard_trace_handler's gzipped traces are picked up too.""" work = tmp_path / "gz" diff --git a/tests/fixtures/dummy_tracelens/TraceLens/Reporting/_dummy.py b/tests/fixtures/dummy_tracelens/TraceLens/Reporting/_dummy.py index 026cbdf7..0f81bb00 100644 --- a/tests/fixtures/dummy_tracelens/TraceLens/Reporting/_dummy.py +++ b/tests/fixtures/dummy_tracelens/TraceLens/Reporting/_dummy.py @@ -94,6 +94,17 @@ def _check_input_file(path: str) -> Optional[int]: return _fail(f"input trace {path} does not exist") if os.path.getsize(path) == 0: return _fail(f"input trace {path} is empty") + if path.endswith(".json"): + # TraceLens loads traces with orjson, which rejects the whole document + # when any byte in it is not valid UTF-8. + try: + with open(path, "rb") as handle: + handle.read().decode("utf-8") + except UnicodeDecodeError: + return _fail( + "orjson.JSONDecodeError: str is not valid UTF-8: surrogates not " + "allowed: line 1 column 1 (char 0)" + ) return None diff --git a/tests/unit/test_tracelens_analyze.py b/tests/unit/test_tracelens_analyze.py index 53d016e3..6a6312ab 100644 --- a/tests/unit/test_tracelens_analyze.py +++ b/tests/unit/test_tracelens_analyze.py @@ -8,6 +8,7 @@ import csv import importlib.util import json +import os import sys from pathlib import Path @@ -163,9 +164,18 @@ def test_pftrace_produces_three_complementary_reports(self, analyzer): assert args[:2] == ["--trace_path", "t.pftrace"] def test_collective_args_carry_rank_regex_and_world_size(self, analyzer): - args = analyzer._collective_args("/root", "/out/coll", 8, []) + traces = [ + "/root/torch_profiler_output/libkineto_trace_rank0_1.json", + "/root/torch_profiler_output/libkineto_trace_rank1_2.json", + ] + args = analyzer._collective_args(traces, "/out/coll", 8, []) assert args[args.index("--world_size") + 1] == "8" assert "rank" in args[args.index("--rank_regex") + 1] + # Scoped to the traces' own directory: a wider glob sweeps up unrelated + # JSON, and rocprofv3 results are hundreds of megabytes each. + trace_glob = args[args.index("--trace_glob") + 1] + assert "torch_profiler_output" in trace_glob + assert trace_glob.endswith(os.path.join("**", "*.json*")) def test_extra_args_are_forwarded(self, analyzer): args = analyzer._pytorch_args("t.json", "/out/t", None, ["--detect_recompute"]) @@ -280,6 +290,35 @@ def test_no_collective_report_for_a_single_rank(self, analyzer, tmp_path, monkey assert len(calls) == 1 assert not any("multi_rank" in part for call in calls for part in call) + def test_collective_report_is_skipped_when_ranks_cannot_be_identified( + self, analyzer, tmp_path, monkeypatch + ): + """dynolog names traces after the pid, and TraceLens needs the rank. + + Attempting the report anyway fails on every multi-process run profiled + through dynolog, which reads as a broken tool rather than a limitation. + """ + for pid in (724, 892): + _write( + tmp_path, f"torch_profiler_output/libkineto_trace_{pid}.json", CHROME_TRACE + ) + calls = [] + monkeypatch.setattr( + analyzer, + "_run", + lambda command, cwd=None: (calls.append(list(command)), (0, ""))[1], + ) + + summary = analyzer.analyze( + root=str(tmp_path), output_dir=str(tmp_path / "out"), python=sys.executable + ) + + assert not any("multi_rank" in part for call in calls for part in call) + skipped = [r for r in summary["results"] if r["status"] == "SKIPPED"] + assert len(skipped) == 1 + assert "rank" in skipped[0]["detail"] + assert "dynolog" in skipped[0]["detail"] + def test_max_traces_caps_work_per_kind(self, analyzer, trace_tree, monkeypatch): calls = [] monkeypatch.setattr(