# mypy: allow-untyped-defs from __future__ import annotations import ctypes import json import logging import threading import time from collections.abc import Callable, Iterable, Mapping from typing import Any, cast, NamedTuple, TYPE_CHECKING import numpy as np from cupti.cupti import ActivityKind # pyrefly: ignore[missing-import] import torch from . import cupti_python from .records import ( Ctype, CudaEvent, FIELD_CTYPE, FIELD_REGISTRY, Kernel, STRING_FIELDS, Sync, ) # A registration request: either a plain iterable of activity kinds (meaning "all # fields"), or a field map {kind: iterable of field ids | "all"} selecting specific # fields per kind. The monitor demuxes the selected fields to columns. ActivitiesSpec = Mapping[ActivityKind, "Iterable[int] | str"] | Iterable[ActivityKind] _PY_PROFILER = torch._C._profiler # The native CUPTI buffer-pool / layout-capture module (C++ side of the monitor). _cupti_monitor_native = _PY_PROFILER._cupti_monitor if TYPE_CHECKING: from typing_extensions import Self logger = logging.getLogger(__name__) # Buffers are a recycling pool bounded by peak concurrent demand, so the count # only keeps climbing if the worker can't drain completed buffers as fast as # CUPTI fills them. Warn once past this many outstanding buffers (1GB at the # default 4MB size) as a sign of that backpressure. _OUTSTANDING_WARN_THRESHOLD = 256 # flush(sync=True) fences at a SYNC point: it enables SYNCHRONIZATION, captures # CUPTI's clock, device-syncs (which produces a SYNCHRONIZATION record at a # timestamp past that point), waits until the native decoder reports a sync record # that recent, then disables SYNCHRONIZATION again. CUPTI delivers buffers in fill # order, so seeing the sync record means everything before it is delivered too. A # device sync -- unlike a tracer kernel -- adds no kernel, no cudaLaunchKernel, and # no dispatcher op to the trace; and enabling SYNCHRONIZATION only for the fence # means the session doesn't record every sync between flushes. KIND + END are the # fields the fence decodes. _FENCE_KIND = ActivityKind.SYNCHRONIZATION _FENCE_END_FIELD = Sync.END.id _FENCE_FIELDS = frozenset({Sync.KIND.id, _FENCE_END_FIELD}) def _has_active_cuda_context() -> bool: try: from cuda.bindings import ( # pyrefly: ignore[missing-import] driver as cuda_driver, ) except ImportError: return False rc, ctx = cuda_driver.cuCtxGetCurrent() if rc == cuda_driver.CUresult.CUDA_SUCCESS: return ctx is not None if rc == cuda_driver.CUresult.CUDA_ERROR_NOT_INITIALIZED: return False raise RuntimeError(f"cuCtxGetCurrent failed with rc={rc}") def _cuda_version_string() -> str: return torch.version.cuda or "" def _deref_cstr(ptr: int) -> str: if not ptr: return "" value = ctypes.cast(ptr, ctypes.c_char_p).value return value.decode(errors="replace") if value is not None else "" # A CUPTI subscriber callback, identified by its (domain, cbid) pair, and the handlers # registered for one. See CuptiMonitor.register_callback_handler / _dispatch_callback. _CallbackKey = tuple[int, int] _CallbackFns = tuple[Callable[..., None], ...] class _CallbackHandler(NamedTuple): """A registered consumer of one CUPTI *subscriber callback*: the ``(domain, cbid)`` it wants and the ``fn(domain, cbid, cbdata)`` to invoke. Subscriber callbacks are a separate axis from activity records: they fire synchronously on the application thread that made the CUDA call, not on the decode worker, and carry a raw ``CUpti_CallbackData`` pointer rather than decoded columns. Returned by ``CuptiMonitor.register_callback_handler``; hand it back to unregister, mirroring ``register`` / ``unregister`` for activity observers.""" domain: int cbid: int fn: Callable[..., None] class _Observer: """A registered consumer of the monitor's records: the activity kinds it requested, its per-kind field selection (``{kind: frozenset(field_ids)}``), and its ``callback(columns)`` -- invoked once per drain (at flush time) with the demuxed columns sliced to its selection (see ``CuptiMonitor.register``).""" def __init__( self, activities: Iterable[ActivityKind], fields: Mapping[ActivityKind, frozenset[int]], callback: Callable[..., None], ) -> None: self.activities: frozenset[ActivityKind] = frozenset(activities) # activity -> the set of field ids wanted for it (the columns to demux). self.fields: dict[ActivityKind, frozenset[int]] = dict(fields) self.callback = callback class _SynchronizedClock: """Maps CUPTI record timestamps to unix-epoch ns on kineto's axis. CUPTI stamps records with cuptiGetTimestamp, which is CLOCK_REALTIME -- the very clock kineto's cpu_ops land on (kineto reads the approx clock for cpu_ops and converts it to unix == CLOCK_REALTIME at serialization). So a record's timestamp already IS its unix time: it is passed through unchanged (identity, offset 0). Records and cpu_ops then share the identical realtime clock, so a runtime call's start is its true realtime -- always >= the start of the cpu_op that issued it, with no clock estimate in the loop. calibrate() requires this and verifies it by bracketing one native read between two clock_gettime(CLOCK_REALTIME) reads; a host whose CUPTI clock is something else raises. The exception is the timestamp callback: when it engages CUPTI stamps records with the profiler's approx clock instead -- kineto's *source* clock -- so records arrive as approx ticks and are mapped to unix via the same _ApproximateClockToUnixTimeConverter kineto uses (a single-slope line -- median rate over 1001 samples -- so two evaluations recover its slope exactly; recovered once so the hot path never calls the converter per record). That path lifts the CLOCK_REALTIME requirement and is the ONLY one that touches the converter. calibrate() reads the native clock through an injected callable, so the conversion math runs (and is tested) without a live CUPTI session. """ def __init__(self) -> None: self._callback_active = False self._native_is_realtime = False self._native_now: Callable[[], int] = lambda: 0 # Session-start anchor. 0 until calibrated -> conversion is identity. self._native_ns = 0 self._unix_ns = 0 self._approx_ns = 0 self._scale = 1.0 def calibrate( self, *, callback_active: bool, native_now: Callable[[], int], converter: Any = None, ) -> None: self._callback_active = callback_active self._native_now = native_now # cuptiGetTimestamp is CLOCK_REALTIME iff it lands inside a clock_gettime bracket. rt_lo = time.clock_gettime_ns(time.CLOCK_REALTIME) cupti_now = native_now() rt_hi = time.clock_gettime_ns(time.CLOCK_REALTIME) self._native_is_realtime = cupti_now != 0 and rt_lo <= cupti_now <= rt_hi if callback_active: # Records are approx-clock ticks. Build the converter first (it only maps approx # reads taken after construction), then bracket the native read -- needed for PM # frames, which stay on the native clock -- between two approx reads and pair at # the midpoint; recover the converter's slope for vectorized column conversion. if converter is None: converter = _PY_PROFILER._ApproximateClockToUnixTimeConverter() approx_before = _PY_PROFILER._get_approximate_time() self._native_ns = ( time.clock_gettime_ns(time.CLOCK_REALTIME) if self._native_is_realtime else native_now() ) approx_after = _PY_PROFILER._get_approximate_time() self._approx_ns = (approx_before + approx_after) // 2 self._unix_ns = converter.to_unix_ns(self._approx_ns) span = 10**12 self._scale = ( converter.to_unix_ns(self._approx_ns + span) - self._unix_ns ) / span else: if not self._native_is_realtime: raise RuntimeError( "cupti_monitor requires cuptiGetTimestamp to be CLOCK_REALTIME so " "record timestamps share kineto's unix clock; this host reports a " "different clock source, unsupported until the CUPTI timestamp " "callback is available" ) # cuptiGetTimestamp == CLOCK_REALTIME == unix: records already carry unix time, # so pass them through (identity). No approx clock or converter needed. self._native_ns = time.clock_gettime_ns(time.CLOCK_REALTIME) self._approx_ns = 0 self._unix_ns = self._native_ns self._scale = 1.0 # A timestamp is in one of two source domains -- approx-clock ticks or CLOCK_REALTIME/unix -- # and both map to unix; these are the two conversion primitives (0, and the uncalibrated # state, always map to itself). Which domain a record is in depends on the timestamp callback, # so the record dispatch lives in CuptiMonitor.convert_time[_array], not here. def convert_approx(self, value: int) -> int: # Approx-clock tick -> unix ns via the recovered converter slope. if value == 0 or self._native_ns == 0: return value return self._unix_ns + int((value - self._approx_ns) * self._scale) def convert_unix(self, value: int) -> int: # Native cuptiGetTimestamp clock (CLOCK_REALTIME) -> unix ns: a constant offset. if value == 0 or self._native_ns == 0: return value return value - self._native_ns + self._unix_ns def convert_approx_array(self, values: np.ndarray) -> np.ndarray: # Vectorized convert_approx. Keep the delta in float (small magnitude) and add the unix # anchor as int64 -- adding at the ~1e18 unix magnitude in float64 would lose ~us. out = values.astype(np.int64) if self._native_ns == 0: return out ticks = (out - self._approx_ns).astype(np.float64) delta = (ticks * self._scale).astype(np.int64) return np.where(out == 0, out, self._unix_ns + delta) def convert_unix_array(self, values: np.ndarray) -> np.ndarray: # Vectorized convert_unix. out = values.astype(np.int64) if self._native_ns == 0: return out offset = self._unix_ns - self._native_ns return np.where(out == 0, out, out + offset) def now_record_ns(self) -> int: # Current record-clock value: the approx clock when the callback is active, else the # native cuptiGetTimestamp clock (read cheaply via clock_gettime when it is realtime). # 0 before calibration. if self._native_ns == 0: return 0 if self._callback_active: return _PY_PROFILER._get_approximate_time() if self._native_is_realtime: return time.clock_gettime_ns(time.CLOCK_REALTIME) return self._native_now() def reset(self) -> None: self._native_ns = 0 self._native_now = lambda: 0 @property def unix_anchor_ns(self) -> int: return self._unix_ns @property def approx_anchor_ns(self) -> int: return self._approx_ns class CuptiMonitor: """Process-wide CUPTI monitor / multiplexer singleton. Like PmSampler, ``CuptiMonitor()`` returns the one instance (constructed on first call); its settings are snapshotted from the class config -- set via ``cupti_monitor.configure()`` before first use -- not from constructor args.""" _instance: CuptiMonitor | None = None _instance_lock = threading.Lock() # Process-wide settings the singleton snapshots when first constructed; set via # configure() (first-come-first-serve, no env var), defaults otherwise. Both cadences # default to -1 (no background flush / drain -- the caller drives flush()). # # use_approx_timestamps defaults OFF; when on, the per-subscriber timestamp callback # re-times HOST and DEVICE records onto the profiler's approx clock (kineto's timebase). # Opt-in via configure(use_approx_timestamps=True). _buffer_size: int = 4 * 1024 * 1024 _background_flush_period_s: float = -1.0 _background_drain_period_s: float = -1.0 _use_approx_timestamps: bool = False _configured: bool = False def __new__(cls) -> Self: with cls._instance_lock: if cls._instance is None: inst = super().__new__(cls) inst._init() cls._instance = inst # the live monitor snapshotted the config; lock it cls._configured = True return cls._instance def _init(self) -> None: cls = type(self) # The monitor is the engine and the multiplexer: it owns the single CUPTI # subscription + buffer pool + native decode worker, which demuxes each # completed buffer into columns; the monitor drains those columns at flush # time and hands every observer the columns it selected. It reaches CUPTI # only through the self._cupti.activity_* wrappers -- no ctypes here. # # It uses CUPTI's v2 user-defined-record API: a subscriber + per-field # selection, decoded columnar against a record layout computed from the # field-size spec (no captured layout needed). This requires libcupti >= 13.2. # # Per-buffer pool size (bytes), default 4 MiB. Bigger buffers complete less often # (fewer worker wakeups, lower overhead) at the cost of more pinned host memory and # coarser delivery. Set process-wide via cupti_monitor.configure(). self.buffer_size = cls._buffer_size # Two independent cadences (seconds); they control different things, so they are # separate knobs. # # background_flush_period_s (default -1) -- how often the native decode thread # self-drives cuptiActivityFlushAll, handing CUPTI's completed buffers to the # decoder. Sign-encoded: # >= 0 -> the decode thread self-flushes on this cadence (0 = continuously). # < 0 -> NO self-flush; the caller must drive flush() itself (e.g. at end of # step). This is the default, and the escape hatch for a libcupti/libnvperf # HES thread-safety bug: cuptiActivityFlushAll drives CUPTI's HW-trace # processing against live collection state and can wild-write the host # heap when it overlaps concurrent host activity (e.g. NCCL collective # setup). The racy op is the flush, NOT the decode -- the decoder keeps # decoding delivered buffers off-thread regardless. Driving flush() only # from the quiescent foreground avoids the race. # # background_drain_period_s (default -1) -- how often the Python thread drains the # decoded columns and dispatches them to observers (GIL work). Same sign-encoding: # >= 0 -> a background thread drains on this cadence (0 = continuously). # < 0 -> NO background drain thread; the caller drives drain via flush(). self.background_flush_period_s = cls._background_flush_period_s self.background_drain_period_s = cls._background_drain_period_s self._cupti = cupti_python.pylibcupti() # The CUPTI subscriber handle. self._subscriber: int | None = None self._latency_enabled = False self._device_ts_enabled = False # Layout state -- a function of registration, recomputed only when the # The fields enabled per kind on the subscriber (a function of the observer # field union, recomputed only on register/deregister, never per buffer). The # record byte layout is NOT tracked here -- each completed buffer carries # CUPTI's own captured layout (ppRecordLayouts) that records.decode reads. self._enabled: dict[int, frozenset[int]] = {} self._lock = threading.Lock() self._started = False self._callbacks_registered = False # --- subscriber callbacks: a second axis, independent of activity records --- # Handlers keyed by (domain, cbid). _dispatch_callback runs on the application # thread inside the CUDA call, so it must not take self._lock (register and # _apply_selection hold that): it reads the immutable tuple values below, which # mutation replaces wholesale under _callback_lock. self._callback_handlers: dict[_CallbackKey, _CallbackFns] = {} # Per-(domain, cbid) arm refcount, so consumers wanting the same callback over # different windows don't disable each other. Enabled on 0->1, disabled on 1->0. self._callback_armed: dict[_CallbackKey, int] = {} self._callback_lock = threading.Lock() # Our single CUpti_CallbackFunc trampoline; must outlive the subscription. self._callback_trampoline: Any = None self._drain_stop = threading.Event() self._drain_and_dispatch_thread: threading.Thread | None = None # Serializes _drain_and_dispatch: the native decoder accumulates columns # GIL-free; Python drains them here. Only ever one driver at a time (the # foreground caller OR the background flush loop, never both), but the lock # keeps a stray concurrent drain from interleaving dispatch. self._drain_lock = threading.Lock() self._observers: list[_Observer] = [] # {external_id: metadata blob} drained alongside the decoded records (see # drain_decoded). Accumulated here until an observer's external-correlation # join consumes it via take_external_metadata(); the blob is attached onto # the kernel event keyed by the same external id. Guarded by _lock. self._external_metadata: dict[int, str] = {} self._next_external_id = 1 # All subsystems push external-correlation ids on ONE CUPTI kind, so CUPTI # inserts a single EXTERNAL_CORRELATION record (that kind's stack top) per op # -- i.e. it tags a kernel with only the *innermost* active id; the enclosing # ids are recovered from our own bookkeeping here. _id_chains[id] is the full # active stack (outermost..id) captured when `id` was pushed, so a consumer # maps a kernel's innermost id to every context active for that op (see # external_id_chain) and picks the one it owns (by membership in its own # name/metadata map). The live LIFO is per-thread (CUPTI's external-correlation # stacks are per-thread) in _push_tls.stack. A popped id's chain is kept a # couple of dispatch cycles -- its activity records arrive after the pop -- then # dropped via the _chains_gc_* generations. _id_chains/_chains_gc_* are guarded # by _lock; the per-thread stack needs no lock. self._push_tls = threading.local() self._id_chains: dict[int, tuple[int, ...]] = {} self._chains_gc_pending: list[int] = [] self._chains_gc_ready: list[int] = [] # Record-timestamp -> unix conversion lives in the clock; the monitor delegates # convert_time / now_record_ns to it and calibrates it in start(). Records normally # arrive on the native (realtime) clock; configure(use_approx_timestamps=True) puts # them directly on the approx clock via CUPTI's per-subscriber timestamp callback # (opt-in, sole-subscriber only). self._timestamp_callback_enabled = cls._use_approx_timestamps self._clock = _SynchronizedClock() self._timestamp_callback_active = False # Snapshot of the native pool size taken before stop() frees it, so # stats() stays meaningful after the monitor has been stopped. self._final_allocated_buffers = 0 self._outstanding_warned = False self._dropped_records = 0 # Opt-in PM sampling (true SM-active % / DRAM-throughput % counters): the monitor registers # each requesting observer as a consumer of the current device's per-device PmSampler # (only one PM session per device is possible), polls it on the flush cadence, and converts # each polled frame's raw CUPTI-clock timestamps into the trace clock before pushing to the # observer's sink (the sampler is clock-agnostic; conversion lives here, where the clock base # does). Each consumer brings its own metrics; the shared session samples their union, starts # on the first consumer, and disables after the last. self._pm_consumers maps an observer's # sink to its sampler handle so poll/release can address it. self._pm_consumers: dict[Callable[[dict[str, Any]], None], Any] = {} self._pm_sampler: Any = None # Serializes PM add/poll/remove so a flush-thread poll never decodes the collector while the # foreground is tearing it down (concurrent decode on one collector is unsafe). self._pm_lock = threading.Lock() def _ensure_subscribed(self) -> int: """Take the process's one CUPTI subscription if we don't already hold it, and return the handle. Shared by the activity path (:meth:`register_callbacks`) and the subscriber-callback path, which needs the subscription but none of the user-defined-record machinery -- so a callback-only consumer pays for neither UDR mode nor the decode worker. """ if self._subscriber is not None: return self._subscriber version = self._cupti.get_version() if version < cupti_python.LIBCUPTI_MIN_VERSION: raise RuntimeError( "CuptiMonitor requires libcupti >= " f"{cupti_python.LIBCUPTI_MIN_VERSION}; loaded " f"{cupti_python.LIBCUPTI_SONAME} reports {version}" ) # CUPTI allows exactly one callback per subscriber, so we always install our own # switchboard and fan out from there. Passing it explicitly -- rather than letting # subscribe() fall back to the module-level no-op -- keeps a consumer that swaps that # global from clobbering our dispatch. self._callback_trampoline = cupti_python._CB_FUNC(self._dispatch_callback) # (A prior consumer that left CUPTI attached -- e.g. Kineto -- can make # cuptiSubscribe_v2 fail with CUPTI_ERROR_MULTIPLE_SUBSCRIBERS; run such consumers # with TEARDOWN_CUPTI=1 so they release CUPTI on teardown rather than us finalizing # global state.) Subscribe solo only when the timestamp callback is opted in: CUPTI # honors it only while multiple subscribers are NOT allowed. Otherwise allow # coexistence (default). try: self._subscriber = self._cupti.subscribe( allow_multiple=not self._timestamp_callback_enabled, callback=self._callback_trampoline, ) except cupti_python.CuptiError as e: self._callback_trampoline = None if self._timestamp_callback_enabled: # We requested sole-subscriber mode only because the approx-clock timestamp # callback needs it; another CUPTI consumer is likely attached. Point the user # at the opt-out so they can coexist (the callback is then off). raise RuntimeError( "cupti_monitor could not subscribe as the sole subscriber, which the " "opt-in approx-clock timestamp callback requires (another CUPTI consumer " "is likely attached). Retry with use_approx_timestamps=False to allow " f"coexisting subscribers: {e}" ) from e raise # Arm the per-subscriber approx-clock timestamp callback right after subscribe, before # any user-defined record can be produced. It stays armed for the life of the # subscription (see _release_subscriber_if_idle), not just the activity session. self._timestamp_callback_active = self._try_arm_approx_timestamp_callback( self._subscriber ) return self._subscriber def register_callbacks(self) -> None: if self._callbacks_registered: return native = _cupti_monitor_native request_addr = native.buffer_request_callback_address() complete_addr = native.buffer_complete_callback_address() # The activity API is subscription-scoped: subscribe, turn on user-defined # records, and register the v2 buffer callbacks. subscriber = self._ensure_subscribed() self._cupti.arm_user_defined_records(subscriber, request_addr, complete_addr) self._callbacks_registered = True # --- subscriber callbacks ---------------------------------------------- def _dispatch_callback( self, _userdata: int, domain: int, cbid: int, cbdata: int ) -> None: """The subscription's single ``CUpti_CallbackFunc``: fan out to the handlers registered for ``(domain, cbid)``. Runs synchronously on the application thread that made the CUDA call, so it takes no lock -- registration swaps in a fresh immutable tuple rather than mutating one in place -- and lets no exception escape: a raise would propagate into CUPTI's C dispatch, and one failing handler must not silence the others. """ for fn in self._callback_handlers.get((domain, cbid), ()): try: fn(domain, cbid, cbdata) except Exception: logger.exception( "CUPTI callback handler failed (domain=%d cbid=%d)", domain, cbid ) def register_callback_handler( self, domain: int, cbid: int, fn: Callable[..., None] ) -> _CallbackHandler: """Register ``fn(domain, cbid, cbdata)`` for one CUPTI ``(domain, cbid)`` subscriber callback; returns a token for :meth:`unregister_callback_handler`. Registering brings the shared subscription up and holds it, exactly as registering an activity observer does, so a callback-only consumer can keep CUPTI open with no observers present. It does NOT enable the callback -- call :meth:`arm_callback` for that, so a consumer can scope delivery to a window (a CUDA-graph capture, say) while staying registered across it. ``fn`` runs on the application thread inside the CUDA call, so it must be cheap and non-blocking; ``cbdata`` is the raw ``CUpti_CallbackData`` pointer as an int. Exceptions are logged and swallowed. """ handler = _CallbackHandler(domain, cbid, fn) self._ensure_subscribed() with self._callback_lock: key = (domain, cbid) self._callback_handlers[key] = self._callback_handlers.get(key, ()) + (fn,) return handler def unregister_callback_handler(self, handler: _CallbackHandler) -> None: """Remove a handler from :meth:`register_callback_handler`. Idempotent. When the last handler for a ``(domain, cbid)`` goes away the callback is disabled outright, whatever its arm refcount was -- delivering to nobody would cost the application thread for nothing. Releases the shared subscription once no handlers and no activity observers remain. """ key = (handler.domain, handler.cbid) with self._callback_lock: remaining = tuple( f for f in self._callback_handlers.get(key, ()) if f is not handler.fn ) if remaining: self._callback_handlers[key] = remaining elif self._callback_handlers.pop(key, None) is not None: self._callback_armed.pop(key, None) if self._subscriber is not None: self._cupti.enable_callback(self._subscriber, *key, False) self._release_subscriber_if_idle() def arm_callback(self, domain: int, cbid: int) -> None: """Enable one ``(domain, cbid)`` callback, refcounted: CUPTI is told to enable on the first arm and to disable on the last :meth:`disarm_callback`, so consumers scoping the same callback to different windows don't disable each other.""" subscriber = self._ensure_subscribed() with self._callback_lock: key = (domain, cbid) count = self._callback_armed.get(key, 0) self._callback_armed[key] = count + 1 if count == 0: self._cupti.enable_callback(subscriber, domain, cbid, True) def disarm_callback(self, domain: int, cbid: int) -> None: """Drop one arm refcount for ``(domain, cbid)``, disabling the callback at zero. No-op when it is not armed.""" with self._callback_lock: key = (domain, cbid) count = self._callback_armed.get(key, 0) if count == 0: return if count > 1: self._callback_armed[key] = count - 1 return del self._callback_armed[key] if self._subscriber is not None: self._cupti.enable_callback(self._subscriber, domain, cbid, False) def _release_subscriber_if_idle(self) -> None: """Unsubscribe once neither activity observers nor callback handlers need CUPTI. The subscription is shared, so it outlives the activity session in both directions: ``stop()`` disarms user-defined records but leaves the handle up while a handler holds it, and a handler registered before any profiling keeps it up for a later ``start()``. """ if self._observers or self._started: return with self._callback_lock: subscriber = self._subscriber if self._callback_handlers or subscriber is None: return for key in self._callback_armed: self._cupti.enable_callback(subscriber, *key, False) self._callback_armed.clear() # Restore CUPTI's default CPU timer before releasing the subscription. if self._timestamp_callback_active: self._cupti.disarm_approx_timestamp_callback(subscriber) self._timestamp_callback_active = False self._cupti.unsubscribe(subscriber) self._subscriber = None self._callback_trampoline = None def start(self) -> None: if self._started: raise RuntimeError("CUPTI monitor is already started") _cupti_monitor_native.reset_buffers() _cupti_monitor_native.configure_buffers(self.buffer_size) self.register_callbacks() # Put activity records on kineto's unix timeline via the clock (see _SynchronizedClock): # normally cuptiGetTimestamp == CLOCK_REALTIME == unix and records pass through, unless # register_callbacks armed the timestamp callback (approx clock). calibrate() reads the # native clock through this callable, which is valid for the life of the subscription. self._clock.calibrate( callback_active=self._timestamp_callback_active, native_now=lambda: self._cupti.get_timestamp(cast(int, self._subscriber)), ) self._drain_stop.clear() # Hand the native decode worker the subscriber + the cuptiActivityGetNextRecord_v2 # (and, for self-flush, cuptiActivityFlushAll) addresses, so it iterates records and # drives the periodic flush without a libcupti link, plus the fence kind/field so it # tracks the SYNCHRONIZATION-END clock for flush(sync). It pulls completed buffers and # decodes them GIL-free; Python drains the accumulated columns at flush time, so # per-buffer decode never contends with the training thread. When # background_flush_period_s >= 0 the decode thread also drives the periodic plain # cuptiActivityFlushAll itself (GIL-free) on that cadence, so there is no separate # flush thread; < 0 disables self-flush (the caller drives flush()). fn_addr = self._cupti.get_next_record_fn_address() if not fn_addr: raise RuntimeError( "libcupti is missing cuptiActivityGetNextRecord_v2 (need >= 13.2); " f"loaded {cupti_python.LIBCUPTI_SONAME}" ) if self.background_flush_period_s >= 0: self_flush = True flush_period_ns = int(self.background_flush_period_s * 1e9) flush_fn = self._cupti.get_flush_fn_address() if not flush_fn: raise RuntimeError( "libcupti is missing cuptiActivityFlushAll; " f"loaded {cupti_python.LIBCUPTI_SONAME}" ) else: self_flush = False flush_period_ns = 0 flush_fn = 0 _cupti_monitor_native.configure_decoder( cast(int, self._subscriber), fn_addr, int(_FENCE_KIND), _FENCE_END_FIELD, self_flush, flush_period_ns, flush_fn, ) _cupti_monitor_native.start_decoder() # The decode thread self-flushes on background_flush_period_s (configured above); this # Python loop only pulls the decoded columns and dispatches them to observers, # which calls Python back and so must hold the GIL. No loop at background_drain_period_s < 0 # -- the caller drives drain via flush() itself. if self.background_drain_period_s >= 0: self._drain_and_dispatch_thread = threading.Thread( target=self._drain_and_dispatch_loop, name="torch-cupti-monitor-drain", daemon=True, ) self._drain_and_dispatch_thread.start() # Kinds/fields are enabled by _apply_selection as observers register. self._started = True def stop(self) -> None: if not self._started: return # Stop the Python drain loop. The decode thread keeps running (and self-flushing # on its cadence) until stop_decoder() below, so it can still decode the fence's # sync record; its plain cadence flush and the fence's foreground flush are both # completed-buffers-only, so overlapping them is harmless. self._drain_stop.set() if self._drain_and_dispatch_thread is not None: self._drain_and_dispatch_thread.join(timeout=5.0) if self._drain_and_dispatch_thread.is_alive(): logger.warning("CUPTI monitor drain thread did not stop within 5s") self._drain_and_dispatch_thread = None # Flush thread is down (no concurrent poll): final tail-drain + disable the PM sessions # while observers are still registered, so their last samples are delivered. self._stop_pm_sampler() # Drain everything in flight (incl. CUPTI's async deliveries) before we tear # the decoder down, so the final window is complete. Then stop the native # decode worker while the subscriber is STILL valid -- it may still decode a # few buffers the fence's trailing flush delivered, and it iterates records # via the subscriber, so it must not outlive the unsubscribe -- and dispatch # the residue. self.flush(sync=True) _cupti_monitor_native.stop_decoder() self._drain_and_dispatch() # Disable everything we enabled, then tear down the activity session. self._disable(self._enabled.keys()) self._enabled = {} if self._subscriber is not None: # Release CUPTI without poisoning it for the next session: turn # user-defined-record mode back off (it changes CUPTI's record layout), # then unsubscribe (in _release_subscriber_if_idle below). Crucially this does # NOT call cuptiFinalize -- on this libcupti a finalize poisons CUPTI for the # rest of the process (a subsequent monitor subscribe stops delivering buffers, # and a classic Kineto session records nothing), so disarm + unsubscribe is the # only clean teardown. This lets the monitor be started and stopped repeatedly # in one process. (Switching to a classic consumer after the monitor is a # separate libcupti limitation -- once the process has used UDR/v2 it # cannot downgrade without the poisonous finalize.) self._cupti.disarm_user_defined_records(self._subscriber) # Force a fresh arm of user-defined records on a subsequent start(). self._callbacks_registered = False self._started = False # The subscription itself is shared, so it is released only once no callback handler # holds it either -- a handler must survive an activity session ending. self._release_subscriber_if_idle() self._final_allocated_buffers = _cupti_monitor_native.allocated_buffers() _cupti_monitor_native.reset_buffers() self._clock.reset() def flush(self, *, sync: bool = False, timeout_s: float = 5.0) -> None: """Flush CUPTI's activity buffers to the processing worker. Both paths issue ``cuptiActivityFlushAll(0)`` -- which hands over COMPLETED records, never in-progress ones -- and end by draining the native decoder's accumulated columns and dispatching them to the observers. The monitor never FORCE-flushes (``CUPTI_ACTIVITY_FLAG_FLUSH_FORCED``): a forced flush hands back a still-running kernel's record with a zero end timestamp and consumes it, so CUPTI never re-delivers the real completion (it would strand a slow-but-healthy collective as a false hang in the comm watchdog), and forcing in-progress buffers over concurrent host activity is the flush race that corrupts the HES heap and freezes the decode worker. Plain (``sync=False``) flushes then drains -- the background flush loop and the per-step foreground driver. With ``sync=True`` it first blocks until the native decoder has processed every record up to the call, so the caller (drain, reconfigure, stop) sees a complete window. CUPTI invokes our buffer-complete callback on its own thread a beat *after* cuptiActivityFlushAll returns, so a single flush + idle-wait can race ahead of that async delivery and miss a just-flushed buffer. To fence deterministically we enable SYNCHRONIZATION just for this call, device-sync (which both completes outstanding GPU work -- so a plain flush now delivers everything -- and produces a SYNCHRONIZATION record past a captured CUPTI timestamp), then flush/poll until the native decoder reports a sync record that recent. CUPTI delivers buffers in fill order, so seeing it means everything before is delivered too -- no timing guess, and concurrent activity only helps. SYNCHRONIZATION is enabled only for the fence so the session doesn't pay to record every sync between flushes.""" if not sync: self._cupti.activity_flush_all() self._account_dropped_records(0, 0) self._drain_and_dispatch() return added = self._begin_fence_kind() try: target = self._fence_sync_point() if target is None: # No CUDA available -> no GPU activity to fence; just flush. self._cupti.activity_flush_all() return # Flush to deliver the sync-point's buffer, then poll until the native # decoder has processed a sync record at/after it (its max-sync clock # reaching target). CUPTI delivers buffers in fill order, so seeing that # sync record means everything before it is delivered and decoded too. # The sync record is guaranteed to exist and be deliverable, so this # terminates; the deadline is only a backstop against an unexpected stall. deadline = time.time() + timeout_s while _cupti_monitor_native.decoder_max_sync_ns() < target: self._cupti.activity_flush_all() if _cupti_monitor_native.decoder_max_sync_ns() >= target: break if time.time() >= deadline: logger.warning("CUPTI monitor flush(sync) did not reach its fence") break time.sleep(0.005) finally: self._end_fence_kind(added) # The fence guarantees everything up to the sync point is decoded; hand # the accumulated window to the observers now. self._drain_and_dispatch() def _begin_fence_kind(self) -> bool: """Enable + make decodable the SYNCHRONIZATION sync-point kind for the duration of a fence. Returns True if this call enabled it (so _end removes it); False if it was already enabled (an observer wanted it -- leave it).""" if _FENCE_KIND in self._enabled or self._subscriber is None: return False # Deliver records pending under the current selection before changing it: # without this, enabling the fence kind drops the still-buffered records # (e.g. kernels/launches) that the fence is about to flush for. self._cupti.activity_flush_all() self._cupti.activity_enable(self._subscriber, _FENCE_KIND, _FENCE_FIELDS) self._enabled = {**self._enabled, _FENCE_KIND: _FENCE_FIELDS} return True def _end_fence_kind(self, added: bool) -> None: """Undo _begin_fence_kind (no-op if the kind was already enabled).""" if not added: return if self._subscriber is not None: # Flush before disabling so the records pending under the current # selection (incl. the fence's own sync record) are delivered rather # than dropped when the kind goes away. self._cupti.activity_flush_all() self._cupti.activity_disable(self._subscriber, _FENCE_KIND) self._enabled = {k: v for k, v in self._enabled.items() if k != _FENCE_KIND} def _fence_sync_point(self) -> int | None: """Establish a deterministic fence point for ``flush(sync=True)``: capture CUPTI's clock, then device-sync. The sync both drains outstanding GPU work and produces a SYNCHRONIZATION record with a timestamp past the captured point; the fence waits until the native decoder reports that record. Unlike a tracer kernel, a sync adds no kernel, no cudaLaunchKernel, and no dispatcher op to the trace. Returns the timestamp, or None if no CUDA device is available. SYNCHRONIZATION is enabled only during the fence, so the decoder's max-sync clock only ever moves for an active fence.""" sub = self._subscriber if sub is None: return None try: # The subscriber-aware _v2 timestamp is required here: plain # cuptiGetTimestamp is CUPTI_ERROR_NOT_COMPATIBLE while the UDR subscriber # is active (13.3), which silently turned this fence into a no-op. target = self._cupti.get_timestamp(sub) torch.cuda.synchronize() return target except Exception: return None def convert_time(self, value: int) -> int: """Convert a record-clock timestamp to unix-epoch ns. Records ride the approx clock while the timestamp callback is engaged, else CLOCK_REALTIME.""" if self._timestamp_callback_active: return self._clock.convert_approx(value) return self._clock.convert_unix(value) def convert_time_array(self, values: np.ndarray) -> np.ndarray: """Vectorized :meth:`convert_time` over a whole record column.""" if self._timestamp_callback_active: return self._clock.convert_approx_array(values) return self._clock.convert_unix_array(values) def now_unix_ns(self) -> int: """Current time on the record clock, converted to unix-epoch ns.""" return self.convert_time(self._clock.now_record_ns()) def now_record_ns(self) -> int: """Current value of the record clock -- the unconverted timebase of decoded record START/END. Use this (not now_unix_ns) to stamp a window boundary compared against raw record timestamps. Returns 0 before the session is calibrated.""" return self._clock.now_record_ns() def _try_arm_approx_timestamp_callback(self, sub_handle: int) -> bool: """Best-effort: hand CUPTI the profiler's approx-clock timestamp callback so it stamps activity records on kineto's exact timebase directly. Opt-in via configure(use_approx_timestamps=True) (and only as the sole subscriber); returns False -- leaving records on the CLOCK_REALTIME pass-through -- when disabled or when CUPTI rejects it. Set as a per-subscriber attribute (CUPTI_ACTIVITY_ATTR_TIMESTAMP_CALLBACK), which coexists with the user-defined-record path -- unlike the global cuptiActivityRegisterTimestampCallback, which returns CUPTI_ERROR_NOT_COMPATIBLE.""" if not self._timestamp_callback_enabled: return False # Older libcupti can't re-time device records when a context predates the subscriber, so # refuse rather than silently drop them; libcupti >= 130303 re-times regardless. if self._cupti.get_version() < 130303 and _has_active_cuda_context(): logger.warning( "CUPTI monitor: use_approx_timestamps requested but a CUDA context already " "exists and this libcupti cannot re-time device records; falling back to the " "CLOCK_REALTIME pass-through." ) return False addr = _cupti_monitor_native.approximate_time_callback_address() if self._cupti.arm_approx_timestamp_callback(sub_handle, addr): logger.info("CUPTI monitor: approx-clock timestamp callback engaged") return True logger.warning( "CUPTI monitor: timestamp callback rejected; using the cuptiGetTimestamp " "(CLOCK_REALTIME) pass-through" ) return False # --- observer registry (this monitor is the multiplexer) --------------- def register( self, activities: ActivitiesSpec, callback: Callable[..., None], ) -> _Observer: """Register an observer. ``activities`` is either an iterable of ``ActivityKind`` (meaning "all fields") or a field map ``{ActivityKind: iterable of field ids | "all"}`` selecting the fields per kind. ``callback(columns)`` fires once per drain (at flush time), with ``columns`` = ``{ActivityKind: {field_id: column}}`` -- the native decoder demuxes every buffer to columns and the drain slices them to this observer's selection (the observer never sees raw bytes or the decode strategy). Recomputes the enabled selection and starts the monitor on first registration.""" kinds, fields = self._normalize_activities(activities) obs = _Observer(kinds, fields, callback) with self._lock: self._observers.append(obs) start_needed = not self._started try: if start_needed: self.start() self._apply_selection() except Exception: # Don't leave a half-registered observer (or a half-started monitor) if # start/selection fails -- e.g. the CUPTI subscribe is rejected. with self._lock: if obs in self._observers: self._observers.remove(obs) raise return obs def unregister(self, obs: _Observer) -> None: """Unregister an observer; drops kinds/fields no longer wanted by anyone, and the monitor stops once the last observer leaves. Idempotent.""" with self._lock: if obs not in self._observers: return self._observers.remove(obs) empty = not self._observers if empty: self.stop() else: self._apply_selection() # --- PM sampling (opt-in GPU utilization counters) ----------------------- def request_pm_sampling( self, sink: Callable[[dict[str, Any]], None], metrics: Iterable[str] ) -> None: """Register ``sink`` as a PM-sampling consumer wanting ``metrics`` on the current device. The shared per-device session samples the union of all consumers' metrics; the first consumer starts it (see PmSampler.configure() for interval/look-back). Frames arrive on the flush thread with ``start_ns`` already converted into the trace clock. No-op if CUDA is unavailable, no metrics are given, or ``sink`` is already registered.""" from torch.profiler._cupti.pm_sampling import PmSampler metrics = list(metrics) if not torch.cuda.is_available() or not metrics: return with self._pm_lock: if sink in self._pm_consumers: return sampler = PmSampler() # per-device singleton for the current device try: handle = sampler.add_consumer(metrics) except Exception as e: logger.warning("PM sampling could not register consumer: %s", e) return self._pm_sampler = sampler self._pm_consumers[sink] = handle def _deliver_pm( self, sink: Callable[[dict[str, Any]], None], frame: dict[str, Any] | None ) -> None: # Convert the sampler's raw CUPTI-clock timestamps into the trace clock (the observer buckets # frames by trace-time start_ns), then push to the observer's sink. No-op on an empty poll. if frame is None: return frame = dict(frame) # PM samples are always on the unix (CLOCK_REALTIME) domain (the timestamp callback only # restamps activity records onto the approx clock), so convert on the unix domain, not the # record path. frame["start_ns"] = self._clock.convert_unix_array(frame["start_ns"]) try: sink(frame) except Exception: logger.exception("PM sampling sink error") def release_pm_sampling(self, sink: Callable[[dict[str, Any]], None]) -> None: """Unregister a PM-sampling consumer; the session disables once the last one leaves. Idempotent. Polls the consumer one last time before removing so its final samples land (removal itself delivers nothing).""" with self._pm_lock: handle = self._pm_consumers.pop(sink, None) if handle is not None: self._deliver_pm(sink, handle.poll()) # final delivery handle.detach() if not self._pm_consumers: self._pm_sampler = None def _poll_pm_sampler(self) -> None: """Poll every PM consumer on the monitor's drain cadence (folded into _drain_and_dispatch) -- pulling the HW ring before it overflows. The final poll at release/stop catches the tail.""" with self._pm_lock: for sink, handle in self._pm_consumers.items(): self._deliver_pm(sink, handle.poll()) def _stop_pm_sampler(self) -> None: """Poll then unregister the monitor's PM consumers (final delivery + disable when the last leaves). Called at monitor stop in case observers have not released yet.""" with self._pm_lock: self._pm_sampler = None entries = list(self._pm_consumers.items()) self._pm_consumers.clear() for sink, handle in entries: self._deliver_pm(sink, handle.poll()) handle.detach() def _normalize_activities( self, activities: ActivitiesSpec ) -> tuple[frozenset[ActivityKind], dict[ActivityKind, frozenset[int]]]: """Resolve a registration request to ``(kinds, fields)``: the ``ActivityKind`` set plus the per-activity field-id selection (``"all"``/``None`` -> the kind's full supported set; ``*_FIELD_KIND`` id 0 is always included).""" if isinstance(activities, Mapping): kinds: list[ActivityKind] = [] fields: dict[ActivityKind, frozenset[int]] = {} for kind, sel in activities.items(): k = ActivityKind(kind) kinds.append(k) fields[k] = self._resolve_fields(k, sel) return frozenset(kinds), fields kind_set = frozenset(ActivityKind(k) for k in activities) # A bare kind list means "all fields of that kind". return kind_set, {k: self._resolve_fields(k, "all") for k in kind_set} @staticmethod def _resolve_fields( kind: ActivityKind, sel: Iterable[int] | str | None ) -> frozenset[int]: if sel is None or sel == "all": resolved = frozenset(f for f in FIELD_REGISTRY.get(kind, frozenset())) else: resolved = frozenset(int(f) for f in sel) # type: ignore[union-attr] return resolved | {0} # FIELD_KIND (0) is required for enable + demux def _apply_selection(self) -> None: """Reconcile CUPTI's enabled per-field selection to the current observer field union. Run only here -- when observers register/deregister -- never per buffer. No demux layout is computed: each completed buffer carries CUPTI's own captured layout (ppRecordLayouts), so this only sets which fields are enabled on the subscriber.""" target = {int(k): frozenset(v) for k, v in self._field_union().items()} # A fence (flush(sync=True)) transiently enables SYNCHRONIZATION outside the # observer union; keep it in the target so a register/deregister mid-fence # doesn't strip it -- otherwise the fence never sees its sync record and # flush(sync) spins until it times out. fence = int(_FENCE_KIND) if fence in self._enabled: target[fence] = self._enabled[fence] if target != self._enabled: self._reconfigure(target) self._enabled = target # queued needs the per-subscriber latency-timestamp attribute (which also gates # submitted, not surfaced here). Enable it once, iff an observer selected the # QUEUED kernel field -- so the always-on timing path pays no latency overhead. if self._subscriber is not None and not self._latency_enabled: if Kernel.QUEUED.id in target.get( int(ActivityKind.CONCURRENT_KERNEL), frozenset() ): self._cupti.enable_kernel_latency_timestamps(self._subscriber, True) self._latency_enabled = True # Device-side CUDA_EVENT timestamps (the deviceTimestamp field, off by default) are # needed to place graph event-record nodes as spans. Enable once, iff an observer # selected the CUDA_EVENT device-timestamp field. if self._subscriber is not None and not self._device_ts_enabled: if CudaEvent.DEVICE_TIMESTAMP.id in target.get( int(ActivityKind.CUDA_EVENT), frozenset() ): self._cupti.enable_cuda_event_device_timestamps(self._subscriber, True) self._device_ts_enabled = True def _reconfigure(self, target: dict[int, frozenset[int]]) -> None: # Reconcile the per-field selection to ``target`` with a minimal diff: only # touch kinds that are being removed or whose field selection changed. Kinds # whose selection is unchanged stay enabled -- toggling them off/on is needless # churn and, for RUNTIME/DRIVER, breaks CUPTI's CUDA-graph kernel tracing (a # graph captured while those kinds were enabled stops emitting per-node kernel # records once they're disabled+re-enabled). Each completed buffer carries # CUPTI's own captured layout, so buffers from before a switch still decode. sub = self._subscriber if sub is None: return removed = [k for k in self._enabled if k not in target] changed = [ k for k in target if k in self._enabled and self._enabled[k] != target[k] ] added = [k for k in target if k not in self._enabled] for kind in (*removed, *changed): self._cupti.activity_disable(sub, kind) # Flush between disabling and (re-)enabling a kind with a new field selection # so records pending under the old selection aren't lost. NON-forced: we only # force-flush while syncing (the fence). Forcing here would push in-progress # buffers concurrently with host activity -- the flush race that freezes the # decode worker. if removed or changed: self._cupti.activity_flush_all() for kind in (*added, *changed): self._cupti.activity_enable(sub, kind, target[kind]) def _disable(self, kinds: Iterable[int]) -> None: sub = self._subscriber if sub is not None: for kind in kinds: self._cupti.activity_disable(sub, kind) def _field_union(self) -> dict[ActivityKind, frozenset[int]]: """The per-activity field selection wanted across all observers.""" union: dict[ActivityKind, frozenset[int]] = {} with self._lock: for obs in self._observers: for kind, fset in obs.fields.items(): union[kind] = union.get(kind, frozenset()) | fset # CUPTI only emits CUDA_EVENT records when SYNCHRONIZATION is also enabled # (the two are joined via cudaEventSyncId), so couple it on whenever any # observer wants CUDA_EVENT. Enable the fence fields (KIND + END) so a # concurrent flush(sync) still finds its decodable sync-point record. if ActivityKind.CUDA_EVENT in union: union[_FENCE_KIND] = union.get(_FENCE_KIND, frozenset()) | _FENCE_FIELDS return union def stats(self) -> dict[str, Any]: with self._lock: return { "started": self._started, "activities": list(self._enabled), "buffers_completed": _cupti_monitor_native.decoder_buffers_decoded(), "buffers_allocated": _cupti_monitor_native.allocated_buffers() if self._started else self._final_allocated_buffers, "buffers_pending": _cupti_monitor_native.pending_buffers(), "valid_total_mb": _cupti_monitor_native.decoder_valid_bytes() / (1024 * 1024), "dropped_records": self._dropped_records, "observers": len(self._observers), } def _thread_push_stack(self) -> list[tuple[int, bool]]: """This thread's live LIFO of active external-correlation frames -- ``(id, cupti_ok)``, where cupti_ok records whether CUPTI accepted that push -- so pop unwinds CUPTI + the native mirror only for frames that actually took. CUPTI's external-correlation stacks are per-thread, so ours is too.""" stack = getattr(self._push_tls, "stack", None) if stack is None: stack = self._push_tls.stack = [] return stack def push_external_correlation_id(self) -> int | None: """Allocate a process-unique external-correlation id, record it as this thread's new innermost active context, and push it onto CUPTI's external- correlation stack. Every CUDA activity recorded until the matching pop gets an EXTERNAL_CORRELATION record linking its correlation_id to this id. All subsystems share ONE CUPTI kind, so CUPTI tags each kernel with only the innermost active id; we snapshot the full active stack here (``_id_chains``) so a consumer recovers every context active for an op from that innermost id (see :meth:`external_id_chain`) -- no per-subsystem kind, no shadowing, no kind-pool ceiling. Returns the id, or None if not started. The frame is recorded even when CUPTI rejects the push, so a matching (and possibly unconditional) pop stays balanced; the frame's cupti_ok flag tells pop not to unwind CUPTI/the mirror for it. A rejected push just leaves CUPTI tagging records with the prior id, so this id's chain goes unreferenced and is GC'd -- never an off-by-one unwind of the enclosing context.""" if not self._started or self._subscriber is None: return None with self._lock: external_id = self._next_external_id self._next_external_id += 1 # Pass the subscriber: the plain push returns NOT_COMPATIBLE under the UDR # subscriber, so the wrapper uses the subscriber-aware _v2 variant. cupti_ok = self._cupti.activity_push_external_correlation_id( external_id, sub_handle=self._subscriber ) stack = self._thread_push_stack() stack.append((external_id, cupti_ok)) chain = tuple(fid for fid, _ in stack) with self._lock: self._id_chains[external_id] = chain if cupti_ok: # Mirror into the native per-thread stack so the current (innermost) id is # readable without a CUPTI peek -- see current_external_correlation_id. _cupti_monitor_native.note_external_push(external_id) return external_id def pop_external_correlation_id(self) -> int | None: """Pop this thread's innermost active external-correlation frame (balances a push). Unwinds CUPTI + the native mirror only when that push was accepted by CUPTI, so a rejected push -- or an over-pop, which no-ops -- never unwinds the enclosing frame. Returns the popped id, or None if not started / nothing active. The id's chain snapshot is retired by a later drain (see :meth:`_gc_external_chains`).""" if not self._started or self._subscriber is None: return None stack = self._thread_push_stack() if not stack: return None external_id, cupti_ok = stack.pop() with self._lock: self._chains_gc_pending.append(external_id) if cupti_ok: self._cupti.activity_pop_external_correlation_id( sub_handle=self._subscriber ) _cupti_monitor_native.note_external_pop() # keep the mirror in sync return external_id def external_id_chain(self, innermost_id: int) -> tuple[int, ...]: """The full active-id stack (outermost..innermost) captured when ``innermost_id`` was pushed -- every external-correlation context active for an op CUPTI tagged with ``innermost_id``. A consumer maps a kernel's (innermost) external id through this and picks the id it owns (by membership in its own name/metadata map), recovering enclosing contexts the single-kind records don't carry. Resolve at parse/dispatch time: the snapshot is dropped a couple of drains after the id is popped. Falls back to ``(innermost_id,)`` when there is no snapshot (already dropped, or an id we didn't push).""" with self._lock: return self._id_chains.get(innermost_id, (innermost_id,)) def _gc_external_chains(self) -> None: """Advance the popped-chain GC one generation: drop chains popped two drains ago (their records are delivered + dispatched by now) and promote this cycle's popped ids to be dropped next. Called once per drain, so a popped id's chain survives the drains that dispatch its trailing records.""" with self._lock: if not (self._chains_gc_ready or self._chains_gc_pending): return for retired in self._chains_gc_ready: self._id_chains.pop(retired, None) self._chains_gc_ready = self._chains_gc_pending self._chains_gc_pending = [] def current_external_correlation_id(self) -> int | None: """The external-correlation id on top of THIS thread's stack (last pushed, not yet popped), or None. Reads the native host-side mirror of CUPTI's stack (CUPTI exposes push/pop but no peek). Lets a consumer on the same thread -- e.g. an in-process NCCL profiler plugin -- associate metadata with the annotation the caller already pushed, instead of pushing its own id.""" cur = _cupti_monitor_native.current_external_id() return cur if cur else None def take_external_metadata(self) -> dict[int, str]: """Move out the {external_id: metadata blob} accumulated from drains since the last call. An observer's external-correlation join consumes this to attach the blob onto its kernel events (keyed by the same external id a producer pushed). Drained-and-reset so blobs aren't re-attached.""" with self._lock: meta = self._external_metadata self._external_metadata = {} return meta def add_collective_metadata(self, **fields: Any) -> None: """Merge extra metadata into the CURRENT collective's entry (the most-recently- pushed external-correlation id on this thread), recursively (nested dicts combine; on a leaf conflict the later value wins). The seam for a backend to contribute schema fields the NCCL profiler plugin doesn't emit (e.g. ``process_group``, ``process_group_ranks``, ``input_sizes``) so the serializer plugins (FlightRecorder/clog) can fill them; the fields ride the same per- collective metadata as the plugin's descriptor. Call inside the comms wrapper's push/pop window (or after :meth:`push_external_correlation_id`). To attach metadata to a specific collective from outside its window, use the native ``metadata_put_external(blob, external_id)`` directly.""" if fields: _cupti_monitor_native.metadata_put_external(json.dumps(fields)) def session_info(self) -> dict[str, Any]: """Monitor/session metadata for consumers that need to describe the capture: versions, clock calibration, and buffer config. Call after start() so the clock fields are populated.""" return { "cupti_version": self._cupti.get_version(), "cuda_version": _cuda_version_string(), "hes_enabled": is_hes_enabled(), "timestamp_mode": "approximate_clock", "session_start_unix_ns": self._clock.unix_anchor_ns, "session_start_approx_ns": self._clock.approx_anchor_ns, "buffer_size": self.buffer_size, "flush_period_ns": int(self.background_flush_period_s * 1e9), "drain_period_ns": int(self.background_drain_period_s * 1e9), "libcupti": cupti_python.LIBCUPTI_SONAME, } def _drain_and_dispatch_loop(self) -> None: # cuptiActivityFlushAll is driven off-thread by the native flusher; this loop # only drains the decoded columns + dispatches them (GIL work) -- the same work # flush(sync=False) does after its flush, minus the flush itself. # Floor the wait so a period of 0 ("continuously") doesn't busy-spin the GIL # (Event.wait(0) returns immediately); 1ms is well below any useful cadence. period = max(self.background_drain_period_s, 0.001) try: while not self._drain_stop.wait(period): if self._started: self._account_dropped_records(0, 0) self._drain_and_dispatch() except BaseException: logger.exception("CUPTI monitor drain thread died") def _drain_and_dispatch(self) -> None: """Drain the column groups the native decoder accumulated and fan them out to the observers. The native worker does the per-buffer decode GIL-free; this only views the drained bytes as their dtype and dispatches, so it is cheap and runs on whichever thread drives flush() (foreground or the flush loop). Native returns a LIST of ``(kind, {field_id: (size, bytes)})`` groups -- one per distinct record layout, so within a group every field column is the same length. Groups are packed into frames (each frame holds at most one group per kind) so a dispatched ``{kind: cols}`` chunk always has length-consistent columns; at steady state every kind has a single layout, so this collapses to one frame -- the same multi-kind chunk as before.""" with self._drain_lock: groups, ext_meta = _cupti_monitor_native.drain_decoded() if ext_meta: with self._lock: self._external_metadata.update(ext_meta) if groups: frames: list[dict[int, dict[int, Any]]] = [] for kind, fields in groups: cols = self._columns_from_native(kind, fields) if not cols: continue frame = next((f for f in frames if kind not in f), None) if frame is None: frame = {} frames.append(frame) frame[kind] = cols with self._lock: observers = list(self._observers) for frame in frames: self._dispatch_observers(frame, observers) # GC popped chains AFTER dispatch: a popped id's chain must survive the # drains that dispatch its trailing records (resolution reads it during # dispatch), so retire it a generation later, never before. self._gc_external_chains() # PM sampling shares the drain cadence: polling the HW ring is GIL work, so it # rides the drain path (foreground flush or the background drain loop) -- the # native self-flush thread must not touch it. self._poll_pm_sampler() self._maybe_warn_backpressure() def _columns_from_native( self, kind: int, fields: Mapping[int, tuple[int, bytes]] ) -> dict[int, Any]: """Turn one native group's ``{field_id: (field_size, bytes)}`` into ``{field_id: column}``: numeric fields are viewed per their :class:`Ctype` (unsigned/signed/float) at the captured width; const char* (string) fields are dereferenced to str.""" str_fields = STRING_FIELDS.get(kind, frozenset()) ctype_by_fid = FIELD_CTYPE.get(kind, {}) cols: dict[int, Any] = {} for fid, (size, raw) in fields.items(): if fid in str_fields and size == 8: ptrs = np.frombuffer(raw, dtype=" 8: # Oversized field (the 20-byte CUpti_ActivityEnvironment union): keep its # first 8 bytes as u8 -- the primary metric pair (power+powerLimit / smClock+ # memoryClock / temperature / fanSpeed), split downstream by ENVIRONMENT_KIND. cols[fid] = ( np.frombuffer(raw, dtype=np.uint8) .reshape(-1, size)[:, :8] .copy() .view(" None: if self._outstanding_warned: return allocated = _cupti_monitor_native.allocated_buffers() if allocated >= _OUTSTANDING_WARN_THRESHOLD: self._outstanding_warned = True logger.warning( "CUPTI monitor allocated %d activity buffers; the processing " "worker is not keeping up with CUPTI, so memory use will grow. " "Reduce traced activity or buffer size.", allocated, ) def _dispatch_observers( self, decoded: dict[int, dict[int, Any]], observers: list[_Observer] ) -> None: """Hand each observer the already-demuxed columns sliced to its selection. Pure fan-out -- no buffer access, so observer callbacks need not finish before the buffer is recycled.""" if not decoded: return for obs in observers: chunk: dict[ActivityKind, dict[int, Any]] = {} for kind, fields in obs.fields.items(): kind_cols = decoded.get(int(kind)) if not kind_cols: continue # A buffer missing any field this observer requested was recorded # before the observer's selection was enabled on the subscriber, so # skip the kind rather than hand over a partial chunk. A correctly # synced reconfigure makes this rare, but an in-flight buffer can # still straddle the field-selection change. if any(f not in kind_cols for f in fields): continue chunk[kind] = {f: kind_cols[f] for f in fields} if chunk: obs.callback(chunk) def _account_dropped_records(self, ctx: int, stream_id: int) -> None: if self._subscriber is None: return self._dropped_records += self._cupti.activity_get_num_dropped_records( self._subscriber, ctx, stream_id ) _hes_enabled = False def enable_hes_early() -> None: global _hes_enabled if _hes_enabled: return if torch.cuda.is_initialized() or _has_active_cuda_context(): raise RuntimeError( "enable_hes_early() must be called before CUDA context creation" ) from cuda.bindings import driver as cuda_driver # pyrefly: ignore[missing-import] rc = cuda_driver.cuInit(0)[0] if rc != cuda_driver.CUresult.CUDA_SUCCESS: raise RuntimeError(f"cuInit failed with rc={rc}") # Use the direct libcupti call (self._cupti.activity_enable_hw_trace), not # cupti-python's activity_enable_hw_trace(): after torch is imported, the # latter makes subsequent cuptiActivityRegisterCallbacks() fail with # CUPTI_ERROR_MULTIPLE_SUBSCRIBERS_NOT_SUPPORTED, while the direct call works. cupti_python.pylibcupti().activity_enable_hw_trace(True) _hes_enabled = True def is_hes_enabled() -> bool: return _hes_enabled def configure( *, buffer_size: int | None = None, background_flush_period_s: float | None = None, background_drain_period_s: float | None = None, use_approx_timestamps: bool | None = None, ) -> None: """Set the process-wide CUPTI monitor settings the singleton snapshots when first constructed (via CuptiMonitor()). First-come-first-serve (like PmSampler.configure): locked once this lands OR the singleton is built, so a later call is ignored with a warning -- pass all settings in one call, before first use. An unset (None) arg keeps its current value. The two cadences default to -1 (caller-driven; a non-negative value opts into background flush/drain at that period).""" with CuptiMonitor._instance_lock: if CuptiMonitor._configured: logger.warning( "cupti_monitor.configure() ignored: already configured " "(first-come-first-serve). Call it once before the first CuptiMonitor()." ) return if buffer_size is not None: CuptiMonitor._buffer_size = buffer_size if background_flush_period_s is not None: CuptiMonitor._background_flush_period_s = background_flush_period_s if background_drain_period_s is not None: CuptiMonitor._background_drain_period_s = background_drain_period_s if use_approx_timestamps is not None: CuptiMonitor._use_approx_timestamps = use_approx_timestamps CuptiMonitor._configured = True def get_config() -> dict[str, Any]: """The process-wide config the singleton will snapshot (or snapshotted): buffer_size, the two cadences, the approx-clock flag, and whether configure()/construction has pinned it (first-come-first-serve).""" return { "buffer_size": CuptiMonitor._buffer_size, "background_flush_period_s": CuptiMonitor._background_flush_period_s, "background_drain_period_s": CuptiMonitor._background_drain_period_s, "use_approx_timestamps": CuptiMonitor._use_approx_timestamps, "configured": CuptiMonitor._configured, # Capability flag, not a setting: lets a consumer detect the shared # subscriber-callback registry (CuptiMonitor.register_callback_handler) and stop # taking a CUPTI subscription of its own, without hasattr-probing. "supports_callback_handlers": True, } def _reset_for_test() -> None: """Test-only: drop the singleton and reset the config to defaults so a test can build a freshly-configured monitor. Callers must have torn down any live session first.""" with CuptiMonitor._instance_lock: CuptiMonitor._instance = None CuptiMonitor._buffer_size = 4 * 1024 * 1024 CuptiMonitor._background_flush_period_s = -1.0 CuptiMonitor._background_drain_period_s = -1.0 CuptiMonitor._use_approx_timestamps = False CuptiMonitor._configured = False