cupti_callback_subscription_for_cuda_api_monitoring.py

1import sys
2from cuda import cuda, cupti
3
4def cupti_callback(userdata, domain, cbid, cbdata):
5    # This function is called when a CUDA API is invoked
6    cbdata_ptr = cupti.CUpti_CallbackData(cbdata)
7    
8    # Check if we are entering or exiting the API call
9    if cbdata_ptr.callbackSite == cupti.CUpti_ApiCallbackSite.CUPTI_API_EXIT:
10        print(f"Callback: Function {cbdata_ptr.functionName} has completed.")
11
12def main():
13    # Initialize CUDA Driver API
14    err, = cuda.cuInit(0)
15    if err != cuda.CUresult.CUDA_SUCCESS:
16        print(f"CUDA Initialization failed with error: {err}")
17        return
18
19    # Subscribe to CUPTI callbacks
20    # This creates a subscriber that listens to Driver and Runtime API domains
21    subscriber = cupti.CUpti_SubscriberHandle()
22    err, subscriber = cupti.cuptiSubscribe(subscriber, cupti_callback, None)
23    
24    # Enable callbacks for the Runtime API domain
25    err = cupti.cuptiEnableDomain(1, subscriber, cupti.CUpti_ApiDomain.CUPTI_CB_DOMAIN_RUNTIME_API)
26
27    print("CUPTI monitoring enabled. Executing CUDA call...")
28
29    # Execute a simple CUDA call to trigger the callback
30    err, device = cuda.cuDeviceGet(0)
31    
32    # Clean up
33    cupti.cuptiUnsubscribe(subscriber)
34    print("Example finished.")
35
36if __name__ == "__main__":
37    main()