Files
2026-07-13 12:24:33 +08:00

103 lines
3.6 KiB
C++

// SPDX-License-Identifier: Apache-2.0
#include "event_recorder.h"
#include <chrono>
#include <utility>
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
// Monotonic wall-clock reader.
//
// Why not just `system_clock::now()`? `system_clock` maps to CLOCK_REALTIME,
// which NTP / chrony can slew backward by tens of microseconds to stay synced
// with upstream. CUDA host callbacks fire at arbitrary wall moments, so two
// consecutive callbacks (e.g. MP_STORE_START then MP_STORE_END on the same
// stream) can straddle a backward slew and land with end_ts < start_ts.
// Jaeger then renders span duration as unsigned 64-bit subtraction, producing
// the characteristic ~213503982d span duration (= 2^64 microseconds).
//
// Fix: anchor once to the (system_clock, steady_clock) pair at first call.
// Every subsequent timestamp is computed as
// epoch_sys + (steady_now - epoch_steady)
// which is monotonic (steady_clock == CLOCK_MONOTONIC on Linux) while
// remaining expressed in Unix-epoch seconds so downstream consumers don't
// notice any change.
static double wall_clock_time() {
static const auto epoch_sys = std::chrono::system_clock::now();
static const auto epoch_steady = std::chrono::steady_clock::now();
auto now_steady = std::chrono::steady_clock::now();
auto since_epoch = epoch_sys.time_since_epoch() + (now_steady - epoch_steady);
return std::chrono::duration<double>(since_epoch).count();
}
// ---------------------------------------------------------------------------
// EventRecorder
// ---------------------------------------------------------------------------
EventRecorder& EventRecorder::instance() {
static EventRecorder recorder;
return recorder;
}
void EventRecorder::push(PendingEvent* event) {
{
std::lock_guard<std::mutex> lock(mutex_);
buffer_.push_back(std::move(*event));
}
delete event;
}
std::vector<PendingEvent> EventRecorder::drain() {
std::lock_guard<std::mutex> lock(mutex_);
std::vector<PendingEvent> result;
result.swap(buffer_);
return result;
}
// ---------------------------------------------------------------------------
// CUDA host callback — runs on a CUDA driver thread, no GIL.
// ---------------------------------------------------------------------------
static void
#ifndef USE_ROCM
CUDART_CB
#endif
event_host_callback(void* data) {
auto* event = static_cast<PendingEvent*>(data);
event->timestamp = wall_clock_time();
EventRecorder::instance().push(event);
}
// ---------------------------------------------------------------------------
// Free functions for pybind11
// ---------------------------------------------------------------------------
void record_event_on_stream(
int64_t cuda_stream_ptr, const std::string& event_type_name,
const std::string& session_id,
const std::unordered_map<std::string, std::string>& str_metadata,
const std::unordered_map<std::string, int64_t>& int_metadata) {
auto* event = new PendingEvent{
event_type_name, session_id, 0.0, str_metadata, int_metadata,
};
auto stream = reinterpret_cast<lmcache_stream_t>(
static_cast<uintptr_t>(cuda_stream_ptr));
LMCACHE_LAUNCH_HOST_FUNC(stream, event_host_callback, event);
}
DrainResult drain_recorded_events() {
auto events = EventRecorder::instance().drain();
DrainResult result;
result.reserve(events.size());
for (auto& e : events) {
result.emplace_back(std::move(e.event_type_name), std::move(e.session_id),
e.timestamp, std::move(e.str_metadata),
std::move(e.int_metadata));
}
return result;
}