Log GPU timing in cudf-polars traces

[TomAugspurger]

Description

The tracing in cudf-polars is currently based around host times. We record a start and stop based on Python's monotonic_ns around a call to IR.do_evaluate. But inside of that IR.do_evaluate we call some non-blocking, asynchronous pylibucdf calls whose runtime extends past the end of stop, running on the GPU.

To measure the actual runtime of GPU operations associated with some IR node, we need to measure when that sequence of GPU operations actually finishes. There's several ways to do this, but I've opted for CUDA Events and cudaLaunchHostFunc.

Using cudaLaunchHostFunc to call a python function can be fraught (deadlocks, often related to the GIL, are apparently a risk)... We deliberately keep the work done inside this function (and so, on the thread calling it) simple: just putting an integer completion token on a Queue.

The actual logging is done by a background thread.

When enabled, we'll emit two traces per IR.do_evaluate:

1. a trace with scope=evaluate_ir_node containing host information
2. a trace with scope=evaluate_ir_node_gpu containing gpu information

We include a trace_event_id (a UUID generated on the fly) so that consumers can correlate GPU traces with host traces.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[TomAugspurger]

Using cudaLaunchHostFunc to call a python function can be fraught

Proof: this is currently hanging on some queries. I don't know why yet, but it's related to recording the completion event on result.stream; An earlier version used kwargs["context"].get_cuda_stream(). This didn't hang, but it also didn't correctly measure the time the GPU operations took (since the stream from get_cuda_stream() might not be downstream of the operations in IR.do_evaluate.

[TomAugspurger]

Using cudaLaunchHostFunc to call a python function can be fraught

Proof 2: even after "fixing" that (by recording the completion event on another stream with)

                timing_stream = get_joined_cuda_stream(
                    exec_ctx.get_cuda_stream,
                    upstreams=[exec_ctx.get_cuda_stream(), result.stream],
                )
 

I'm getting intermittent hangs. Maybe this suggests the approach of using cuda.bindings plus a "simple" host callback function is more fragile than I thought. Or maybe I'm missing something about how we're supposed to record the events.

[TomAugspurger]

I think query_id should be included through context variables?

[TomAugspurger]

ctypes probably isn't the way to do this long term, but might be OK as a POC.

[TomAugspurger]

/ok to test 1eb332d