Add Quent Resource and IR-node tracing to cudf-polars - #23179
Add Quent Resource and IR-node tracing to cudf-polars#23179TomAugspurger wants to merge 82 commits into
Conversation
|
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. |
github-actions
Bot
added
Python
This PR builds on NVIDIA#22337 by adding Quent Resources to cudf-polars.
86782f1 to
0efdf65
Compare
TomAugspurger
added
improvement
|
/ok to test 2fc4e74 |
|
/ok to test 9821ef7 |
- write a zipfile - include in ndjson output - add quent role
this gives us the timeline animation.
| # do this just once | ||
| from cudf_polars.quent._types import Task | ||
|
|
||
| pynvml.nvmlInit() | ||
| maybe_handle = get_device_handle() |
There was a problem hiding this comment.
It should yeah. I think the original intent was to call each of these once (per IR node type) on import and then use the returned value on each call to IR.do_evaluate. I'll move these outside def wrapper(...) to achieve that.
| execute_persisted_query, | ||
| ) | ||
| from cudf_polars.quent._context import LocalQuentContext | ||
| from cudf_polars.quent._context import ( |
There was a problem hiding this comment.
Remind me why we have these "internal" packages in the quent subsystem. We don't do this anywhere else.
There was a problem hiding this comment.
I want to be extra careful with the public-facing API here: We hope to delete a bunch of what's under cudf_polars.quent, but we do have some user-facing pieces. Having cudf_polars/quent/__init__.py contain just user-facing things makes it clear what we can just delete and what we need to provide compatibility wrappers for.
| # Now that ``comm`` exists, declare the engine-scoped inter-rank network | ||
| # topology once (a no-op for single-rank runs). |
There was a problem hiding this comment.
This seems to suggest that if there's only one GPU we wouldn't have quent events. But I think that is not right.
| def _integer_variant(value: int) -> str: | ||
| """Return the narrowest Quent integer variant that can hold ``value``.""" | ||
| for variant, lo, hi in _INT_VARIANTS: | ||
| if lo <= value <= hi: | ||
| return variant | ||
| raise ValueError(f"Integer value {value} does not fit any Quent integer type.") | ||
|
|
||
|
|
||
| def _common_integer_variant(values: list[int]) -> str: | ||
| """Return the narrowest Quent integer variant that can hold all ``values``.""" | ||
| for variant, lo, hi in _INT_VARIANTS: | ||
| if all(lo <= value <= hi for value in values): | ||
| return variant | ||
| raise ValueError("Integer list values do not fit any Quent integer type.") |
There was a problem hiding this comment.
Why do we care about this and not just encode as i64 for everything?
There was a problem hiding this comment.
Mmm, yeah let's just use u64 for unsigned values and i64 for signed values.
There was a problem hiding this comment.
I guess some part of the rust backend for Quent does deserialize these when you're using the UI, and presumably it uses the size associated with the value. If it were just Python it wouldn't matter, but given that there's some runtime benefit to using smaller sizes here, maybe this code is worth keeping.
| # ids are derived per-collect and ``get_stable_plan_id`` is a deterministic | ||
| # function of the IR structure, so an un-namespaced plan id would collide | ||
| # across the two identical collects. | ||
| q = pl.LazyFrame({"x": [1, 2, 3]}).filter(pl.col("x") > 1) |
There was a problem hiding this comment.
Why not just have a test-file-scoped engine fixture?
There was a problem hiding this comment.
I think because the test itself wants to assert some things that happen during engine initialization and shutdown. Previously, this was increasing the duration of the test quite a bit, but I'll see whether that's still true. Perhaps I was using the fixtures giving us cluster resources improperly.
There was a problem hiding this comment.
There's also some state that's cached on the Engine (_quent_events) and QuentContext (_query_group_cache) that makes it a pain to try to reuse any of those resources across tests. That state would need to be cleared for the assertions that depend on them to pass.
|
8b14654 is one small change bringing us to Quent I'm testing out some additional changes to better capture data as is flows through our DAG. One major gap right now is (pinned) host memory: we don't see it at all currently. We can easily define another |
This reverts commit 15153e7.
4 participants
Description
This PR builds on #22337 by
IR.do_evaluatelevel tracing to our emitted events{run_id}/{type}/{uuid}.ndjson+ a metadata file).I've added resource definitions / usage for:
These resources are used in the new
Task, which wraps our individualIR.do_evaluateSee this gist for some traces captured with
Structurally, the biggest change is to
IRExecutionContext. We'd like to emit some start/end traces in<IR>.do_evaluate. Those traces need to know what Quent IR node they're executing for, so I've addedIRExecutionContext.quent_ir_execution_context, which contains the operator ID. When we build theIRExecutionContextfor any given IR node, we make sure to set the Quent Operator this IR node is working on (seeir_context_for_node.Aside from that, the biggest changes by lines of code are probably:
Checklist