Full structlog-based tracing

[rjzamora]

Description

• Closes [FEA] Add IR-node profiling for the rapidsmpf runtime #21244
• Another Alternative to Add basic runtime profiling in cudf-polars #21277 and Add coarse-grained streaming-node tracing #21309
• Part of [FEA][Story] Dynamic planning in cuDF-Polars #20482

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[rjzamora]

@TomAugspurger - If you think we can post-process structlog output to produce annotated query plans (like in #21309), something like this may be a bit less intrusive.

[rjzamora]

I think think this alternative is looking best. I am able to reconstruct the logs to print the information I need. E.g.

============================================================
Query 1, Iteration 0 (duration: 1.497s)
============================================================
SORT [rows=4, chunks=1, bytes_out=1.1KB, time=29ms, workers=4, duplicated]
  REPARTITION [rows=4, chunks=1, workers=4, duplicated]
    SELECT [rows=4, chunks=1, bytes_out=1.1KB, time=1.3e+02ms, workers=4, duplicated]
      SELECT [rows=4, chunks=1, bytes_out=1.2KB, time=19ms, workers=4, duplicated]
        GROUPBY [rows=4, chunks=1, bytes_out=1.6KB, time=3.4ms, workers=4, duplicated]
          REPARTITION [rows=44, chunks=1, workers=4, duplicated]
            GROUPBY [rows=44, chunks=11, bytes_out=3.6KB, time=3.3e+02ms, workers=4]
              HSTACK [rows=1.8e+02M, chunks=11, bytes_out=8.9GB, time=99ms, workers=4]
                SCAN [rows=1.8e+02M, chunks=11, bytes_out=7.6GB, time=6.5s, workers=4]

However, we can definitely add the post-processing step in a follow-up.

[wence-]

I think think this alternative is looking best. I am able to reconstruct the logs to print the information I need. E.g.

============================================================
Query 1, Iteration 0 (duration: 1.497s)
============================================================
SORT [rows=4, chunks=1, bytes_out=1.1KB, time=29ms, workers=4, duplicated]
  REPARTITION [rows=4, chunks=1, workers=4, duplicated]
    SELECT [rows=4, chunks=1, bytes_out=1.1KB, time=1.3e+02ms, workers=4, duplicated]
      SELECT [rows=4, chunks=1, bytes_out=1.2KB, time=19ms, workers=4, duplicated]
        GROUPBY [rows=4, chunks=1, bytes_out=1.6KB, time=3.4ms, workers=4, duplicated]
          REPARTITION [rows=44, chunks=1, workers=4, duplicated]
            GROUPBY [rows=44, chunks=11, bytes_out=3.6KB, time=3.3e+02ms, workers=4]
              HSTACK [rows=1.8e+02M, chunks=11, bytes_out=8.9GB, time=99ms, workers=4]
                SCAN [rows=1.8e+02M, chunks=11, bytes_out=7.6GB, time=6.5s, workers=4]

However, we can definitely add the post-processing step in a follow-up.

I think think this alternative is looking best. I am able to reconstruct the logs to print the information I need. E.g.

============================================================
Query 1, Iteration 0 (duration: 1.497s)
============================================================

[...]

            SCAN [rows=1.8e+02M, chunks=11, bytes_out=7.6GB, time=6.5s, workers=4]

Not sure how a scan can take 6.5s, but the query 1.5? Is this the sum of the times on the workers? But then there's still something weird because 4 * 1.5 is 6, not 6.5

[rjzamora]

Not sure how a scan can take 6.5s, but the query 1.5? Is this the sum of the times on the workers? But then there's still something weird because 4 * 1.5 is 6, not 6.5

The toy script is just adding up times for every do_evaluate event that executes within the "context" of that Scan node, and then those times are summed over the 4 ranks. We are also using 4 producers on each rank, so I suppose the total amount of "measured" time can be high? (Note that I just threw in the runtime for fun, the chunk count and row-count were my immediate focus)

[TomAugspurger]

Left a few quick thoughts. Looking good though.

[TomAugspurger]

Not a blocker, but this is related to #21315. Depending on which goes first, we'll want to have just one way of serializing query plans.

I think we'll want a env var / flag CUDF_POLARS_LOG_QUERY_PLAN similar to CUDF_POLARS_LOG_TRACES_MEMORY to control whether or not to log this. I'm not worried about the runtime overhead of this, but maybe some query plans can get largish, and people would want to exclude it from their logs? For this PR, it probably doesn't matter since it's just an ID per IR node. My PR is capturing more information though.

[TomAugspurger]

https://github.com/TomAugspurger/pygdf/blob/05443333fa3d3b1b4e9475cfe4a9d7fbaa0c6665/python/cudf_polars/cudf_polars/dsl/tracing.py#L215-L227 adds a bound_contextvars contextmanager that respects our LOG_TRACES. Maybe copy use that instead of unbinding in the finally?

[TomAugspurger]

And I'd recommend also binding ir_type, so that we can easily tell what type of IR node lead to us executing this specific IR.do_evaluate.

[rjzamora]

I think the bound_contextvars utility makes this code much messier. The finally will ensure we unbind, so I don't think we benefit from refactoring to use another (non-asyn) context manager within this context manager.

And I'd recommend also binding ir_type, so that we can easily tell what type of IR node lead to us executing this specific IR.do_evaluate.

Good idea. This actually made me realize that the ir_* naming is confusing, because the do_evaluate may belong to a different type of IR node. I decided to change these names to actor_ir_id and actor_ir_type to make the distinction more clear.

[TomAugspurger]

It'd be helpful to bind some kind of query plan ID here. This (combined with my suggestion to bind the ir_node_id below) will make our lowest-level logs easier to understand, since each record will have

plan_id, ir_node_id, type, ...

We'll need to make sure to unbind it at the end of the query. Hopefully we have a good spot to do it...

How to compute the query ID though? Really, it could just be a str(uuid.uuid4()). In fact I think we might want that, rather than something deterministic based on the query plan, so that two executions of the same plan gets different ids.

[rjzamora]

I like this idea, but I'm wondering if it can/should wait for a follow-up PR?

Since we are logging the query plan on the client, we would need to create a random ID here, and then pass that ID into evaluate_pipeline somehow so that the same value can be bound/unbound on the workers. I do think this is doable, but I'm wincing a bit in my head.

I suppose we could add a trace_token field to ConfigOptions so we can just piggy-back on the fact that those options are already plumbed through everywhere. (not sure if that makes sense though)

[rjzamora]

/merge