[ROCm] Fix teardown memory corruption and truncated traces on the rocprofiler-sdk backend

[haishuok0525]

The bug

On the rocprofiler-sdk backend, stopping a trace corrupts the heap, and the trace it publishes is usually missing most of the GPU work without reporting anything.

Both symptoms have the same cause. rocprofiler-sdk does not emplace a kernel dispatch record at the moment a flush runs. The record is produced on the SDK's own signal handler thread (rocprofiler::hsa::AsyncSignalHandler) once the dispatch completion signal fires, and rocprofiler_stop_context() does not wait for sessions already in flight. So records keep arriving for tens of milliseconds after collection has stopped.

RocprofLogger::stopLogging() flushed, stopped the context and returned, so the rest of teardown ran while that thread was still appending to rows_ and externalCorrelations_:

• Truncation. Whatever had not arrived when processActivities() ran was dropped, with no error or warning. The loss is FIFO, so the trace is a plausible-looking prefix of the real one: on a 32-node graph replayed 200 times, every node shows up 54 to 58 times instead of 200.
• Heap corruption. processActivities() walks those containers and clearLogs() clears them without the mutex the appends hold, so they race the late appends. glibc catches it when it next looks at the heap (malloc(): unsorted double linked list corrupted), or the walk faults in processActivities(), or the process wedges.

The fix

Flush, stop the context, then keep flushing until the row count has been stable for three polls 10ms apart, bounded by a 5s total timeout that logs a warning and proceeds.

Once the in-flight records have arrived, nothing is appending any more, so the rest of teardown no longer races the callback thread. That is why this fixes the corruption as well as the truncation without touching the locking.

This is deliberately the smallest change that fixes both symptoms. It is a workaround for rocprofiler-sdk offering no way to wait for in-flight records at a mid-run stop (it guarantees delivery at finalize only), and should be revisited once the SDK does.

Why a poll

rocprofiler-sdk has no equivalent of CUPTI's forced flush. rocprofiler_flush_buffer() hands over only the records already emplaced in the buffer, and there is no call that waits for in-flight dispatch records to be emplaced. The alternative that gives a hard guarantee is to track correlation ids — note the highest one issued at stop and wait for its record — which is a considerably larger change. A behaviour change on the SDK side is being raised with the ROCm profiler team.

A poll rather than a fixed sleep because a sleep has to be sized for the worst case: every stop pays the full duration, and when it turns out too short the trace is silently truncated again. The poll returns as soon as the rows stop growing, about 20ms on the runs below, and the 5s bound is only reached if records are still arriving, in which case it warns.

Validation

HIP graph replay dispatching 6400 kernels (32 nodes x 200 replays) through libkineto's public API, MI355X, ROCm 10.0.0, rocprofiler-sdk 1.3.5. Arms interleaved run by run, 60s timeout per run so a wedged process is counted rather than waited on:

	runs	clean exit	segfault	hung	heap messages	complete trace	kernels (median)
main	40	13	10	17	17	0	~1400
this change	40	40	0	0	0	39	6400

One run with this change reported 6399 of 6400, without a timeout warning; main's plain-stream control arm showed the same single-kernel shortfall once, so it looks like a trace-window boundary effect rather than an undrained record. At 64000 dispatches this change is complete on 10 of 10 runs. The poll adds about 20ms to teardown (median 516ms vs 496ms end to end).

Alternatives measured

approach	failures	complete traces	kernels (median)
main, flush then stop	15 / 30	0 of 15	1874
ordering corrected only, stop then flush once	13 / 30	0 of 17	1595
hipDeviceSynchronize(), then stop and flush	5 / 20	0 of 15	1894

Correcting the ordering alone moves neither symptom (p = 0.80 on failures). hipDeviceSynchronize() guarantees the completion signals fired, not that the SDK's handler has emplaced the records.

Not in this PR

• Locking. processActivities() and clearLogs() still touch rows_ and externalCorrelations_ without their mutexes. With the wait in place that is no longer reachable in practice, but if the timeout ever expires while records are still arriving, it becomes a heap corruption again rather than a short trace.
• rows_ ownership. Rows are new'd and clearLogs() never deletes them, because the emitted activities point into them. That needs its own redesign.

Repro

Plain torch.profiler around a HIP graph replay, no ROCm-specific options needed. The standalone C++ reproducer used for the numbers above is in a comment below.

[pytorch-bot]

The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:

• ciflow/rocm

Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows.

[haishuok0525]

Attaching the standalone reproducer referenced in the description, in case it is useful for reviewing this. It drives libkineto directly through its public API, so it needs no torch and no ROCm-specific profiler options — just a GPU and a ROCm build of this repo.

It captures a HIP graph, replays it inside a trace window, and saves a chrome trace. Since the number of dispatches is known exactly (E2E_NODES * E2E_REPLAYS, 6400 by default), both symptoms are countable: whether the process survived stopTrace(), and how many kernel events the trace actually contains.

E2E_STREAM=1 issues the same number of dispatches as plain stream launches instead, which is what separates "graphs are involved" from "records are still arriving during teardown".

gdb is not in the image I was testing in, so it installs a signal handler and prints its own frames. One caveat that turned out to matter: on a corrupted heap, that handler can itself wedge, because backtrace_symbols_fd allocates. That is why the table in the description counts hangs separately from segfaults — the underlying event is the same, and a 60s timeout per run is enough to tell them apart.

Build

# libkineto
cmake -DKINETO_BACKEND=rocm -DROCM_SOURCE_DIR=/opt/rocm -DCMAKE_BUILD_TYPE=Release \
      /path/to/kineto/libkineto
make -j kineto

# the repro, against the static lib just built
K=/path/to/kineto/libkineto
hipcc -std=c++20 -O2 -c e2e_graph.cpp -I$K/include -I$K/third_party/fmt/include -o repro.o
hipcc repro.o libkineto.a -L/opt/rocm/lib -lrocprofiler-sdk -ldl -lpthread -o repro

Run

export HIP_VISIBLE_DEVICES=0
for i in $(seq 20); do
  rm -f out.json
  timeout 60 env E2E_OUT=$PWD/out.json ./repro > run$i.log 2>&1
  rc=$?
  k=$(python3 -c 'import json;t=json.load(open("out.json"));
print(sum(1 for e in t.get("traceEvents",[]) if e.get("cat")=="kernel"))' 2>/dev/null || echo 0)
  echo "run $i: rc=$rc kernels=$k/6400"
done
grep -h 'malloc()\|free():\|corrupted' run*.log | sort | uniq -c

rc=0 is a clean run, rc=139 a segfault, rc=124 a process that had to be timed out.

On current main expect a mix of clean runs, segfaults and timeouts, and kernels well under 6400 even on the clean ones — 10 of 20 clean and a mean of 1361 kernels on the last batch I ran. With this PR applied, 20 of 20 clean and 6400 every time.

e2e_graph.cpp

// End-to-end check of kineto PRs #1564 and #1567 using the real libkineto,
// not a transcription of it.
//
// Profiles a HIP graph replay through kineto's own public API and writes a
// chrome trace. Whether the PRs work is then a question about the trace file:
// does every graph kernel carry "graph id" and "graph node id", are there as
// many distinct node ids as the graph has nodes, and did the trace keep all
// the dispatches instead of being truncated at teardown.
//
// E2E_NODES   nodes in the captured graph
// E2E_REPLAYS replays inside the trace window

#include <hip/hip_runtime.h>

#include <libkineto.h>

#include <csignal>
#include <cstdio>
#include <cstdlib>
#include <execinfo.h>
#include <set>
#include <string>
#include <unistd.h>

// gdb is not available in the test image, so catch the fault in-process and
// print the frames. The point is to see which component is on the stack when
// an unpatched build dies at teardown.
extern "C" void segv_handler(int sig)
{
    void*  frames[32];
    int    n = backtrace(frames, 32);
    fprintf(stderr, "\n=== caught signal %d, %d frames ===\n", sig, n);
    fflush(stderr);
    backtrace_symbols_fd(frames, n, STDERR_FILENO);
    _exit(139);
}

#define HC(x)                                                                            \
    do                                                                                   \
    {                                                                                    \
        hipError_t e = (x);                                                              \
        if(e != hipSuccess)                                                              \
        {                                                                                \
            fprintf(stderr, "HIP error %s at %d: %s\n", #x, __LINE__,                    \
                    hipGetErrorString(e));                                               \
            exit(1);                                                                     \
        }                                                                                \
    } while(0)

__global__ void e2ekernel(float* p, int n)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if(i < n) p[i] += 1.0f;
}

static long envLong(const char* k, long d)
{
    const char* v = getenv(k);
    if(!v || !*v) return d;
    char* e = nullptr;
    long  x = strtol(v, &e, 0);
    return x > 0 ? x : d;
}

int main()
{
    signal(SIGSEGV, segv_handler);
    signal(SIGABRT, segv_handler);

    const long nodes   = envLong("E2E_NODES", 32);
    const long replays = envLong("E2E_REPLAYS", 200);
    const int  n       = 1024;

    float* buf = nullptr;
    HC(hipMalloc(&buf, n * sizeof(float)));
    HC(hipMemset(buf, 0, n * sizeof(float)));
    hipStream_t s;
    HC(hipStreamCreate(&s));

    // Capture the graph before tracing starts, the way a framework would.
    hipStream_t cap;
    HC(hipStreamCreate(&cap));
    HC(hipStreamBeginCapture(cap, hipStreamCaptureModeGlobal));
    for(long i = 0; i < nodes; ++i)
        hipLaunchKernelGGL(e2ekernel, dim3(1), dim3(64), 0, cap, buf, n);
    hipGraph_t g = nullptr;
    HC(hipStreamEndCapture(cap, &g));
    hipGraphExec_t exec = nullptr;
    HC(hipGraphInstantiate(&exec, g, nullptr, nullptr, 0));

    // Warm up outside the trace window.
    for(int i = 0; i < 5; ++i) HC(hipGraphLaunch(exec, s));
    HC(hipStreamSynchronize(s));

    libkineto_init(false, true);
    std::set<libkineto::ActivityType> types;
    auto&                             profiler = libkineto::api().activityProfiler();
    libkineto::api().initProfilerIfRegistered();
    profiler.prepareTrace(types);

    for(int i = 0; i < 5; ++i) HC(hipGraphLaunch(exec, s));
    HC(hipStreamSynchronize(s));

    // E2E_STREAM=1 issues the same number of dispatches as plain launches, to
    // separate "graphs are involved" from "records arrive during stopTrace".
    const bool plain = getenv("E2E_STREAM") != nullptr;
    printf("start trace: %ld dispatches expected (%s)\n",
           nodes * replays,
           plain ? "plain stream launches" : "graph replays");
    profiler.startTrace();

    if(plain)
        for(long i = 0; i < nodes * replays; ++i)
            hipLaunchKernelGGL(e2ekernel, dim3(1), dim3(64), 0, s, buf, n);
    else
        for(long i = 0; i < replays; ++i) HC(hipGraphLaunch(exec, s));
    HC(hipDeviceSynchronize());

    auto trace = profiler.stopTrace();
    printf("activities in trace: %zu\n", trace->activities()->size());

    const char* out = getenv("E2E_OUT");
    std::string path = (out && *out) ? out : "/tmp/e2e_trace.json";
    trace->save(path);
    printf("saved %s\n", path.c_str());
    return 0;
}

[ryanzhang22]

Is continuously polling the only way to wait for rocprofiler-sdk to potentially finish draining? In the CUDA paths we have a CUPTI API which blocks and forces a flush, even for incomplete buffer records.

I see that the other changes have made concurrent modification on the rows vector safer, which is a good change, but I'm wondering if there's any way we can have a stronger guarantee on completeness here?

[haishuok0525]

Thanks, and sorry for the slow reply.

As far as I can tell, rocprofiler-sdk has no equivalent of CUPTI's forced flush. rocprofiler_flush_buffer() only hands over records already emplaced in the buffer, and a dispatch record is emplaced later, by the SDK's own signal handler thread once the completion signal fires. The SDK guarantees delivery at finalize, but not at a mid-run stop, which is the case we need here.

The option with a hard guarantee would be tracking correlation ids: note the highest one issued at stop and wait until its record comes through. That's a considerably bigger change. @mwootton is raising a behaviour change with the SDK team instead.

In the meantime I've cut this PR down to just the bounded wait: flush, stop, then flush until the row count has been stable for 30ms, with a 5s total timeout that warns. That alone fixes both the truncation and the crash on the repro (numbers are in the updated description). The locking changes are out of this PR and can follow separately if wanted.

[ryanzhang22]

If we want to simplify this, wouldn't the safest thing be to keep the changes that guard concurrent modification instead? The sleep loop can still let stuff past if we're unlucky, and I'd much rather live with the consequences of dropping a few events than segfaulting a training job.

[haishuok0525]

That's fair, and some history would help here.

The first version of this PR did both: stop the context, flush until the row count stopped growing, and take rowsMutex_ / externalCorrelationsMutex_ in processActivities() and clearLogs().

After discussing it with @mwootton, I narrowed it for two reasons. First, the root cause is on the SDK side: rocprofiler-sdk has no reliable drain at a mid-run stop, and that is being raised with the SDK team. We didn't want the mutex to read as the fix for that, when what it really does is make a racy teardown survivable. Second, once the wait drains the in-flight records, nothing is appending to rows_ when processActivities() and clearLogs() run, so the normal path doesn't need the lock. I checked that before dropping it: over 40 runs with the wait alone there were 0 crashes, against 27 of 40 crashed or hung on main.

You're right that the loop is a heuristic, though. If the 5s timeout expires, or a long-running kernel completes after the row count already looked stable, a record can still land while those functions are running, and without the lock that's heap corruption again rather than a few missing events. I agree with your trade-off: a slightly short trace is much better than taking down a training job. I'm happy to put the locking back in this PR. @mwootton, any objection?

[mwootton]

So the 816 and 817 lines are swapped in the original file. Trivial mistake.

rocprofiler_stop_context()
rocprofiler_flush_buffer()

is the correct order.

I thought that should be sufficient. I had an agent dig through the rocprofiler_flush_buffer() implementation and it appears not (see the TLDR above). This is the same dumb situation as roctracer; and it probably needs the same workaround. The RoctracerLogger used to track the highest outstanding correlation id at the time of the 'stop' and waited for that record to come through. Rocprofiler-sdk can guaranty record delivery on finalize, but not mid flight stop; which is what we need here.

The heap corruption if from the late callbacks happening after tracing is stopped; those should not be happening if the flushing is done correctly. So the stop -> flush -> readout phases SHOULD be exclusive without the mutex. We can add the mutex for safety. But in this case it is also "papering over" the bigger problem.

It's a pretty big change to adopt the correlation tracking. I will inquire about a behavior change in the sdk. There are other backends that don't have this drain problem; so this is a choice.

[haishuok0525]

So the 816 and 817 lines are swapped in the original file.

Thanks Michael. Updated as you suggested: this PR now only fixes the ordering and adds the wait. It flushes, stops the context, then keeps flushing until the row count is stable, with a 5s total timeout that logs a warning. The locking changes are gone.

The wait alone is enough on the repro. Over 40 runs, main segfaulted or hung on 27 and never produced a complete trace. With this change all 40 exited cleanly and 39 were complete; the one exception was short by a single kernel. Details are in the description.

Agreed that correlation tracking or an SDK-side drain is the real fix. Thanks for raising it with the SDK team.

[ryanzhang22]

isn't this still wrong order (should be stop, then flush)? tbh we can probably just rocprofiler_stop_context and then let the loop below handle flush

[haishuok0525]

The flush before stop came from the same discussion: it follows the flush / disable / flush pattern in the SDK samples, so whatever is already buffered gets delivered before the context stops. You're right that it's redundant here, since the loop flushes after the stop anyway. I'll drop it, so it becomes stop, then the flush loop.