diff --git a/src/cupti.cpp b/src/cupti.cpp index 1e1dd8d..a2c7f87 100644 --- a/src/cupti.cpp +++ b/src/cupti.cpp @@ -37,9 +37,27 @@ static std::atomic g_bufferCycle{0}; // runtime calls) thread_local uint32_t runtimeEnterCorrelationId = 0; -// Thread-local rate limiter for callback probes (default 100/sec, -// configurable via PARCAGPU_RATE_LIMIT). +// Configured probe rates (tokens/sec), read from env once in init_debug. +// Globals because the buckets below are thread_local: setRate from init_debug +// would only reach the single thread that ran it. +std::atomic g_eagerRateLimit{100.0}; +std::atomic g_graphRateLimit{20.0}; + +// Thread-local rate limiters. Graph launches get a separate, lower bucket +// (PARCAGPU_GRAPH_RATE_LIMIT) so graph replays don't flood cupti_events and +// eager bursts can't starve graph sampling. thread_local TokenBucket callbackLimiter(100.0); +thread_local TokenBucket graphCallbackLimiter(20.0); + +// Copy the configured rates into this thread's buckets, once per thread. +static inline void applyRateLimitsOnce() { + thread_local bool applied = false; + if (!applied) { + callbackLimiter.setRate(g_eagerRateLimit.load(std::memory_order_relaxed)); + graphCallbackLimiter.setRate(g_graphRateLimit.load(std::memory_order_relaxed)); + applied = true; + } +} // --------------------------------------------------------------------------- // PC sampling probabilistic control. @@ -179,7 +197,15 @@ void init_debug() { if (rateEnv != nullptr) { double rate = atof(rateEnv); if (rate > 0) { - callbackLimiter.setRate(rate); + g_eagerRateLimit.store(rate, std::memory_order_relaxed); + } + } + + const char *graphRateEnv = getenv("PARCAGPU_GRAPH_RATE_LIMIT"); + if (graphRateEnv != nullptr) { + double rate = atof(graphRateEnv); + if (rate > 0) { + g_graphRateLimit.store(rate, std::memory_order_relaxed); } } @@ -642,7 +668,7 @@ class CuptiProfiler : public proton::Singleton { return; } - // Check if this is a graph launch (never rate limit these) + // Check if this is a graph launch (rate-limited via a separate bucket). bool isGraphLaunch = false; if (signedCbid < 0) { // Driver API: cuGraphLaunch = 514, cuGraphLaunch_ptsz = 515 @@ -658,15 +684,17 @@ class CuptiProfiler : public proton::Singleton { CUPTI_RUNTIME_TRACE_CBID_cudaGraphLaunch_ptsz_v10000); } - // Rate limit probes using token bucket. Skip rate limiting for graph - // launches (they share one correlation ID across many kernels) and when - // PC sampling is active (every kernel needs its correlation callback so - // PC samples can be matched with CPU stacks on the agent side). - if (!isGraphLaunch && !g_pcSamplingState.active) { - if (!callbackLimiter.tryAcquire()) { - DEBUG_PRINTF( - "[PARCAGPU] Rate limited: skipping probe for correlationId=%u\n", - correlationId); + // Rate limit via per-class buckets (eager vs graph). Skip when PC + // sampling is active: every kernel needs its correlation callback to + // match PC samples to CPU stacks on the agent side. + if (!g_pcSamplingState.active) { + applyRateLimitsOnce(); + TokenBucket &limiter = + isGraphLaunch ? graphCallbackLimiter : callbackLimiter; + if (!limiter.tryAcquire()) { + DEBUG_PRINTF("[PARCAGPU] Rate limited: skipping probe for " + "correlationId=%u (graph=%d)\n", + correlationId, isGraphLaunch); return; } }