103 lines
3.6 KiB
C++
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;
|
|
}
|