Add coarse-grained streaming-node tracing - #21309
Conversation
rjzamora
added
2 - In Progress
github-actions
Bot
added
Python
| # Remap partitioning if schema has changed | ||
| partitioning = remap_partitioning( | ||
| metadata_in.partitioning, ir.children[0].schema, ir.schema | ||
| ) |
There was a problem hiding this comment.
This matches up child columns 1-1 with our columns. Is this always valid if we say preserve_partitioning?
| except (IndexError, KeyError): | ||
| return None # Column missing in old or new schema |
There was a problem hiding this comment.
Hmm, I think it should be impossible to get IndexError because if some column_index is not in old_names then something has gone catastrophically wrong.
There was a problem hiding this comment.
If we don't have the relevant old name in the new schema what does that mean? We've dropped some columns or (I think?) renamed some columns. Is it bad if we rename and don't notice?
There was a problem hiding this comment.
We get an index error if the column name we are "partitioned" on is not in our new schema. This can happen if we rename the column(s) we are partitioned on. Right now we just "forget" that we are partitioned in this case, but it won't be hard to handle this later.
| """Metadata payload for a channel.""" | ||
| # If inter_rank partitioning was invalidated, the whole partitioning is invalid | ||
| if isinstance(partitioning.inter_rank, HashScheme) and new_inter_rank is None: | ||
| return None |
There was a problem hiding this comment.
OK, so if we were not previously inter-rank partitioned we can still potentially preserve local?
There was a problem hiding this comment.
This check isn't necessary, and should technically check the rank count - The node deciding to shuffle will check this kinds of detail anyway.
| def _repr_trace_tree( | ||
| ir: IR, | ||
| partition_info: MutableMapping[IR, PartitionInfo], | ||
| tracer: StreamingQueryTracer, | ||
| *, | ||
| offset: str = "", | ||
| ) -> str: | ||
| """Recursively build a tree representation with tracer data.""" | ||
| header = _repr_ir(ir, offset=offset) | ||
| header = header.rstrip("\n") | ||
|
|
||
| # Get node tracer if it exists | ||
| if (node_tracer := tracer.node_tracers.get(ir)) is not None: | ||
| # Add actual row count if available | ||
| if node_tracer.row_count is not None: | ||
| header += f" rows={_fmt_row_count(node_tracer.row_count)}" | ||
|
|
||
| # Add decision if present | ||
| if node_tracer.decision is not None: | ||
| header += f" decision={node_tracer.decision}" | ||
|
|
||
| # Add actual chunk count | ||
| header += f" chunks={node_tracer.chunk_count}" | ||
|
|
||
| children_strs = [ | ||
| _repr_trace_tree(child, partition_info, tracer, offset=offset + " ") | ||
| for child in ir.children | ||
| ] | ||
|
|
||
| header += "\n" | ||
| return header + "".join( | ||
| f"{line}{offset} (repeated {count} times)\n" | ||
| if (count := sum(1 for _ in group)) > 1 | ||
| else line | ||
| for line, group in groupby(children_strs) | ||
| ) |
There was a problem hiding this comment.
Can we also have an option to dump this tree in a structured form (json would be fine)? That way automated tools have a chance of reading it rather than needing to parse.
There was a problem hiding this comment.
I'm working on that and will have something later today.
There was a problem hiding this comment.
This PR will already dump the information in structlog form when CUDF_POLARS_LOG_TRACES=1. Tom is working on the code needed to post-process that output to provide something similar to the above. I suppose we could drop some of this logic if we feel like the dump + post-process workflow works well enough for multiple ranks.
| from cudf_polars.dsl.ir import IR | ||
|
|
||
|
|
||
| def _stable_ir_id(ir_node: IR) -> int: |
There was a problem hiding this comment.
This seems like a reasonable method to have on the IR class itself.
There was a problem hiding this comment.
Yeah, that may make sense. This hash just uses the repr, so it isn't as "robust" as hash(), but it's useful for purposes like this.
| def _repr_trace_tree( | ||
| ir: IR, | ||
| partition_info: MutableMapping[IR, PartitionInfo], | ||
| tracer: StreamingQueryTracer, | ||
| *, | ||
| offset: str = "", | ||
| ) -> str: | ||
| """Recursively build a tree representation with tracer data.""" | ||
| header = _repr_ir(ir, offset=offset) | ||
| header = header.rstrip("\n") | ||
|
|
||
| # Get node tracer if it exists | ||
| if (node_tracer := tracer.node_tracers.get(ir)) is not None: | ||
| # Add actual row count if available | ||
| if node_tracer.row_count is not None: | ||
| header += f" rows={_fmt_row_count(node_tracer.row_count)}" | ||
|
|
||
| # Add decision if present | ||
| if node_tracer.decision is not None: | ||
| header += f" decision={node_tracer.decision}" | ||
|
|
||
| # Add actual chunk count | ||
| header += f" chunks={node_tracer.chunk_count}" | ||
|
|
||
| children_strs = [ | ||
| _repr_trace_tree(child, partition_info, tracer, offset=offset + " ") | ||
| for child in ir.children | ||
| ] | ||
|
|
||
| header += "\n" | ||
| return header + "".join( | ||
| f"{line}{offset} (repeated {count} times)\n" | ||
| if (count := sum(1 for _ in group)) > 1 | ||
| else line | ||
| for line, group in groupby(children_strs) | ||
| ) |
There was a problem hiding this comment.
I'm working on that and will have something later today.
rjzamora
added
0 - Waiting on Author
|
Closing in favor of #21316 |
- Closes #21244 - **Another** Alternative to #21277 and #21309 - Part of #20482 Authors: - Richard (Rick) Zamora (https://github.com/rjzamora) Approvers: - Tom Augspurger (https://github.com/TomAugspurger) URL: #21316
4 participants
Description
cc @TomAugspurger (who suggested something along this general path) - This is hybrid between #21277 and the logging approach suggested offline:
When tracing is enabled (e.g.
CUDF_POLARS_LOG_TRACES=1), we will collect coarse-grained information about streaming nodes, and log"Streaming Node"events. These events includeir_idinformation that should match their_idinfo added to the existing"Execute IR"events. WhenCUDF_POLARS_LOG_TRACES=0, we will still collect this information if/whenTracingOptionsare passed in by the user, but we will not emit the structlog events in that case.My motivation for preserving the "in-memory" tracing infrastructure was simply personal convenience. I struggled a bit to get what I wanted out of a "structlog-only" approach.
Checklist