Thanks @RSchwan, this looks good to me.

I mentioned this pull request to @nvlukasz over dinner and he pointed out that this won't work because we need to care about the cases in which another framework initiates the graph capture. Apparently he's considered this topic for quite some time.

That's exactly why I added @nvlukasz as a reviewer 😁 Thanks!

But is it actually a problem in this particular case? In case runtime.captures has no elements, the calls to runtime.core.cuda_stream_is_capturing(stream.cuda_stream) and runtime.core.cuda_stream_get_capture_id(stream.cuda_stream) are not used since runtime.captures.get(capture_id) will always return None. Hence, I'm just short-circuiting useless API calls (in a logical sense).

But is it actually a problem in this particular case? In case runtime.captures has no elements, the calls to runtime.core.cuda_stream_is_capturing(stream.cuda_stream) and runtime.core.cuda_stream_get_capture_id(stream.cuda_stream) are not used since runtime.captures.get(capture_id) will always return None. Hence, I'm just short-circuiting useless API calls (in a logical sense).

Good point! What you said makes sense to me, I'll ask @nvlukasz to take another look to confirm.

Sorry for the confusion, everything's good and merged now. Thanks again @RSchwan!