From 0efdf65e17480b57b65b5f46c7b76818505d4dc6 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 7 Jul 2026 14:34:28 -0700 Subject: [PATCH 01/60] [WIP]: Add Quent Resource tracing to cudf-polars This PR builds on https://github.com/rapidsai/cudf/pull/22337 by adding Quent Resources to cudf-polars. --- python/cudf_polars/cudf_polars/dsl/ir.py | 4 + python/cudf_polars/cudf_polars/dsl/tracing.py | 137 +++++- python/cudf_polars/cudf_polars/engine/core.py | 67 ++- python/cudf_polars/cudf_polars/engine/dask.py | 41 +- python/cudf_polars/cudf_polars/engine/ray.py | 34 +- python/cudf_polars/cudf_polars/engine/spmd.py | 25 +- .../cudf_polars/cudf_polars/quent/__init__.py | 20 +- .../cudf_polars/cudf_polars/quent/_context.py | 118 ++++- python/cudf_polars/cudf_polars/quent/_plan.py | 14 + .../cudf_polars/cudf_polars/quent/_types.py | 461 +++++++++++++++++- .../actor_graph/collectives/shuffle.py | 34 +- .../streaming/actor_graph/collectives/sort.py | 10 +- .../cudf_polars/streaming/actor_graph/core.py | 11 + .../streaming/actor_graph/dispatch.py | 29 +- .../streaming/actor_graph/groupby.py | 4 +- .../cudf_polars/streaming/actor_graph/io.py | 13 +- .../cudf_polars/streaming/actor_graph/join.py | 8 +- .../streaming/actor_graph/nodes.py | 13 +- .../cudf_polars/streaming/actor_graph/over.py | 8 +- .../streaming/actor_graph/repartition.py | 8 +- .../streaming/actor_graph/union.py | 4 +- .../streaming/actor_graph/utils.py | 103 +++- python/cudf_polars/tests/quent/test_quent.py | 46 ++ .../tests/quent/test_quent_integration.py | 14 + 24 files changed, 1156 insertions(+), 70 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index f61a738d0bbd..4228edc4545e 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -84,6 +84,7 @@ from cudf_polars.containers.dataframe import NamedColumn from cudf_polars.dsl.utils.io import CachedParquetInfo + from cudf_polars.quent._context import QuentIRExecutionContext from cudf_polars.streaming.rank_aware_source import RankAwareSource from cudf_polars.typing import CSECache, ClosedInterval, Schema, Slice as Zlice from cudf_polars.utils.config import ParquetOptions @@ -138,11 +139,14 @@ class IRExecutionContext: A zero-argument callable that returns a CUDA stream. query_id Identifier for the query being executed. + quent_ir_execution_context + Optional Quent tracing context bound to a physical operator. """ py_executor: concurrent.futures.ThreadPoolExecutor | None = field(default=None) get_cuda_stream: Callable[[], Stream] = field(default=get_cuda_stream) query_id: uuid.UUID = field(default_factory=uuid.uuid4) + quent_ir_execution_context: QuentIRExecutionContext | None = None async def to_thread( self, func: Callable[P, T], /, *args: P.args, **kwargs: P.kwargs diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index 2c1dd1f26f0c..ef11dfa14d06 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -9,7 +9,9 @@ import enum import functools import os +import threading import time +import uuid from typing import TYPE_CHECKING, Any, Concatenate, Literal, ParamSpec import nvtx @@ -49,6 +51,89 @@ import cudf_polars.containers from cudf_polars.dsl import ir + from cudf_polars.dsl.ir import IRExecutionContext + + +def _dataframe_size_bytes(frame: cudf_polars.containers.DataFrame) -> int: + return sum(col.device_buffer_size() for col in frame.table.columns()) + + +def _begin_quent_do_evaluate_events( + cls: type[ir.IR], + ir_execution_context: IRExecutionContext, +) -> tuple[Any, Any, bool] | None: + import cudf_polars.quent + from cudf_polars.dsl.ir import DataFrameScan, Scan + + quent_ir_execution_context = ir_execution_context.quent_ir_execution_context + if quent_ir_execution_context is None: + return None + + token = uuid.uuid4() + quent_task = cudf_polars.quent.Task( + instance_name=( + f"{cls.__name__}-{quent_ir_execution_context.quent_operator.id.hex[:8]}-" + f"{token.hex[:8]}" + ), + operator_id=quent_ir_execution_context.quent_operator.id, + ) + quent_processor = quent_ir_execution_context.context.get_or_declare_processor( + quent_ir_execution_context.logger, + thread_ident=threading.get_ident(), + pool_id=quent_ir_execution_context.thread_pool_id, + ) + is_io_node = issubclass(cls, (Scan, DataFrameScan)) + quent_ir_execution_context.logger.emit(quent_task.queueing()) + if not is_io_node: + quent_ir_execution_context.logger.emit( + quent_task.allocating(resource_id=quent_processor.id) + ) + return quent_task, quent_processor, is_io_node + + +def _end_quent_do_evaluate_events( + cls: type[ir.IR], + frames: Sequence[cudf_polars.containers.DataFrame], + result: cudf_polars.containers.DataFrame, + ir_execution_context: IRExecutionContext, + quent_state: tuple[Any, Any, bool], +) -> None: + import cudf_polars.quent + + quent_task, quent_processor, is_io_node = quent_state + quent_ir_execution_context = ir_execution_context.quent_ir_execution_context + if quent_ir_execution_context is None: + return + + output_capacity_bytes = _dataframe_size_bytes(result) + if is_io_node: + quent_ir_execution_context.logger.emit( + quent_task.loading( + use_thread=quent_processor, + use_channel=quent_ir_execution_context.disk_to_device_channel, + channel_capacity_bytes=output_capacity_bytes, + use_memory=quent_ir_execution_context.device_memory, + memory_capacity_bytes=output_capacity_bytes, + ) + ) + else: + quent_ir_execution_context.logger.emit( + quent_task.computing( + use_thread=quent_processor, + use_memory=quent_ir_execution_context.device_memory, + memory_capacity_bytes=output_capacity_bytes, + ) + ) + quent_ir_execution_context.logger.emit( + quent_ir_execution_context.quent_operator.statistics( + statistics=cudf_polars.quent.Statistics( + input_bytes=sum(_dataframe_size_bytes(frame) for frame in frames), + output_bytes=output_capacity_bytes, + output_rows=result.num_rows, + ) + ) + ) + quent_ir_execution_context.logger.emit(quent_task.exit()) class Scope(enum.StrEnum): @@ -158,42 +243,37 @@ def log_do_evaluate( func The ``IR.do_evaluate`` method to wrap. """ - if not LOG_TRACES: - return func - else: # pragma: no cover; requires CUDF_POLARS_LOG_TRACES=1 - - @functools.wraps(func) - def wrapper( - cls: type[ir.IR], - *args: P.args, - **kwargs: P.kwargs, - ) -> cudf_polars.containers.DataFrame: - # do this just once + + @functools.wraps(func) + def wrapper( + cls: type[ir.IR], + *args: P.args, + **kwargs: P.kwargs, + ) -> cudf_polars.containers.DataFrame: + ir_execution_context: IRExecutionContext | None = kwargs.get("context") # type: ignore[assignment] + + frames: list[cudf_polars.containers.DataFrame] = ( + list(args) + [v for k, v in kwargs.items() if k != "context"] + )[cls._n_non_child_args :] # type: ignore[assignment] + + quent_state = None + if ir_execution_context is not None: + quent_state = _begin_quent_do_evaluate_events(cls, ir_execution_context) + + if LOG_TRACES: # pragma: no cover; requires CUDF_POLARS_LOG_TRACES=1 pynvml.nvmlInit() maybe_handle = get_device_handle() pid = _getpid() log = structlog.get_logger() - # By convention, all non-dataframe arguments (non-child) come first. - # Anything remaining is a dataframe, except for 'context' kwarg. - frames: list[cudf_polars.containers.DataFrame] = ( - list(args) + [v for k, v in kwargs.items() if k != "context"] - )[cls._n_non_child_args :] # type: ignore[assignment] - before_start = time.monotonic_ns() before = make_snapshot( cls, frames, phase="input", device_handle=maybe_handle, pid=pid ) before_end = time.monotonic_ns() - - # The decorator preserves the exact signature of the original do_evaluate method. - # Each IR.do_evaluate method is a classmethod that takes the IR class as first - # argument, followed by the method-specific arguments, and returns a DataFrame. - start = time.monotonic_ns() result = func(cls, *args, **kwargs) stop = time.monotonic_ns() - after_start = time.monotonic_ns() after = make_snapshot( cls, @@ -214,10 +294,17 @@ def wrapper( } ) log.info("Execute IR", **record) + else: + result = func(cls, *args, **kwargs) + + if ir_execution_context is not None and quent_state is not None: + _end_quent_do_evaluate_events( + cls, frames, result, ir_execution_context, quent_state + ) - return result + return result - return wrapper + return wrapper @contextlib.contextmanager diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index c2aea367f683..d8f484c01503 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -31,7 +31,7 @@ attach_cached_parquet_metadata, prefetch_parquet_file_metadata_for_ir, ) -from cudf_polars.quent._plan import build_plan +from cudf_polars.quent._plan import build_plan, build_quent_operator_map from cudf_polars.streaming.actor_graph.collectives import ReserveOpIDs from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id from cudf_polars.streaming.actor_graph.core import generate_network @@ -421,6 +421,43 @@ def _find_memory_error(exc: BaseException) -> MemoryError | None: return None +def _declare_network_channels( + comm: Communicator, + local_quent_context: LocalQuentContext, +) -> None: + """ + Declare network link channels for inter-rank communication. + + Creates a Network resource group and one Channel per remote rank, + emitting their lifecycle events to the quent logger. + """ + if comm.nranks <= 1: + return + + from cudf_polars.quent._types import Channel, Network + + network = Network(engine_id=local_quent_context.context.engine.id) + local_quent_context.logger.emit(network.declare()) + local_quent_context.network = network + + link_channels: dict[int, Channel] = {} + for target_rank in range(comm.nranks): + if target_rank == comm.rank: + continue + link = Channel( + instance_name=f"rank-{comm.rank} -> rank-{target_rank}", + resource_type_name="Link", + parent_group_id=network.id, + source=local_quent_context.device_memory, + target=local_quent_context.device_memory, + ) + local_quent_context.logger.emit(link.initializing()) + local_quent_context.logger.emit(link.operating()) + link_channels[target_rank] = link + + local_quent_context.link_channels = link_channels + + def execute_ir_on_rank( ctx: Context, comm: Communicator, @@ -430,6 +467,9 @@ def execute_ir_on_rank( config_options: ConfigOptions[StreamingExecutor], stats: StatsCollector, collective_id_map: dict[IR, list[int]], + *, + quent_operator_map: dict[IR, cudf_polars.quent.Operator] | None = None, + local_quent_context: LocalQuentContext | None = None, ) -> tuple[pl.DataFrame, list[ChannelMetadata]]: """ Execute a Polars IR query on a single rank's GPU. @@ -456,6 +496,12 @@ def execute_ir_on_rank( Statistics collector. collective_id_map Mapping from IR nodes to their pre-allocated collective operation IDs. + quent_operator_map + Mapping from IR nodes to their Quent operators, or ``None`` when tracing + is disabled. + local_quent_context + The local Quent context for this rank, or ``None`` when tracing is + disabled. Returns ------- @@ -476,6 +522,8 @@ def execute_ir_on_rank( ir_context=ir_context, collective_id_map=collective_id_map, metadata_collector=metadata_collector, + quent_operator_map=quent_operator_map, + local_quent_context=local_quent_context, ) try: @@ -703,8 +751,11 @@ def evaluate_on_rank( """ stats = allgather_stats(comm, ctx.br(), ir, config_options, py_executor) + physical_op_by_id: dict[str, cudf_polars.quent.Operator] | None = None + quent_operator_map: dict[IR, cudf_polars.quent.Operator] | None = None if config_options.executor.quent_context is not None: assert local_quent_context is not None + _declare_network_channels(comm, local_quent_context) logical_plan_id = ir.get_stable_plan_id() plan, ops, ports, logical_op_by_id = build_plan( ir, @@ -731,7 +782,7 @@ def evaluate_on_rank( if config_options.executor.quent_context is not None: assert local_quent_context is not None physical_plan_id = uuid.uuid4() - local_quent_context.context._emit_physical_plan_events( + physical_op_by_id = local_quent_context.context._emit_physical_plan_events( local_quent_context.logger, ir, config_options, @@ -741,6 +792,7 @@ def evaluate_on_rank( node_map=node_map, logical_op_by_id=logical_op_by_id, ) + quent_operator_map = build_quent_operator_map(ir, physical_op_by_id) ir_context = IRExecutionContext( py_executor, get_cuda_stream=ctx.br().stream_pool.get_stream, query_id=query_id ) @@ -754,7 +806,7 @@ def evaluate_on_rank( attach_cached_parquet_metadata(ir, cached_parquet_info_map) with ReserveOpIDs(ir, config_options) as collective_id_map: - return execute_ir_on_rank( + result = execute_ir_on_rank( ctx, comm, ir, @@ -763,4 +815,13 @@ def evaluate_on_rank( config_options, stats, collective_id_map, + quent_operator_map=quent_operator_map, + local_quent_context=local_quent_context, ) + + if local_quent_context is not None: + for link in local_quent_context.link_channels.values(): + local_quent_context.logger.emit(link.finalizing()) + local_quent_context.logger.emit(link.exit()) + + return result diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index fb5c7eed3ea3..2c2c4124300a 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -41,7 +41,11 @@ HardwareBindingPolicy, bind_to_gpu, ) -from cudf_polars.quent._context import LocalQuentContext +from cudf_polars.quent._context import ( + LocalQuentContext, + declare_worker_resources, + finalize_worker_resources, +) from cudf_polars.utils.config import DaskContext, MemoryResourceConfig if TYPE_CHECKING: @@ -119,7 +123,10 @@ class _WorkerContext: base_mr: rmm.mr.DeviceMemoryResource | None quent_logger: cudf_polars.quent._logging.QuentLogger | None quent_worker: cudf_polars.quent._types.Worker - mr: RmmResourceAdaptor | None = None # set after `Context` is built (below). + mr: RmmResourceAdaptor | None = None + device_memory: cudf_polars.quent._types.Memory | None = None + disk_to_device_channel: cudf_polars.quent._types.Channel | None = None + thread_pool: cudf_polars.quent._types.ThreadPool | None = None def _setup_root( @@ -305,12 +312,21 @@ def _setup_worker( ) if quent_context is not None: - quent_logger: cudf_polars.quent._logging.QuentLogger | None = ( - cudf_polars.quent._logging.QuentLogger() - ) + quent_logger = cudf_polars.quent._logging.QuentLogger() else: quent_logger = None + device_memory = None + disk_to_device_channel = None + thread_pool = None + if quent_logger is not None: + device_memory, disk_to_device_channel, thread_pool = declare_worker_resources( + quent_logger, + instance_suffix=f"rank-{comm.rank}", + engine_id=engine_id, + worker_id=worker_id, + ) + mp_ctx = _WorkerContext( comm=comm, ctx=ctx, @@ -319,6 +335,9 @@ def _setup_worker( mr=mr, quent_worker=quent_worker, quent_logger=quent_logger, + device_memory=device_memory, + disk_to_device_channel=disk_to_device_channel, + thread_pool=thread_pool, ) setattr(dask_worker, attr, mp_ctx) if mp_ctx.quent_logger is not None: @@ -347,6 +366,12 @@ def _teardown_worker( traces = [] if mp_ctx is not None: if mp_ctx.quent_worker is not None and mp_ctx.quent_logger is not None: + if mp_ctx.device_memory is not None: + finalize_worker_resources( + mp_ctx.quent_logger, + device_memory=mp_ctx.device_memory, + disk_to_device_channel=mp_ctx.disk_to_device_channel, + ) mp_ctx.quent_logger.emit(mp_ctx.quent_worker._exit()) traces = mp_ctx.quent_logger.drain() @@ -497,10 +522,15 @@ def _worker_evaluate( local_quent_context: LocalQuentContext | None = None if quent_context is not None: assert mp_ctx.quent_logger is not None + assert mp_ctx.device_memory is not None + assert mp_ctx.thread_pool is not None local_quent_context = LocalQuentContext( context=quent_context, worker=mp_ctx.quent_worker, logger=mp_ctx.quent_logger, + thread_pool_id=mp_ctx.thread_pool.id, + device_memory=mp_ctx.device_memory, + disk_to_device_channel=mp_ctx.disk_to_device_channel, ) # evaluate_on_rank always collects metadata internally so we can read # metadata[-1].duplicated to decide whether to suppress this rank's output. @@ -1023,6 +1053,7 @@ def shutdown(self) -> None: if quent_context is not None: assert self._quent_logger is not None quent_context._emit_engine_exit_events(self._quent_logger) + quent_context.emit_resource_exit_events(self._quent_logger) if ctx.owned_client is not None: ctx.owned_client.close() if ctx.owned_cluster is not None: diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index 63e7e15e7d1f..2b9ba237e5f6 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -39,7 +39,11 @@ HardwareBindingPolicy, bind_to_gpu, ) -from cudf_polars.quent._context import LocalQuentContext +from cudf_polars.quent._context import ( + LocalQuentContext, + declare_worker_resources, + finalize_worker_resources, +) from cudf_polars.quent._types import Worker from cudf_polars.utils.config import MemoryResourceConfig, RayContext @@ -231,8 +235,21 @@ def __init__( engine=engine, instance_name=f"RankActor-{worker_id.hex[:8]}", ) + self._device_memory = None + self._disk_to_device_channel = None + self._quent_thread_pool = None if self._quent_logger is not None: self._quent_logger.emit(self._quent_worker._init()) + ( + self._device_memory, + self._disk_to_device_channel, + self._quent_thread_pool, + ) = declare_worker_resources( + self._quent_logger, + instance_suffix=f"RankActor-{worker_id.hex[:8]}", + engine_id=engine.id, + worker_id=worker_id, + ) def setup_root(self) -> bytes: """ @@ -335,6 +352,15 @@ def _exit(self) -> list[dict[str, Any]]: # Maybe generalize this to all application-level things, # followed by framework (ray) level things. if self._quent_worker is not None and self._quent_logger is not None: + if ( + self._device_memory is not None + and self._disk_to_device_channel is not None + ): + finalize_worker_resources( + self._quent_logger, + device_memory=self._device_memory, + disk_to_device_channel=self._disk_to_device_channel, + ) self._quent_logger.emit(self._quent_worker._exit()) return self._drain_quent_events() return [] @@ -446,10 +472,15 @@ def evaluate_polars_ir( local_quent_context: LocalQuentContext | None = None if quent_context is not None: assert self._quent_logger is not None + assert self._device_memory is not None + assert self._quent_thread_pool is not None local_quent_context = LocalQuentContext( context=quent_context, worker=self._quent_worker, logger=self._quent_logger, + thread_pool_id=self._quent_thread_pool.id, + device_memory=self._device_memory, + disk_to_device_channel=self._disk_to_device_channel, ) # evaluate_on_rank always collects metadata internally so we can read # metadata[-1].duplicated to decide whether to suppress this rank's @@ -916,6 +947,7 @@ def shutdown(self) -> None: if quent_context is not None: assert self._quent_logger is not None quent_context._emit_engine_exit_events(self._quent_logger) + quent_context.emit_resource_exit_events(self._quent_logger) self._rank_actors = None super().shutdown() diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index 17b4c7b5ba5c..eb8d16385258 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -43,7 +43,11 @@ HardwareBindingPolicy, bind_to_gpu, ) -from cudf_polars.quent._context import LocalQuentContext +from cudf_polars.quent._context import ( + LocalQuentContext, + declare_worker_resources, + finalize_worker_resources, +) from cudf_polars.quent._types import Worker from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id from cudf_polars.streaming.actor_graph.utils import set_memory_resource @@ -121,14 +125,24 @@ def evaluate_pipeline_spmd_mode( assert quent_logger is not None quent_context._emit_query_group_events(quent_logger) quent_context._emit_query_events(quent_logger) + worker_id = config_options.executor.spmd_context.worker_id + device_memory, disk_to_device_channel, thread_pool = declare_worker_resources( + quent_logger, + instance_suffix=f"rank-{comm.rank}", + engine_id=quent_context.engine.id, + worker_id=worker_id, + ) local_quent_context = LocalQuentContext( context=quent_context, worker=Worker( - id=config_options.executor.spmd_context.worker_id, + id=worker_id, engine=quent_context.engine, instance_name=f"rank-{comm.rank}", ), logger=quent_logger, + thread_pool_id=thread_pool.id, + device_memory=device_memory, + disk_to_device_channel=disk_to_device_channel, ) df, metadata = evaluate_on_rank( @@ -142,6 +156,12 @@ def evaluate_pipeline_spmd_mode( ) if quent_context is not None: assert config_options.executor.spmd_context.quent_logger is not None + assert local_quent_context is not None + finalize_worker_resources( + config_options.executor.spmd_context.quent_logger, + device_memory=local_quent_context.device_memory, + disk_to_device_channel=local_quent_context.disk_to_device_channel, + ) quent_context._emit_query_exit_events( config_options.executor.spmd_context.quent_logger ) @@ -782,6 +802,7 @@ def shutdown(self) -> None: if quent_context is not None: assert self._quent_logger is not None quent_context._emit_engine_exit_events(self._quent_logger) + quent_context.emit_resource_exit_events(self._quent_logger) super().shutdown() diff --git a/python/cudf_polars/cudf_polars/quent/__init__.py b/python/cudf_polars/cudf_polars/quent/__init__.py index 4656f8a67101..877f303b0b57 100644 --- a/python/cudf_polars/cudf_polars/quent/__init__.py +++ b/python/cudf_polars/cudf_polars/quent/__init__.py @@ -5,26 +5,44 @@ from __future__ import annotations -from cudf_polars.quent._context import QuentContext +from cudf_polars.quent._context import ( + LocalQuentContext, + QuentContext, + QuentIRExecutionContext, +) from cudf_polars.quent._types import ( Attribute, + Channel, Engine, HomogeneousListValue, Implementation, + Network, + Operator, Query, QueryGroup, ScalarValue, + Statistics, + Task, Value, + Worker, ) __all__ = [ "Attribute", + "Channel", "Engine", "HomogeneousListValue", "Implementation", + "LocalQuentContext", + "Network", + "Operator", "QuentContext", + "QuentIRExecutionContext", "Query", "QueryGroup", "ScalarValue", + "Statistics", + "Task", "Value", + "Worker", ] diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index 627d864cfa69..4e4281a8c4f3 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -16,10 +16,14 @@ ) from cudf_polars.quent._types import ( Attribute, + Channel, Engine, Implementation, + Memory, + Processor, Query, QueryGroup, + ThreadPool, ) if TYPE_CHECKING: @@ -28,6 +32,7 @@ from cudf_polars.dsl.ir import IR from cudf_polars.quent._logging import QuentLogger from cudf_polars.quent._types import ( + Network, Operator, Plan, Port, @@ -36,7 +41,9 @@ from cudf_polars.utils.config import ConfigOptions, StreamingExecutor __all__ = [ + "LocalQuentContext", "QuentContext", + "QuentIRExecutionContext", ] @@ -66,6 +73,11 @@ class QuentContext: def __post_init__(self) -> None: object.__setattr__(self, "_query_group_cache_", set()) + object.__setattr__(self, "_processor_map_", {}) + + @property + def _processor_map(self) -> dict[int, Processor]: + return self._processor_map_ # type: ignore[attr-defined] def serialize(self) -> bytes: """ @@ -304,8 +316,83 @@ def _emit_physical_plan_events( parent_operators_by_node_id=parent_operators_by_node_id, ) + def get_or_declare_processor( + self, quent_logger: QuentLogger, thread_ident: int, pool_id: uuid.UUID + ) -> Processor: + """Get (or declare a new) Quent Processor for a CPU thread.""" + if thread_ident in self._processor_map: + return self._processor_map[thread_ident] + + processor = Processor(pool_id=pool_id) + self._processor_map[thread_ident] = processor + quent_logger.emit(processor.initializing()) + quent_logger.emit(processor.operating()) + return processor + + def emit_resource_exit_events(self, quent_logger: QuentLogger) -> None: + """Emit finalizing/exit events for declared processors.""" + for processor in self._processor_map.values(): + quent_logger.emit(processor.finalizing()) + quent_logger.emit(processor.exit()) + + +def declare_worker_resources( + logger: QuentLogger, + *, + instance_suffix: str, + engine_id: uuid.UUID, + worker_id: uuid.UUID, +) -> tuple[Memory, Channel, ThreadPool]: + """ + Declare per-worker Quent resources and emit their lifecycle events. -@dataclasses.dataclass(frozen=True, kw_only=True) + Returns device memory, disk-to-device channel, and thread pool handles. + """ + device_memory = Memory( + instance_name=f"{instance_suffix} device memory", + resource_type_name="memory", + parent_group_id=engine_id, + ) + filesystem = Memory( + instance_name=f"{instance_suffix} filesystem", + resource_type_name="filesystem", + parent_group_id=worker_id, + ) + disk_to_device_channel = Channel( + instance_name=f"{instance_suffix} disk -> device", + resource_type_name="DiskToDevice", + parent_group_id=worker_id, + source=filesystem, + target=device_memory, + ) + thread_pool = ThreadPool(worker_id=worker_id) + logger.emit(device_memory.initializing()) + logger.emit(device_memory.operating(0)) + logger.emit(filesystem.initializing()) + logger.emit(filesystem.operating(0)) + logger.emit(disk_to_device_channel.initializing()) + logger.emit(disk_to_device_channel.operating()) + logger.emit(thread_pool.declare()) + return device_memory, disk_to_device_channel, thread_pool + + +def finalize_worker_resources( + logger: QuentLogger, + *, + device_memory: Memory, + disk_to_device_channel: Channel | None, +) -> None: + """Emit finalizing/exit events for per-worker Quent resources.""" + if disk_to_device_channel is not None: + logger.emit(disk_to_device_channel.finalizing()) + logger.emit(disk_to_device_channel.exit()) + logger.emit(disk_to_device_channel.source.finalizing()) + logger.emit(disk_to_device_channel.source.exit()) + logger.emit(device_memory.finalizing()) + logger.emit(device_memory.exit()) + + +@dataclasses.dataclass(kw_only=True) class LocalQuentContext: """ A Quent Context that is only ever used on the local worker rank. @@ -317,3 +404,32 @@ class LocalQuentContext: context: QuentContext worker: Worker logger: QuentLogger + thread_pool_id: uuid.UUID + device_memory: Memory + disk_to_device_channel: Channel | None = None + network: Network | None = None + link_channels: dict[int, Channel] = dataclasses.field(default_factory=dict) + + +@dataclasses.dataclass(kw_only=True) +class QuentIRExecutionContext(LocalQuentContext): + """Like ``LocalQuentContext``, but with a Quent Operator bound too.""" + + quent_operator: Operator + + @classmethod + def from_execution_context( + cls, execution_context: LocalQuentContext, quent_operator: Operator + ) -> Self: + """Create a ``QuentIRExecutionContext`` from a ``LocalQuentContext``.""" + return cls( + quent_operator=quent_operator, + context=execution_context.context, + worker=execution_context.worker, + logger=execution_context.logger, + thread_pool_id=execution_context.thread_pool_id, + device_memory=execution_context.device_memory, + disk_to_device_channel=execution_context.disk_to_device_channel, + network=execution_context.network, + link_channels=execution_context.link_channels, + ) diff --git a/python/cudf_polars/cudf_polars/quent/_plan.py b/python/cudf_polars/cudf_polars/quent/_plan.py index edbb41d1efb9..542b525c2f0d 100644 --- a/python/cudf_polars/cudf_polars/quent/_plan.py +++ b/python/cudf_polars/cudf_polars/quent/_plan.py @@ -8,6 +8,7 @@ import functools from typing import TYPE_CHECKING +from cudf_polars.dsl.traversal import traversal from cudf_polars.quent._types import ( Attribute, Edge, @@ -177,3 +178,16 @@ def build_parent_operators_map( ] for physical_sid, logical_sids in node_map.items() } + + +def build_quent_operator_map( + ir: IR, + physical_op_by_id: dict[str, Operator], +) -> dict[IR, Operator]: + """Build a map from IR nodes to their physical-plan Quent operators.""" + result: dict[IR, Operator] = {} + for node in traversal([ir]): + stable_id = str(node.get_stable_id()) + if stable_id in physical_op_by_id: + result[node] = physical_op_by_id[stable_id] + return result diff --git a/python/cudf_polars/cudf_polars/quent/_types.py b/python/cudf_polars/cudf_polars/quent/_types.py index 37299162a05b..cb6e2d068080 100644 --- a/python/cudf_polars/cudf_polars/quent/_types.py +++ b/python/cudf_polars/cudf_polars/quent/_types.py @@ -12,7 +12,7 @@ import sys import time import uuid -from typing import Any, TypeAlias +from typing import Any, Literal, TypeAlias from cudf_polars import __version__ @@ -30,6 +30,10 @@ class EventName(enum.Enum): OPERATOR = "Operator" PORT = "Port" TASK = "Task" + MEMORY = "Memory" + CHANNEL = "Channel" + THREAD_POOL = "ThreadPool" + PROCESSOR = "Processor" if sys.version_info >= (3, 14): # pragma: no cover; requires Python 3.14+ @@ -72,6 +76,50 @@ def to_dict(self) -> dict[str, Any]: } +@dataclasses.dataclass(frozen=True, slots=True) +class StatisticsAttribute: + """Typed key/value pair for Quent statistics custom attributes.""" + + key: str + value_type: Literal["U64", "F64", "String"] + value: int | float | str + + def to_dict(self) -> dict[str, Any]: + return {"key": self.key, "value": {self.value_type: self.value}} + + +@dataclasses.dataclass(frozen=True, slots=True) +class Statistics: + """Operator statistics payload.""" + + input_bytes: int + output_bytes: int + output_rows: int + custom_attributes: list[StatisticsAttribute] = dataclasses.field( + default_factory=list + ) + + def to_dict(self) -> dict[str, Any]: + """Serialize to Quent's custom attributes format.""" + base_attributes: list[StatisticsAttribute] = [ + StatisticsAttribute( + key="input_bytes", value_type="U64", value=self.input_bytes + ), + StatisticsAttribute( + key="output_bytes", value_type="U64", value=self.output_bytes + ), + StatisticsAttribute( + key="output_rows", value_type="U64", value=self.output_rows + ), + ] + return { + "custom_attributes": [ + *(attribute.to_dict() for attribute in base_attributes), + *(attribute.to_dict() for attribute in self.custom_attributes), + ] + } + + @dataclasses.dataclass(frozen=True, slots=True) class Operator: """ @@ -123,6 +171,14 @@ def declare(self, timestamp: int | None = None) -> Event: data={EventName.OPERATOR.value: {"Declaration": self.to_dict()}}, ) + def statistics(self, statistics: Statistics, timestamp: int | None = None) -> Event: + """Emit post-execution operator statistics.""" + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={EventName.OPERATOR.value: {"Statistics": statistics.to_dict()}}, + ) + @dataclasses.dataclass(frozen=True, slots=True) class Engine: @@ -493,3 +549,406 @@ def _deserialize_value(value: dict[str, Any] | None) -> Value | None: if variant == "String": return str(deserialized) raise ValueError(f"Unsupported Quent custom attribute variant: '{variant}'") + + +# Resource types + + +@dataclasses.dataclass(frozen=True, slots=True, kw_only=True) +class Memory: + """A Quent Memory resource.""" + + id: uuid.UUID = dataclasses.field(default_factory=new_quent_id) + instance_name: str + resource_type_name: str + parent_group_id: uuid.UUID + + def initializing(self, timestamp: int | None = None) -> Event: + """Build a Quent Memory Initializing event.""" + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={ + EventName.MEMORY.value: { + "seq": 0, + "state": { + "MemoryInitializing": { + "instance_name": self.instance_name, + "parent_group_id": str(self.parent_group_id), + "resource_type_name": self.resource_type_name, + } + }, + } + }, + ) + + def operating(self, capacity_bytes: int, timestamp: int | None = None) -> Event: + """Build a Quent Memory Operating event.""" + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={ + EventName.MEMORY.value: { + "seq": 1, + "state": {"MemoryOperating": {"capacity_bytes": capacity_bytes}}, + } + }, + ) + + def finalizing(self, timestamp: int | None = None) -> Event: + """Build a Quent Memory Finalizing event.""" + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={ + EventName.MEMORY.value: {"seq": 2, "state": {"MemoryFinalizing": None}} + }, + ) + + def exit(self, timestamp: int | None = None) -> Event: + """Build a Quent Memory Exit event.""" + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={EventName.MEMORY.value: {"seq": 3, "state": "Exit"}}, + ) + + +@dataclasses.dataclass(frozen=True, slots=True, kw_only=True) +class Channel: + """ + A Quent Channel resource. + + A Channel is a unidirectional data-transfer resource between two entities. + Examples include disk-to-device I/O channels and inter-rank network links. + """ + + id: uuid.UUID = dataclasses.field(default_factory=new_quent_id) + instance_name: str + resource_type_name: str + parent_group_id: uuid.UUID + source: Memory + target: Memory + + def initializing(self, timestamp: int | None = None) -> Event: + """Build a Quent Channel Initializing event.""" + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={ + EventName.CHANNEL.value: { + "seq": 0, + "state": { + "ChannelInitializing": { + "instance_name": self.instance_name, + "parent_group_id": str(self.parent_group_id), + "resource_type_name": self.resource_type_name, + "source_id": str(self.source.id), + "target_id": str(self.target.id), + } + }, + } + }, + ) + + def operating( + self, capacity_bytes: int | None = None, timestamp: int | None = None + ) -> Event: + """Build a Quent Channel Operating event.""" + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={ + EventName.CHANNEL.value: { + "seq": 1, + "state": {"ChannelOperating": {"capacity_bytes": capacity_bytes}}, + } + }, + ) + + def finalizing(self, timestamp: int | None = None) -> Event: + """Build a Quent Channel Finalizing event.""" + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={ + EventName.CHANNEL.value: { + "seq": 2, + "state": {"ChannelFinalizing": None}, + } + }, + ) + + def exit(self, timestamp: int | None = None) -> Event: + """Build a Quent Channel Exit event.""" + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={EventName.CHANNEL.value: {"seq": 3, "state": "Exit"}}, + ) + + +@dataclasses.dataclass(frozen=True, slots=True, kw_only=True) +class Network: + """A Quent Network resource group.""" + + id: uuid.UUID = dataclasses.field(default_factory=new_quent_id) + engine_id: uuid.UUID + + def declare(self, timestamp: int | None = None) -> Event: + """Build a Network declaration event.""" + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={ + "Network": { + "Declaration": { + "instance_name": "Network", + "parent_group_id": str(self.engine_id), + } + } + }, + ) + + +@dataclasses.dataclass(frozen=True, slots=True, kw_only=True) +class ThreadPool: + """A Quent ThreadPool resource group.""" + + id: uuid.UUID = dataclasses.field(default_factory=new_quent_id) + worker_id: uuid.UUID + + def declare(self, timestamp: int | None = None) -> Event: + """Build a ThreadPool declaration event.""" + instance_name = f"Thread Pool {self.id.hex[:8]}" + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={ + EventName.THREAD_POOL.value: { + "Declaration": { + "instance_name": instance_name, + "parent_group_id": str(self.worker_id), + } + }, + }, + ) + + +@dataclasses.dataclass(frozen=True, slots=True, kw_only=True) +class Processor: + """A Quent Processor resource representing a CPU thread.""" + + id: uuid.UUID = dataclasses.field(default_factory=new_quent_id) + pool_id: uuid.UUID + + def initializing(self, timestamp: int | None = None) -> Event: + """Build a Quent Processor Initializing event.""" + instance_name = f"Thread {self.id.hex[:8]}" + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={ + EventName.PROCESSOR.value: { + "seq": 0, + "state": { + "ProcessorInitializing": { + "instance_name": instance_name, + "parent_group_id": str(self.pool_id), + "resource_type_name": "processor", + } + }, + } + }, + ) + + def operating(self, timestamp: int | None = None) -> Event: + """Build a Quent Processor Operating event.""" + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={ + EventName.PROCESSOR.value: { + "seq": 1, + "state": {"ProcessorOperating": None}, + } + }, + ) + + def finalizing(self, timestamp: int | None = None) -> Event: + """Build a Quent Processor Finalizing event.""" + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={ + EventName.PROCESSOR.value: { + "seq": 2, + "state": {"ProcessorFinalizing": None}, + } + }, + ) + + def exit(self, timestamp: int | None = None) -> Event: + """Build a Quent Processor Exit event.""" + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={EventName.PROCESSOR.value: {"seq": 3, "state": "Exit"}}, + ) + + +@dataclasses.dataclass(frozen=True, slots=True, kw_only=True) +class Task: + """A Quent Task representing a unit of work on an operator.""" + + id: uuid.UUID = dataclasses.field(default_factory=new_quent_id) + operator_id: uuid.UUID + instance_name: str | None = None + + def queueing(self, timestamp: int | None = None) -> Event: + """Build a Quent Task Queueing event.""" + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={ + EventName.TASK.value: { + "seq": 0, + "state": { + "Queueing": { + "instance_name": self.instance_name or self.id.hex[:8], + "operator_id": str(self.operator_id), + } + }, + } + }, + ) + + def allocating( + self, + resource_id: uuid.UUID, + capacity: int | None = None, + timestamp: int | None = None, + ) -> Event: + """Build a Quent Task Allocating event.""" + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={ + EventName.TASK.value: { + "seq": 1, + "state": { + "Allocating": { + "use_thread": { + "resource_id": str(resource_id), + "capacity": capacity, + } + } + }, + } + }, + ) + + def loading( + self, + use_thread: Processor | None = None, + use_channel: Channel | None = None, + channel_capacity_bytes: int = 0, + use_memory: Memory | None = None, + memory_capacity_bytes: int = 0, + timestamp: int | None = None, + ) -> Event: + """Build a Quent Task Loading event.""" + loading_data: dict[str, dict[str, Any]] = {} + if use_thread is not None: + loading_data["use_thread"] = { + "resource_id": str(use_thread.id), + "capacity": None, + } + if use_channel is not None: + loading_data["use_fs_to_mem"] = { + "resource_id": str(use_channel.id), + "capacity": {"capacity_bytes": channel_capacity_bytes}, + } + if use_memory is not None: + loading_data["use_memory"] = { + "resource_id": str(use_memory.id), + "capacity": {"capacity_bytes": memory_capacity_bytes}, + } + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={ + EventName.TASK.value: { + "seq": 2, + "state": {"Loading": loading_data}, + } + }, + ) + + def computing( + self, + use_thread: Processor | None = None, + use_memory: Memory | None = None, + memory_capacity_bytes: int = 0, + timestamp: int | None = None, + ) -> Event: + """Build a Quent Task Computing event.""" + computing_data: dict[str, dict[str, Any]] = {} + if use_thread is not None: + computing_data["use_thread"] = { + "resource_id": str(use_thread.id), + "capacity": None, + } + if use_memory is not None: + computing_data["use_memory"] = { + "resource_id": str(use_memory.id), + "capacity": {"capacity_bytes": memory_capacity_bytes}, + } + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={ + EventName.TASK.value: { + "seq": 3, + "state": {"Computing": computing_data}, + } + }, + ) + + def sending( + self, + use_thread: Processor | None = None, + use_link: Channel | None = None, + link_capacity_bytes: int = 0, + timestamp: int | None = None, + ) -> Event: + """Build a Quent Task Sending event.""" + sending_data: dict[str, dict[str, Any]] = {} + if use_thread is not None: + sending_data["use_thread"] = { + "resource_id": str(use_thread.id), + "capacity": None, + } + if use_link is not None: + sending_data["use_link"] = { + "resource_id": str(use_link.id), + "capacity": {"capacity_bytes": link_capacity_bytes}, + } + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={ + EventName.TASK.value: { + "seq": 4, + "state": {"Sending": sending_data}, + } + }, + ) + + def exit(self, timestamp: int | None = None) -> Event: + """Build a Quent Task Exit event.""" + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={EventName.TASK.value: {"seq": 5, "state": "Exit"}}, + ) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/shuffle.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/shuffle.py index 577db3928913..c1436ac52405 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/shuffle.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/shuffle.py @@ -27,9 +27,11 @@ from rapidsmpf.streaming.core.context import Context from rapidsmpf.streaming.core.message import Message +import cudf_polars.quent from cudf_polars.dsl.expr import Col from cudf_polars.streaming.actor_graph.dispatch import ( generate_ir_sub_network, + ir_context_for_node, ) from cudf_polars.streaming.actor_graph.nodes import shutdown_on_error from cudf_polars.streaming.actor_graph.utils import ( @@ -398,15 +400,32 @@ async def _global_shuffle( skip_insert = metadata_in.duplicated and comm.rank != 0 shuffle = ShuffleManager(context, comm, num_partitions, collective_id) + total_bytes_inserted: int = 0 async with shuffle.inserting() as inserter: while (msg := await ch_in.recv(context)) is not None: if not skip_insert: - inserter.insert_hash( - TableChunk.from_message( - msg, br=context.br() - ).make_available_and_spill(context.br(), allow_overbooking=True), - columns_to_hash, - ) + chunk = TableChunk.from_message( + msg, br=context.br() + ).make_available_and_spill(context.br(), allow_overbooking=True) + total_bytes_inserted += chunk.data_alloc_size() + inserter.insert_hash(chunk, columns_to_hash) + + if ( + (quent_ctx := ir_context.quent_ir_execution_context) is not None + and quent_ctx.link_channels + and total_bytes_inserted > 0 + ): + quent_task = cudf_polars.quent.Task( + instance_name=f"Shuffle-send-{quent_ctx.quent_operator.id.hex[:8]}", + operator_id=quent_ctx.quent_operator.id, + ) + quent_ctx.logger.emit(quent_task.queueing()) + bytes_per_rank = total_bytes_inserted // comm.nranks + for link in quent_ctx.link_channels.values(): + quent_ctx.logger.emit( + quent_task.sending(use_link=link, link_capacity_bytes=bytes_per_rank) + ) + quent_ctx.logger.emit(quent_task.exit()) for partition_id in shuffle.local_partitions(): stream = ir_context.get_cuda_stream() @@ -505,6 +524,7 @@ def _( # Create output ChannelManager channels[ir] = ChannelManager(rec.state["context"]) + ir_context = ir_context_for_node(rec, ir) # Complete shuffle node nodes[ir] = [ @@ -512,7 +532,7 @@ def _( context, rec.state["comm"], ir, - rec.state["ir_context"], + ir_context, ch_in=channels[child].reserve_output_slot(), ch_out=channels[ir].reserve_input_slot(), columns_to_hash=columns_to_hash, diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py index 8bfb1c8a86ab..2999f80bb9ce 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py @@ -28,7 +28,10 @@ from cudf_polars.dsl.utils.naming import names_to_indices, unique_names from cudf_polars.streaming.actor_graph.collectives.allgather import AllGatherManager from cudf_polars.streaming.actor_graph.collectives.shuffle import ShuffleManager -from cudf_polars.streaming.actor_graph.dispatch import generate_ir_sub_network +from cudf_polars.streaming.actor_graph.dispatch import ( + generate_ir_sub_network, + ir_context_for_node, +) from cudf_polars.streaming.actor_graph.nodes import ( default_node_single, shutdown_on_error, @@ -685,6 +688,7 @@ def _sort_rapidsmpf_network(ir: Sort, rec: SubNetGenerator) -> tuple[dict, dict] executor = rec.state["config_options"].executor partition_info = rec.state["partition_info"] dynamic = executor.dynamic_planning is not None + ir_context = ir_context_for_node(rec, ir) if partition_info[ir].count == 1 and ( not dynamic or isinstance(ir.children[0], Repartition) @@ -695,7 +699,7 @@ def _sort_rapidsmpf_network(ir: Sort, rec: SubNetGenerator) -> tuple[dict, dict] default_node_single( rec.state["context"], ir, - rec.state["ir_context"], + ir_context, channels[ir].reserve_input_slot(), channels[ir.children[0]].reserve_output_slot(), ) @@ -718,7 +722,7 @@ def _sort_rapidsmpf_network(ir: Sort, rec: SubNetGenerator) -> tuple[dict, dict] rec.state["context"], rec.state["comm"], ir, - rec.state["ir_context"], + ir_context, ch_in=channels[child].reserve_output_slot(), ch_out=channels[ir].reserve_input_slot(), by=by, diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py index c923d342db99..2da64d30fa36 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py @@ -12,6 +12,7 @@ from rapidsmpf.streaming.core.leaf_actor import pull_from_channel import cudf_polars.dsl.tracing +import cudf_polars.quent from cudf_polars.dsl.ir import ( DataFrameScan, Join, @@ -218,6 +219,8 @@ def generate_network( ir_context: IRExecutionContext, collective_id_map: dict[IR, list[int]], metadata_collector: list[ChannelMetadata] | None, + quent_operator_map: dict[IR, cudf_polars.quent.Operator] | None = None, + local_quent_context: cudf_polars.quent.LocalQuentContext | None = None, ) -> tuple[list[Any], DeferredMessages]: """ Translate the IR graph to a RapidsMPF streaming network. @@ -244,6 +247,12 @@ def generate_network( The list to collect the final metadata. This list will be mutated when the network is executed. If None, metadata will not be collected. + quent_operator_map + Mapping from IR nodes to their Quent operators, or ``None`` when tracing + is disabled. + local_quent_context + The local Quent context for this rank, or ``None`` when tracing is + disabled. Returns ------- @@ -276,6 +285,8 @@ def generate_network( "max_io_threads": max_io_threads_local, "stats": stats, "collective_id_map": collective_id_map, + "quent_operator_map": quent_operator_map, + "quent_execution_context": local_quent_context, } mapper: SubNetGenerator = CachingVisitor( generate_ir_sub_network_wrapper, state=state diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py index 2554d95fe750..57d420950e0a 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py @@ -1,9 +1,10 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Dispatching for the RapidsMPF streaming runtime.""" from __future__ import annotations +import dataclasses from functools import singledispatch from typing import TYPE_CHECKING, Any, NamedTuple, TypeAlias, TypedDict @@ -15,6 +16,7 @@ from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.streaming.core.context import Context + import cudf_polars.quent from cudf_polars.dsl.ir import IR, IRExecutionContext from cudf_polars.streaming.actor_graph.utils import ChannelManager from cudf_polars.streaming.base import ( @@ -58,6 +60,10 @@ class GenState(TypedDict): Statistics collector. collective_id_map The mapping of IR nodes to lists of collective IDs. + quent_operator_map + Mapping from IR nodes to physical-plan Quent operators. + quent_execution_context + Rank-local Quent execution context. """ context: Context @@ -69,6 +75,27 @@ class GenState(TypedDict): max_io_threads: int stats: StatsCollector collective_id_map: dict[IR, list[int]] + quent_operator_map: dict[IR, cudf_polars.quent.Operator] | None + quent_execution_context: cudf_polars.quent.LocalQuentContext | None + + +def ir_context_for_node(rec: SubNetGenerator, ir: IR) -> IRExecutionContext: + """Return ``ir_context`` with the physical Quent operator bound when tracing.""" + import cudf_polars.quent + + ir_context = rec.state["ir_context"] + quent_operator_map = rec.state["quent_operator_map"] + quent_execution_context = rec.state["quent_execution_context"] + if quent_operator_map is not None and quent_execution_context is not None: + quent_operator = quent_operator_map[ir] + return dataclasses.replace( + ir_context, + quent_ir_execution_context=cudf_polars.quent.QuentIRExecutionContext.from_execution_context( + execution_context=quent_execution_context, + quent_operator=quent_operator, + ), + ) + return ir_context SubNetGenerator: TypeAlias = GenericTransformer[ diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py index 10b8d673d06d..c10e7cc4a087 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py @@ -25,6 +25,7 @@ from cudf_polars.streaming.actor_graph.collectives.shuffle import ShuffleManager from cudf_polars.streaming.actor_graph.dispatch import ( generate_ir_sub_network, + ir_context_for_node, ) from cudf_polars.streaming.actor_graph.tracing import send_chunk from cudf_polars.streaming.actor_graph.utils import ( @@ -773,6 +774,7 @@ def _( actors, channels = process_children(ir, rec) channels[ir] = ChannelManager(rec.state["context"]) collective_ids = list(rec.state["collective_id_map"].get(ir, [])) + ir_context = ir_context_for_node(rec, ir) assert len(collective_ids) == 2, ( f"{type(ir).__name__} requires 2 collective IDs, got {len(collective_ids)}" ) @@ -781,7 +783,7 @@ def _( rec.state["context"], rec.state["comm"], ir, - rec.state["ir_context"], + ir_context, channels[ir].reserve_input_slot(), channels[ir.children[0]].reserve_output_slot(), config_options.executor.target_partition_size, diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py index 4c5773874aa5..567c60036f2b 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -32,6 +32,7 @@ from cudf_polars.dsl.to_ast import to_parquet_filter from cudf_polars.streaming.actor_graph.dispatch import ( generate_ir_sub_network, + ir_context_for_node, ) from cudf_polars.streaming.actor_graph.nodes import ( define_actor, @@ -302,7 +303,7 @@ def _( estimated_chunk_bytes = config_options.executor.target_partition_size context = rec.state["context"] - ir_context = rec.state["ir_context"] + ir_context = ir_context_for_node(rec, ir) channels: dict[IR, ChannelManager] = {ir: ChannelManager(rec.state["context"])} nodes: dict[IR, list[Any]] = { ir: [ @@ -484,7 +485,7 @@ def _( ir: PythonScan, rec: SubNetGenerator ) -> tuple[dict[IR, list[Any]], dict[IR, ChannelManager]]: context = rec.state["context"] - ir_context = rec.state["ir_context"] + ir_context = ir_context_for_node(rec, ir) channels: dict[IR, ChannelManager] = {ir: ChannelManager(context)} nodes: dict[IR, list[Any]] = { ir: [ @@ -766,6 +767,7 @@ def _( parquet_options = config_options.parquet_options partition_info = rec.state["partition_info"][ir] num_producers = rec.state["max_io_threads"] + ir_context = ir_context_for_node(rec, ir) channels: dict[IR, ChannelManager] = {ir: ChannelManager(rec.state["context"])} assert partition_info.io_plan is not None, "Scan node must have a partition plan" @@ -811,7 +813,7 @@ def _( # Just estimate the local count as well. local_count=math.ceil(partition_info.count / rec.state["comm"].nranks), ), - rec.state["ir_context"], + ir_context, ) nodes[ir] = [native_node, metadata_node] else: @@ -819,7 +821,7 @@ def _( scan_node( rec.state["context"], ir, - rec.state["ir_context"], + ir_context, ch_out, num_producers=num_producers, estimated_chunk_bytes=( @@ -946,12 +948,13 @@ def _( """Generate network for StreamingSink node.""" nodes, channels = process_children(ir, rec) channels[ir] = ChannelManager(rec.state["context"]) + ir_context = ir_context_for_node(rec, ir) nodes[ir] = [ sink_node( rec.state["context"], rec.state["comm"], ir, - rec.state["ir_context"], + ir_context, channels[ir.children[0]].reserve_output_slot(), channels[ir].reserve_input_slot(), rec.state["partition_info"][ir], diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py index 1a3f8efc63c3..76f895b0e5d2 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -33,6 +33,7 @@ from cudf_polars.streaming.actor_graph.collectives.shuffle import _global_shuffle from cudf_polars.streaming.actor_graph.dispatch import ( generate_ir_sub_network, + ir_context_for_node, ) from cudf_polars.streaming.actor_graph.nodes import default_node_multi from cudf_polars.streaming.actor_graph.tracing import send_chunk @@ -1477,6 +1478,7 @@ def _( # Create output ChannelManager channels[ir] = ChannelManager(rec.state["context"]) + ir_context = ir_context_for_node(rec, ir) if pwise_join: # Partition-wise join (use default_node_multi) @@ -1485,7 +1487,7 @@ def _( default_node_multi( rec.state["context"], ir, - rec.state["ir_context"], + ir_context, channels[ir].reserve_input_slot(), ( channels[left].reserve_output_slot(), @@ -1517,7 +1519,7 @@ def _( rec.state["context"], rec.state["comm"], ir, - rec.state["ir_context"], + ir_context, channels[ir].reserve_input_slot(), channels[left].reserve_output_slot(), channels[right].reserve_output_slot(), @@ -1540,7 +1542,7 @@ def _( rec.state["context"], rec.state["comm"], ir, - rec.state["ir_context"], + ir_context, channels[ir].reserve_input_slot(), channels[left].reserve_output_slot(), channels[right].reserve_output_slot(), diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py index 8066f56ddaea..24ad977da026 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py @@ -22,6 +22,7 @@ from cudf_polars.dsl.ir import IR, Empty from cudf_polars.streaming.actor_graph.dispatch import ( generate_ir_sub_network, + ir_context_for_node, ) from cudf_polars.streaming.actor_graph.tracing import send_chunk from cudf_polars.streaming.actor_graph.utils import ( @@ -520,6 +521,7 @@ def _( # Create output ChannelManager channels[ir] = ChannelManager(rec.state["context"]) + ir_context = ir_context_for_node(rec, ir) if len(ir.children) == 1: # Single-channel default node @@ -527,7 +529,7 @@ def _( default_node_single( rec.state["context"], ir, - rec.state["ir_context"], + ir_context, channels[ir].reserve_input_slot(), channels[ir.children[0]].reserve_output_slot(), ) @@ -538,7 +540,7 @@ def _( default_node_multi( rec.state["context"], ir, - rec.state["ir_context"], + ir_context, channels[ir].reserve_input_slot(), tuple(channels[c].reserve_output_slot() for c in ir.children), ) @@ -595,7 +597,7 @@ def _( ) -> tuple[dict[IR, list[Any]], dict[IR, ChannelManager]]: """Generate network for Empty node - produces one empty chunk.""" context = rec.state["context"] - ir_context = rec.state["ir_context"] + ir_context = ir_context_for_node(rec, ir) channels: dict[IR, ChannelManager] = {ir: ChannelManager(rec.state["context"])} nodes: dict[IR, list[Any]] = { ir: [empty_node(context, ir, ir_context, channels[ir].reserve_input_slot())] @@ -630,6 +632,7 @@ def generate_ir_sub_network_wrapper( if (fanout_info := rec.state["fanout_nodes"].get(ir)) is not None: count = fanout_info.num_consumers manager = ChannelManager(rec.state["context"], count=count) + ir_context = ir_context_for_node(rec, ir) fanout_node: Any if fanout_info.unbounded: fanout_node = fanout_node_unbounded( @@ -637,7 +640,7 @@ def generate_ir_sub_network_wrapper( channels[ir].reserve_output_slot(), *[manager.reserve_input_slot() for _ in range(count)], trace_ir=ir, - ir_context=rec.state["ir_context"], + ir_context=ir_context, ) else: # "bounded" fanout_node = fanout_node_bounded( @@ -645,7 +648,7 @@ def generate_ir_sub_network_wrapper( channels[ir].reserve_output_slot(), *[manager.reserve_input_slot() for _ in range(count)], trace_ir=ir, - ir_context=rec.state["ir_context"], + ir_context=ir_context, ) nodes[ir].append(fanout_node) channels[ir] = manager diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py index 0322645d02d1..7ea80b12877e 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py @@ -62,7 +62,10 @@ LocalRepartitioner, ShuffleManager, ) -from cudf_polars.streaming.actor_graph.dispatch import generate_ir_sub_network +from cudf_polars.streaming.actor_graph.dispatch import ( + generate_ir_sub_network, + ir_context_for_node, +) from cudf_polars.streaming.actor_graph.tracing import send_chunk from cudf_polars.streaming.actor_graph.utils import ( ChannelManager, @@ -788,12 +791,13 @@ def _( else 0 ) scalar_plan = _build_scalar_over_plan(ir) if ir.is_scalar else None + ir_context = ir_context_for_node(rec, ir) actors[ir] = [ over_actor( rec.state["context"], rec.state["comm"], ir, - rec.state["ir_context"], + ir_context, channels[ir].reserve_input_slot(), channels[ir.children[0]].reserve_output_slot(), collective_ids, diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/repartition.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/repartition.py index 3688b48f0c09..d2f4572fa644 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/repartition.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/repartition.py @@ -17,7 +17,10 @@ from cudf_polars.containers import DataFrame from cudf_polars.streaming.actor_graph.collectives.allgather import AllGatherManager -from cudf_polars.streaming.actor_graph.dispatch import generate_ir_sub_network +from cudf_polars.streaming.actor_graph.dispatch import ( + generate_ir_sub_network, + ir_context_for_node, +) from cudf_polars.streaming.actor_graph.nodes import shutdown_on_error from cudf_polars.streaming.actor_graph.tracing import send_chunk from cudf_polars.streaming.actor_graph.utils import ( @@ -258,6 +261,7 @@ def _( # Look up the reserved shuffle ID for this operation collective_id = rec.state["collective_id_map"][ir][0] + ir_context = ir_context_for_node(rec, ir) # Add python node nodes[ir] = [ @@ -265,7 +269,7 @@ def _( rec.state["context"], rec.state["comm"], ir, - rec.state["ir_context"], + ir_context, channels[ir].reserve_input_slot(), channels[ir.children[0]].reserve_output_slot(), output_count=partition_info[ir].count, diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/union.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/union.py index 5d2c891cb12e..f8a1cfee0a62 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/union.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/union.py @@ -13,6 +13,7 @@ from cudf_polars.dsl.ir import Union from cudf_polars.streaming.actor_graph.dispatch import ( generate_ir_sub_network, + ir_context_for_node, ) from cudf_polars.streaming.actor_graph.nodes import define_actor, shutdown_on_error from cudf_polars.streaming.actor_graph.utils import ( @@ -122,6 +123,7 @@ def _( # Create output ChannelManager channels[ir] = ChannelManager(rec.state["context"]) + ir_context = ir_context_for_node(rec, ir) # Add simple python node nodes[ir] = [ @@ -129,7 +131,7 @@ def _( rec.state["context"], rec.state["comm"], ir, - rec.state["ir_context"], + ir_context, channels[ir].reserve_input_slot(), *[channels[c].reserve_output_slot() for c in ir.children], ) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py index c7728056c085..a0015b956168 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -9,7 +9,9 @@ import itertools import operator import struct +import threading import time +import uuid from collections import deque from contextlib import asynccontextmanager from dataclasses import dataclass @@ -35,9 +37,20 @@ from rapidsmpf.streaming.core.message import Message import cudf_polars.dsl.tracing +import cudf_polars.quent from cudf_polars.containers import DataFrame from cudf_polars.dsl.expr import Col, NamedExpr -from cudf_polars.dsl.ir import Cache, Filter, GroupBy, HStack, Join, Projection, Select +from cudf_polars.dsl.ir import ( + Cache, + DataFrameScan, + Filter, + GroupBy, + HStack, + Join, + Projection, + Scan, + Select, +) from cudf_polars.dsl.tracing import Scope from cudf_polars.dsl.utils.naming import names_to_indices from cudf_polars.streaming.actor_graph.collectives.allgather import AllGatherManager @@ -265,6 +278,36 @@ async def shutdown_on_error( if ir_context is not None: contextvars["cudf_polars_query_id"] = str(ir_context.query_id) + quent_ir_execution_context = None + quent_task = None + quent_processor = None + is_io_node = False + if ir_context is not None: + quent_ir_execution_context = ir_context.quent_ir_execution_context + if quent_ir_execution_context is not None: + token = uuid.uuid4() + quent_task = cudf_polars.quent.Task( + instance_name=( + f"Actor-{type(trace_ir).__name__}-" + f"{quent_ir_execution_context.quent_operator.id.hex[:8]}-" + f"{token.hex[:8]}" + ), + operator_id=quent_ir_execution_context.quent_operator.id, + ) + quent_processor = ( + quent_ir_execution_context.context.get_or_declare_processor( + quent_ir_execution_context.logger, + thread_ident=threading.get_ident(), + pool_id=quent_ir_execution_context.thread_pool_id, + ) + ) + is_io_node = issubclass(type(trace_ir), (Scan, DataFrameScan)) + quent_ir_execution_context.logger.emit(quent_task.queueing()) + if not is_io_node: + quent_ir_execution_context.logger.emit( + quent_task.allocating(resource_id=quent_processor.id) + ) + with cudf_polars.dsl.tracing.bound_contextvars(**contextvars): start = time.monotonic_ns() try: @@ -298,6 +341,64 @@ async def shutdown_on_error( "Streaming Actor", start=start, stop=stop, **record ) + if quent_ir_execution_context is not None and quent_task is not None: + if is_io_node: + quent_ir_execution_context.logger.emit( + quent_task.loading( + use_thread=quent_processor, + use_channel=quent_ir_execution_context.disk_to_device_channel, + use_memory=quent_ir_execution_context.device_memory, + ) + ) + else: + quent_ir_execution_context.logger.emit( + quent_task.computing( + use_thread=quent_processor, + use_memory=quent_ir_execution_context.device_memory, + ) + ) + custom_attributes = [] + if tracer is not None and tracer.chunk_count is not None: + custom_attributes.append( + cudf_polars.quent._types.StatisticsAttribute( + key="chunk_count", + value_type="U64", + value=tracer.chunk_count, + ) + ) + if tracer is not None and tracer.duplicated is not None: + custom_attributes.append( + cudf_polars.quent._types.StatisticsAttribute( + key="duplicated", + value_type="U64", + value=1 if tracer.duplicated else 0, + ) + ) + if tracer is not None and tracer.decision is not None: + custom_attributes.append( + cudf_polars.quent._types.StatisticsAttribute( + key="decision", + value_type="String", + value=tracer.decision, + ) + ) + if tracer is None or tracer.row_count is None: + # TODO: See if `output_rows` is nullable. + output_rows = 0 + else: + output_rows = tracer.row_count + quent_ir_execution_context.logger.emit( + quent_ir_execution_context.quent_operator.statistics( + statistics=cudf_polars.quent.Statistics( + output_rows=output_rows, + input_bytes=0, + output_bytes=0, + custom_attributes=custom_attributes, + ) + ) + ) + quent_ir_execution_context.logger.emit(quent_task.exit()) + def _update_ordering_indices( ordering: Ordering, new_indices: tuple[int, ...] diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index be1aedc770b8..501c1d4ea6e3 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -21,10 +21,13 @@ Attribute, Engine, Implementation, + Memory, Operator, Plan, Port, Query, + Statistics, + Task, Worker, _deserialize_value, ) @@ -237,6 +240,49 @@ def test_operator_declare_serialization( assert decl["type_name"] == op.type_name +def test_operator_statistics_serialization( + ir_and_config: tuple[IR, ConfigOptions[StreamingExecutor]], +) -> None: + ir, config_options = ir_and_config + _, operators, _, _ = build_plan( + ir, config_options, Query(), uuid.uuid4(), _make_worker() + ) + op = operators[0] + stats = Statistics(input_bytes=123, output_bytes=456, output_rows=7) + + event = op.statistics(stats, timestamp=101) + d = event.to_dict() + + assert d["id"] == str(op.id) + payload = d["data"]["Operator"]["Statistics"]["custom_attributes"] + assert payload == [ + {"key": "input_bytes", "value": {"U64": 123}}, + {"key": "output_bytes", "value": {"U64": 456}}, + {"key": "output_rows", "value": {"U64": 7}}, + ] + + +def test_memory_lifecycle_events() -> None: + memory = Memory( + instance_name="device", + resource_type_name="memory", + parent_group_id=uuid.uuid4(), + ) + assert memory.initializing().to_dict()["data"]["Memory"]["seq"] == 0 + assert memory.operating(1024).to_dict()["data"]["Memory"]["seq"] == 1 + assert memory.finalizing().to_dict()["data"]["Memory"]["seq"] == 2 + assert memory.exit().to_dict()["data"]["Memory"]["seq"] == 3 + + +def test_task_lifecycle_events() -> None: + operator_id = uuid.uuid4() + task = Task(operator_id=operator_id, instance_name="task-0") + queue = task.queueing().to_dict() + assert queue["data"]["Task"]["state"]["Queueing"]["operator_id"] == str(operator_id) + assert task.allocating(uuid.uuid4()).to_dict()["data"]["Task"]["seq"] == 1 + assert task.exit().to_dict()["data"]["Task"]["seq"] == 5 + + def test_port_declare_serialization( ir_and_config: tuple[IR, ConfigOptions[StreamingExecutor]], ) -> None: diff --git a/python/cudf_polars/tests/quent/test_quent_integration.py b/python/cudf_polars/tests/quent/test_quent_integration.py index d6fb79b30185..f5f05fed65c5 100644 --- a/python/cudf_polars/tests/quent/test_quent_integration.py +++ b/python/cudf_polars/tests/quent/test_quent_integration.py @@ -156,3 +156,17 @@ def check_quent_events(engine: StreamingEngine, quent_context: QuentContext) -> assert query_executing["data"]["Query"]["seq"] == 2 assert query_exit["id"] == str(quent_context.query.id) assert query_exit["data"]["Query"]["seq"] == 3 + + +def test_quent_events_include_resources( + engine_with_quent_context: StreamingEngine, quent_context: QuentContext +) -> None: + q = pl.LazyFrame({"x": [1, 2, 3, 4]}).filter(pl.col("x") > 1) + with engine_with_quent_context: + q.collect(engine=engine_with_quent_context) + + quent_events = engine_with_quent_context._quent_events + memory_events = [x for x in quent_events if "Memory" in x["data"]] + task_events = [x for x in quent_events if "Task" in x["data"]] + assert len(memory_events) > 0 + assert len(task_events) > 0 From 43ba6e50fff464dac7d218cbb60a8a3723e0cc3f Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 8 Jul 2026 14:15:54 -0700 Subject: [PATCH 02/60] many things --- python/cudf_polars/cudf_polars/dsl/tracing.py | 4 +- python/cudf_polars/cudf_polars/engine/core.py | 13 +- python/cudf_polars/cudf_polars/engine/dask.py | 18 +- python/cudf_polars/cudf_polars/engine/ray.py | 14 +- python/cudf_polars/cudf_polars/engine/spmd.py | 89 ++++++--- .../cudf_polars/cudf_polars/quent/_context.py | 122 +++++++++--- .../cudf_polars/streaming/actor_graph/core.py | 11 ++ .../streaming/actor_graph/utils.py | 8 +- .../cudf_polars/cudf_polars/utils/config.py | 15 +- python/cudf_polars/tests/quent/test_quent.py | 134 +++++++++++++ .../tests/quent/test_quent_integration.py | 177 +++++++++++++++++- 11 files changed, 531 insertions(+), 74 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index ef11dfa14d06..94cded4f1563 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -77,10 +77,8 @@ def _begin_quent_do_evaluate_events( ), operator_id=quent_ir_execution_context.quent_operator.id, ) - quent_processor = quent_ir_execution_context.context.get_or_declare_processor( - quent_ir_execution_context.logger, + quent_processor = quent_ir_execution_context.get_or_declare_processor( thread_ident=threading.get_ident(), - pool_id=quent_ir_execution_context.thread_pool_id, ) is_io_node = issubclass(cls, (Scan, DataFrameScan)) quent_ir_execution_context.logger.emit(quent_task.queueing()) diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index d8f484c01503..9a3b3ff3d10b 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -756,11 +756,20 @@ def evaluate_on_rank( if config_options.executor.quent_context is not None: assert local_quent_context is not None _declare_network_channels(comm, local_quent_context) - logical_plan_id = ir.get_stable_plan_id() + # ``get_stable_plan_id`` is a deterministic function of the IR + # structure, so every rank derives the same logical plan ID for a + # given query (only rank 0 emits the declaration, but physical plans + # on every rank reference it as their parent). It is *not* unique + # across collects, though: re-running an identical query would reuse + # the same plan ID under a different parent query. Namespacing by the + # per-collect ``query_id`` (which is identical across ranks but unique + # per collect) keeps the cross-rank agreement while making the plan ID + # unique per collect. + logical_plan_id = uuid.uuid5(query_id, str(ir.get_stable_plan_id())) plan, ops, ports, logical_op_by_id = build_plan( ir, config_options, - query=local_quent_context.context.query, + query=local_quent_context.query, plan_id=logical_plan_id, worker=local_quent_context.worker, instance_name="logical", diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 2c2c4124300a..122cbfecb0ab 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -43,6 +43,7 @@ ) from cudf_polars.quent._context import ( LocalQuentContext, + ProcessorRegistry, declare_worker_resources, finalize_worker_resources, ) @@ -127,6 +128,7 @@ class _WorkerContext: device_memory: cudf_polars.quent._types.Memory | None = None disk_to_device_channel: cudf_polars.quent._types.Channel | None = None thread_pool: cudf_polars.quent._types.ThreadPool | None = None + processor_registry: ProcessorRegistry | None = None def _setup_root( @@ -319,7 +321,9 @@ def _setup_worker( device_memory = None disk_to_device_channel = None thread_pool = None + processor_registry = None if quent_logger is not None: + processor_registry = ProcessorRegistry() device_memory, disk_to_device_channel, thread_pool = declare_worker_resources( quent_logger, instance_suffix=f"rank-{comm.rank}", @@ -338,6 +342,7 @@ def _setup_worker( device_memory=device_memory, disk_to_device_channel=disk_to_device_channel, thread_pool=thread_pool, + processor_registry=processor_registry, ) setattr(dask_worker, attr, mp_ctx) if mp_ctx.quent_logger is not None: @@ -366,6 +371,10 @@ def _teardown_worker( traces = [] if mp_ctx is not None: if mp_ctx.quent_worker is not None and mp_ctx.quent_logger is not None: + if mp_ctx.processor_registry is not None: + mp_ctx.processor_registry._emit_processor_exit_events( + mp_ctx.quent_logger + ) if mp_ctx.device_memory is not None: finalize_worker_resources( mp_ctx.quent_logger, @@ -524,11 +533,14 @@ def _worker_evaluate( assert mp_ctx.quent_logger is not None assert mp_ctx.device_memory is not None assert mp_ctx.thread_pool is not None + assert mp_ctx.processor_registry is not None local_quent_context = LocalQuentContext( context=quent_context, + query=quent_context.query_for(query_id), worker=mp_ctx.quent_worker, logger=mp_ctx.quent_logger, thread_pool_id=mp_ctx.thread_pool.id, + processor_registry=mp_ctx.processor_registry, device_memory=mp_ctx.device_memory, disk_to_device_channel=mp_ctx.disk_to_device_channel, ) @@ -616,8 +628,9 @@ def evaluate_pipeline_dask_mode( if quent_context is not None: quent_logger = dask_context.quent_logger assert quent_logger is not None + query = quent_context.query_for(query_id) quent_context._emit_query_group_events(quent_logger) - quent_context._emit_query_events(quent_logger) + quent_context._emit_query_events(quent_logger, query) # Strip dask_context before pickling config_options for remote calls. worker_config = dataclasses.replace( @@ -644,7 +657,7 @@ def evaluate_pipeline_dask_mode( if quent_context is not None: quent_logger = dask_context.quent_logger assert quent_logger is not None - quent_context._emit_query_exit_events(quent_logger) + quent_context._emit_query_exit_events(quent_logger, query) return pl.concat(dfs), metadata_collector or None @@ -1053,7 +1066,6 @@ def shutdown(self) -> None: if quent_context is not None: assert self._quent_logger is not None quent_context._emit_engine_exit_events(self._quent_logger) - quent_context.emit_resource_exit_events(self._quent_logger) if ctx.owned_client is not None: ctx.owned_client.close() if ctx.owned_cluster is not None: diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index 2b9ba237e5f6..ac9cf7971dfe 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -41,6 +41,7 @@ ) from cudf_polars.quent._context import ( LocalQuentContext, + ProcessorRegistry, declare_worker_resources, finalize_worker_resources, ) @@ -120,8 +121,9 @@ def evaluate_pipeline_ray_mode( if quent_context is not None: quent_logger = config_options.executor.ray_context.quent_logger assert quent_logger is not None + query = quent_context.query_for(query_id) quent_context._emit_query_group_events(quent_logger) - quent_context._emit_query_events(quent_logger) + quent_context._emit_query_events(quent_logger, query) # Serialize the IR into the Ray object store so actors fetch by reference # instead of receiving N copies. @@ -150,7 +152,7 @@ def evaluate_pipeline_ray_mode( if quent_context is not None: quent_logger = config_options.executor.ray_context.quent_logger assert quent_logger is not None - quent_context._emit_query_exit_events(quent_logger) + quent_context._emit_query_exit_events(quent_logger, query) return pl.concat(dfs), metadata_collector or None @@ -238,7 +240,9 @@ def __init__( self._device_memory = None self._disk_to_device_channel = None self._quent_thread_pool = None + self._processor_registry: ProcessorRegistry | None = None if self._quent_logger is not None: + self._processor_registry = ProcessorRegistry() self._quent_logger.emit(self._quent_worker._init()) ( self._device_memory, @@ -352,6 +356,8 @@ def _exit(self) -> list[dict[str, Any]]: # Maybe generalize this to all application-level things, # followed by framework (ray) level things. if self._quent_worker is not None and self._quent_logger is not None: + if self._processor_registry is not None: + self._processor_registry._emit_processor_exit_events(self._quent_logger) if ( self._device_memory is not None and self._disk_to_device_channel is not None @@ -474,11 +480,14 @@ def evaluate_polars_ir( assert self._quent_logger is not None assert self._device_memory is not None assert self._quent_thread_pool is not None + assert self._processor_registry is not None local_quent_context = LocalQuentContext( context=quent_context, + query=quent_context.query_for(query_id), worker=self._quent_worker, logger=self._quent_logger, thread_pool_id=self._quent_thread_pool.id, + processor_registry=self._processor_registry, device_memory=self._device_memory, disk_to_device_channel=self._disk_to_device_channel, ) @@ -947,7 +956,6 @@ def shutdown(self) -> None: if quent_context is not None: assert self._quent_logger is not None quent_context._emit_engine_exit_events(self._quent_logger) - quent_context.emit_resource_exit_events(self._quent_logger) self._rank_actors = None super().shutdown() diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index eb8d16385258..bdcfae430a01 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -45,6 +45,7 @@ ) from cudf_polars.quent._context import ( LocalQuentContext, + ProcessorRegistry, declare_worker_resources, finalize_worker_resources, ) @@ -70,6 +71,7 @@ from cudf_polars.dsl.ir import IR from cudf_polars.engine.core import T from cudf_polars.engine.options import StreamingOptions + from cudf_polars.quent._types import Channel, Memory, ThreadPool from cudf_polars.streaming.parallel import ConfigOptions from cudf_polars.utils.config import StreamingExecutor @@ -117,32 +119,33 @@ def evaluate_pipeline_spmd_mode( comm = config_options.executor.spmd_context.comm context = config_options.executor.spmd_context.context py_executor = config_options.executor.spmd_context.py_executor + spmd_context = config_options.executor.spmd_context quent_context = config_options.executor.quent_context local_quent_context: LocalQuentContext | None = None if quent_context is not None: quent_logger = config_options.executor.spmd_context.quent_logger assert quent_logger is not None + query = quent_context.query_for(query_id) quent_context._emit_query_group_events(quent_logger) - quent_context._emit_query_events(quent_logger) + quent_context._emit_query_events(quent_logger, query) worker_id = config_options.executor.spmd_context.worker_id - device_memory, disk_to_device_channel, thread_pool = declare_worker_resources( - quent_logger, - instance_suffix=f"rank-{comm.rank}", - engine_id=quent_context.engine.id, - worker_id=worker_id, - ) + assert spmd_context.processor_registry is not None + assert spmd_context.thread_pool_id is not None + assert spmd_context.device_memory is not None local_quent_context = LocalQuentContext( context=quent_context, + query=query, worker=Worker( id=worker_id, engine=quent_context.engine, instance_name=f"rank-{comm.rank}", ), logger=quent_logger, - thread_pool_id=thread_pool.id, - device_memory=device_memory, - disk_to_device_channel=disk_to_device_channel, + thread_pool_id=spmd_context.thread_pool_id, + processor_registry=spmd_context.processor_registry, + device_memory=spmd_context.device_memory, + disk_to_device_channel=spmd_context.disk_to_device_channel, ) df, metadata = evaluate_on_rank( @@ -157,13 +160,11 @@ def evaluate_pipeline_spmd_mode( if quent_context is not None: assert config_options.executor.spmd_context.quent_logger is not None assert local_quent_context is not None - finalize_worker_resources( - config_options.executor.spmd_context.quent_logger, - device_memory=local_quent_context.device_memory, - disk_to_device_channel=local_quent_context.disk_to_device_channel, - ) + # Device memory and the disk->device channel are engine-scoped and are + # finalized once at engine shutdown, not per query. quent_context._emit_query_exit_events( - config_options.executor.spmd_context.quent_logger + config_options.executor.spmd_context.quent_logger, + local_quent_context.query, ) return df, metadata if collect_metadata else None @@ -454,6 +455,10 @@ def __init__( self._comm: Communicator | None = comm self._ctx: Context | None = None self._py_executor: ThreadPoolExecutor | None = None + self._processor_registry: ProcessorRegistry | None = None + self._quent_thread_pool: ThreadPool | None = None + self._device_memory: Memory | None = None + self._disk_to_device_channel: Channel | None = None exit_stack = contextlib.ExitStack() @@ -492,6 +497,26 @@ def __init__( instance_name=f"rank-{self.rank}", # relies on self.comm ) + if quent_context is not None: + assert self._quent_logger is not None + self._quent_logger.emit(self._quent_worker._init()) + self._processor_registry = ProcessorRegistry() + # Device memory, the disk->device channel, and the thread pool + # are all engine-scoped: they live for the lifetime of the + # engine's `py_executor`, not any single query. Declaring them + # once here (rather than per collect) means a single + # init/finalize pair per engine. + ( + self._device_memory, + self._disk_to_device_channel, + self._quent_thread_pool, + ) = declare_worker_resources( + self._quent_logger, + instance_suffix=f"rank-{self.rank}", + engine_id=engine_id, + worker_id=self._quent_worker.id, + ) + # Register after `_cleanup_ctx` so on teardown (LIFO) the # executor shuts down first. `wait=True` is safe because # rapidsmpf's `run_actor_network` awaits its only submitted @@ -517,6 +542,14 @@ def __init__( quent_logger=self._quent_logger, context=self._ctx, py_executor=self._py_executor, + processor_registry=self._processor_registry, + thread_pool_id=( + self._quent_thread_pool.id + if self._quent_thread_pool is not None + else None + ), + device_memory=self._device_memory, + disk_to_device_channel=self._disk_to_device_channel, ), }, engine_options={ @@ -525,9 +558,6 @@ def __init__( }, exit_stack=exit_stack, ) - - if self._quent_logger is not None: - self._quent_logger.emit(self._quent_worker._init()) except Exception: exit_stack.close() raise @@ -659,6 +689,12 @@ def _reset( engine_id=engine_id, worker_id=self._quent_worker.id, quent_logger=self._quent_logger, + processor_registry=self._processor_registry, + thread_pool_id=( + self._quent_thread_pool.id + if self._quent_thread_pool is not None + else None + ), ), }, engine_options={ @@ -794,15 +830,24 @@ def shutdown(self) -> None: # quent traces before that. # Clear the references only after shutdown completes. - if self._quent_logger is not None: - self._quent_logger.emit(self._quent_worker._exit()) quent_context: cudf_polars.quent.QuentContext | None = self.config[ "executor_options" ].get("quent_context") + + if self._quent_logger is not None: + if quent_context is not None: + assert self._processor_registry is not None + self._processor_registry._emit_processor_exit_events(self._quent_logger) + if self._device_memory is not None: + finalize_worker_resources( + self._quent_logger, + device_memory=self._device_memory, + disk_to_device_channel=self._disk_to_device_channel, + ) + self._quent_logger.emit(self._quent_worker._exit()) if quent_context is not None: assert self._quent_logger is not None quent_context._emit_engine_exit_events(self._quent_logger) - quent_context.emit_resource_exit_events(self._quent_logger) super().shutdown() diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index 4e4281a8c4f3..cb55185ca73a 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -7,6 +7,7 @@ import dataclasses import json +import threading import uuid from typing import TYPE_CHECKING @@ -42,11 +43,59 @@ __all__ = [ "LocalQuentContext", + "ProcessorRegistry", "QuentContext", "QuentIRExecutionContext", ] +class ProcessorRegistry: + """ + Engine/worker-scoped registry of dynamically declared Quent Processors. + + One registry is owned by the object that owns the Python + :class:`~concurrent.futures.ThreadPoolExecutor` (e.g. ``SPMDEngine``, + a Dask worker, or a Ray actor). Processors are declared lazily on first + use by a thread-pool worker and finalized once at executor shutdown. + """ + + def __init__(self) -> None: + self._processors: dict[int, Processor] = {} + self._lock = threading.Lock() + self._closed = False + + def get_or_declare_processor( + self, logger: QuentLogger, thread_ident: int, pool_id: uuid.UUID + ) -> Processor: + """Get (or declare a new) Quent Processor for a CPU thread.""" + with self._lock: + if self._closed: + raise RuntimeError( + "Cannot declare processors after registry has been closed" + ) + if thread_ident in self._processors: + return self._processors[thread_ident] + + processor = Processor(pool_id=pool_id) + self._processors[thread_ident] = processor + + logger.emit(processor.initializing()) + logger.emit(processor.operating()) + return processor + + def _emit_processor_exit_events(self, logger: QuentLogger) -> None: + """Emit finalizing/exit events for all declared processors.""" + with self._lock: + if self._closed: + return + self._closed = True + processors = list(self._processors.values()) + + for processor in processors: + logger.emit(processor.finalizing()) + logger.emit(processor.exit()) + + @dataclasses.dataclass(frozen=True, kw_only=True) class QuentContext: """ @@ -73,11 +122,6 @@ class QuentContext: def __post_init__(self) -> None: object.__setattr__(self, "_query_group_cache_", set()) - object.__setattr__(self, "_processor_map_", {}) - - @property - def _processor_map(self) -> dict[int, Processor]: - return self._processor_map_ # type: ignore[attr-defined] def serialize(self) -> bytes: """ @@ -169,19 +213,32 @@ def _emit_query_group_events(self, logger: QuentLogger) -> None: self._query_group_cache.add(self.query_group.id) logger.emit(self.query_group._declare(engine=self.engine)) - def _emit_query_events(self, logger: QuentLogger) -> None: + def query_for(self, query_id: uuid.UUID) -> Query: + """ + Build a per-collect Quent Query with a unique id. + + The engine-scoped ``QuentContext`` is reused across many + ``.collect()`` calls, so each collect must derive its own + :class:`Query` (identified by the per-collect ``query_id``) rather + than reusing the shared ``self.query``. The ``instance_name`` from + the template ``self.query`` is preserved. + """ + return Query(id=query_id, instance_name=self.query.instance_name) + + def _emit_query_events(self, logger: QuentLogger, query: Query) -> None: """ Emit Quent Query events. This includes events for 'Declare', 'Init', and 'Planning'. """ - logger.emit(self.query._init(query_group=self.query_group)) - logger.emit(self.query._planning()) - logger.emit(self.query._executing()) + print(f"Query int: {query.id}") + logger.emit(query._init(query_group=self.query_group)) + logger.emit(query._planning()) + logger.emit(query._executing()) - def _emit_query_exit_events(self, logger: QuentLogger) -> None: + def _emit_query_exit_events(self, logger: QuentLogger, query: Query) -> None: """Emit a Quent Query exit event.""" - logger.emit(self.query._exit()) + logger.emit(query._exit()) def _emit_plan_declarations( self, @@ -316,25 +373,6 @@ def _emit_physical_plan_events( parent_operators_by_node_id=parent_operators_by_node_id, ) - def get_or_declare_processor( - self, quent_logger: QuentLogger, thread_ident: int, pool_id: uuid.UUID - ) -> Processor: - """Get (or declare a new) Quent Processor for a CPU thread.""" - if thread_ident in self._processor_map: - return self._processor_map[thread_ident] - - processor = Processor(pool_id=pool_id) - self._processor_map[thread_ident] = processor - quent_logger.emit(processor.initializing()) - quent_logger.emit(processor.operating()) - return processor - - def emit_resource_exit_events(self, quent_logger: QuentLogger) -> None: - """Emit finalizing/exit events for declared processors.""" - for processor in self._processor_map.values(): - quent_logger.emit(processor.finalizing()) - quent_logger.emit(processor.exit()) - def declare_worker_resources( logger: QuentLogger, @@ -399,17 +437,39 @@ class LocalQuentContext: This can contain non-serializable objects (like a ``QuentLogger``) and entities that are only valid on the local rank. + + The ``processor_registry`` is engine/worker-scoped and outlives + individual queries. It is injected by the backend that owns the + ``ThreadPoolExecutor``. + + The ``query`` is per-collect: each ``.collect()`` derives a fresh + :class:`Query` from its unique ``query_id`` (see + :meth:`QuentContext.query_for`), rather than reusing the shared + ``context.query``. """ context: QuentContext + query: Query worker: Worker logger: QuentLogger thread_pool_id: uuid.UUID + processor_registry: ProcessorRegistry device_memory: Memory disk_to_device_channel: Channel | None = None network: Network | None = None link_channels: dict[int, Channel] = dataclasses.field(default_factory=dict) + def get_or_declare_processor( + self, + thread_ident: int, + ) -> Processor: + """Get (or declare a new) Quent Processor for a CPU thread.""" + return self.processor_registry.get_or_declare_processor( + self.logger, + thread_ident=thread_ident, + pool_id=self.thread_pool_id, + ) + @dataclasses.dataclass(kw_only=True) class QuentIRExecutionContext(LocalQuentContext): @@ -425,9 +485,11 @@ def from_execution_context( return cls( quent_operator=quent_operator, context=execution_context.context, + query=execution_context.query, worker=execution_context.worker, logger=execution_context.logger, thread_pool_id=execution_context.thread_pool_id, + processor_registry=execution_context.processor_registry, device_memory=execution_context.device_memory, disk_to_device_channel=execution_context.disk_to_device_channel, network=execution_context.network, diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py index 2da64d30fa36..b8b565559e6b 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py @@ -95,6 +95,17 @@ def evaluate_logical_plan( engine_id=engine_id, worker_id=engine._quent_worker.id, quent_logger=engine._quent_logger, + # The engine owns these engine-scoped Quent resources; the + # per-collect config must forward them so query evaluation + # reuses them instead of re-declaring per collect. + processor_registry=engine._processor_registry, + thread_pool_id=( + engine._quent_thread_pool.id + if engine._quent_thread_pool is not None + else None + ), + device_memory=engine._device_memory, + disk_to_device_channel=engine._disk_to_device_channel, ), ), ) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py index a0015b956168..72fa7be2fe93 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -294,12 +294,8 @@ async def shutdown_on_error( ), operator_id=quent_ir_execution_context.quent_operator.id, ) - quent_processor = ( - quent_ir_execution_context.context.get_or_declare_processor( - quent_ir_execution_context.logger, - thread_ident=threading.get_ident(), - pool_id=quent_ir_execution_context.thread_pool_id, - ) + quent_processor = quent_ir_execution_context.get_or_declare_processor( + thread_ident=threading.get_ident(), ) is_io_node = issubclass(type(trace_ir), (Scan, DataFrameScan)) quent_ir_execution_context.logger.emit(quent_task.queueing()) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index e81355411f85..8e19b9120d8e 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -44,8 +44,9 @@ from rapidsmpf.streaming.core.context import Context from cudf_polars.engine.ray import RankActor - from cudf_polars.quent._context import QuentContext + from cudf_polars.quent._context import ProcessorRegistry, QuentContext from cudf_polars.quent._logging import QuentLogger + from cudf_polars.quent._types import Channel, Memory __all__ = [ @@ -570,6 +571,14 @@ class SPMDContext: The active RapidsMPF context. py_executor Thread-pool executor used to drive the actor network on each rank. + processor_registry + Engine/worker-scoped registry of dynamically declared Processors. + thread_pool_id + ID of the engine-scoped Quent ThreadPool resource. + device_memory + The engine-scoped Quent device Memory resource. + disk_to_device_channel + The engine-scoped Quent disk-to-device Channel resource. """ comm: Communicator @@ -578,6 +587,10 @@ class SPMDContext: engine_id: uuid.UUID worker_id: uuid.UUID quent_logger: QuentLogger | None + processor_registry: ProcessorRegistry | None = None + thread_pool_id: uuid.UUID | None = None + device_memory: Memory | None = None + disk_to_device_channel: Channel | None = None @dataclasses.dataclass(frozen=True) diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index 501c1d4ea6e3..29f82dc105d5 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -16,6 +16,7 @@ import cudf_polars.quent import cudf_polars.quent._logging from cudf_polars.dsl.translate import Translator +from cudf_polars.quent._context import ProcessorRegistry from cudf_polars.quent._plan import build_plan, port_names_for_node from cudf_polars.quent._types import ( Attribute, @@ -36,6 +37,7 @@ if TYPE_CHECKING: from cudf_polars.dsl.ir import IR from cudf_polars.quent import QuentContext + from cudf_polars.quent._types import Processor from cudf_polars.utils.config import StreamingExecutor @@ -604,6 +606,138 @@ def test_emit_query_group_events_idempotent(quent_context: QuentContext): assert len(logger._buffer) == 1 +def test_processor_registry_declares_once_per_thread() -> None: + pytest.importorskip("structlog") + + logger = cudf_polars.quent._logging.QuentLogger() + registry = ProcessorRegistry() + pool_id = uuid.uuid4() + thread_ident = 42 + + processor_a = registry.get_or_declare_processor( + logger, thread_ident=thread_ident, pool_id=pool_id + ) + processor_b = registry.get_or_declare_processor( + logger, thread_ident=thread_ident, pool_id=pool_id + ) + + assert processor_a is processor_b + processor_events = [x for x in _drained_events(logger) if "Processor" in x["data"]] + assert len(processor_events) == 2 + assert processor_events[0]["data"]["Processor"]["state"] == { + "ProcessorInitializing": { + "instance_name": f"Thread {processor_a.id.hex[:8]}", + "parent_group_id": str(pool_id), + "resource_type_name": "processor", + } + } + assert processor_events[1]["data"]["Processor"]["state"] == { + "ProcessorOperating": None + } + + +def test_processor_registry_concurrent_first_use_declares_once() -> None: + pytest.importorskip("structlog") + + logger = cudf_polars.quent._logging.QuentLogger() + registry = ProcessorRegistry() + pool_id = uuid.uuid4() + thread_ident = 123 + + def get_processor(_: int) -> Processor: + return registry.get_or_declare_processor( + logger, thread_ident=thread_ident, pool_id=pool_id + ) + + with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor: + processors = list(executor.map(get_processor, range(32))) + + assert len({processor.id for processor in processors}) == 1 + processor_events = [x for x in _drained_events(logger) if "Processor" in x["data"]] + assert len(processor_events) == 2 + + +def test_processor_registry_reused_across_quent_contexts() -> None: + pytest.importorskip("structlog") + logger = cudf_polars.quent._logging.QuentLogger() + registry = ProcessorRegistry() + pool_id = uuid.uuid4() + thread_ident = 99 + + context_a = cudf_polars.quent.QuentContext() + context_b = cudf_polars.quent.QuentContext() + local_a = cudf_polars.quent.LocalQuentContext( + context=context_a, + query=context_a.query_for(uuid.uuid4()), + worker=Worker(id=uuid.uuid4(), engine=context_a.engine, instance_name="rank-0"), + logger=logger, + thread_pool_id=pool_id, + processor_registry=registry, + device_memory=Memory( + instance_name="device", + resource_type_name="memory", + parent_group_id=context_a.engine.id, + ), + ) + local_b = cudf_polars.quent.LocalQuentContext( + context=context_b, + query=context_b.query_for(uuid.uuid4()), + worker=Worker(id=uuid.uuid4(), engine=context_b.engine, instance_name="rank-0"), + logger=logger, + thread_pool_id=pool_id, + processor_registry=registry, + device_memory=Memory( + instance_name="device", + resource_type_name="memory", + parent_group_id=context_b.engine.id, + ), + ) + + processor_a = local_a.get_or_declare_processor(thread_ident=thread_ident) + processor_b = local_b.get_or_declare_processor(thread_ident=thread_ident) + + assert processor_a is processor_b + processor_events = [x for x in _drained_events(logger) if "Processor" in x["data"]] + assert len(processor_events) == 2 + + +def test_processor_registry_exit_events_idempotent() -> None: + pytest.importorskip("structlog") + from cudf_polars.quent._context import ProcessorRegistry + + logger = cudf_polars.quent._logging.QuentLogger() + registry = ProcessorRegistry() + pool_id = uuid.uuid4() + + registry.get_or_declare_processor(logger, thread_ident=1, pool_id=pool_id) + registry.get_or_declare_processor(logger, thread_ident=2, pool_id=pool_id) + + registry._emit_processor_exit_events(logger) + registry._emit_processor_exit_events(logger) + + events = _drained_events(logger) + finalizing_events = [ + x + for x in events + if "Processor" in x["data"] + and x["data"]["Processor"]["state"] == {"ProcessorFinalizing": None} + ] + exit_events = [ + x + for x in events + if "Processor" in x["data"] and x["data"]["Processor"]["state"] == "Exit" + ] + assert len(finalizing_events) == 2 + assert len(exit_events) == 2 + + +def _drained_events( + logger: cudf_polars.quent._logging.QuentLogger, +) -> list[dict]: + """Drain Quent logger events into the same shape as engine._quent_events.""" + return [x["event"] for x in logger.drain()] + + def test_serialize_list_raises(): with pytest.raises(NotImplementedError, match="not supported yet"): Attribute("list", [1, 2]).serialize() diff --git a/python/cudf_polars/tests/quent/test_quent_integration.py b/python/cudf_polars/tests/quent/test_quent_integration.py index f5f05fed65c5..ac81f1e92ab3 100644 --- a/python/cudf_polars/tests/quent/test_quent_integration.py +++ b/python/cudf_polars/tests/quent/test_quent_integration.py @@ -144,17 +144,21 @@ def check_quent_events(engine: StreamingEngine, quent_context: QuentContext) -> assert len(query_events) == 4 query_init, query_planning, query_executing, query_exit = query_events - assert query_init["id"] == str(quent_context.query.id) + # Each ``.collect()`` derives a fresh per-collect query id, so the emitted + # id must be unique to this collect rather than the engine-scoped template + # ``quent_context.query`` id. + query_id = query_init["id"] + assert query_id != str(quent_context.query.id) assert ( query_init["data"]["Query"]["state"]["Init"]["query_group_id"] == query_group_declaration["id"] ) assert query_init["data"]["Query"]["seq"] == 0 - assert query_planning["id"] == str(quent_context.query.id) + assert query_planning["id"] == query_id assert query_planning["data"]["Query"]["seq"] == 1 - assert query_executing["id"] == str(quent_context.query.id) + assert query_executing["id"] == query_id assert query_executing["data"]["Query"]["seq"] == 2 - assert query_exit["id"] == str(quent_context.query.id) + assert query_exit["id"] == query_id assert query_exit["data"]["Query"]["seq"] == 3 @@ -170,3 +174,168 @@ def test_quent_events_include_resources( task_events = [x for x in quent_events if "Task" in x["data"]] assert len(memory_events) > 0 assert len(task_events) > 0 + + +def test_quent_device_memory_declared_once_per_engine( + engine_with_quent_context: StreamingEngine, quent_context: QuentContext +) -> None: + # Device memory is an engine-scoped resource: it must be initialized and + # finalized exactly once per engine, regardless of how many collects run. + q1 = pl.LazyFrame({"x": [1, 2, 3]}).filter(pl.col("x") > 1) + q2 = pl.LazyFrame({"y": [4, 5, 6]}).filter(pl.col("y") > 4) + with engine_with_quent_context: + q1.collect(engine=engine_with_quent_context) + q2.collect(engine=engine_with_quent_context) + + quent_events = engine_with_quent_context._quent_events + memory_events = [x for x in quent_events if "Memory" in x["data"]] + + device_init_events = [ + x + for x in memory_events + if isinstance(x["data"]["Memory"]["state"], dict) + and "MemoryInitializing" in x["data"]["Memory"]["state"] + and "device memory" + in x["data"]["Memory"]["state"]["MemoryInitializing"]["instance_name"] + ] + device_exit_events = [ + x + for x in memory_events + if x["data"]["Memory"]["state"] == "Exit" + and x["id"] in {e["id"] for e in device_init_events} + ] + + # Device memory is engine/worker-scoped: it is initialized and finalized + # exactly once per worker, matching the number of engine-scoped ThreadPool + # declarations. Critically, running two collects must NOT re-declare it + # (the per-query bug would produce a fresh device memory per collect, i.e. + # twice as many inits as thread pools). + thread_pool_decls = [ + x + for x in quent_events + if "ThreadPool" in x["data"] and "Declaration" in x["data"]["ThreadPool"] + ] + assert len(thread_pool_decls) >= 1 + assert len(device_init_events) == len(thread_pool_decls) + + # Every device memory id is initialized once and exited once. + init_ids = [x["id"] for x in device_init_events] + assert len(set(init_ids)) == len(init_ids) + assert {x["id"] for x in device_exit_events} == set(init_ids) + + +def test_quent_processor_lifecycle_balanced( + engine_with_quent_context: StreamingEngine, quent_context: QuentContext +) -> None: + q = pl.LazyFrame({"x": list(range(100))}).filter(pl.col("x") > 1) + with engine_with_quent_context: + q.collect(engine=engine_with_quent_context) + + quent_events = engine_with_quent_context._quent_events + check_processor_lifecycle(quent_events) + + +def test_quent_processor_lifecycle_across_multiple_collects( + engine_with_quent_context: StreamingEngine, quent_context: QuentContext +) -> None: + q1 = pl.LazyFrame({"x": [1, 2, 3]}).filter(pl.col("x") > 1) + q2 = pl.LazyFrame({"y": [4, 5, 6]}).filter(pl.col("y") > 4) + with engine_with_quent_context: + q1.collect(engine=engine_with_quent_context) + q2.collect(engine=engine_with_quent_context) + + quent_events = engine_with_quent_context._quent_events + check_processor_lifecycle(quent_events) + + +def test_quent_query_id_unique_per_collect( + engine_with_quent_context: StreamingEngine, quent_context: QuentContext +) -> None: + q1 = pl.LazyFrame({"x": [1, 2, 3]}).filter(pl.col("x") > 1) + q2 = pl.LazyFrame({"y": [4, 5, 6]}).filter(pl.col("y") > 4) + with engine_with_quent_context: + q1.collect(engine=engine_with_quent_context) + q2.collect(engine=engine_with_quent_context) + + quent_events = engine_with_quent_context._quent_events + query_init_ids = [ + x["id"] + for x in quent_events + if "Query" in x["data"] and "Init" in x["data"]["Query"].get("state", {}) + ] + assert len(query_init_ids) == 2 + # Each collect reuses the engine-scoped QuentContext but must emit a + # distinct query id. + assert len(set(query_init_ids)) == 2 + assert str(quent_context.query.id) not in query_init_ids + + +def test_quent_plan_id_unique_per_collect( + engine_with_quent_context: StreamingEngine, quent_context: QuentContext +) -> None: + # Run the *same* query twice. ``get_stable_plan_id`` is a deterministic + # function of the IR structure, so without namespacing by the per-collect + # query id both collects would emit the same logical plan id under + # different parent queries. + q = pl.LazyFrame({"x": [1, 2, 3]}).filter(pl.col("x") > 1) + with engine_with_quent_context: + q.collect(engine=engine_with_quent_context) + q.collect(engine=engine_with_quent_context) + + quent_events = engine_with_quent_context._quent_events + logical_plan_decls = [ + x + for x in quent_events + if "Plan" in x["data"] + and "Declaration" in x["data"]["Plan"] + and x["data"]["Plan"]["Declaration"]["instance_name"] == "logical" + ] + assert len(logical_plan_decls) == 2 + + plan_ids = [x["id"] for x in logical_plan_decls] + assert len(set(plan_ids)) == 2 + + # Each logical plan must hang off the distinct per-collect query id. + parent_query_ids = [ + x["data"]["Plan"]["Declaration"]["parent"]["query_id"] + for x in logical_plan_decls + ] + assert len(set(parent_query_ids)) == 2 + + +def check_processor_lifecycle(quent_events: list[dict]) -> None: + thread_pool_ids = { + x["id"] + for x in quent_events + if "ThreadPool" in x["data"] and "Declaration" in x["data"]["ThreadPool"] + } + assert len(thread_pool_ids) >= 1 + + processor_events = [x for x in quent_events if "Processor" in x["data"]] + init_events = [ + x + for x in processor_events + if "ProcessorInitializing" in x["data"]["Processor"]["state"] + ] + finalizing_events = [ + x + for x in processor_events + if x["data"]["Processor"]["state"] == {"ProcessorFinalizing": None} + ] + exit_events = [ + x for x in processor_events if x["data"]["Processor"]["state"] == "Exit" + ] + + assert len(init_events) == len(finalizing_events) == len(exit_events) + assert len(init_events) > 0 + + init_by_id = {x["id"]: x for x in init_events} + finalizing_by_id = {x["id"]: x for x in finalizing_events} + exit_by_id = {x["id"]: x for x in exit_events} + assert init_by_id.keys() == finalizing_by_id.keys() == exit_by_id.keys() + + for init_event in init_by_id.values(): + parent_group_id = init_event["data"]["Processor"]["state"][ + "ProcessorInitializing" + ]["parent_group_id"] + assert parent_group_id in thread_pool_ids From 43a2df94915f3eede66b9521d76276a03c6fa54a Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 13 Jul 2026 12:33:17 -0700 Subject: [PATCH 03/60] is_io_node --- python/cudf_polars/cudf_polars/dsl/ir.py | 7 +++++++ python/cudf_polars/cudf_polars/dsl/tracing.py | 6 ++---- python/cudf_polars/cudf_polars/streaming/io.py | 6 ++++++ 3 files changed, 15 insertions(+), 4 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 4228edc4545e..b67c19119a38 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -248,6 +248,9 @@ class IR(Node["IR"]): schema: Schema """Mapping from column names to their data types.""" + is_io_node: bool = False + """Whether the node is an IO node.""" + def get_hashable(self) -> Hashable: """ Hashable representation of node, treating schema dictionary. @@ -701,6 +704,8 @@ class Scan(IR): PARQUET_DEFAULT_CHUNK_SIZE: int = 0 # unlimited PARQUET_DEFAULT_PASS_LIMIT: int = 16 * 1024**3 # 16GiB + is_io_node: bool = True + def __init__( self, schema: Schema, @@ -1635,6 +1640,8 @@ class DataFrameScan(IR): projection: tuple[str, ...] | None """List of columns to project out.""" + is_io_node: bool = True + def __init__( self, schema: Schema, diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index 94cded4f1563..86c52fc91e4d 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -63,7 +63,6 @@ def _begin_quent_do_evaluate_events( ir_execution_context: IRExecutionContext, ) -> tuple[Any, Any, bool] | None: import cudf_polars.quent - from cudf_polars.dsl.ir import DataFrameScan, Scan quent_ir_execution_context = ir_execution_context.quent_ir_execution_context if quent_ir_execution_context is None: @@ -80,13 +79,12 @@ def _begin_quent_do_evaluate_events( quent_processor = quent_ir_execution_context.get_or_declare_processor( thread_ident=threading.get_ident(), ) - is_io_node = issubclass(cls, (Scan, DataFrameScan)) quent_ir_execution_context.logger.emit(quent_task.queueing()) - if not is_io_node: + if not cls.is_io_node: quent_ir_execution_context.logger.emit( quent_task.allocating(resource_id=quent_processor.id) ) - return quent_task, quent_processor, is_io_node + return quent_task, quent_processor, cls.is_io_node def _end_quent_do_evaluate_events( diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index b3b4438812c5..5ba0f5a2e4a2 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -192,6 +192,8 @@ class SplitScan(IR): (skip_rows and n_rows) is calculated at IO time. """ + is_io_node: bool = True + __slots__ = ( "base_scan", "cached_parquet_info", @@ -379,6 +381,8 @@ class FusedScan(IR): SINGLE_FILE (N = 1). """ + is_io_node: bool = True + __slots__ = ( "base_scan", "cached_parquet_info", @@ -614,6 +618,8 @@ def _( class StreamingScan(IR): """A streaming scan node.""" + is_io_node: bool = True + __slots__ = ( "base_scan", "scan_type", From 8745a5bf3c66bb14e98e6597555a3c6c99e0c99b Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 13 Jul 2026 13:33:42 -0700 Subject: [PATCH 04/60] remove print --- python/cudf_polars/cudf_polars/quent/_context.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index cb55185ca73a..8c9c37fcf2e2 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -231,7 +231,6 @@ def _emit_query_events(self, logger: QuentLogger, query: Query) -> None: This includes events for 'Declare', 'Init', and 'Planning'. """ - print(f"Query int: {query.id}") logger.emit(query._init(query_group=self.query_group)) logger.emit(query._planning()) logger.emit(query._executing()) From 0c5e50557317c33af2589b86e02366ace8af70c1 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 13 Jul 2026 14:14:09 -0700 Subject: [PATCH 05/60] change emission times --- python/cudf_polars/cudf_polars/dsl/tracing.py | 52 +++++++++---------- 1 file changed, 26 insertions(+), 26 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index 86c52fc91e4d..006db1101523 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -52,6 +52,7 @@ import cudf_polars.containers from cudf_polars.dsl import ir from cudf_polars.dsl.ir import IRExecutionContext + from cudf_polars.quent import Task def _dataframe_size_bytes(frame: cudf_polars.containers.DataFrame) -> int: @@ -61,7 +62,7 @@ def _dataframe_size_bytes(frame: cudf_polars.containers.DataFrame) -> int: def _begin_quent_do_evaluate_events( cls: type[ir.IR], ir_execution_context: IRExecutionContext, -) -> tuple[Any, Any, bool] | None: +) -> Task | None: import cudf_polars.quent quent_ir_execution_context = ir_execution_context.quent_ir_execution_context @@ -84,7 +85,25 @@ def _begin_quent_do_evaluate_events( quent_ir_execution_context.logger.emit( quent_task.allocating(resource_id=quent_processor.id) ) - return quent_task, quent_processor, cls.is_io_node + quent_ir_execution_context.logger.emit( + quent_task.computing( + use_thread=quent_processor, + use_memory=quent_ir_execution_context.device_memory, + # memory_capacity_bytes=output_capacity_bytes, + ) + ) + else: + quent_ir_execution_context.logger.emit( + quent_task.loading( + use_thread=quent_processor, + use_channel=quent_ir_execution_context.disk_to_device_channel, + # channel_capacity_bytes=output_capacity_bytes, + use_memory=quent_ir_execution_context.device_memory, + # memory_capacity_bytes=output_capacity_bytes, + ) + ) + + return quent_task def _end_quent_do_evaluate_events( @@ -92,34 +111,15 @@ def _end_quent_do_evaluate_events( frames: Sequence[cudf_polars.containers.DataFrame], result: cudf_polars.containers.DataFrame, ir_execution_context: IRExecutionContext, - quent_state: tuple[Any, Any, bool], + quent_task: Task, ) -> None: import cudf_polars.quent - quent_task, quent_processor, is_io_node = quent_state quent_ir_execution_context = ir_execution_context.quent_ir_execution_context if quent_ir_execution_context is None: return output_capacity_bytes = _dataframe_size_bytes(result) - if is_io_node: - quent_ir_execution_context.logger.emit( - quent_task.loading( - use_thread=quent_processor, - use_channel=quent_ir_execution_context.disk_to_device_channel, - channel_capacity_bytes=output_capacity_bytes, - use_memory=quent_ir_execution_context.device_memory, - memory_capacity_bytes=output_capacity_bytes, - ) - ) - else: - quent_ir_execution_context.logger.emit( - quent_task.computing( - use_thread=quent_processor, - use_memory=quent_ir_execution_context.device_memory, - memory_capacity_bytes=output_capacity_bytes, - ) - ) quent_ir_execution_context.logger.emit( quent_ir_execution_context.quent_operator.statistics( statistics=cudf_polars.quent.Statistics( @@ -252,9 +252,9 @@ def wrapper( list(args) + [v for k, v in kwargs.items() if k != "context"] )[cls._n_non_child_args :] # type: ignore[assignment] - quent_state = None + quent_task = None if ir_execution_context is not None: - quent_state = _begin_quent_do_evaluate_events(cls, ir_execution_context) + quent_task = _begin_quent_do_evaluate_events(cls, ir_execution_context) if LOG_TRACES: # pragma: no cover; requires CUDF_POLARS_LOG_TRACES=1 pynvml.nvmlInit() @@ -293,9 +293,9 @@ def wrapper( else: result = func(cls, *args, **kwargs) - if ir_execution_context is not None and quent_state is not None: + if ir_execution_context is not None and quent_task is not None: _end_quent_do_evaluate_events( - cls, frames, result, ir_execution_context, quent_state + cls, frames, result, ir_execution_context, quent_task ) return result From 323b6aa40f842dd332db1aec8a72e18db90fd812 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 13 Jul 2026 14:42:20 -0700 Subject: [PATCH 06/60] Remove Actor-scoped tasks --- .../streaming/actor_graph/utils.py | 54 +++---------------- 1 file changed, 7 insertions(+), 47 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py index ddefb18165cd..0afafefad246 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -9,9 +9,7 @@ import itertools import operator import struct -import threading import time -import uuid from collections import defaultdict, deque from contextlib import asynccontextmanager from dataclasses import dataclass @@ -43,13 +41,11 @@ from cudf_polars.dsl.expr import Cast, Col, NamedExpr, TemporalFunction from cudf_polars.dsl.ir import ( Cache, - DataFrameScan, Filter, GroupBy, HStack, Join, Projection, - Scan, Select, ) from cudf_polars.dsl.tracing import Scope @@ -279,32 +275,6 @@ async def shutdown_on_error( if ir_context is not None: contextvars["cudf_polars_query_id"] = str(ir_context.query_id) - quent_ir_execution_context = None - quent_task = None - quent_processor = None - is_io_node = False - if ir_context is not None: - quent_ir_execution_context = ir_context.quent_ir_execution_context - if quent_ir_execution_context is not None: - token = uuid.uuid4() - quent_task = cudf_polars.quent.Task( - instance_name=( - f"Actor-{type(trace_ir).__name__}-" - f"{quent_ir_execution_context.quent_operator.id.hex[:8]}-" - f"{token.hex[:8]}" - ), - operator_id=quent_ir_execution_context.quent_operator.id, - ) - quent_processor = quent_ir_execution_context.get_or_declare_processor( - thread_ident=threading.get_ident(), - ) - is_io_node = issubclass(type(trace_ir), (Scan, DataFrameScan)) - quent_ir_execution_context.logger.emit(quent_task.queueing()) - if not is_io_node: - quent_ir_execution_context.logger.emit( - quent_task.allocating(resource_id=quent_processor.id) - ) - with cudf_polars.dsl.tracing.bound_contextvars(**contextvars): start = time.monotonic_ns() try: @@ -338,22 +308,13 @@ async def shutdown_on_error( "Streaming Actor", start=start, stop=stop, **record ) - if quent_ir_execution_context is not None and quent_task is not None: - if is_io_node: - quent_ir_execution_context.logger.emit( - quent_task.loading( - use_thread=quent_processor, - use_channel=quent_ir_execution_context.disk_to_device_channel, - use_memory=quent_ir_execution_context.device_memory, - ) - ) - else: - quent_ir_execution_context.logger.emit( - quent_task.computing( - use_thread=quent_processor, - use_memory=quent_ir_execution_context.device_memory, - ) - ) + if ( + ir_context is not None + and ( + quent_ir_execution_context := ir_context.quent_ir_execution_context + ) + is not None + ): custom_attributes = [] if tracer is not None and tracer.chunk_count is not None: custom_attributes.append( @@ -394,7 +355,6 @@ async def shutdown_on_error( ) ) ) - quent_ir_execution_context.logger.emit(quent_task.exit()) def _update_ordering_indices( From bbc7966b7e968e5ec6dc6491048f76f20147d332 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 13 Jul 2026 14:55:53 -0700 Subject: [PATCH 07/60] Revert collective trace --- .../actor_graph/collectives/shuffle.py | 30 ++++--------------- 1 file changed, 6 insertions(+), 24 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/shuffle.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/shuffle.py index c1436ac52405..391edf0bd55d 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/shuffle.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/shuffle.py @@ -27,7 +27,6 @@ from rapidsmpf.streaming.core.context import Context from rapidsmpf.streaming.core.message import Message -import cudf_polars.quent from cudf_polars.dsl.expr import Col from cudf_polars.streaming.actor_graph.dispatch import ( generate_ir_sub_network, @@ -400,32 +399,15 @@ async def _global_shuffle( skip_insert = metadata_in.duplicated and comm.rank != 0 shuffle = ShuffleManager(context, comm, num_partitions, collective_id) - total_bytes_inserted: int = 0 async with shuffle.inserting() as inserter: while (msg := await ch_in.recv(context)) is not None: if not skip_insert: - chunk = TableChunk.from_message( - msg, br=context.br() - ).make_available_and_spill(context.br(), allow_overbooking=True) - total_bytes_inserted += chunk.data_alloc_size() - inserter.insert_hash(chunk, columns_to_hash) - - if ( - (quent_ctx := ir_context.quent_ir_execution_context) is not None - and quent_ctx.link_channels - and total_bytes_inserted > 0 - ): - quent_task = cudf_polars.quent.Task( - instance_name=f"Shuffle-send-{quent_ctx.quent_operator.id.hex[:8]}", - operator_id=quent_ctx.quent_operator.id, - ) - quent_ctx.logger.emit(quent_task.queueing()) - bytes_per_rank = total_bytes_inserted // comm.nranks - for link in quent_ctx.link_channels.values(): - quent_ctx.logger.emit( - quent_task.sending(use_link=link, link_capacity_bytes=bytes_per_rank) - ) - quent_ctx.logger.emit(quent_task.exit()) + inserter.insert_hash( + TableChunk.from_message( + msg, br=context.br() + ).make_available_and_spill(context.br(), allow_overbooking=True), + columns_to_hash, + ) for partition_id in shuffle.local_partitions(): stream = ir_context.get_cuda_stream() From ba03d7c54e334aca4a64720020b46ff2fc1cdd4a Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 13 Jul 2026 14:56:07 -0700 Subject: [PATCH 08/60] simplify tests --- .../tests/quent/test_quent_integration.py | 104 +++++------------- 1 file changed, 27 insertions(+), 77 deletions(-) diff --git a/python/cudf_polars/tests/quent/test_quent_integration.py b/python/cudf_polars/tests/quent/test_quent_integration.py index ac81f1e92ab3..804d031508ab 100644 --- a/python/cudf_polars/tests/quent/test_quent_integration.py +++ b/python/cudf_polars/tests/quent/test_quent_integration.py @@ -161,35 +161,41 @@ def check_quent_events(engine: StreamingEngine, quent_context: QuentContext) -> assert query_exit["id"] == query_id assert query_exit["data"]["Query"]["seq"] == 3 - -def test_quent_events_include_resources( - engine_with_quent_context: StreamingEngine, quent_context: QuentContext -) -> None: - q = pl.LazyFrame({"x": [1, 2, 3, 4]}).filter(pl.col("x") > 1) - with engine_with_quent_context: - q.collect(engine=engine_with_quent_context) - - quent_events = engine_with_quent_context._quent_events memory_events = [x for x in quent_events if "Memory" in x["data"]] task_events = [x for x in quent_events if "Task" in x["data"]] assert len(memory_events) > 0 assert len(task_events) > 0 + # A single collect exercises the full processor lifecycle, so fold that + # check in here rather than paying for a dedicated engine startup. + check_processor_lifecycle(quent_events) + -def test_quent_device_memory_declared_once_per_engine( +def test_quent_events_multiple_collects( engine_with_quent_context: StreamingEngine, quent_context: QuentContext ) -> None: - # Device memory is an engine-scoped resource: it must be initialized and - # finalized exactly once per engine, regardless of how many collects run. - q1 = pl.LazyFrame({"x": [1, 2, 3]}).filter(pl.col("x") > 1) - q2 = pl.LazyFrame({"y": [4, 5, 6]}).filter(pl.col("y") > 4) + # Everything that depends on running more than one collect against the same + # engine is folded into this single test to avoid paying for extra engine + # startups. Running the *same* query twice is the strongest scenario: query + # 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) with engine_with_quent_context: - q1.collect(engine=engine_with_quent_context) - q2.collect(engine=engine_with_quent_context) + q.collect(engine=engine_with_quent_context) + q.collect(engine=engine_with_quent_context) quent_events = engine_with_quent_context._quent_events - memory_events = [x for x in quent_events if "Memory" in x["data"]] + # The processor lifecycle stays balanced across multiple collects. + check_processor_lifecycle(quent_events) + + # Device memory is engine/worker-scoped: it is initialized and finalized + # exactly once per worker, matching the number of engine-scoped ThreadPool + # declarations. Critically, running two collects must NOT re-declare it + # (the per-query bug would produce a fresh device memory per collect, i.e. + # twice as many inits as thread pools). + memory_events = [x for x in quent_events if "Memory" in x["data"]] device_init_events = [ x for x in memory_events @@ -204,12 +210,6 @@ def test_quent_device_memory_declared_once_per_engine( if x["data"]["Memory"]["state"] == "Exit" and x["id"] in {e["id"] for e in device_init_events} ] - - # Device memory is engine/worker-scoped: it is initialized and finalized - # exactly once per worker, matching the number of engine-scoped ThreadPool - # declarations. Critically, running two collects must NOT re-declare it - # (the per-query bug would produce a fresh device memory per collect, i.e. - # twice as many inits as thread pools). thread_pool_decls = [ x for x in quent_events @@ -217,72 +217,24 @@ def test_quent_device_memory_declared_once_per_engine( ] assert len(thread_pool_decls) >= 1 assert len(device_init_events) == len(thread_pool_decls) - # Every device memory id is initialized once and exited once. init_ids = [x["id"] for x in device_init_events] assert len(set(init_ids)) == len(init_ids) assert {x["id"] for x in device_exit_events} == set(init_ids) - -def test_quent_processor_lifecycle_balanced( - engine_with_quent_context: StreamingEngine, quent_context: QuentContext -) -> None: - q = pl.LazyFrame({"x": list(range(100))}).filter(pl.col("x") > 1) - with engine_with_quent_context: - q.collect(engine=engine_with_quent_context) - - quent_events = engine_with_quent_context._quent_events - check_processor_lifecycle(quent_events) - - -def test_quent_processor_lifecycle_across_multiple_collects( - engine_with_quent_context: StreamingEngine, quent_context: QuentContext -) -> None: - q1 = pl.LazyFrame({"x": [1, 2, 3]}).filter(pl.col("x") > 1) - q2 = pl.LazyFrame({"y": [4, 5, 6]}).filter(pl.col("y") > 4) - with engine_with_quent_context: - q1.collect(engine=engine_with_quent_context) - q2.collect(engine=engine_with_quent_context) - - quent_events = engine_with_quent_context._quent_events - check_processor_lifecycle(quent_events) - - -def test_quent_query_id_unique_per_collect( - engine_with_quent_context: StreamingEngine, quent_context: QuentContext -) -> None: - q1 = pl.LazyFrame({"x": [1, 2, 3]}).filter(pl.col("x") > 1) - q2 = pl.LazyFrame({"y": [4, 5, 6]}).filter(pl.col("y") > 4) - with engine_with_quent_context: - q1.collect(engine=engine_with_quent_context) - q2.collect(engine=engine_with_quent_context) - - quent_events = engine_with_quent_context._quent_events + # Each collect reuses the engine-scoped QuentContext but must emit a + # distinct query id. query_init_ids = [ x["id"] for x in quent_events if "Query" in x["data"] and "Init" in x["data"]["Query"].get("state", {}) ] assert len(query_init_ids) == 2 - # Each collect reuses the engine-scoped QuentContext but must emit a - # distinct query id. assert len(set(query_init_ids)) == 2 assert str(quent_context.query.id) not in query_init_ids - -def test_quent_plan_id_unique_per_collect( - engine_with_quent_context: StreamingEngine, quent_context: QuentContext -) -> None: - # Run the *same* query twice. ``get_stable_plan_id`` is a deterministic - # function of the IR structure, so without namespacing by the per-collect - # query id both collects would emit the same logical plan id under - # different parent queries. - q = pl.LazyFrame({"x": [1, 2, 3]}).filter(pl.col("x") > 1) - with engine_with_quent_context: - q.collect(engine=engine_with_quent_context) - q.collect(engine=engine_with_quent_context) - - quent_events = engine_with_quent_context._quent_events + # Without namespacing by the per-collect query id, both identical collects + # would emit the same logical plan id under different parent queries. logical_plan_decls = [ x for x in quent_events @@ -291,10 +243,8 @@ def test_quent_plan_id_unique_per_collect( and x["data"]["Plan"]["Declaration"]["instance_name"] == "logical" ] assert len(logical_plan_decls) == 2 - plan_ids = [x["id"] for x in logical_plan_decls] assert len(set(plan_ids)) == 2 - # Each logical plan must hang off the distinct per-collect query id. parent_query_ids = [ x["data"]["Plan"]["Declaration"]["parent"]["query_id"] From a20d3b920539bc4d3f90fa1f5d70a9382f5148b9 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 13 Jul 2026 15:07:49 -0700 Subject: [PATCH 09/60] docs --- python/cudf_polars/cudf_polars/dsl/tracing.py | 67 ++++++++++++++++++- 1 file changed, 64 insertions(+), 3 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index 006db1101523..b958657c4356 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -63,6 +63,37 @@ def _begin_quent_do_evaluate_events( cls: type[ir.IR], ir_execution_context: IRExecutionContext, ) -> Task | None: + """ + Build and emit Quent events for the start of an IR node's evaluation. + + Parameters + ---------- + cls + The IR node class. + ir_execution_context + The IR execution context. To emit any events, this must have a + quent_ir_execution_context bound. + + Returns + ------- + The Quent task instance, or None if no Quent IR execution context is bound. + + Notes + ----- + The following events are emitted: + + - Queueing + - Loading (I/O nodes only) + - Allocating (non-I/O nodes only) + - Computing (non-I/O nodes only) + + The Loading, Allocating, and Computing events will indicate the host CPU thread + and device memory that they're using. + + See Also + -------- + _end_quent_do_evaluate_events + """ import cudf_polars.quent quent_ir_execution_context = ir_execution_context.quent_ir_execution_context @@ -89,7 +120,6 @@ def _begin_quent_do_evaluate_events( quent_task.computing( use_thread=quent_processor, use_memory=quent_ir_execution_context.device_memory, - # memory_capacity_bytes=output_capacity_bytes, ) ) else: @@ -97,9 +127,7 @@ def _begin_quent_do_evaluate_events( quent_task.loading( use_thread=quent_processor, use_channel=quent_ir_execution_context.disk_to_device_channel, - # channel_capacity_bytes=output_capacity_bytes, use_memory=quent_ir_execution_context.device_memory, - # memory_capacity_bytes=output_capacity_bytes, ) ) @@ -113,6 +141,39 @@ def _end_quent_do_evaluate_events( ir_execution_context: IRExecutionContext, quent_task: Task, ) -> None: + """ + Build and emit Quent events for the end of an IR node's evaluation. + + Parameters + ---------- + cls + The IR node class. + frames + The input dataframes passed to the IR node. + result + The output dataframe returned from the IR node. + ir_execution_context + The IR execution context. To emit any events, this must have a + quent_ir_execution_context bound. + quent_task + The Quent task instance created by ``_begin_quent_do_evaluate_events``. + + Notes + ----- + This method emits an ``Exit`` event for the Quent Task, whose timestamp represents + when the IR node completed host-side processing. + + A ``Statistics`` record, associated with the Quent Operator bound to the IR execution context, + is also emitted. It includes + + - input bytes: the total size of the input dataframes. + - output bytes: the size of the output dataframe. + - output rows: the number of rows in the output dataframe. + + See Also + -------- + _begin_quent_do_evaluate_events + """ import cudf_polars.quent quent_ir_execution_context = ir_execution_context.quent_ir_execution_context From 3b7e6afd3a3cfce40cdc438f30e7f815e1f2f0a7 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 13 Jul 2026 15:18:27 -0700 Subject: [PATCH 10/60] restore old tracing flow --- python/cudf_polars/cudf_polars/dsl/tracing.py | 57 +++++++++++-------- 1 file changed, 32 insertions(+), 25 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index b958657c4356..0d454dcddb29 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -300,34 +300,41 @@ def log_do_evaluate( func The ``IR.do_evaluate`` method to wrap. """ - - @functools.wraps(func) - def wrapper( - cls: type[ir.IR], - *args: P.args, - **kwargs: P.kwargs, - ) -> cudf_polars.containers.DataFrame: - ir_execution_context: IRExecutionContext | None = kwargs.get("context") # type: ignore[assignment] - - frames: list[cudf_polars.containers.DataFrame] = ( - list(args) + [v for k, v in kwargs.items() if k != "context"] - )[cls._n_non_child_args :] # type: ignore[assignment] - - quent_task = None - if ir_execution_context is not None: - quent_task = _begin_quent_do_evaluate_events(cls, ir_execution_context) - - if LOG_TRACES: # pragma: no cover; requires CUDF_POLARS_LOG_TRACES=1 + if not LOG_TRACES: + return func + else: # pragma: no cover; requires CUDF_POLARS_LOG_TRACES=1 + + @functools.wraps(func) + def wrapper( + cls: type[ir.IR], + *args: P.args, + **kwargs: P.kwargs, + ) -> cudf_polars.containers.DataFrame: pynvml.nvmlInit() maybe_handle = get_device_handle() pid = _getpid() log = structlog.get_logger() + # By convention, all non-dataframe arguments (non-child) come first. + # Anything remaining is a dataframe, except for 'context' kwarg. + frames: list[cudf_polars.containers.DataFrame] = ( + list(args) + [v for k, v in kwargs.items() if k != "context"] + )[cls._n_non_child_args :] # type: ignore[assignment] + + # And the kwonly 'context' argument has the IR execution context. + ir_execution_context: IRExecutionContext = kwargs["context"] # type: ignore[assignment] + quent_task = _begin_quent_do_evaluate_events(cls, ir_execution_context) + before_start = time.monotonic_ns() before = make_snapshot( cls, frames, phase="input", device_handle=maybe_handle, pid=pid ) before_end = time.monotonic_ns() + + # The decorator preserves the exact signature of the original do_evaluate method. + # Each IR.do_evaluate method is a classmethod that takes the IR class as first + # argument, followed by the method-specific arguments, and returns a DataFrame. + start = time.monotonic_ns() result = func(cls, *args, **kwargs) stop = time.monotonic_ns() @@ -351,17 +358,17 @@ def wrapper( } ) log.info("Execute IR", **record) - else: + result = func(cls, *args, **kwargs) - if ir_execution_context is not None and quent_task is not None: - _end_quent_do_evaluate_events( - cls, frames, result, ir_execution_context, quent_task - ) + if quent_task is not None: + _end_quent_do_evaluate_events( + cls, frames, result, ir_execution_context, quent_task + ) - return result + return result - return wrapper + return wrapper @contextlib.contextmanager From 2fc4e742963d27093d565ac2d2d10e8c66265497 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 13 Jul 2026 15:30:56 -0700 Subject: [PATCH 11/60] imports --- python/cudf_polars/cudf_polars/dsl/tracing.py | 10 ++++---- python/cudf_polars/cudf_polars/engine/core.py | 8 +++--- python/cudf_polars/cudf_polars/engine/dask.py | 17 ++++++------- .../cudf_polars/cudf_polars/engine/options.py | 2 +- python/cudf_polars/cudf_polars/engine/ray.py | 13 +++++----- python/cudf_polars/cudf_polars/engine/spmd.py | 23 ++++++++--------- .../cudf_polars/cudf_polars/quent/__init__.py | 20 --------------- .../cudf_polars/streaming/actor_graph/core.py | 7 +++--- .../streaming/actor_graph/dispatch.py | 11 ++++---- .../streaming/actor_graph/utils.py | 4 +-- .../cudf_polars/streaming/benchmarks/utils.py | 3 ++- python/cudf_polars/tests/quent/conftest.py | 5 ++-- python/cudf_polars/tests/quent/test_quent.py | 25 +++++++++++-------- .../tests/quent/test_quent_integration.py | 2 +- .../cudf_polars/tests/streaming/test_spmd.py | 3 ++- python/cudf_polars/tests/test_config.py | 6 +++-- 16 files changed, 72 insertions(+), 87 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index 0d454dcddb29..915bf44c94c5 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -52,7 +52,7 @@ import cudf_polars.containers from cudf_polars.dsl import ir from cudf_polars.dsl.ir import IRExecutionContext - from cudf_polars.quent import Task + from cudf_polars.quent._types import Task def _dataframe_size_bytes(frame: cudf_polars.containers.DataFrame) -> int: @@ -94,14 +94,14 @@ def _begin_quent_do_evaluate_events( -------- _end_quent_do_evaluate_events """ - import cudf_polars.quent + import cudf_polars.quent._types quent_ir_execution_context = ir_execution_context.quent_ir_execution_context if quent_ir_execution_context is None: return None token = uuid.uuid4() - quent_task = cudf_polars.quent.Task( + quent_task = cudf_polars.quent._types.Task( instance_name=( f"{cls.__name__}-{quent_ir_execution_context.quent_operator.id.hex[:8]}-" f"{token.hex[:8]}" @@ -174,7 +174,7 @@ def _end_quent_do_evaluate_events( -------- _begin_quent_do_evaluate_events """ - import cudf_polars.quent + import cudf_polars.quent._types quent_ir_execution_context = ir_execution_context.quent_ir_execution_context if quent_ir_execution_context is None: @@ -183,7 +183,7 @@ def _end_quent_do_evaluate_events( output_capacity_bytes = _dataframe_size_bytes(result) quent_ir_execution_context.logger.emit( quent_ir_execution_context.quent_operator.statistics( - statistics=cudf_polars.quent.Statistics( + statistics=cudf_polars.quent._types.Statistics( input_bytes=sum(_dataframe_size_bytes(frame) for frame in frames), output_bytes=output_capacity_bytes, output_rows=result.num_rows, diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index a9a2dc4df3b5..96ce3decd38c 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -53,8 +53,8 @@ from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.streaming.core.context import Context - import cudf_polars.quent import cudf_polars.quent._logging + import cudf_polars.quent._types from cudf_polars.dsl.ir import IR from cudf_polars.quent._context import LocalQuentContext from cudf_polars.streaming.base import PartitionInfo @@ -497,7 +497,7 @@ def execute_ir_on_rank( stats: StatsCollector, collective_id_map: dict[IR, list[int]], *, - quent_operator_map: dict[IR, cudf_polars.quent.Operator] | None = None, + quent_operator_map: dict[IR, cudf_polars.quent._types.Operator] | None = None, local_quent_context: LocalQuentContext | None = None, ) -> tuple[pl.DataFrame, list[ChannelMetadata]]: """ @@ -780,8 +780,8 @@ def evaluate_on_rank( """ stats = allgather_stats(comm, ctx.br(), ir, config_options, py_executor) - physical_op_by_id: dict[str, cudf_polars.quent.Operator] | None = None - quent_operator_map: dict[IR, cudf_polars.quent.Operator] | None = None + physical_op_by_id: dict[str, cudf_polars.quent._types.Operator] | None = None + quent_operator_map: dict[IR, cudf_polars.quent._types.Operator] | None = None if config_options.executor.quent_context is not None: assert local_quent_context is not None _declare_network_channels(comm, local_quent_context) diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index d1068c0ea5f8..45dcda871624 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -60,6 +60,7 @@ from cudf_polars.dsl.ir import IR from cudf_polars.engine.core import T from cudf_polars.engine.options import StreamingOptions + from cudf_polars.quent._context import QuentContext from cudf_polars.streaming.parallel import ConfigOptions from cudf_polars.utils.config import StreamingExecutor @@ -143,7 +144,7 @@ def _setup_root( dask_worker: distributed.Worker | None = None, engine_id: uuid.UUID, worker_id: uuid.UUID, - quent_context: cudf_polars.quent.QuentContext | None, + quent_context: QuentContext | None, ) -> bytes: """ Initialize the root rank on one Dask worker. @@ -233,7 +234,7 @@ def _setup_worker( worker_ids: list[uuid.UUID], engine_id: uuid.UUID, num_py_executors: int, - quent_context: cudf_polars.quent.QuentContext | None, + quent_context: QuentContext | None, dask_worker: distributed.Worker | None = None, ) -> None: """ @@ -495,7 +496,7 @@ def _worker_evaluate( uid: str, collect_metadata: bool = False, query_id: uuid.UUID, - quent_context: cudf_polars.quent.QuentContext | None = None, + quent_context: QuentContext | None = None, dask_worker: distributed.Worker | None = None, ) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]: """ @@ -768,9 +769,7 @@ def __init__( executor_options = executor_options or {} engine_options = engine_options or {} - quent_context: cudf_polars.quent.QuentContext | None = executor_options.get( - "quent_context" - ) + quent_context: QuentContext | None = executor_options.get("quent_context") if quent_context is not None: self._quent_logger = cudf_polars.quent._logging.QuentLogger() else: @@ -1054,9 +1053,9 @@ def shutdown(self) -> None: ctx = self._dask_context self._dask_context = None exceptions: list[Exception] = [] - quent_context: cudf_polars.quent.QuentContext | None = self.config[ - "executor_options" - ].get("quent_context") + quent_context: QuentContext | None = self.config["executor_options"].get( + "quent_context" + ) try: # Teardown emits Worker.exit, then we drain all buffered events # (including the exit event) from workers. diff --git a/python/cudf_polars/cudf_polars/engine/options.py b/python/cudf_polars/cudf_polars/engine/options.py index c959f24f98d9..5a70f0dd23e1 100644 --- a/python/cudf_polars/cudf_polars/engine/options.py +++ b/python/cudf_polars/cudf_polars/engine/options.py @@ -23,7 +23,7 @@ if TYPE_CHECKING: from collections.abc import Callable - from cudf_polars.quent import QuentContext + from cudf_polars.quent._context import QuentContext from cudf_polars.utils.config import ( DynamicPlanningOptions, ParquetOptions, diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index a42258a672bf..128a9c7cd2c5 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -61,6 +61,7 @@ from cudf_polars.dsl.ir import IR from cudf_polars.engine.core import T from cudf_polars.engine.options import StreamingOptions + from cudf_polars.quent._context import QuentContext from cudf_polars.streaming.parallel import ConfigOptions from cudf_polars.utils.config import StreamingExecutor @@ -437,7 +438,7 @@ def evaluate_polars_ir( config_options: ConfigOptions[StreamingExecutor], *, collect_metadata: bool, - quent_context: cudf_polars.quent.QuentContext | None, + quent_context: QuentContext | None, query_id: uuid.UUID, ) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]: """ @@ -670,9 +671,7 @@ def __init__( check_reserved_keys(executor_options, engine_options) - quent_context: cudf_polars.quent.QuentContext | None = executor_options.get( - "quent_context" - ) + quent_context: QuentContext | None = executor_options.get("quent_context") if quent_context is not None: self._quent_logger = cudf_polars.quent._logging.QuentLogger() else: @@ -924,9 +923,9 @@ def shutdown(self) -> None: if self._rank_actors is None: return # already shut down; idempotent exceptions: list[Exception] = [] - quent_context: cudf_polars.quent.QuentContext | None = self.config[ - "executor_options" - ].get("quent_context") + quent_context: QuentContext | None = self.config["executor_options"].get( + "quent_context" + ) try: # If Ray is no longer initialized (for example, if ``ray.shutdown()`` was # called before ``RayEngine.shutdown()``), the actors are gone as well. diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index c39fe310415d..4b356605d049 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -47,6 +47,7 @@ from cudf_polars.quent._context import ( LocalQuentContext, ProcessorRegistry, + QuentContext, declare_worker_resources, finalize_worker_resources, ) @@ -242,7 +243,7 @@ def synchronize_quent_context( *, comm: Communicator, context: Context, -) -> cudf_polars.quent.QuentContext: +) -> QuentContext: """ Ensure all ranks use the same Quent engine ID. @@ -250,19 +251,19 @@ def synchronize_quent_context( ranks participate in an AllGather so every process converges on that value. """ if comm.rank == 0: - quent_context = cudf_polars.quent.QuentContext() + quent_context = QuentContext() data = quent_context.serialize() else: data = b"" if comm.nranks == 1: # skip the collective - return cudf_polars.quent.QuentContext() + return QuentContext() with reserve_op_id() as op_id: all_data = all_gather_host_data(comm, context.br(), op_id, data) - return cudf_polars.quent.QuentContext.deserialize(all_data[0]) + return QuentContext.deserialize(all_data[0]) class SPMDEngine(StreamingEngine): @@ -421,9 +422,7 @@ def __init__( ) -> None: executor_options = executor_options or {} engine_options = engine_options or {} - quent_context: cudf_polars.quent.QuentContext | None = executor_options.get( - "quent_context" - ) + quent_context: QuentContext | None = executor_options.get("quent_context") if quent_context is not None: self._quent_logger = cudf_polars.quent._logging.QuentLogger() else: @@ -638,9 +637,7 @@ def _reset( if existing_quent_context is not None: executor_options.setdefault("quent_context", existing_quent_context) engine_options = engine_options or {} - quent_context: cudf_polars.quent.QuentContext | None = executor_options.get( - "quent_context" - ) + quent_context: QuentContext | None = executor_options.get("quent_context") rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) # Collective: synchronize all ranks before tearing down the Context. @@ -835,9 +832,9 @@ def shutdown(self) -> None: # quent traces before that. # Clear the references only after shutdown completes. - quent_context: cudf_polars.quent.QuentContext | None = self.config[ - "executor_options" - ].get("quent_context") + quent_context: QuentContext | None = self.config["executor_options"].get( + "quent_context" + ) if self._quent_logger is not None: if quent_context is not None: diff --git a/python/cudf_polars/cudf_polars/quent/__init__.py b/python/cudf_polars/cudf_polars/quent/__init__.py index 877f303b0b57..3a4e5b991f0a 100644 --- a/python/cudf_polars/cudf_polars/quent/__init__.py +++ b/python/cudf_polars/cudf_polars/quent/__init__.py @@ -5,44 +5,24 @@ from __future__ import annotations -from cudf_polars.quent._context import ( - LocalQuentContext, - QuentContext, - QuentIRExecutionContext, -) from cudf_polars.quent._types import ( Attribute, - Channel, Engine, HomogeneousListValue, Implementation, - Network, - Operator, Query, QueryGroup, ScalarValue, - Statistics, - Task, Value, - Worker, ) __all__ = [ "Attribute", - "Channel", "Engine", "HomogeneousListValue", "Implementation", - "LocalQuentContext", - "Network", - "Operator", - "QuentContext", - "QuentIRExecutionContext", "Query", "QueryGroup", "ScalarValue", - "Statistics", - "Task", "Value", - "Worker", ] diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py index b8b565559e6b..148c1132161e 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py @@ -12,7 +12,8 @@ from rapidsmpf.streaming.core.leaf_actor import pull_from_channel import cudf_polars.dsl.tracing -import cudf_polars.quent +import cudf_polars.quent._context +import cudf_polars.quent._types from cudf_polars.dsl.ir import ( DataFrameScan, Join, @@ -230,8 +231,8 @@ def generate_network( ir_context: IRExecutionContext, collective_id_map: dict[IR, list[int]], metadata_collector: list[ChannelMetadata] | None, - quent_operator_map: dict[IR, cudf_polars.quent.Operator] | None = None, - local_quent_context: cudf_polars.quent.LocalQuentContext | None = None, + quent_operator_map: dict[IR, cudf_polars.quent._types.Operator] | None = None, + local_quent_context: cudf_polars.quent._context.LocalQuentContext | None = None, ) -> tuple[list[Any], DeferredMessages]: """ Translate the IR graph to a RapidsMPF streaming network. diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py index 57d420950e0a..adad3c6f89a4 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py @@ -16,7 +16,8 @@ from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.streaming.core.context import Context - import cudf_polars.quent + import cudf_polars.quent._context + import cudf_polars.quent._types from cudf_polars.dsl.ir import IR, IRExecutionContext from cudf_polars.streaming.actor_graph.utils import ChannelManager from cudf_polars.streaming.base import ( @@ -75,13 +76,13 @@ class GenState(TypedDict): max_io_threads: int stats: StatsCollector collective_id_map: dict[IR, list[int]] - quent_operator_map: dict[IR, cudf_polars.quent.Operator] | None - quent_execution_context: cudf_polars.quent.LocalQuentContext | None + quent_operator_map: dict[IR, cudf_polars.quent._types.Operator] | None + quent_execution_context: cudf_polars.quent._context.LocalQuentContext | None def ir_context_for_node(rec: SubNetGenerator, ir: IR) -> IRExecutionContext: """Return ``ir_context`` with the physical Quent operator bound when tracing.""" - import cudf_polars.quent + import cudf_polars.quent._context ir_context = rec.state["ir_context"] quent_operator_map = rec.state["quent_operator_map"] @@ -90,7 +91,7 @@ def ir_context_for_node(rec: SubNetGenerator, ir: IR) -> IRExecutionContext: quent_operator = quent_operator_map[ir] return dataclasses.replace( ir_context, - quent_ir_execution_context=cudf_polars.quent.QuentIRExecutionContext.from_execution_context( + quent_ir_execution_context=cudf_polars.quent._context.QuentIRExecutionContext.from_execution_context( execution_context=quent_execution_context, quent_operator=quent_operator, ), diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py index 0afafefad246..99ad2ae07421 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -36,7 +36,7 @@ from rapidsmpf.streaming.core.message import Message import cudf_polars.dsl.tracing -import cudf_polars.quent +import cudf_polars.quent._types from cudf_polars.containers import DataFrame from cudf_polars.dsl.expr import Cast, Col, NamedExpr, TemporalFunction from cudf_polars.dsl.ir import ( @@ -347,7 +347,7 @@ async def shutdown_on_error( output_rows = tracer.row_count quent_ir_execution_context.logger.emit( quent_ir_execution_context.quent_operator.statistics( - statistics=cudf_polars.quent.Statistics( + statistics=cudf_polars.quent._types.Statistics( output_rows=output_rows, input_bytes=0, output_bytes=0, diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index bfbe25649cfc..64428628a057 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -68,6 +68,7 @@ try: import cudf_polars.dsl.tracing import cudf_polars.quent + import cudf_polars.quent._context from cudf_polars.dsl.ir import IRExecutionContext from cudf_polars.dsl.tracing import Scope from cudf_polars.dsl.translate import Translator @@ -754,7 +755,7 @@ def get_executor_options( run_config.streaming_options.to_executor_options() ) executor_options["max_io_threads"] = run_config.max_io_threads - executor_options["quent_context"] = cudf_polars.quent.QuentContext( + executor_options["quent_context"] = cudf_polars.quent._context.QuentContext( engine=cudf_polars.quent.Engine(id=run_config.run_id) ) diff --git a/python/cudf_polars/tests/quent/conftest.py b/python/cudf_polars/tests/quent/conftest.py index 8d90ced657fa..849364e2f32a 100644 --- a/python/cudf_polars/tests/quent/conftest.py +++ b/python/cudf_polars/tests/quent/conftest.py @@ -10,15 +10,16 @@ import pytest import cudf_polars.quent +import cudf_polars.quent._context if TYPE_CHECKING: - from cudf_polars.quent import QuentContext + from cudf_polars.quent._context import QuentContext @pytest.fixture def quent_context() -> QuentContext: """A Quent Context with a QueryGroup and Query set.""" - return cudf_polars.quent.QuentContext( + return cudf_polars.quent._context.QuentContext( query_group=cudf_polars.quent.QueryGroup(instance_name="test_query_group"), query=cudf_polars.quent.Query(instance_name="test_query"), ) diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index 29f82dc105d5..3b5d7bbd01ba 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -16,7 +16,11 @@ import cudf_polars.quent import cudf_polars.quent._logging from cudf_polars.dsl.translate import Translator -from cudf_polars.quent._context import ProcessorRegistry +from cudf_polars.quent._context import ( + LocalQuentContext, + ProcessorRegistry, + QuentContext, +) from cudf_polars.quent._plan import build_plan, port_names_for_node from cudf_polars.quent._types import ( Attribute, @@ -36,7 +40,6 @@ if TYPE_CHECKING: from cudf_polars.dsl.ir import IR - from cudf_polars.quent import QuentContext from cudf_polars.quent._types import Processor from cudf_polars.utils.config import StreamingExecutor @@ -556,20 +559,20 @@ def test_query_lifecycle() -> None: @pytest.fixture def quent_context() -> QuentContext: - return cudf_polars.quent.QuentContext( + return QuentContext( query_group=cudf_polars.quent.QueryGroup(instance_name="test_query_group"), query=cudf_polars.quent.Query(instance_name="test_query"), ) def test_quent_context_serialization() -> None: - quent_context = cudf_polars.quent.QuentContext( + quent_context = QuentContext( query_group=cudf_polars.quent.QueryGroup(instance_name="test_query_group"), query=cudf_polars.quent.Query(instance_name="test_query"), ) data = quent_context.serialize() - new = cudf_polars.quent.QuentContext.deserialize(data) + new = QuentContext.deserialize(data) assert new == quent_context @@ -586,14 +589,14 @@ def test_quent_context_serialization_with_custom_attributes() -> None: ], ) ) - quent_context = cudf_polars.quent.QuentContext( + quent_context = QuentContext( engine=engine, query_group=cudf_polars.quent.QueryGroup(instance_name="test_query_group"), query=cudf_polars.quent.Query(instance_name="test_query"), ) data = quent_context.serialize() - new = cudf_polars.quent.QuentContext.deserialize(data) + new = QuentContext.deserialize(data) assert new == quent_context @@ -664,9 +667,9 @@ def test_processor_registry_reused_across_quent_contexts() -> None: pool_id = uuid.uuid4() thread_ident = 99 - context_a = cudf_polars.quent.QuentContext() - context_b = cudf_polars.quent.QuentContext() - local_a = cudf_polars.quent.LocalQuentContext( + context_a = QuentContext() + context_b = QuentContext() + local_a = LocalQuentContext( context=context_a, query=context_a.query_for(uuid.uuid4()), worker=Worker(id=uuid.uuid4(), engine=context_a.engine, instance_name="rank-0"), @@ -679,7 +682,7 @@ def test_processor_registry_reused_across_quent_contexts() -> None: parent_group_id=context_a.engine.id, ), ) - local_b = cudf_polars.quent.LocalQuentContext( + local_b = LocalQuentContext( context=context_b, query=context_b.query_for(uuid.uuid4()), worker=Worker(id=uuid.uuid4(), engine=context_b.engine, instance_name="rank-0"), diff --git a/python/cudf_polars/tests/quent/test_quent_integration.py b/python/cudf_polars/tests/quent/test_quent_integration.py index 804d031508ab..acdf891a7f4b 100644 --- a/python/cudf_polars/tests/quent/test_quent_integration.py +++ b/python/cudf_polars/tests/quent/test_quent_integration.py @@ -15,7 +15,7 @@ from collections.abc import Iterator from cudf_polars.engine.core import StreamingEngine - from cudf_polars.quent import QuentContext + from cudf_polars.quent._context import QuentContext # Quent tracing requires structlog to emit events. Skip the whole module when # it is unavailable so the engine fixture below is never even constructed. diff --git a/python/cudf_polars/tests/streaming/test_spmd.py b/python/cudf_polars/tests/streaming/test_spmd.py index e38684a3f0e6..774462b7610a 100644 --- a/python/cudf_polars/tests/streaming/test_spmd.py +++ b/python/cudf_polars/tests/streaming/test_spmd.py @@ -18,6 +18,7 @@ from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor import cudf_polars.quent +import cudf_polars.quent._context from cudf_polars.engine.core import _find_memory_error from cudf_polars.engine.hardware_binding import HardwareBindingPolicy from cudf_polars.engine.options import StreamingOptions @@ -397,7 +398,7 @@ def test_reset_rejects_construction_time_engine_options( def test_quent_context_user_provided(spmd_engine: SPMDEngine) -> None: # Ensure that the user-provided quent context is used if provided - quent_context = cudf_polars.quent.QuentContext( + quent_context = cudf_polars.quent._context.QuentContext( engine=cudf_polars.quent.Engine( id=uuid.uuid4(), implementation=cudf_polars.quent.Implementation( diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index e5deb48cb14d..8672d61442de 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -14,7 +14,7 @@ from rmm._cuda import gpu import cudf_polars.callback -import cudf_polars.quent +import cudf_polars.quent._context import cudf_polars.utils.config from cudf_polars.callback import ( _is_concurrent_managed_access_supported, @@ -395,7 +395,9 @@ def test_hash_streaming_executor() -> None: config = ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", - executor_options={"quent_context": cudf_polars.quent.QuentContext()}, + executor_options={ + "quent_context": cudf_polars.quent._context.QuentContext() + }, ) ) assert hash(config.executor) == hash(config.executor) From f038b597f2f017452c42570b1df2e05845db8091 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 14 Jul 2026 05:19:31 -0700 Subject: [PATCH 12/60] test skip --- python/cudf_polars/tests/quent/test_quent_integration.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/cudf_polars/tests/quent/test_quent_integration.py b/python/cudf_polars/tests/quent/test_quent_integration.py index acdf891a7f4b..4f4c85da12a7 100644 --- a/python/cudf_polars/tests/quent/test_quent_integration.py +++ b/python/cudf_polars/tests/quent/test_quent_integration.py @@ -11,6 +11,8 @@ import polars as pl +from cudf_polars.dsl.tracing import LOG_TRACES + if TYPE_CHECKING: from collections.abc import Iterator @@ -171,6 +173,7 @@ def check_quent_events(engine: StreamingEngine, quent_context: QuentContext) -> check_processor_lifecycle(quent_events) +@pytest.mark.skipif(not LOG_TRACES, reason="requires CUDF_POLARS_LOG_TRACES=1") def test_quent_events_multiple_collects( engine_with_quent_context: StreamingEngine, quent_context: QuentContext ) -> None: From 681638e906cd88cd3a221eda86e68cb0a8a7c7ed Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 14 Jul 2026 05:48:56 -0700 Subject: [PATCH 13/60] tweak the test --- .../tests/quent/test_quent_integration.py | 31 ++++++++++--------- 1 file changed, 17 insertions(+), 14 deletions(-) diff --git a/python/cudf_polars/tests/quent/test_quent_integration.py b/python/cudf_polars/tests/quent/test_quent_integration.py index 4f4c85da12a7..5ce2c5736307 100644 --- a/python/cudf_polars/tests/quent/test_quent_integration.py +++ b/python/cudf_polars/tests/quent/test_quent_integration.py @@ -166,14 +166,15 @@ def check_quent_events(engine: StreamingEngine, quent_context: QuentContext) -> memory_events = [x for x in quent_events if "Memory" in x["data"]] task_events = [x for x in quent_events if "Task" in x["data"]] assert len(memory_events) > 0 - assert len(task_events) > 0 + + if LOG_TRACES: + assert len(task_events) > 0 # A single collect exercises the full processor lifecycle, so fold that # check in here rather than paying for a dedicated engine startup. check_processor_lifecycle(quent_events) -@pytest.mark.skipif(not LOG_TRACES, reason="requires CUDF_POLARS_LOG_TRACES=1") def test_quent_events_multiple_collects( engine_with_quent_context: StreamingEngine, quent_context: QuentContext ) -> None: @@ -280,15 +281,17 @@ def check_processor_lifecycle(quent_events: list[dict]) -> None: ] assert len(init_events) == len(finalizing_events) == len(exit_events) - assert len(init_events) > 0 - - init_by_id = {x["id"]: x for x in init_events} - finalizing_by_id = {x["id"]: x for x in finalizing_events} - exit_by_id = {x["id"]: x for x in exit_events} - assert init_by_id.keys() == finalizing_by_id.keys() == exit_by_id.keys() - - for init_event in init_by_id.values(): - parent_group_id = init_event["data"]["Processor"]["state"][ - "ProcessorInitializing" - ]["parent_group_id"] - assert parent_group_id in thread_pool_ids + + if LOG_TRACES: + assert len(init_events) > 0 + + init_by_id = {x["id"]: x for x in init_events} + finalizing_by_id = {x["id"]: x for x in finalizing_events} + exit_by_id = {x["id"]: x for x in exit_events} + assert init_by_id.keys() == finalizing_by_id.keys() == exit_by_id.keys() + + for init_event in init_by_id.values(): + parent_group_id = init_event["data"]["Processor"]["state"][ + "ProcessorInitializing" + ]["parent_group_id"] + assert parent_group_id in thread_pool_ids From 9652021cefc28d14cf6d734aef63f54fdd66b420 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 14 Jul 2026 05:51:46 -0700 Subject: [PATCH 14/60] docs --- python/cudf_polars/cudf_polars/quent/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/cudf_polars/cudf_polars/quent/__init__.py b/python/cudf_polars/cudf_polars/quent/__init__.py index 3a4e5b991f0a..4656f8a67101 100644 --- a/python/cudf_polars/cudf_polars/quent/__init__.py +++ b/python/cudf_polars/cudf_polars/quent/__init__.py @@ -5,6 +5,7 @@ from __future__ import annotations +from cudf_polars.quent._context import QuentContext from cudf_polars.quent._types import ( Attribute, Engine, @@ -21,6 +22,7 @@ "Engine", "HomogeneousListValue", "Implementation", + "QuentContext", "Query", "QueryGroup", "ScalarValue", From e02620cd9c749922d25684f419131c0efcda21f0 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 14 Jul 2026 08:48:36 -0700 Subject: [PATCH 15/60] Move network declaration --- python/cudf_polars/cudf_polars/dsl/tracing.py | 3 +- python/cudf_polars/cudf_polars/engine/core.py | 39 +------------------ .../cudf_polars/cudf_polars/quent/_context.py | 38 ++++++++++++++++++ 3 files changed, 40 insertions(+), 40 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index 915bf44c94c5..f2ebfe8d6ebe 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -338,6 +338,7 @@ def wrapper( start = time.monotonic_ns() result = func(cls, *args, **kwargs) stop = time.monotonic_ns() + after_start = time.monotonic_ns() after = make_snapshot( cls, @@ -359,8 +360,6 @@ def wrapper( ) log.info("Execute IR", **record) - result = func(cls, *args, **kwargs) - if quent_task is not None: _end_quent_do_evaluate_events( cls, frames, result, ir_execution_context, quent_task diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 96ce3decd38c..f0ac804ec2ca 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -450,43 +450,6 @@ def _find_memory_error(exc: BaseException) -> MemoryError | None: return None -def _declare_network_channels( - comm: Communicator, - local_quent_context: LocalQuentContext, -) -> None: - """ - Declare network link channels for inter-rank communication. - - Creates a Network resource group and one Channel per remote rank, - emitting their lifecycle events to the quent logger. - """ - if comm.nranks <= 1: - return - - from cudf_polars.quent._types import Channel, Network - - network = Network(engine_id=local_quent_context.context.engine.id) - local_quent_context.logger.emit(network.declare()) - local_quent_context.network = network - - link_channels: dict[int, Channel] = {} - for target_rank in range(comm.nranks): - if target_rank == comm.rank: - continue - link = Channel( - instance_name=f"rank-{comm.rank} -> rank-{target_rank}", - resource_type_name="Link", - parent_group_id=network.id, - source=local_quent_context.device_memory, - target=local_quent_context.device_memory, - ) - local_quent_context.logger.emit(link.initializing()) - local_quent_context.logger.emit(link.operating()) - link_channels[target_rank] = link - - local_quent_context.link_channels = link_channels - - def execute_ir_on_rank( ctx: Context, comm: Communicator, @@ -784,7 +747,7 @@ def evaluate_on_rank( quent_operator_map: dict[IR, cudf_polars.quent._types.Operator] | None = None if config_options.executor.quent_context is not None: assert local_quent_context is not None - _declare_network_channels(comm, local_quent_context) + local_quent_context._declare_network_channels(comm) # ``get_stable_plan_id`` is a deterministic function of the IR # structure, so every rank derives the same logical plan ID for a # given query (only rank 0 emits the declaration, but physical plans diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index 8c9c37fcf2e2..9415cde07650 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -30,6 +30,8 @@ if TYPE_CHECKING: from typing import Self + from rapidsmpf.communicator.communicator import Communicator + from cudf_polars.dsl.ir import IR from cudf_polars.quent._logging import QuentLogger from cudf_polars.quent._types import ( @@ -469,6 +471,42 @@ def get_or_declare_processor( pool_id=self.thread_pool_id, ) + def _declare_network_channels( + self, + comm: Communicator, + ) -> None: + """ + Declare network link channels for inter-rank communication. + + Creates a Network resource group and one Channel per remote rank, + emitting their lifecycle events to the quent logger. + """ + if comm.nranks <= 1: + return + + from cudf_polars.quent._types import Channel, Network + + network = Network(engine_id=self.context.engine.id) + self.logger.emit(network.declare()) + self.network = network + + link_channels: dict[int, Channel] = {} + for target_rank in range(comm.nranks): + if target_rank == comm.rank: + continue + link = Channel( + instance_name=f"rank-{comm.rank} -> rank-{target_rank}", + resource_type_name="Link", + parent_group_id=network.id, + source=self.device_memory, + target=self.device_memory, + ) + self.logger.emit(link.initializing()) + self.logger.emit(link.operating()) + link_channels[target_rank] = link + + self.link_channels = link_channels + @dataclasses.dataclass(kw_only=True) class QuentIRExecutionContext(LocalQuentContext): From c78bfc823655013427b80760d2ca142c122c6850 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 14 Jul 2026 08:57:14 -0700 Subject: [PATCH 16/60] Engine-scoped network, channels --- python/cudf_polars/cudf_polars/engine/core.py | 10 +- python/cudf_polars/cudf_polars/engine/dask.py | 23 +++++ python/cudf_polars/cudf_polars/engine/ray.py | 21 ++++ python/cudf_polars/cudf_polars/engine/spmd.py | 19 +++- .../cudf_polars/cudf_polars/quent/_context.py | 97 ++++++++++++------- .../cudf_polars/cudf_polars/utils/config.py | 9 +- 6 files changed, 131 insertions(+), 48 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index f0ac804ec2ca..7d6f61fba2a4 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -747,7 +747,6 @@ def evaluate_on_rank( quent_operator_map: dict[IR, cudf_polars.quent._types.Operator] | None = None if config_options.executor.quent_context is not None: assert local_quent_context is not None - local_quent_context._declare_network_channels(comm) # ``get_stable_plan_id`` is a deterministic function of the IR # structure, so every rank derives the same logical plan ID for a # given query (only rank 0 emits the declaration, but physical plans @@ -807,7 +806,7 @@ def evaluate_on_rank( attach_cached_parquet_metadata(ir, cached_parquet_info_map) with ReserveOpIDs(ir, config_options) as collective_id_map: - result = execute_ir_on_rank( + return execute_ir_on_rank( ctx, comm, ir, @@ -819,10 +818,3 @@ def evaluate_on_rank( quent_operator_map=quent_operator_map, local_quent_context=local_quent_context, ) - - if local_quent_context is not None: - for link in local_quent_context.link_channels.values(): - local_quent_context.logger.emit(link.finalizing()) - local_quent_context.logger.emit(link.exit()) - - return result diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 45dcda871624..4eee4e6f8289 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -45,7 +45,9 @@ from cudf_polars.quent._context import ( LocalQuentContext, ProcessorRegistry, + declare_network_channels, declare_worker_resources, + finalize_network_channels, finalize_worker_resources, ) from cudf_polars.utils.config import DaskContext, MemoryResourceConfig @@ -132,6 +134,10 @@ class _WorkerContext: disk_to_device_channel: cudf_polars.quent._types.Channel | None = None thread_pool: cudf_polars.quent._types.ThreadPool | None = None processor_registry: ProcessorRegistry | None = None + network: cudf_polars.quent._types.Network | None = None + link_channels: dict[int, cudf_polars.quent._types.Channel] = dataclasses.field( + default_factory=dict + ) def _setup_root( @@ -328,6 +334,8 @@ def _setup_worker( disk_to_device_channel = None thread_pool = None processor_registry = None + network = None + link_channels: dict[int, cudf_polars.quent._types.Channel] = {} if quent_logger is not None: processor_registry = ProcessorRegistry() device_memory, disk_to_device_channel, thread_pool = declare_worker_resources( @@ -336,6 +344,14 @@ def _setup_worker( engine_id=engine_id, worker_id=worker_id, ) + # Inter-rank network topology is engine-scoped: declare it once here + # alongside the other worker resources (a no-op for single-rank runs). + network, link_channels = declare_network_channels( + quent_logger, + comm=comm, + engine_id=engine_id, + device_memory=device_memory, + ) mp_ctx = _WorkerContext( comm=comm, @@ -349,6 +365,8 @@ def _setup_worker( disk_to_device_channel=disk_to_device_channel, thread_pool=thread_pool, processor_registry=processor_registry, + network=network, + link_channels=link_channels, statistics=statistics, ) setattr(dask_worker, attr, mp_ctx) @@ -382,6 +400,9 @@ def _teardown_worker( mp_ctx.processor_registry._emit_processor_exit_events( mp_ctx.quent_logger ) + finalize_network_channels( + mp_ctx.quent_logger, link_channels=mp_ctx.link_channels + ) if mp_ctx.device_memory is not None: finalize_worker_resources( mp_ctx.quent_logger, @@ -551,6 +572,8 @@ def _worker_evaluate( processor_registry=mp_ctx.processor_registry, device_memory=mp_ctx.device_memory, disk_to_device_channel=mp_ctx.disk_to_device_channel, + network=mp_ctx.network, + link_channels=mp_ctx.link_channels, ) # evaluate_on_rank always collects metadata internally so we can read # metadata[-1].duplicated to decide whether to suppress this rank's output. diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index 128a9c7cd2c5..4ba9ce732fe1 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -43,7 +43,9 @@ from cudf_polars.quent._context import ( LocalQuentContext, ProcessorRegistry, + declare_network_channels, declare_worker_resources, + finalize_network_channels, finalize_worker_resources, ) from cudf_polars.quent._types import Worker @@ -62,6 +64,7 @@ from cudf_polars.engine.core import T from cudf_polars.engine.options import StreamingOptions from cudf_polars.quent._context import QuentContext + from cudf_polars.quent._types import Channel from cudf_polars.streaming.parallel import ConfigOptions from cudf_polars.utils.config import StreamingExecutor @@ -242,6 +245,10 @@ def __init__( self._device_memory = None self._disk_to_device_channel = None self._quent_thread_pool = None + # Network topology is engine-scoped, but ``comm`` is not available until + # ``setup_worker``; declare it there and store the results here. + self._network: cudf_polars.quent._types.Network | None = None + self._link_channels: dict[int, Channel] = {} self._processor_registry: ProcessorRegistry | None = None if self._quent_logger is not None: self._processor_registry = ProcessorRegistry() @@ -304,6 +311,15 @@ def setup_worker(self, root_ucxx_address_as_bytes: bytes) -> None: progress_thread=ProgressThread(self._rapidsmpf_statistics), ) barrier(self._comm) + # Now that ``comm`` exists, declare the engine-scoped inter-rank network + # topology once (a no-op for single-rank runs). + if self._quent_logger is not None and self._device_memory is not None: + self._network, self._link_channels = declare_network_channels( + self._quent_logger, + comm=self._comm, + engine_id=self._quent_worker.engine.id, + device_memory=self._device_memory, + ) assert self._base_mr is not None self._ctx = Context.from_options( self._comm.logger, @@ -364,6 +380,9 @@ def _exit(self) -> list[dict[str, Any]]: if self._quent_worker is not None and self._quent_logger is not None: if self._processor_registry is not None: self._processor_registry._emit_processor_exit_events(self._quent_logger) + finalize_network_channels( + self._quent_logger, link_channels=self._link_channels + ) if ( self._device_memory is not None and self._disk_to_device_channel is not None @@ -496,6 +515,8 @@ def evaluate_polars_ir( processor_registry=self._processor_registry, device_memory=self._device_memory, disk_to_device_channel=self._disk_to_device_channel, + network=self._network, + link_channels=self._link_channels, ) # evaluate_on_rank always collects metadata internally so we can read # metadata[-1].duplicated to decide whether to suppress this rank's diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index 4b356605d049..a1f1896aa3a4 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -48,7 +48,9 @@ LocalQuentContext, ProcessorRegistry, QuentContext, + declare_network_channels, declare_worker_resources, + finalize_network_channels, finalize_worker_resources, ) from cudf_polars.quent._types import Worker @@ -73,7 +75,7 @@ from cudf_polars.dsl.ir import IR from cudf_polars.engine.core import T from cudf_polars.engine.options import StreamingOptions - from cudf_polars.quent._types import Channel, Memory, ThreadPool + from cudf_polars.quent._types import Channel, Memory, Network, ThreadPool from cudf_polars.streaming.parallel import ConfigOptions from cudf_polars.utils.config import StreamingExecutor @@ -148,6 +150,8 @@ def evaluate_pipeline_spmd_mode( processor_registry=spmd_context.processor_registry, device_memory=spmd_context.device_memory, disk_to_device_channel=spmd_context.disk_to_device_channel, + network=spmd_context.network, + link_channels=spmd_context.link_channels, ) df, metadata = evaluate_on_rank( @@ -468,6 +472,8 @@ def __init__( self._quent_thread_pool: ThreadPool | None = None self._device_memory: Memory | None = None self._disk_to_device_channel: Channel | None = None + self._network: Network | None = None + self._link_channels: dict[int, Channel] = {} exit_stack = contextlib.ExitStack() @@ -520,6 +526,12 @@ def __init__( engine_id=engine_id, worker_id=self._quent_worker.id, ) + self._network, self._link_channels = declare_network_channels( + self._quent_logger, + comm=comm, + engine_id=engine_id, + device_memory=self._device_memory, + ) # Register after `_cleanup_ctx` so on teardown (LIFO) the # executor shuts down first. `wait=True` is safe because @@ -554,6 +566,8 @@ def __init__( ), device_memory=self._device_memory, disk_to_device_channel=self._disk_to_device_channel, + network=self._network, + link_channels=self._link_channels, ), }, engine_options={ @@ -840,6 +854,9 @@ def shutdown(self) -> None: if quent_context is not None: assert self._processor_registry is not None self._processor_registry._emit_processor_exit_events(self._quent_logger) + finalize_network_channels( + self._quent_logger, link_channels=self._link_channels + ) if self._device_memory is not None: finalize_worker_resources( self._quent_logger, diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index 9415cde07650..78722a18b83c 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -21,6 +21,7 @@ Engine, Implementation, Memory, + Network, Processor, Query, QueryGroup, @@ -35,7 +36,6 @@ from cudf_polars.dsl.ir import IR from cudf_polars.quent._logging import QuentLogger from cudf_polars.quent._types import ( - Network, Operator, Plan, Port, @@ -431,6 +431,59 @@ def finalize_worker_resources( logger.emit(device_memory.exit()) +def declare_network_channels( + logger: QuentLogger, + *, + comm: Communicator, + engine_id: uuid.UUID, + device_memory: Memory, +) -> tuple[Network | None, dict[int, Channel]]: + """ + Declare engine-scoped network link channels for inter-rank communication. + + Creates a Network resource group and one Channel per remote rank, emitting + their lifecycle events to the quent logger. This is engine/worker-scoped: + the inter-rank topology is fixed for the lifetime of the engine, so it is + declared once at worker setup rather than per query. + + Returns ``(None, {})`` for single-rank runs, which have no inter-rank + communication. + """ + if comm.nranks <= 1: + return None, {} + + network = Network(engine_id=engine_id) + logger.emit(network.declare()) + + link_channels: dict[int, Channel] = {} + for target_rank in range(comm.nranks): + if target_rank == comm.rank: + continue + link = Channel( + instance_name=f"rank-{comm.rank} -> rank-{target_rank}", + resource_type_name="Link", + parent_group_id=network.id, + source=device_memory, + target=device_memory, + ) + logger.emit(link.initializing()) + logger.emit(link.operating()) + link_channels[target_rank] = link + + return network, link_channels + + +def finalize_network_channels( + logger: QuentLogger, + *, + link_channels: dict[int, Channel], +) -> None: + """Emit finalizing/exit events for engine-scoped network link channels.""" + for link in link_channels.values(): + logger.emit(link.finalizing()) + logger.emit(link.exit()) + + @dataclasses.dataclass(kw_only=True) class LocalQuentContext: """ @@ -447,6 +500,12 @@ class LocalQuentContext: :class:`Query` from its unique ``query_id`` (see :meth:`QuentContext.query_for`), rather than reusing the shared ``context.query``. + + The ``device_memory``, ``disk_to_device_channel``, ``network``, and + ``link_channels`` resources are all engine/worker-scoped: they are declared + once at worker setup (see :func:`declare_worker_resources` and + :func:`declare_network_channels`) and injected into each per-collect + context, rather than being declared per query. """ context: QuentContext @@ -471,42 +530,6 @@ def get_or_declare_processor( pool_id=self.thread_pool_id, ) - def _declare_network_channels( - self, - comm: Communicator, - ) -> None: - """ - Declare network link channels for inter-rank communication. - - Creates a Network resource group and one Channel per remote rank, - emitting their lifecycle events to the quent logger. - """ - if comm.nranks <= 1: - return - - from cudf_polars.quent._types import Channel, Network - - network = Network(engine_id=self.context.engine.id) - self.logger.emit(network.declare()) - self.network = network - - link_channels: dict[int, Channel] = {} - for target_rank in range(comm.nranks): - if target_rank == comm.rank: - continue - link = Channel( - instance_name=f"rank-{comm.rank} -> rank-{target_rank}", - resource_type_name="Link", - parent_group_id=network.id, - source=self.device_memory, - target=self.device_memory, - ) - self.logger.emit(link.initializing()) - self.logger.emit(link.operating()) - link_channels[target_rank] = link - - self.link_channels = link_channels - @dataclasses.dataclass(kw_only=True) class QuentIRExecutionContext(LocalQuentContext): diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 8e19b9120d8e..540a74837c4b 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -46,7 +46,7 @@ from cudf_polars.engine.ray import RankActor from cudf_polars.quent._context import ProcessorRegistry, QuentContext from cudf_polars.quent._logging import QuentLogger - from cudf_polars.quent._types import Channel, Memory + from cudf_polars.quent._types import Channel, Memory, Network __all__ = [ @@ -579,6 +579,11 @@ class SPMDContext: The engine-scoped Quent device Memory resource. disk_to_device_channel The engine-scoped Quent disk-to-device Channel resource. + network + The engine-scoped Quent Network resource group (``None`` for + single-rank runs). + link_channels + The engine-scoped Quent inter-rank Link channels, keyed by target rank. """ comm: Communicator @@ -591,6 +596,8 @@ class SPMDContext: thread_pool_id: uuid.UUID | None = None device_memory: Memory | None = None disk_to_device_channel: Channel | None = None + network: Network | None = None + link_channels: dict[int, Channel] = dataclasses.field(default_factory=dict) @dataclasses.dataclass(frozen=True) From 22bef5a9e5347ce6e39824f7ecd45e4119485b4b Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 14 Jul 2026 11:22:11 -0700 Subject: [PATCH 17/60] Fixes - always emit Task exit events - Handle everything in _reset --- python/cudf_polars/cudf_polars/dsl/tracing.py | 34 +++++++++++++------ python/cudf_polars/cudf_polars/engine/spmd.py | 4 +++ 2 files changed, 28 insertions(+), 10 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index f2ebfe8d6ebe..285eef030627 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -137,7 +137,7 @@ def _begin_quent_do_evaluate_events( def _end_quent_do_evaluate_events( cls: type[ir.IR], frames: Sequence[cudf_polars.containers.DataFrame], - result: cudf_polars.containers.DataFrame, + result: cudf_polars.containers.DataFrame | None, ir_execution_context: IRExecutionContext, quent_task: Task, ) -> None: @@ -151,7 +151,8 @@ def _end_quent_do_evaluate_events( frames The input dataframes passed to the IR node. result - The output dataframe returned from the IR node. + The output dataframe returned from the IR node. This will be ``None`` + if an exception was raised while evaluating the IR node. ir_execution_context The IR execution context. To emit any events, this must have a quent_ir_execution_context bound. @@ -180,13 +181,18 @@ def _end_quent_do_evaluate_events( if quent_ir_execution_context is None: return - output_capacity_bytes = _dataframe_size_bytes(result) + if result is not None: + output_rows = result.num_rows + output_capacity_bytes = _dataframe_size_bytes(result) + else: + output_rows = 0 + output_capacity_bytes = 0 quent_ir_execution_context.logger.emit( quent_ir_execution_context.quent_operator.statistics( statistics=cudf_polars.quent._types.Statistics( input_bytes=sum(_dataframe_size_bytes(frame) for frame in frames), output_bytes=output_capacity_bytes, - output_rows=result.num_rows, + output_rows=output_rows, ) ) ) @@ -336,7 +342,20 @@ def wrapper( # argument, followed by the method-specific arguments, and returns a DataFrame. start = time.monotonic_ns() - result = func(cls, *args, **kwargs) + try: + result = func(cls, *args, **kwargs) + except Exception: + if quent_task is not None: + _end_quent_do_evaluate_events( + cls, frames, None, ir_execution_context, quent_task + ) + raise + else: + if quent_task is not None: + _end_quent_do_evaluate_events( + cls, frames, result, ir_execution_context, quent_task + ) + stop = time.monotonic_ns() after_start = time.monotonic_ns() @@ -360,11 +379,6 @@ def wrapper( ) log.info("Execute IR", **record) - if quent_task is not None: - _end_quent_do_evaluate_events( - cls, frames, result, ir_execution_context, quent_task - ) - return result return wrapper diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index a1f1896aa3a4..ec8b7d7deaaa 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -711,6 +711,10 @@ def _reset( if self._quent_thread_pool is not None else None ), + device_memory=self._device_memory, + disk_to_device_channel=self._disk_to_device_channel, + network=self._network, + link_channels=self._link_channels, ), }, engine_options={ From cd3fefc384bb23f6149e183fbcdac887d19669df Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 14 Jul 2026 11:24:58 -0700 Subject: [PATCH 18/60] docs --- .../streaming/actor_graph/dispatch.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py index adad3c6f89a4..e5c7e8a3507d 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py @@ -81,7 +81,22 @@ class GenState(TypedDict): def ir_context_for_node(rec: SubNetGenerator, ir: IR) -> IRExecutionContext: - """Return ``ir_context`` with the physical Quent operator bound when tracing.""" + """ + Return ``ir_context`` with the physical Quent operator bound when tracing. + + Parameters + ---------- + rec + The recursive SubNetGenerator callable. + ir + The IR node to return the execution context for. + + Returns + ------- + ir_context + A clone of rec.state["ir_context"] with ``quent_ir_execution_context`` + bound to the physical Quent operator for the given IR node. + """ import cudf_polars.quent._context ir_context = rec.state["ir_context"] From 23107f371190010e1803697d80749144da2c26a3 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 14 Jul 2026 11:28:53 -0700 Subject: [PATCH 19/60] fixup --- python/cudf_polars/cudf_polars/quent/_context.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index 78722a18b83c..4caa9fab5be2 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -71,12 +71,12 @@ def get_or_declare_processor( ) -> Processor: """Get (or declare a new) Quent Processor for a CPU thread.""" with self._lock: + if thread_ident in self._processors: + return self._processors[thread_ident] if self._closed: raise RuntimeError( "Cannot declare processors after registry has been closed" ) - if thread_ident in self._processors: - return self._processors[thread_ident] processor = Processor(pool_id=pool_id) self._processors[thread_ident] = processor From 5ddc59051dca42676dd0a6ed7fc77a86da30a6ae Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 14 Jul 2026 12:00:17 -0700 Subject: [PATCH 20/60] Refactor stuff out of _tracing and into _quent --- .../cudf_polars/containers/dataframe.py | 4 + python/cudf_polars/cudf_polars/dsl/tracing.py | 187 +++--------------- .../cudf_polars/cudf_polars/quent/_context.py | 114 +++++++++++ .../cudf_polars/cudf_polars/quent/_types.py | 38 +++- 4 files changed, 186 insertions(+), 157 deletions(-) diff --git a/python/cudf_polars/cudf_polars/containers/dataframe.py b/python/cudf_polars/cudf_polars/containers/dataframe.py index ad75a8fde151..2fd6c850021d 100644 --- a/python/cudf_polars/cudf_polars/containers/dataframe.py +++ b/python/cudf_polars/cudf_polars/containers/dataframe.py @@ -108,6 +108,10 @@ def __init__( self.stream = stream self._num_rows_override = num_rows + def _size_bytes(self) -> int: + """Return the size of the dataframe in bytes.""" + return sum(c.device_buffer_size() for c in self.table.columns()) + def copy(self) -> Self: """Return a shallow copy of self.""" return type(self)( diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index 285eef030627..af6b876e9839 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -9,9 +9,7 @@ import enum import functools import os -import threading import time -import uuid from typing import TYPE_CHECKING, Any, Concatenate, Literal, ParamSpec import nvtx @@ -52,151 +50,6 @@ import cudf_polars.containers from cudf_polars.dsl import ir from cudf_polars.dsl.ir import IRExecutionContext - from cudf_polars.quent._types import Task - - -def _dataframe_size_bytes(frame: cudf_polars.containers.DataFrame) -> int: - return sum(col.device_buffer_size() for col in frame.table.columns()) - - -def _begin_quent_do_evaluate_events( - cls: type[ir.IR], - ir_execution_context: IRExecutionContext, -) -> Task | None: - """ - Build and emit Quent events for the start of an IR node's evaluation. - - Parameters - ---------- - cls - The IR node class. - ir_execution_context - The IR execution context. To emit any events, this must have a - quent_ir_execution_context bound. - - Returns - ------- - The Quent task instance, or None if no Quent IR execution context is bound. - - Notes - ----- - The following events are emitted: - - - Queueing - - Loading (I/O nodes only) - - Allocating (non-I/O nodes only) - - Computing (non-I/O nodes only) - - The Loading, Allocating, and Computing events will indicate the host CPU thread - and device memory that they're using. - - See Also - -------- - _end_quent_do_evaluate_events - """ - import cudf_polars.quent._types - - quent_ir_execution_context = ir_execution_context.quent_ir_execution_context - if quent_ir_execution_context is None: - return None - - token = uuid.uuid4() - quent_task = cudf_polars.quent._types.Task( - instance_name=( - f"{cls.__name__}-{quent_ir_execution_context.quent_operator.id.hex[:8]}-" - f"{token.hex[:8]}" - ), - operator_id=quent_ir_execution_context.quent_operator.id, - ) - quent_processor = quent_ir_execution_context.get_or_declare_processor( - thread_ident=threading.get_ident(), - ) - quent_ir_execution_context.logger.emit(quent_task.queueing()) - if not cls.is_io_node: - quent_ir_execution_context.logger.emit( - quent_task.allocating(resource_id=quent_processor.id) - ) - quent_ir_execution_context.logger.emit( - quent_task.computing( - use_thread=quent_processor, - use_memory=quent_ir_execution_context.device_memory, - ) - ) - else: - quent_ir_execution_context.logger.emit( - quent_task.loading( - use_thread=quent_processor, - use_channel=quent_ir_execution_context.disk_to_device_channel, - use_memory=quent_ir_execution_context.device_memory, - ) - ) - - return quent_task - - -def _end_quent_do_evaluate_events( - cls: type[ir.IR], - frames: Sequence[cudf_polars.containers.DataFrame], - result: cudf_polars.containers.DataFrame | None, - ir_execution_context: IRExecutionContext, - quent_task: Task, -) -> None: - """ - Build and emit Quent events for the end of an IR node's evaluation. - - Parameters - ---------- - cls - The IR node class. - frames - The input dataframes passed to the IR node. - result - The output dataframe returned from the IR node. This will be ``None`` - if an exception was raised while evaluating the IR node. - ir_execution_context - The IR execution context. To emit any events, this must have a - quent_ir_execution_context bound. - quent_task - The Quent task instance created by ``_begin_quent_do_evaluate_events``. - - Notes - ----- - This method emits an ``Exit`` event for the Quent Task, whose timestamp represents - when the IR node completed host-side processing. - - A ``Statistics`` record, associated with the Quent Operator bound to the IR execution context, - is also emitted. It includes - - - input bytes: the total size of the input dataframes. - - output bytes: the size of the output dataframe. - - output rows: the number of rows in the output dataframe. - - See Also - -------- - _begin_quent_do_evaluate_events - """ - import cudf_polars.quent._types - - quent_ir_execution_context = ir_execution_context.quent_ir_execution_context - if quent_ir_execution_context is None: - return - - if result is not None: - output_rows = result.num_rows - output_capacity_bytes = _dataframe_size_bytes(result) - else: - output_rows = 0 - output_capacity_bytes = 0 - quent_ir_execution_context.logger.emit( - quent_ir_execution_context.quent_operator.statistics( - statistics=cudf_polars.quent._types.Statistics( - input_bytes=sum(_dataframe_size_bytes(frame) for frame in frames), - output_bytes=output_capacity_bytes, - output_rows=output_rows, - ) - ) - ) - quent_ir_execution_context.logger.emit(quent_task.exit()) class Scope(enum.StrEnum): @@ -316,6 +169,8 @@ def wrapper( *args: P.args, **kwargs: P.kwargs, ) -> cudf_polars.containers.DataFrame: + from cudf_polars.quent._types import Task + pynvml.nvmlInit() maybe_handle = get_device_handle() pid = _getpid() @@ -329,7 +184,16 @@ def wrapper( # And the kwonly 'context' argument has the IR execution context. ir_execution_context: IRExecutionContext = kwargs["context"] # type: ignore[assignment] - quent_task = _begin_quent_do_evaluate_events(cls, ir_execution_context) + quent_task = Task.from_ir(cls, ir_execution_context) + if ( + quent_task is not None + and ir_execution_context.quent_ir_execution_context is not None + ): + ir_execution_context.quent_ir_execution_context.context._emit_task_begin_events( + cls, + quent_task, + ir_execution_context.quent_ir_execution_context, + ) before_start = time.monotonic_ns() before = make_snapshot( @@ -344,18 +208,29 @@ def wrapper( start = time.monotonic_ns() try: result = func(cls, *args, **kwargs) - except Exception: - if quent_task is not None: - _end_quent_do_evaluate_events( - cls, frames, None, ir_execution_context, quent_task + except Exception: # pragma: no cover; + if ( + quent_task is not None + and ir_execution_context.quent_ir_execution_context is not None + ): + ir_execution_context.quent_ir_execution_context.context._emit_task_end_events( + quent_task, + ir_execution_context.quent_ir_execution_context, + frames, + None, ) raise else: - if quent_task is not None: - _end_quent_do_evaluate_events( - cls, frames, result, ir_execution_context, quent_task + if ( + quent_task is not None + and ir_execution_context.quent_ir_execution_context is not None + ): + ir_execution_context.quent_ir_execution_context.context._emit_task_end_events( + quent_task, + ir_execution_context.quent_ir_execution_context, + frames, + result, ) - stop = time.monotonic_ns() after_start = time.monotonic_ns() diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index 4caa9fab5be2..c0abc1638b94 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -25,20 +25,24 @@ Processor, Query, QueryGroup, + Statistics, ThreadPool, ) if TYPE_CHECKING: + from collections.abc import Sequence from typing import Self from rapidsmpf.communicator.communicator import Communicator + from cudf_polars.containers import DataFrame from cudf_polars.dsl.ir import IR from cudf_polars.quent._logging import QuentLogger from cudf_polars.quent._types import ( Operator, Plan, Port, + Task, Worker, ) from cudf_polars.utils.config import ConfigOptions, StreamingExecutor @@ -374,6 +378,116 @@ def _emit_physical_plan_events( parent_operators_by_node_id=parent_operators_by_node_id, ) + def _emit_task_begin_events( + self, + ir_type: type[IR], + quent_task: Task, + quent_ir_execution_context: QuentIRExecutionContext, + ) -> None: + """ + Emit begin events for a Quent Task. + + Parameters + ---------- + ir_type: type[IR] + The IR type of the operator. + quent_task: Task + The Quent Task to emit events for. + quent_ir_execution_context: QuentIRExecutionContext + The Quent IR execution context. + + Notes + ----- + The following events are emitted: + + - Queueing + - Loading (I/O nodes only) + - Allocating (non-I/O nodes only) + - Computing (non-I/O nodes only) + + The Loading, Allocating, and Computing events will indicate the host CPU thread + and device memory that they're using. + """ + quent_processor = quent_ir_execution_context.get_or_declare_processor( + thread_ident=threading.get_ident(), + ) + quent_ir_execution_context.logger.emit(quent_task.queueing()) + if not ir_type.is_io_node: + quent_ir_execution_context.logger.emit( + quent_task.allocating(resource_id=quent_processor.id) + ) + quent_ir_execution_context.logger.emit( + quent_task.computing( + use_thread=quent_processor, + use_memory=quent_ir_execution_context.device_memory, + ) + ) + else: + quent_ir_execution_context.logger.emit( + quent_task.loading( + use_thread=quent_processor, + use_channel=quent_ir_execution_context.disk_to_device_channel, + use_memory=quent_ir_execution_context.device_memory, + ) + ) + + def _emit_task_end_events( + self, + quent_task: Task, + quent_ir_execution_context: QuentIRExecutionContext, + frames: Sequence[DataFrame], + result: DataFrame | None, + ) -> None: + """ + Build and emit Quent events for the end of an IR node's evaluation. + + Parameters + ---------- + quent_task: Task + The Quent Task to emit events for. + quent_ir_execution_context: QuentIRExecutionContext + The Quent IR execution context. + frames + The input dataframes passed to the IR node. + result + The output dataframe returned from the IR node. This will be ``None`` + if an exception was raised while evaluating the IR node. + ir_execution_context + The IR execution context. To emit any events, this must have a + quent_ir_execution_context bound. + + Notes + ----- + This method emits an ``Exit`` event for the Quent Task, whose timestamp represents + when the IR node completed host-side processing. + + A ``Statistics`` record, associated with the Quent Operator bound to the IR execution context, + is also emitted. It includes + + - input bytes: the total size of the input dataframes. + - output bytes: the size of the output dataframe. + - output rows: the number of rows in the output dataframe. + """ + if quent_ir_execution_context is None: + return + + if result is not None: + output_rows = result.num_rows + output_capacity_bytes = result._size_bytes() + else: + output_rows = 0 + output_capacity_bytes = 0 + quent_ir_execution_context.logger.emit( + quent_ir_execution_context.quent_operator.statistics( + statistics=Statistics( + input_bytes=sum(frame._size_bytes() for frame in frames), + output_bytes=output_capacity_bytes, + output_rows=output_rows, + ) + ) + ) + quent_ir_execution_context.logger.emit(quent_task.exit()) + def declare_worker_resources( logger: QuentLogger, diff --git a/python/cudf_polars/cudf_polars/quent/_types.py b/python/cudf_polars/cudf_polars/quent/_types.py index cb6e2d068080..4ef05698d9e6 100644 --- a/python/cudf_polars/cudf_polars/quent/_types.py +++ b/python/cudf_polars/cudf_polars/quent/_types.py @@ -12,10 +12,13 @@ import sys import time import uuid -from typing import Any, Literal, TypeAlias +from typing import TYPE_CHECKING, Any, Literal, Self, TypeAlias from cudf_polars import __version__ +if TYPE_CHECKING: + from cudf_polars.dsl.ir import IR, IRExecutionContext + QUENT_SCOPE = "QUENT" @@ -805,6 +808,39 @@ class Task: operator_id: uuid.UUID instance_name: str | None = None + @classmethod + def from_ir( + cls, ir_type: type[IR], ir_execution_context: IRExecutionContext + ) -> Self | None: + """ + Maybe build an operator-scoped Quent Task from an IR execution context. + + Parameters + ---------- + ir_type + The IR type of the operator. + ir_execution_context + The IR execution context. + + Returns + ------- + Task | None + The operator-scoped Quent Task, or ``None`` if the IR execution context + is not bound to a Quent operator. + """ + quent_ir_execution_context = ir_execution_context.quent_ir_execution_context + if quent_ir_execution_context is None: + return None + + token = uuid.uuid4() + return cls( + instance_name=( + f"{ir_type.__name__}-{quent_ir_execution_context.quent_operator.id.hex[:8]}-" + f"{token.hex[:8]}" + ), + operator_id=quent_ir_execution_context.quent_operator.id, + ) + def queueing(self, timestamp: int | None = None) -> Event: """Build a Quent Task Queueing event.""" return Event( From 94428263e150c1bcab02252876bff7a90ad91603 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 14 Jul 2026 12:05:14 -0700 Subject: [PATCH 21/60] more refactor --- python/cudf_polars/cudf_polars/dsl/tracing.py | 8 +++++++- python/cudf_polars/cudf_polars/quent/_types.py | 15 ++++++--------- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index af6b876e9839..2c55a3305f40 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -184,7 +184,13 @@ def wrapper( # And the kwonly 'context' argument has the IR execution context. ir_execution_context: IRExecutionContext = kwargs["context"] # type: ignore[assignment] - quent_task = Task.from_ir(cls, ir_execution_context) + if ir_execution_context.quent_ir_execution_context is not None: + quent_task = Task.from_ir( + cls, ir_execution_context.quent_ir_execution_context + ) + else: + quent_task = None + if ( quent_task is not None and ir_execution_context.quent_ir_execution_context is not None diff --git a/python/cudf_polars/cudf_polars/quent/_types.py b/python/cudf_polars/cudf_polars/quent/_types.py index 4ef05698d9e6..da27e3f85448 100644 --- a/python/cudf_polars/cudf_polars/quent/_types.py +++ b/python/cudf_polars/cudf_polars/quent/_types.py @@ -17,7 +17,8 @@ from cudf_polars import __version__ if TYPE_CHECKING: - from cudf_polars.dsl.ir import IR, IRExecutionContext + from cudf_polars.dsl.ir import IR + from cudf_polars.quent._context import QuentIRExecutionContext QUENT_SCOPE = "QUENT" @@ -810,17 +811,17 @@ class Task: @classmethod def from_ir( - cls, ir_type: type[IR], ir_execution_context: IRExecutionContext + cls, ir_type: type[IR], quent_ir_execution_context: QuentIRExecutionContext ) -> Self | None: """ - Maybe build an operator-scoped Quent Task from an IR execution context. + Build an operator-scoped Quent Task from an IR execution context. Parameters ---------- ir_type The IR type of the operator. - ir_execution_context - The IR execution context. + quent_ir_execution_context + The Quent IR execution context, which is used to get the operator ID. Returns ------- @@ -828,10 +829,6 @@ def from_ir( The operator-scoped Quent Task, or ``None`` if the IR execution context is not bound to a Quent operator. """ - quent_ir_execution_context = ir_execution_context.quent_ir_execution_context - if quent_ir_execution_context is None: - return None - token = uuid.uuid4() return cls( instance_name=( From 5b61041532a02166fbe43ae98955012043f1948d Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 14 Jul 2026 14:38:24 -0700 Subject: [PATCH 22/60] test size_bytes --- python/cudf_polars/tests/containers/test_dataframe.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/tests/containers/test_dataframe.py b/python/cudf_polars/tests/containers/test_dataframe.py index 5006fcaa9b4a..f49530d2953f 100644 --- a/python/cudf_polars/tests/containers/test_dataframe.py +++ b/python/cudf_polars/tests/containers/test_dataframe.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -205,3 +205,10 @@ def test_serialization_roundtrip(polars_tbl): res = DataFrame.deserialize(header, frames, stream=stream) assert_frame_equal(df.to_polars(), res.to_polars()) + + +def test_size_bytes(): + stream = get_cuda_stream() + df = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]}) + df = DataFrame.from_polars(df, stream=stream) + assert df._size_bytes() == df.to_polars().estimated_size() From faafe6dd139812927d8e4f196c5cc3261322e23a Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 14 Jul 2026 15:04:45 -0700 Subject: [PATCH 23/60] testing --- python/cudf_polars/cudf_polars/engine/dask.py | 3 +- python/cudf_polars/cudf_polars/engine/ray.py | 3 +- python/cudf_polars/cudf_polars/engine/spmd.py | 3 +- .../cudf_polars/cudf_polars/quent/_context.py | 28 +- python/cudf_polars/tests/quent/test_quent.py | 420 +++++++++++++++++- 5 files changed, 434 insertions(+), 23 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 4eee4e6f8289..dcc7ea7d952c 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -348,7 +348,8 @@ def _setup_worker( # alongside the other worker resources (a no-op for single-rank runs). network, link_channels = declare_network_channels( quent_logger, - comm=comm, + rank=comm.rank, + nranks=comm.nranks, engine_id=engine_id, device_memory=device_memory, ) diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index 4ba9ce732fe1..f840cb99db1e 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -316,7 +316,8 @@ def setup_worker(self, root_ucxx_address_as_bytes: bytes) -> None: if self._quent_logger is not None and self._device_memory is not None: self._network, self._link_channels = declare_network_channels( self._quent_logger, - comm=self._comm, + rank=self._comm.rank, + nranks=self._comm.nranks, engine_id=self._quent_worker.engine.id, device_memory=self._device_memory, ) diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index ec8b7d7deaaa..519fa0a221fe 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -528,7 +528,8 @@ def __init__( ) self._network, self._link_channels = declare_network_channels( self._quent_logger, - comm=comm, + rank=comm.rank, + nranks=comm.nranks, engine_id=engine_id, device_memory=self._device_memory, ) diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index c0abc1638b94..f0c84d896487 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -33,8 +33,6 @@ from collections.abc import Sequence from typing import Self - from rapidsmpf.communicator.communicator import Communicator - from cudf_polars.containers import DataFrame from cudf_polars.dsl.ir import IR from cudf_polars.quent._logging import QuentLogger @@ -61,8 +59,11 @@ class ProcessorRegistry: One registry is owned by the object that owns the Python :class:`~concurrent.futures.ThreadPoolExecutor` (e.g. ``SPMDEngine``, - a Dask worker, or a Ray actor). Processors are declared lazily on first - use by a thread-pool worker and finalized once at executor shutdown. + a Dask worker, or a Ray actor). + + Processors (thread resources) are declared on-demand in ``get_or_declare_processor``. + Call ``_emit_processor_exit_events`` on engine shutdown to emit finalizing/exit events + for all declared processors. """ def __init__(self) -> None: @@ -77,11 +78,6 @@ def get_or_declare_processor( with self._lock: if thread_ident in self._processors: return self._processors[thread_ident] - if self._closed: - raise RuntimeError( - "Cannot declare processors after registry has been closed" - ) - processor = Processor(pool_id=pool_id) self._processors[thread_ident] = processor @@ -92,9 +88,6 @@ def get_or_declare_processor( def _emit_processor_exit_events(self, logger: QuentLogger) -> None: """Emit finalizing/exit events for all declared processors.""" with self._lock: - if self._closed: - return - self._closed = True processors = list(self._processors.values()) for processor in processors: @@ -548,7 +541,8 @@ def finalize_worker_resources( def declare_network_channels( logger: QuentLogger, *, - comm: Communicator, + rank: int, + nranks: int, engine_id: uuid.UUID, device_memory: Memory, ) -> tuple[Network | None, dict[int, Channel]]: @@ -563,18 +557,18 @@ def declare_network_channels( Returns ``(None, {})`` for single-rank runs, which have no inter-rank communication. """ - if comm.nranks <= 1: + if nranks <= 1: return None, {} network = Network(engine_id=engine_id) logger.emit(network.declare()) link_channels: dict[int, Channel] = {} - for target_rank in range(comm.nranks): - if target_rank == comm.rank: + for target_rank in range(nranks): + if target_rank == rank: continue link = Channel( - instance_name=f"rank-{comm.rank} -> rank-{target_rank}", + instance_name=f"rank-{rank} -> rank-{target_rank}", resource_type_name="Link", parent_group_id=network.id, source=device_memory, diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index 3b5d7bbd01ba..c168c3f56ae0 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -15,21 +15,29 @@ import cudf_polars.quent import cudf_polars.quent._logging +from cudf_polars.containers import DataFrame +from cudf_polars.dsl.ir import DataFrameScan, Filter from cudf_polars.dsl.translate import Translator from cudf_polars.quent._context import ( LocalQuentContext, ProcessorRegistry, QuentContext, + QuentIRExecutionContext, + declare_network_channels, + finalize_network_channels, ) from cudf_polars.quent._plan import build_plan, port_names_for_node from cudf_polars.quent._types import ( Attribute, + Channel, Engine, Implementation, Memory, + Network, Operator, Plan, Port, + Processor, Query, Statistics, Task, @@ -37,10 +45,10 @@ _deserialize_value, ) from cudf_polars.utils.config import ConfigOptions +from cudf_polars.utils.cuda_stream import get_cuda_stream if TYPE_CHECKING: from cudf_polars.dsl.ir import IR - from cudf_polars.quent._types import Processor from cudf_polars.utils.config import StreamingExecutor @@ -52,6 +60,58 @@ def _make_worker() -> Worker: ) +def _make_dataframe(pl_df: pl.DataFrame) -> DataFrame: + return DataFrame.from_polars(pl_df, get_cuda_stream()) + + +def _make_quent_ir_execution_context( + *, + operator_id: uuid.UUID | None = None, + disk_to_device_channel: Channel | None = None, +) -> tuple[cudf_polars.quent._logging.QuentLogger, QuentIRExecutionContext]: + pytest.importorskip("structlog") + logger = cudf_polars.quent._logging.QuentLogger() + context = QuentContext() + engine_id = context.engine.id + worker_id = uuid.uuid4() + pool_id = uuid.uuid4() + device_memory = Memory( + instance_name="device", + resource_type_name="memory", + parent_group_id=engine_id, + ) + operator_id = operator_id or uuid.uuid4() + plan = Plan( + id=uuid.uuid4(), + query=Query(), + parent_plan=None, + instance_name="logical", + edges=[], + worker=None, + ) + operator = Operator( + id=operator_id, + plan=plan, + parent_operators=[], + instance_name="Filter-NodeIndex(0)", + type_name="Filter", + ) + local_context = LocalQuentContext( + context=context, + query=context.query_for(uuid.uuid4()), + worker=Worker(id=worker_id, engine=context.engine, instance_name="rank-0"), + logger=logger, + thread_pool_id=pool_id, + processor_registry=ProcessorRegistry(), + device_memory=device_memory, + disk_to_device_channel=disk_to_device_channel, + ) + quent_ir_execution_context = QuentIRExecutionContext.from_execution_context( + local_context, operator + ) + return logger, quent_ir_execution_context + + @pytest.mark.parametrize( "value,expected_variant", [ @@ -704,7 +764,7 @@ def test_processor_registry_reused_across_quent_contexts() -> None: assert len(processor_events) == 2 -def test_processor_registry_exit_events_idempotent() -> None: +def test_processor_registry_exit_events() -> None: pytest.importorskip("structlog") from cudf_polars.quent._context import ProcessorRegistry @@ -715,7 +775,6 @@ def test_processor_registry_exit_events_idempotent() -> None: registry.get_or_declare_processor(logger, thread_ident=1, pool_id=pool_id) registry.get_or_declare_processor(logger, thread_ident=2, pool_id=pool_id) - registry._emit_processor_exit_events(logger) registry._emit_processor_exit_events(logger) events = _drained_events(logger) @@ -756,3 +815,358 @@ def test_quent_serialize_none(): "key": "none", "value": None, } + + +def test_task_from_ir() -> None: + operator_id = uuid.uuid4() + _logger, quent_ir_execution_context = _make_quent_ir_execution_context( + operator_id=operator_id + ) + + task = Task.from_ir(Filter, quent_ir_execution_context) + + assert task is not None + assert task.operator_id == operator_id + assert task.instance_name is not None + assert task.instance_name.startswith("Filter-") + assert operator_id.hex[:8] in task.instance_name + + +def test_task_loading_serialization() -> None: + operator_id = uuid.uuid4() + processor = Processor(pool_id=uuid.uuid4()) + channel = Channel( + instance_name="disk -> device", + resource_type_name="DiskToDevice", + parent_group_id=uuid.uuid4(), + source=Memory( + instance_name="filesystem", + resource_type_name="filesystem", + parent_group_id=uuid.uuid4(), + ), + target=Memory( + instance_name="device", + resource_type_name="memory", + parent_group_id=uuid.uuid4(), + ), + ) + memory = Memory( + instance_name="device", + resource_type_name="memory", + parent_group_id=uuid.uuid4(), + ) + task = Task(operator_id=operator_id, instance_name="scan-task") + + event = task.loading( + use_thread=processor, + use_channel=channel, + channel_capacity_bytes=4096, + use_memory=memory, + memory_capacity_bytes=8192, + timestamp=100, + ) + d = event.to_dict() + + assert d["id"] == str(task.id) + loading = d["data"]["Task"]["state"]["Loading"] + assert loading["use_thread"] == { + "resource_id": str(processor.id), + "capacity": None, + } + assert loading["use_fs_to_mem"] == { + "resource_id": str(channel.id), + "capacity": {"capacity_bytes": 4096}, + } + assert loading["use_memory"] == { + "resource_id": str(memory.id), + "capacity": {"capacity_bytes": 8192}, + } + + +def test_task_computing_serialization() -> None: + operator_id = uuid.uuid4() + processor = Processor(pool_id=uuid.uuid4()) + memory = Memory( + instance_name="device", + resource_type_name="memory", + parent_group_id=uuid.uuid4(), + ) + task = Task(operator_id=operator_id, instance_name="filter-task") + + event = task.computing( + use_thread=processor, + use_memory=memory, + memory_capacity_bytes=16384, + timestamp=101, + ) + d = event.to_dict() + + computing = d["data"]["Task"]["state"]["Computing"] + assert computing["use_thread"] == { + "resource_id": str(processor.id), + "capacity": None, + } + assert computing["use_memory"] == { + "resource_id": str(memory.id), + "capacity": {"capacity_bytes": 16384}, + } + + +def test_task_sending_serialization() -> None: + operator_id = uuid.uuid4() + processor = Processor(pool_id=uuid.uuid4()) + link = Channel( + instance_name="rank-0 -> rank-1", + resource_type_name="Link", + parent_group_id=uuid.uuid4(), + source=Memory( + instance_name="device", + resource_type_name="memory", + parent_group_id=uuid.uuid4(), + ), + target=Memory( + instance_name="device", + resource_type_name="memory", + parent_group_id=uuid.uuid4(), + ), + ) + task = Task(operator_id=operator_id, instance_name="shuffle-task") + + event = task.sending( + use_thread=processor, + use_link=link, + link_capacity_bytes=2048, + timestamp=102, + ) + d = event.to_dict() + + sending = d["data"]["Task"]["state"]["Sending"] + assert sending["use_thread"] == { + "resource_id": str(processor.id), + "capacity": None, + } + assert sending["use_link"] == { + "resource_id": str(link.id), + "capacity": {"capacity_bytes": 2048}, + } + + +def test_network_declare_serialization() -> None: + engine_id = uuid.uuid4() + network = Network(engine_id=engine_id) + + event = network.declare(timestamp=555) + d = event.to_dict() + + assert d["id"] == str(network.id) + assert d["timestamp"] == 555 + assert d["data"]["Network"]["Declaration"] == { + "instance_name": "Network", + "parent_group_id": str(engine_id), + } + + +def test_declare_network_channels_single_rank() -> None: + pytest.importorskip("structlog") + logger = cudf_polars.quent._logging.QuentLogger() + device_memory = Memory( + instance_name="device", + resource_type_name="memory", + parent_group_id=uuid.uuid4(), + ) + + network, link_channels = declare_network_channels( + logger, + rank=0, + nranks=1, + engine_id=uuid.uuid4(), + device_memory=device_memory, + ) + + assert network is None + assert link_channels == {} + assert _drained_events(logger) == [] + + +@pytest.mark.parametrize( + "rank,nranks,expected_targets", [(0, 3, [1, 2]), (1, 3, [0, 2])] +) +def test_declare_network_channels_multi_rank( + rank: int, nranks: int, expected_targets: list[int] +) -> None: + pytest.importorskip("structlog") + logger = cudf_polars.quent._logging.QuentLogger() + engine_id = uuid.uuid4() + device_memory = Memory( + instance_name="device", + resource_type_name="memory", + parent_group_id=engine_id, + ) + + network, link_channels = declare_network_channels( + logger, + rank=rank, + nranks=nranks, + engine_id=engine_id, + device_memory=device_memory, + ) + + assert network is not None + assert set(link_channels) == set(expected_targets) + for target_rank, link in link_channels.items(): + assert link.instance_name == f"rank-{rank} -> rank-{target_rank}" + assert link.resource_type_name == "Link" + assert link.parent_group_id == network.id + assert link.source is device_memory + assert link.target is device_memory + + events = _drained_events(logger) + network_events = [event for event in events if "Network" in event["data"]] + channel_events = [event for event in events if "Channel" in event["data"]] + assert len(network_events) == 1 + assert network_events[0]["data"]["Network"]["Declaration"][ + "parent_group_id" + ] == str(engine_id) + assert len(channel_events) == len(expected_targets) * 2 + + +def test_finalize_network_channels() -> None: + pytest.importorskip("structlog") + logger = cudf_polars.quent._logging.QuentLogger() + device_memory = Memory( + instance_name="device", + resource_type_name="memory", + parent_group_id=uuid.uuid4(), + ) + link_channels = { + target_rank: Channel( + instance_name=f"rank-0 -> rank-{target_rank}", + resource_type_name="Link", + parent_group_id=uuid.uuid4(), + source=device_memory, + target=device_memory, + ) + for target_rank in (1, 2) + } + + finalize_network_channels(logger, link_channels=link_channels) + + events = _drained_events(logger) + finalizing_events = [ + event + for event in events + if event["data"]["Channel"]["state"] == {"ChannelFinalizing": None} + ] + exit_events = [ + event for event in events if event["data"]["Channel"]["state"] == "Exit" + ] + assert len(finalizing_events) == 2 + assert len(exit_events) == 2 + + +def test_emit_task_begin_events_computing_node() -> None: + logger, quent_ir_execution_context = _make_quent_ir_execution_context() + task = Task.from_ir(Filter, quent_ir_execution_context) + assert task is not None + + quent_ir_execution_context.context._emit_task_begin_events( + Filter, + task, + quent_ir_execution_context, + ) + + events = _drained_events(logger) + task_events = [event for event in events if "Task" in event["data"]] + assert [event["data"]["Task"]["seq"] for event in task_events] == [0, 1, 3] + assert "Queueing" in task_events[0]["data"]["Task"]["state"] + assert "Allocating" in task_events[1]["data"]["Task"]["state"] + assert "Computing" in task_events[2]["data"]["Task"]["state"] + processor_events = [event for event in events if "Processor" in event["data"]] + assert len(processor_events) == 2 + + +def test_emit_task_begin_events_io_node() -> None: + channel = Channel( + instance_name="disk -> device", + resource_type_name="DiskToDevice", + parent_group_id=uuid.uuid4(), + source=Memory( + instance_name="filesystem", + resource_type_name="filesystem", + parent_group_id=uuid.uuid4(), + ), + target=Memory( + instance_name="device", + resource_type_name="memory", + parent_group_id=uuid.uuid4(), + ), + ) + logger, quent_ir_execution_context = _make_quent_ir_execution_context( + disk_to_device_channel=channel + ) + task = Task.from_ir(DataFrameScan, quent_ir_execution_context) + assert task is not None + + quent_ir_execution_context.context._emit_task_begin_events( + DataFrameScan, + task, + quent_ir_execution_context, + ) + + events = _drained_events(logger) + task_events = [event for event in events if "Task" in event["data"]] + assert [event["data"]["Task"]["seq"] for event in task_events] == [0, 2] + assert "Queueing" in task_events[0]["data"]["Task"]["state"] + loading = task_events[1]["data"]["Task"]["state"]["Loading"] + assert loading["use_fs_to_mem"]["resource_id"] == str(channel.id) + assert loading["use_memory"]["resource_id"] == str( + quent_ir_execution_context.device_memory.id + ) + + +def test_emit_task_end_events() -> None: + logger, quent_ir_execution_context = _make_quent_ir_execution_context() + task = Task(operator_id=quent_ir_execution_context.quent_operator.id) + input_frame = _make_dataframe(pl.DataFrame({"x": [1, 2, 3]})) + result = _make_dataframe(pl.DataFrame({"y": list(range(7))})) + + quent_ir_execution_context.context._emit_task_end_events( + task, + quent_ir_execution_context, + [input_frame], + result, + ) + + events = _drained_events(logger) + operator_events = [event for event in events if "Operator" in event["data"]] + task_events = [event for event in events if "Task" in event["data"]] + assert len(operator_events) == 1 + stats = operator_events[0]["data"]["Operator"]["Statistics"]["custom_attributes"] + assert stats == [ + {"key": "input_bytes", "value": {"U64": input_frame._size_bytes()}}, + {"key": "output_bytes", "value": {"U64": result._size_bytes()}}, + {"key": "output_rows", "value": {"U64": result.num_rows}}, + ] + assert task_events[0]["data"]["Task"]["state"] == "Exit" + + +def test_emit_task_end_events_on_failure() -> None: + logger, quent_ir_execution_context = _make_quent_ir_execution_context() + task = Task(operator_id=quent_ir_execution_context.quent_operator.id) + input_frame = _make_dataframe(pl.DataFrame({"x": [10, 20]})) + + quent_ir_execution_context.context._emit_task_end_events( + task, + quent_ir_execution_context, + [input_frame], + None, + ) + + events = _drained_events(logger) + operator_events = [event for event in events if "Operator" in event["data"]] + stats = operator_events[0]["data"]["Operator"]["Statistics"]["custom_attributes"] + assert stats == [ + {"key": "input_bytes", "value": {"U64": input_frame._size_bytes()}}, + {"key": "output_bytes", "value": {"U64": 0}}, + {"key": "output_rows", "value": {"U64": 0}}, + ] From 6186ae0e73acd276c8558ec954c74fd645d5c7ba Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 14 Jul 2026 15:13:43 -0700 Subject: [PATCH 24/60] test coverage --- .../cudf_polars/cudf_polars/quent/_context.py | 2 +- python/cudf_polars/tests/quent/conftest.py | 32 +++++ python/cudf_polars/tests/quent/test_quent.py | 117 +++++------------- 3 files changed, 65 insertions(+), 86 deletions(-) diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index f0c84d896487..2d0c17e67dd5 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -461,7 +461,7 @@ def _emit_task_end_events( - output bytes: the size of the output dataframe. - output rows: the number of rows in the output dataframe. """ - if quent_ir_execution_context is None: + if quent_ir_execution_context is None: # pragma: no cover; return if result is not None: diff --git a/python/cudf_polars/tests/quent/conftest.py b/python/cudf_polars/tests/quent/conftest.py index 849364e2f32a..ee94d385f66c 100644 --- a/python/cudf_polars/tests/quent/conftest.py +++ b/python/cudf_polars/tests/quent/conftest.py @@ -5,17 +5,49 @@ from __future__ import annotations +import uuid from typing import TYPE_CHECKING import pytest import cudf_polars.quent import cudf_polars.quent._context +from cudf_polars.quent._types import Channel, Memory, Processor if TYPE_CHECKING: from cudf_polars.quent._context import QuentContext +@pytest.fixture +def processor() -> Processor: + return Processor(pool_id=uuid.uuid4()) + + +@pytest.fixture +def device_memory() -> Memory: + return Memory( + instance_name="device", + resource_type_name="memory", + parent_group_id=uuid.uuid4(), + ) + + +@pytest.fixture +def disk_to_device_channel(device_memory: Memory) -> Channel: + filesystem = Memory( + instance_name="filesystem", + resource_type_name="filesystem", + parent_group_id=uuid.uuid4(), + ) + return Channel( + instance_name="disk -> device", + resource_type_name="DiskToDevice", + parent_group_id=uuid.uuid4(), + source=filesystem, + target=device_memory, + ) + + @pytest.fixture def quent_context() -> QuentContext: """A Quent Context with a QueryGroup and Query set.""" diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index c168c3f56ae0..4323574354cd 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -37,7 +37,6 @@ Operator, Plan, Port, - Processor, Query, Statistics, Task, @@ -49,6 +48,7 @@ if TYPE_CHECKING: from cudf_polars.dsl.ir import IR + from cudf_polars.quent._types import Processor from cudf_polars.utils.config import StreamingExecutor @@ -832,36 +832,19 @@ def test_task_from_ir() -> None: assert operator_id.hex[:8] in task.instance_name -def test_task_loading_serialization() -> None: +def test_task_loading_serialization( + processor: Processor, + device_memory: Memory, + disk_to_device_channel: Channel, +) -> None: operator_id = uuid.uuid4() - processor = Processor(pool_id=uuid.uuid4()) - channel = Channel( - instance_name="disk -> device", - resource_type_name="DiskToDevice", - parent_group_id=uuid.uuid4(), - source=Memory( - instance_name="filesystem", - resource_type_name="filesystem", - parent_group_id=uuid.uuid4(), - ), - target=Memory( - instance_name="device", - resource_type_name="memory", - parent_group_id=uuid.uuid4(), - ), - ) - memory = Memory( - instance_name="device", - resource_type_name="memory", - parent_group_id=uuid.uuid4(), - ) task = Task(operator_id=operator_id, instance_name="scan-task") event = task.loading( use_thread=processor, - use_channel=channel, + use_channel=disk_to_device_channel, channel_capacity_bytes=4096, - use_memory=memory, + use_memory=device_memory, memory_capacity_bytes=8192, timestamp=100, ) @@ -874,28 +857,25 @@ def test_task_loading_serialization() -> None: "capacity": None, } assert loading["use_fs_to_mem"] == { - "resource_id": str(channel.id), + "resource_id": str(disk_to_device_channel.id), "capacity": {"capacity_bytes": 4096}, } assert loading["use_memory"] == { - "resource_id": str(memory.id), + "resource_id": str(device_memory.id), "capacity": {"capacity_bytes": 8192}, } -def test_task_computing_serialization() -> None: +def test_task_computing_serialization( + processor: Processor, + device_memory: Memory, +) -> None: operator_id = uuid.uuid4() - processor = Processor(pool_id=uuid.uuid4()) - memory = Memory( - instance_name="device", - resource_type_name="memory", - parent_group_id=uuid.uuid4(), - ) task = Task(operator_id=operator_id, instance_name="filter-task") event = task.computing( use_thread=processor, - use_memory=memory, + use_memory=device_memory, memory_capacity_bytes=16384, timestamp=101, ) @@ -907,30 +887,24 @@ def test_task_computing_serialization() -> None: "capacity": None, } assert computing["use_memory"] == { - "resource_id": str(memory.id), + "resource_id": str(device_memory.id), "capacity": {"capacity_bytes": 16384}, } -def test_task_sending_serialization() -> None: +def test_task_sending_serialization( + processor: Processor, + device_memory: Memory, +) -> None: operator_id = uuid.uuid4() - processor = Processor(pool_id=uuid.uuid4()) + task = Task(operator_id=operator_id, instance_name="shuffle-task") link = Channel( instance_name="rank-0 -> rank-1", resource_type_name="Link", parent_group_id=uuid.uuid4(), - source=Memory( - instance_name="device", - resource_type_name="memory", - parent_group_id=uuid.uuid4(), - ), - target=Memory( - instance_name="device", - resource_type_name="memory", - parent_group_id=uuid.uuid4(), - ), + source=device_memory, + target=device_memory, ) - task = Task(operator_id=operator_id, instance_name="shuffle-task") event = task.sending( use_thread=processor, @@ -966,14 +940,9 @@ def test_network_declare_serialization() -> None: } -def test_declare_network_channels_single_rank() -> None: +def test_declare_network_channels_single_rank(device_memory: Memory) -> None: pytest.importorskip("structlog") logger = cudf_polars.quent._logging.QuentLogger() - device_memory = Memory( - instance_name="device", - resource_type_name="memory", - parent_group_id=uuid.uuid4(), - ) network, link_channels = declare_network_channels( logger, @@ -992,16 +961,14 @@ def test_declare_network_channels_single_rank() -> None: "rank,nranks,expected_targets", [(0, 3, [1, 2]), (1, 3, [0, 2])] ) def test_declare_network_channels_multi_rank( - rank: int, nranks: int, expected_targets: list[int] + device_memory: Memory, + rank: int, + nranks: int, + expected_targets: list[int], ) -> None: pytest.importorskip("structlog") logger = cudf_polars.quent._logging.QuentLogger() engine_id = uuid.uuid4() - device_memory = Memory( - instance_name="device", - resource_type_name="memory", - parent_group_id=engine_id, - ) network, link_channels = declare_network_channels( logger, @@ -1030,14 +997,9 @@ def test_declare_network_channels_multi_rank( assert len(channel_events) == len(expected_targets) * 2 -def test_finalize_network_channels() -> None: +def test_finalize_network_channels(device_memory: Memory) -> None: pytest.importorskip("structlog") logger = cudf_polars.quent._logging.QuentLogger() - device_memory = Memory( - instance_name="device", - resource_type_name="memory", - parent_group_id=uuid.uuid4(), - ) link_channels = { target_rank: Channel( instance_name=f"rank-0 -> rank-{target_rank}", @@ -1085,24 +1047,9 @@ def test_emit_task_begin_events_computing_node() -> None: assert len(processor_events) == 2 -def test_emit_task_begin_events_io_node() -> None: - channel = Channel( - instance_name="disk -> device", - resource_type_name="DiskToDevice", - parent_group_id=uuid.uuid4(), - source=Memory( - instance_name="filesystem", - resource_type_name="filesystem", - parent_group_id=uuid.uuid4(), - ), - target=Memory( - instance_name="device", - resource_type_name="memory", - parent_group_id=uuid.uuid4(), - ), - ) +def test_emit_task_begin_events_io_node(disk_to_device_channel: Channel) -> None: logger, quent_ir_execution_context = _make_quent_ir_execution_context( - disk_to_device_channel=channel + disk_to_device_channel=disk_to_device_channel ) task = Task.from_ir(DataFrameScan, quent_ir_execution_context) assert task is not None @@ -1118,7 +1065,7 @@ def test_emit_task_begin_events_io_node() -> None: assert [event["data"]["Task"]["seq"] for event in task_events] == [0, 2] assert "Queueing" in task_events[0]["data"]["Task"]["state"] loading = task_events[1]["data"]["Task"]["state"]["Loading"] - assert loading["use_fs_to_mem"]["resource_id"] == str(channel.id) + assert loading["use_fs_to_mem"]["resource_id"] == str(disk_to_device_channel.id) assert loading["use_memory"]["resource_id"] == str( quent_ir_execution_context.device_memory.id ) From 0b2a8f08fd161851f5d57f34a8fcbb8a2c024b62 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 15 Jul 2026 05:19:06 -0700 Subject: [PATCH 25/60] fixup --- python/cudf_polars/tests/quent/test_quent.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index 4323574354cd..8b1933cbc992 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -75,15 +75,19 @@ def _make_quent_ir_execution_context( engine_id = context.engine.id worker_id = uuid.uuid4() pool_id = uuid.uuid4() - device_memory = Memory( - instance_name="device", - resource_type_name="memory", - parent_group_id=engine_id, - ) + if disk_to_device_channel is not None: + device_memory = disk_to_device_channel.target + else: + device_memory = Memory( + instance_name="device", + resource_type_name="memory", + parent_group_id=engine_id, + ) operator_id = operator_id or uuid.uuid4() + query = context.query_for(uuid.uuid4()) plan = Plan( id=uuid.uuid4(), - query=Query(), + query=query, parent_plan=None, instance_name="logical", edges=[], @@ -98,7 +102,7 @@ def _make_quent_ir_execution_context( ) local_context = LocalQuentContext( context=context, - query=context.query_for(uuid.uuid4()), + query=query, worker=Worker(id=worker_id, engine=context.engine, instance_name="rank-0"), logger=logger, thread_pool_id=pool_id, From c6f38ace084682bedebd28cd25a22e898ada78d9 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Thu, 16 Jul 2026 09:38:20 -0700 Subject: [PATCH 26/60] Fixup --- python/cudf_polars/cudf_polars/dsl/tracing.py | 3 + .../cudf_polars/cudf_polars/quent/_context.py | 80 ++++++-- .../cudf_polars/cudf_polars/quent/_export.py | 144 ++++++++++++++ .../cudf_polars/cudf_polars/quent/_types.py | 62 ++++-- .../cudf_polars/streaming/benchmarks/utils.py | 18 +- python/cudf_polars/tests/quent/test_export.py | 177 ++++++++++++++++++ python/cudf_polars/tests/quent/test_quent.py | 63 ++++--- 7 files changed, 474 insertions(+), 73 deletions(-) create mode 100644 python/cudf_polars/cudf_polars/quent/_export.py create mode 100644 python/cudf_polars/tests/quent/test_export.py diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index 2c55a3305f40..4486c6a9093f 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -199,6 +199,7 @@ def wrapper( cls, quent_task, ir_execution_context.quent_ir_execution_context, + input_frames_bytes=sum(frame._size_bytes() for frame in frames), ) before_start = time.monotonic_ns() @@ -220,6 +221,7 @@ def wrapper( and ir_execution_context.quent_ir_execution_context is not None ): ir_execution_context.quent_ir_execution_context.context._emit_task_end_events( + cls, quent_task, ir_execution_context.quent_ir_execution_context, frames, @@ -232,6 +234,7 @@ def wrapper( and ir_execution_context.quent_ir_execution_context is not None ): ir_execution_context.quent_ir_execution_context.context._emit_task_end_events( + cls, quent_task, ir_execution_context.quent_ir_execution_context, frames, diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index 2d0c17e67dd5..b9c692ace1ab 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -121,6 +121,10 @@ class QuentContext: def __post_init__(self) -> None: object.__setattr__(self, "_query_group_cache_", set()) + # Maps task id -> unix timestamp when an I/O task entered Allocating. + # Used to backdate the Loading transition so channel rate spans cover + # the full disk-to-device transfer duration. + object.__setattr__(self, "_io_loading_started_at_", {}) def serialize(self) -> bytes: """ @@ -376,6 +380,7 @@ def _emit_task_begin_events( ir_type: type[IR], quent_task: Task, quent_ir_execution_context: QuentIRExecutionContext, + input_frames_bytes: int, ) -> None: """ Emit begin events for a Quent Task. @@ -388,44 +393,47 @@ def _emit_task_begin_events( The Quent Task to emit events for. quent_ir_execution_context: QuentIRExecutionContext The Quent IR execution context. + input_frames_bytes: int + The total size of the input dataframes in bytes. Notes ----- - The following events are emitted: + This emits the following events: - - Queueing - - Loading (I/O nodes only) - - Allocating (non-I/O nodes only) - - Computing (non-I/O nodes only) - - The Loading, Allocating, and Computing events will indicate the host CPU thread - and device memory that they're using. + - queueing + - allocating + - loading (I/O nodes only) + - computing (non-I/O nodes only) """ quent_processor = quent_ir_execution_context.get_or_declare_processor( thread_ident=threading.get_ident(), ) quent_ir_execution_context.logger.emit(quent_task.queueing()) - if not ir_type.is_io_node: - quent_ir_execution_context.logger.emit( - quent_task.allocating(resource_id=quent_processor.id) - ) + quent_ir_execution_context.logger.emit( + quent_task.allocating(resource_id=quent_processor.id) + ) + if ir_type.is_io_node: quent_ir_execution_context.logger.emit( - quent_task.computing( + quent_task.loading( use_thread=quent_processor, + use_channel=quent_ir_execution_context.disk_to_device_channel, + channel_capacity_bytes=input_frames_bytes, use_memory=quent_ir_execution_context.device_memory, + memory_capacity_bytes=input_frames_bytes, ) ) else: quent_ir_execution_context.logger.emit( - quent_task.loading( + quent_task.computing( use_thread=quent_processor, - use_channel=quent_ir_execution_context.disk_to_device_channel, use_memory=quent_ir_execution_context.device_memory, + memory_capacity_bytes=input_frames_bytes, ) ) def _emit_task_end_events( self, + ir_type: type[IR], quent_task: Task, quent_ir_execution_context: QuentIRExecutionContext, frames: Sequence[DataFrame], @@ -436,6 +444,8 @@ def _emit_task_end_events( Parameters ---------- + ir_type: type[IR] + The IR type of the operator. quent_task: Task The Quent Task to emit events for. quent_ir_execution_context: QuentIRExecutionContext @@ -454,6 +464,12 @@ def _emit_task_end_events( This method emits an ``Exit`` event for the Quent Task, whose timestamp represents when the IR node completed host-side processing. + The Quent Task FSM only permits a task to exit from the ``Computing`` + state. Non-I/O nodes already entered ``Computing`` in + :meth:`_emit_task_begin_events`. I/O nodes were left in ``Allocating`` + during the load; here they transition through ``Loading`` (with byte + counts), ``Computing``, and ``Exit``. + A ``Statistics`` record, associated with the Quent Operator bound to the IR execution context, is also emitted. It includes @@ -479,9 +495,32 @@ def _emit_task_end_events( ) ) ) + if ir_type.is_io_node: + quent_processor = quent_ir_execution_context.get_or_declare_processor( + thread_ident=threading.get_ident(), + ) + quent_ir_execution_context.logger.emit( + quent_task.computing( + use_thread=quent_processor, + use_memory=quent_ir_execution_context.device_memory, + memory_capacity_bytes=output_capacity_bytes, + ) + ) quent_ir_execution_context.logger.emit(quent_task.exit()) +def _device_memory_capacity_bytes() -> int | None: + """Return total device memory in bytes, or ``None`` if unavailable.""" + try: + import pynvml + + pynvml.nvmlInit() + handle = pynvml.nvmlDeviceGetHandleByIndex(0) + return int(pynvml.nvmlDeviceGetMemoryInfo(handle).total) + except Exception: # pragma: no cover - no GPU / pynvml unavailable + return None + + def declare_worker_resources( logger: QuentLogger, *, @@ -512,12 +551,15 @@ def declare_worker_resources( target=device_memory, ) thread_pool = ThreadPool(worker_id=worker_id) + device_memory_capacity = _device_memory_capacity_bytes() logger.emit(device_memory.initializing()) - logger.emit(device_memory.operating(0)) + logger.emit(device_memory.operating(device_memory_capacity)) logger.emit(filesystem.initializing()) - logger.emit(filesystem.operating(0)) + # Filesystem capacity is unknown; declare as unbounded. + logger.emit(filesystem.operating(None)) logger.emit(disk_to_device_channel.initializing()) - logger.emit(disk_to_device_channel.operating()) + # Channel capacity is a rate bound; unbounded when unknown. + logger.emit(disk_to_device_channel.operating(None)) logger.emit(thread_pool.declare()) return device_memory, disk_to_device_channel, thread_pool @@ -575,7 +617,7 @@ def declare_network_channels( target=device_memory, ) logger.emit(link.initializing()) - logger.emit(link.operating()) + logger.emit(link.operating(None)) link_channels[target_rank] = link return network, link_channels diff --git a/python/cudf_polars/cudf_polars/quent/_export.py b/python/cudf_polars/cudf_polars/quent/_export.py new file mode 100644 index 000000000000..d652a4d36c64 --- /dev/null +++ b/python/cudf_polars/cudf_polars/quent/_export.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Export Quent events to the filesystem directory layout.""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, Any + +from cudf_polars.quent._types import EventName, new_quent_id + +if TYPE_CHECKING: + import uuid + from pathlib import Path + + +SIDECAR_FILE_NAME = "model.qmi" +EXTENSION = "ndjson" + +MODEL_QMI: dict[str, Any] = { + "quent": { + "version": "0.1.0", + "commit": "9c2924510756d9eca288819e125330805053051f", + "remote": "https://github.com/rapidsai/quent", + }, + "model": { + "name": "Simulator", + "package": "quent-simulator-instrumentation", + "type_path": "quent_simulator_instrumentation::SimulatorEvent", + "source": { + "version": "0.1.0", + "commit": "9c2924510756d9eca288819e125330805053051f", + "remote": "https://github.com/rapidsai/quent", + }, + "analyzer_package": "quent-simulator-analyzer", + }, +} + +ENTITY_DIRECTORIES: dict[str, str] = { + EventName.ENGINE.value: "engine", + EventName.WORKER.value: "worker", + EventName.QUERY_GROUP.value: "query_group", + EventName.QUERY.value: "query", + EventName.PLAN.value: "plan", + EventName.OPERATOR.value: "operator", + EventName.PORT.value: "port", + EventName.TASK.value: "task", + EventName.MEMORY.value: "memory", + EventName.CHANNEL.value: "channel", + EventName.THREAD_POOL.value: "thread_pool", + EventName.PROCESSOR.value: "processor", + EventName.NETWORK.value: "network", +} + + +def unwrap_event_data(data: dict[str, Any]) -> tuple[str, Any]: + """ + Extract the entity name and unwrapped payload from a buffered event. + + Buffered events wrap payloads as ``{"Engine": {...}}``; directory export + stores the payload directly because the entity type is implied by the + subdirectory name. + """ + if len(data) != 1: + msg = ( + "Expected event data with exactly one entity wrapper, " + f"got {len(data)} keys: {sorted(data)}" + ) + raise ValueError(msg) + entity_name, payload = next(iter(data.items())) + if entity_name not in ENTITY_DIRECTORIES: + msg = f"Unknown Quent entity type: {entity_name!r}" + raise ValueError(msg) + return entity_name, payload + + +def to_export_line(event: dict[str, Any]) -> tuple[str, dict[str, Any]]: + """Convert a buffered event envelope to directory export line format.""" + entity_name, payload = unwrap_event_data(event["data"]) + directory = ENTITY_DIRECTORIES[entity_name] + export_line = { + "id": event["id"], + "timestamp": event["timestamp"], + "data": payload, + } + return directory, export_line + + +def write_sidecar(context_dir: Path, sidecar: dict[str, Any]) -> None: + """Atomically write the ``model.qmi`` provenance sidecar.""" + tmp_path = context_dir / f".{SIDECAR_FILE_NAME}.tmp" + final_path = context_dir / SIDECAR_FILE_NAME + tmp_path.write_text(json.dumps(sidecar, indent=2) + "\n", encoding="utf-8") + tmp_path.replace(final_path) + + +def write_quent_export( + events: list[dict[str, Any]], + export_root: Path, + context_id: uuid.UUID, + *, + sidecar: dict[str, Any] | None = None, +) -> Path: + """ + Write Quent events to the filesystem export layout. + + Parameters + ---------- + events + Buffered Quent event envelopes from ``engine._quent_events``. + export_root + Root directory for exported contexts (e.g. ``logs``). + context_id + Context UUID, typically the engine/run id. + sidecar + Optional provenance payload for ``model.qmi``. Defaults to + :data:`MODEL_QMI`. + + Returns + ------- + Path + The context directory ``export_root//``. + """ + context_dir = export_root / str(context_id) + context_dir.mkdir(parents=True, exist_ok=True) + + write_sidecar(context_dir, sidecar or MODEL_QMI) + + grouped: dict[str, list[dict[str, Any]]] = {} + for event in events: + directory, export_line = to_export_line(event) + grouped.setdefault(directory, []).append(export_line) + + for directory, lines in grouped.items(): + entity_dir = context_dir / directory + entity_dir.mkdir(parents=True, exist_ok=True) + stream_path = entity_dir / f"{new_quent_id()}.{EXTENSION}" + with stream_path.open("w", encoding="utf-8") as stream_file: + for line in lines: + stream_file.write(json.dumps(line, separators=(",", ":"))) + stream_file.write("\n") + + return context_dir diff --git a/python/cudf_polars/cudf_polars/quent/_types.py b/python/cudf_polars/cudf_polars/quent/_types.py index da27e3f85448..144ee126ef6d 100644 --- a/python/cudf_polars/cudf_polars/quent/_types.py +++ b/python/cudf_polars/cudf_polars/quent/_types.py @@ -9,6 +9,7 @@ import dataclasses import enum +import itertools import sys import time import uuid @@ -17,6 +18,8 @@ from cudf_polars import __version__ if TYPE_CHECKING: + from collections.abc import Iterator + from cudf_polars.dsl.ir import IR from cudf_polars.quent._context import QuentIRExecutionContext @@ -38,6 +41,7 @@ class EventName(enum.Enum): CHANNEL = "Channel" THREAD_POOL = "ThreadPool" PROCESSOR = "Processor" + NETWORK = "Network" if sys.version_info >= (3, 14): # pragma: no cover; requires Python 3.14+ @@ -555,6 +559,27 @@ def _deserialize_value(value: dict[str, Any] | None) -> Value | None: raise ValueError(f"Unsupported Quent custom attribute variant: '{variant}'") +# Resource capacity helpers +# +# Quent distinguishes unit resources (Processor/thread), occupancy capacities +# (Memory), and rate capacities (Channel). See quent/docs/modeling/resource.md. + + +def occupancy_usage_capacity_bytes(capacity_bytes: int) -> dict[str, int]: + """Usage capacity for a Memory resource (occupancy over the usage span).""" + return {"capacity_bytes": capacity_bytes} + + +def rate_usage_capacity_bytes(capacity_bytes: int) -> dict[str, int]: + """ + Usage capacity for a Channel resource (total bytes over the usage span). + + Rate capacity values represent the total quantity transferred during the + span, not bytes per second. + """ + return {"capacity_bytes": capacity_bytes} + + # Resource types @@ -586,7 +611,9 @@ def initializing(self, timestamp: int | None = None) -> Event: }, ) - def operating(self, capacity_bytes: int, timestamp: int | None = None) -> Event: + def operating( + self, capacity_bytes: int | None = None, timestamp: int | None = None + ) -> Event: """Build a Quent Memory Operating event.""" return Event( id=self.id, @@ -705,7 +732,7 @@ def declare(self, timestamp: int | None = None) -> Event: id=self.id, timestamp=timestamp if timestamp is not None else time.time_ns(), data={ - "Network": { + EventName.NETWORK.value: { "Declaration": { "instance_name": "Network", "parent_group_id": str(self.engine_id), @@ -808,6 +835,9 @@ class Task: id: uuid.UUID = dataclasses.field(default_factory=new_quent_id) operator_id: uuid.UUID instance_name: str | None = None + _seq: Iterator[int] = dataclasses.field( + default_factory=itertools.count, compare=False, repr=False + ) @classmethod def from_ir( @@ -845,7 +875,7 @@ def queueing(self, timestamp: int | None = None) -> Event: timestamp=timestamp if timestamp is not None else time.time_ns(), data={ EventName.TASK.value: { - "seq": 0, + "seq": next(self._seq), "state": { "Queueing": { "instance_name": self.instance_name or self.id.hex[:8], @@ -859,7 +889,6 @@ def queueing(self, timestamp: int | None = None) -> Event: def allocating( self, resource_id: uuid.UUID, - capacity: int | None = None, timestamp: int | None = None, ) -> Event: """Build a Quent Task Allocating event.""" @@ -868,12 +897,12 @@ def allocating( timestamp=timestamp if timestamp is not None else time.time_ns(), data={ EventName.TASK.value: { - "seq": 1, + "seq": next(self._seq), "state": { "Allocating": { "use_thread": { "resource_id": str(resource_id), - "capacity": capacity, + "capacity": None, } } }, @@ -900,19 +929,19 @@ def loading( if use_channel is not None: loading_data["use_fs_to_mem"] = { "resource_id": str(use_channel.id), - "capacity": {"capacity_bytes": channel_capacity_bytes}, + "capacity": rate_usage_capacity_bytes(channel_capacity_bytes), } if use_memory is not None: loading_data["use_memory"] = { "resource_id": str(use_memory.id), - "capacity": {"capacity_bytes": memory_capacity_bytes}, + "capacity": occupancy_usage_capacity_bytes(memory_capacity_bytes), } return Event( id=self.id, timestamp=timestamp if timestamp is not None else time.time_ns(), data={ EventName.TASK.value: { - "seq": 2, + "seq": next(self._seq), "state": {"Loading": loading_data}, } }, @@ -922,11 +951,14 @@ def computing( self, use_thread: Processor | None = None, use_memory: Memory | None = None, + input_bytes: int = 0, memory_capacity_bytes: int = 0, timestamp: int | None = None, ) -> Event: """Build a Quent Task Computing event.""" - computing_data: dict[str, dict[str, Any]] = {} + computing_data: dict[str, Any] = {} + computing_data["instance_name"] = "" + computing_data["input_bytes"] = input_bytes if use_thread is not None: computing_data["use_thread"] = { "resource_id": str(use_thread.id), @@ -935,14 +967,14 @@ def computing( if use_memory is not None: computing_data["use_memory"] = { "resource_id": str(use_memory.id), - "capacity": {"capacity_bytes": memory_capacity_bytes}, + "capacity": occupancy_usage_capacity_bytes(memory_capacity_bytes), } return Event( id=self.id, timestamp=timestamp if timestamp is not None else time.time_ns(), data={ EventName.TASK.value: { - "seq": 3, + "seq": next(self._seq), "state": {"Computing": computing_data}, } }, @@ -965,14 +997,14 @@ def sending( if use_link is not None: sending_data["use_link"] = { "resource_id": str(use_link.id), - "capacity": {"capacity_bytes": link_capacity_bytes}, + "capacity": rate_usage_capacity_bytes(link_capacity_bytes), } return Event( id=self.id, timestamp=timestamp if timestamp is not None else time.time_ns(), data={ EventName.TASK.value: { - "seq": 4, + "seq": next(self._seq), "state": {"Sending": sending_data}, } }, @@ -983,5 +1015,5 @@ def exit(self, timestamp: int | None = None) -> Event: return Event( id=self.id, timestamp=timestamp if timestamp is not None else time.time_ns(), - data={EventName.TASK.value: {"seq": 5, "state": "Exit"}}, + data={EventName.TASK.value: {"seq": next(self._seq), "state": "Exit"}}, ) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index 64428628a057..3c44a561096c 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -1562,13 +1562,15 @@ def inject( def _write_quent_traces( engine: StreamingEngine, run_id: uuid.UUID, *, collect_traces: bool ) -> None: - """Write collected Quent events to logs/{run_id}.ndjson.""" + """Write collected Quent events to logs// directory export layout.""" if not (_HAS_STRUCTLOG or collect_traces): return + from cudf_polars.quent._export import write_quent_export + quent_logs = list(engine._quent_events) - # The quent UI currently requires the filename to match the engine's ID. + # The quent UI currently requires the context directory to match the engine's ID. for log in quent_logs: if log.get("data", {}).get("Engine", {}).get("Init") and log.get("id") != str( run_id @@ -1580,13 +1582,11 @@ def _write_quent_traces( warnings.warn(msg, stacklevel=2) logs_dir = Path("logs") - logs_dir.mkdir(parents=True, exist_ok=True) - output_path = logs_dir / f"{run_id}.ndjson" - with output_path.open("w") as f: - for log in quent_logs: - f.write(json.dumps(log)) - f.write("\n") - print(f"Wrote {len(quent_logs)} Quent trace events to {output_path}") + output_path = write_quent_export(quent_logs, logs_dir, run_id) + print( + f"Wrote {len(quent_logs)} Quent trace events to {output_path} " + f"(directory export layout)" + ) def _consolidate_logs( diff --git a/python/cudf_polars/tests/quent/test_export.py b/python/cudf_polars/tests/quent/test_export.py new file mode 100644 index 000000000000..ed5a599c305c --- /dev/null +++ b/python/cudf_polars/tests/quent/test_export.py @@ -0,0 +1,177 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for Quent filesystem export.""" + +from __future__ import annotations + +import json +import uuid +from typing import TYPE_CHECKING, Any + +import pytest + +from cudf_polars.quent._export import ( + EXTENSION, + MODEL_QMI, + SIDECAR_FILE_NAME, + to_export_line, + unwrap_event_data, + write_quent_export, +) +from cudf_polars.quent._types import Engine, Network, Query, QueryGroup + +if TYPE_CHECKING: + from pathlib import Path + + +def _buffered_events() -> list[dict[str, Any]]: + engine = Engine(id=uuid.UUID("019dd571-105a-7c53-a15b-713cbdd7666b")) + query_group = QueryGroup( + id=uuid.UUID("019dd571-1062-77c2-9803-62a66b6e0c5f"), + instance_name="test-group", + ) + query = Query( + id=uuid.UUID("019dd571-1062-77c2-9803-62bd37658144"), + instance_name="test-query", + ) + network = Network(engine_id=engine.id) + return [ + engine._init().to_dict(), + query_group._declare(engine).to_dict(), + query._init(query_group).to_dict(), + network.declare().to_dict(), + engine._exit().to_dict(), + ] + + +def _read_ndjson_lines(path: Path) -> list[dict[str, Any]]: + return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines()] + + +def test_unwrap_event_data() -> None: + entity_name, payload = unwrap_event_data( + {"Engine": {"Init": {"instance_name": "x"}}} + ) + assert entity_name == "Engine" + assert payload == {"Init": {"instance_name": "x"}} + + +def test_unwrap_event_data_rejects_multiple_wrappers() -> None: + with pytest.raises(ValueError, match="exactly one entity wrapper"): + unwrap_event_data({"Engine": {}, "Query": {}}) + + +def test_unwrap_event_data_rejects_unknown_entity() -> None: + with pytest.raises(ValueError, match="Unknown Quent entity type"): + unwrap_event_data({"UnknownEntity": {}}) + + +def test_to_export_line_unwraps_payload() -> None: + event = { + "id": "019dd571-105a-7c53-a15b-713cbdd7666b", + "timestamp": 1777402450018164995, + "data": {"Engine": {"Init": {"instance_name": "test"}}}, + } + directory, export_line = to_export_line(event) + assert directory == "engine" + assert export_line == { + "id": "019dd571-105a-7c53-a15b-713cbdd7666b", + "timestamp": 1777402450018164995, + "data": {"Init": {"instance_name": "test"}}, + } + + +def test_write_quent_export_creates_context_layout(tmp_path: Path) -> None: + context_id = uuid.UUID("019dd571-105a-7c53-a15b-713cbdd7666b") + events = _buffered_events() + + context_dir = write_quent_export(events, tmp_path, context_id) + + assert context_dir == tmp_path / str(context_id) + assert (context_dir / SIDECAR_FILE_NAME).exists() + assert ( + json.loads((context_dir / SIDECAR_FILE_NAME).read_text(encoding="utf-8")) + == MODEL_QMI + ) + + expected_dirs = {"engine", "query_group", "query", "network"} + created_dirs = {path.name for path in context_dir.iterdir() if path.is_dir()} + assert created_dirs == expected_dirs + + for entity_dir in expected_dirs: + stream_files = list((context_dir / entity_dir).glob(f"*.{EXTENSION}")) + assert len(stream_files) == 1 + lines = _read_ndjson_lines(stream_files[0]) + assert lines + for line in lines: + assert "id" in line + assert "timestamp" in line + assert isinstance(line["data"], dict) + assert len(line["data"]) == 1 or "seq" in line["data"] + + +def test_write_quent_export_unwraps_buffered_envelopes(tmp_path: Path) -> None: + context_id = uuid.UUID("019dd571-105a-7c53-a15b-713cbdd7666b") + events = _buffered_events() + + write_quent_export(events, tmp_path, context_id) + + engine_stream = next((tmp_path / str(context_id) / "engine").glob(f"*.{EXTENSION}")) + engine_lines = _read_ndjson_lines(engine_stream) + assert engine_lines[0]["data"] == { + "Init": { + "implementation": { + "name": "cudf-polars", + "version": engine_lines[0]["data"]["Init"]["implementation"]["version"], + "custom_attributes": [], + }, + "instance_name": "cudf-polars-019dd571", + } + } + assert engine_lines[1]["data"] == {"Exit": None} + + network_stream = next( + (tmp_path / str(context_id) / "network").glob(f"*.{EXTENSION}") + ) + network_lines = _read_ndjson_lines(network_stream) + assert network_lines[0]["data"] == { + "Declaration": { + "instance_name": "Network", + "parent_group_id": str(context_id), + } + } + + +def test_write_quent_export_rejects_malformed_event(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="exactly one entity wrapper"): + write_quent_export( + [{"id": "x", "timestamp": 1, "data": {"Engine": {}, "Query": {}}}], + tmp_path, + uuid.uuid4(), + ) + + +def test_write_quent_traces_benchmark_writer( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + pytest.importorskip("structlog") + from cudf_polars.streaming.benchmarks import utils as benchmark_utils + + run_id = uuid.UUID("019dd571-105a-7c53-a15b-713cbdd7666b") + events = _buffered_events() + + class FakeEngine: + _quent_events = events + + monkeypatch.chdir(tmp_path) + benchmark_utils._write_quent_traces( + FakeEngine(), # type: ignore[arg-type] + run_id, + collect_traces=True, + ) + + context_dir = tmp_path / "logs" / str(run_id) + assert (context_dir / SIDECAR_FILE_NAME).exists() + assert (context_dir / "engine").is_dir() + assert (context_dir / "query").is_dir() diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index 8b1933cbc992..187c965a9536 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -348,8 +348,11 @@ def test_task_lifecycle_events() -> None: task = Task(operator_id=operator_id, instance_name="task-0") queue = task.queueing().to_dict() assert queue["data"]["Task"]["state"]["Queueing"]["operator_id"] == str(operator_id) + assert queue["data"]["Task"]["seq"] == 0 + # ``seq`` is a per-instance counter that increments by one on each + # transition, in emission order (queueing == 0, allocating == 1, exit == 2). assert task.allocating(uuid.uuid4()).to_dict()["data"]["Task"]["seq"] == 1 - assert task.exit().to_dict()["data"]["Task"]["seq"] == 5 + assert task.exit().to_dict()["data"]["Task"]["seq"] == 2 def test_port_declare_serialization( @@ -1030,7 +1033,7 @@ def test_finalize_network_channels(device_memory: Memory) -> None: assert len(exit_events) == 2 -def test_emit_task_begin_events_computing_node() -> None: +def test_emit_task_events_computing_node() -> None: logger, quent_ir_execution_context = _make_quent_ir_execution_context() task = Task.from_ir(Filter, quent_ir_execution_context) assert task is not None @@ -1039,19 +1042,33 @@ def test_emit_task_begin_events_computing_node() -> None: Filter, task, quent_ir_execution_context, + input_frames_bytes=0, + ) + + # Simulate the result + result = _make_dataframe(pl.DataFrame({"y": list(range(7))})) + + quent_ir_execution_context.context._emit_task_end_events( + Filter, + task, + quent_ir_execution_context, + [], + result, ) events = _drained_events(logger) task_events = [event for event in events if "Task" in event["data"]] - assert [event["data"]["Task"]["seq"] for event in task_events] == [0, 1, 3] + # queueing -> allocating -> computing -> exit + assert [event["data"]["Task"]["seq"] for event in task_events] == [0, 1, 2, 3] assert "Queueing" in task_events[0]["data"]["Task"]["state"] assert "Allocating" in task_events[1]["data"]["Task"]["state"] assert "Computing" in task_events[2]["data"]["Task"]["state"] + assert "Exit" in task_events[3]["data"]["Task"]["state"] processor_events = [event for event in events if "Processor" in event["data"]] assert len(processor_events) == 2 -def test_emit_task_begin_events_io_node(disk_to_device_channel: Channel) -> None: +def test_emit_task_events_io_node(disk_to_device_channel: Channel) -> None: logger, quent_ir_execution_context = _make_quent_ir_execution_context( disk_to_device_channel=disk_to_device_channel ) @@ -1062,43 +1079,28 @@ def test_emit_task_begin_events_io_node(disk_to_device_channel: Channel) -> None DataFrameScan, task, quent_ir_execution_context, + input_frames_bytes=0, ) - events = _drained_events(logger) - task_events = [event for event in events if "Task" in event["data"]] - assert [event["data"]["Task"]["seq"] for event in task_events] == [0, 2] - assert "Queueing" in task_events[0]["data"]["Task"]["state"] - loading = task_events[1]["data"]["Task"]["state"]["Loading"] - assert loading["use_fs_to_mem"]["resource_id"] == str(disk_to_device_channel.id) - assert loading["use_memory"]["resource_id"] == str( - quent_ir_execution_context.device_memory.id - ) - - -def test_emit_task_end_events() -> None: - logger, quent_ir_execution_context = _make_quent_ir_execution_context() - task = Task(operator_id=quent_ir_execution_context.quent_operator.id) - input_frame = _make_dataframe(pl.DataFrame({"x": [1, 2, 3]})) + # Simulate the result result = _make_dataframe(pl.DataFrame({"y": list(range(7))})) - quent_ir_execution_context.context._emit_task_end_events( + DataFrameScan, task, quent_ir_execution_context, - [input_frame], + [], result, ) events = _drained_events(logger) - operator_events = [event for event in events if "Operator" in event["data"]] + # queueing -> allocating -> loading -> computing -> exit task_events = [event for event in events if "Task" in event["data"]] - assert len(operator_events) == 1 - stats = operator_events[0]["data"]["Operator"]["Statistics"]["custom_attributes"] - assert stats == [ - {"key": "input_bytes", "value": {"U64": input_frame._size_bytes()}}, - {"key": "output_bytes", "value": {"U64": result._size_bytes()}}, - {"key": "output_rows", "value": {"U64": result.num_rows}}, - ] - assert task_events[0]["data"]["Task"]["state"] == "Exit" + assert [event["data"]["Task"]["seq"] for event in task_events] == [0, 1, 2, 3, 4] + assert "Queueing" in task_events[0]["data"]["Task"]["state"] + assert "Allocating" in task_events[1]["data"]["Task"]["state"] + assert "Loading" in task_events[2]["data"]["Task"]["state"] + assert "Computing" in task_events[3]["data"]["Task"]["state"] + assert "Exit" in task_events[4]["data"]["Task"]["state"] def test_emit_task_end_events_on_failure() -> None: @@ -1107,6 +1109,7 @@ def test_emit_task_end_events_on_failure() -> None: input_frame = _make_dataframe(pl.DataFrame({"x": [10, 20]})) quent_ir_execution_context.context._emit_task_end_events( + Filter, task, quent_ir_execution_context, [input_frame], From 740abaf3389ce45fb6d4792376bc6eb39ca41740 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Thu, 16 Jul 2026 13:50:15 -0700 Subject: [PATCH 27/60] pass through --- python/cudf_polars/cudf_polars/quent/_context.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index b9c692ace1ab..fe80846b35de 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -427,6 +427,7 @@ def _emit_task_begin_events( quent_task.computing( use_thread=quent_processor, use_memory=quent_ir_execution_context.device_memory, + input_bytes=input_frames_bytes, memory_capacity_bytes=input_frames_bytes, ) ) From 615b082d5f505cc439b27e2065835304fc3b9562 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 17 Jul 2026 09:40:51 -0700 Subject: [PATCH 28/60] Track Operator Statistics via Tracer --- python/cudf_polars/cudf_polars/dsl/ir.py | 4 +++ python/cudf_polars/cudf_polars/dsl/tracing.py | 9 ++++-- .../cudf_polars/cudf_polars/quent/_context.py | 21 -------------- .../streaming/actor_graph/tracing.py | 28 +++++++------------ .../streaming/actor_graph/utils.py | 14 ++++++++-- python/cudf_polars/tests/quent/test_quent.py | 25 ----------------- 6 files changed, 32 insertions(+), 69 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index b67c19119a38..70412785b299 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -85,6 +85,7 @@ from cudf_polars.containers.dataframe import NamedColumn from cudf_polars.dsl.utils.io import CachedParquetInfo from cudf_polars.quent._context import QuentIRExecutionContext + from cudf_polars.streaming.actor_graph.tracing import ActorTracer from cudf_polars.streaming.rank_aware_source import RankAwareSource from cudf_polars.typing import CSECache, ClosedInterval, Schema, Slice as Zlice from cudf_polars.utils.config import ParquetOptions @@ -141,12 +142,15 @@ class IRExecutionContext: Identifier for the query being executed. quent_ir_execution_context Optional Quent tracing context bound to a physical operator. + tracer + The actor tracer. """ py_executor: concurrent.futures.ThreadPoolExecutor | None = field(default=None) get_cuda_stream: Callable[[], Stream] = field(default=get_cuda_stream) query_id: uuid.UUID = field(default_factory=uuid.uuid4) quent_ir_execution_context: QuentIRExecutionContext | None = None + tracer: ActorTracer | None = None async def to_thread( self, func: Callable[P, T], /, *args: P.args, **kwargs: P.kwargs diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index 4486c6a9093f..a2a941fe13d7 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -224,7 +224,6 @@ def wrapper( cls, quent_task, ir_execution_context.quent_ir_execution_context, - frames, None, ) raise @@ -233,11 +232,11 @@ def wrapper( quent_task is not None and ir_execution_context.quent_ir_execution_context is not None ): + # TODO: This should emit some Chunk-level statistics (duration, rows, bytes, schema, etc.) ir_execution_context.quent_ir_execution_context.context._emit_task_end_events( cls, quent_task, ir_execution_context.quent_ir_execution_context, - frames, result, ) stop = time.monotonic_ns() @@ -263,6 +262,12 @@ def wrapper( ) log.info("Execute IR", **record) + # Update the tracer with all the relevant information. + if (tracer := ir_execution_context.tracer) is not None: + # ActorTracer.send updates row_count and chunk_count + tracer.input_bytes += sum(frame._size_bytes() for frame in frames) + tracer.output_bytes += result._size_bytes() + return result return wrapper diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index fe80846b35de..fd8ebca9d95a 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -25,12 +25,10 @@ Processor, Query, QueryGroup, - Statistics, ThreadPool, ) if TYPE_CHECKING: - from collections.abc import Sequence from typing import Self from cudf_polars.containers import DataFrame @@ -437,7 +435,6 @@ def _emit_task_end_events( ir_type: type[IR], quent_task: Task, quent_ir_execution_context: QuentIRExecutionContext, - frames: Sequence[DataFrame], result: DataFrame | None, ) -> None: """ @@ -470,32 +467,14 @@ def _emit_task_end_events( :meth:`_emit_task_begin_events`. I/O nodes were left in ``Allocating`` during the load; here they transition through ``Loading`` (with byte counts), ``Computing``, and ``Exit``. - - A ``Statistics`` record, associated with the Quent Operator bound to the IR execution context, - is also emitted. It includes - - - input bytes: the total size of the input dataframes. - - output bytes: the size of the output dataframe. - - output rows: the number of rows in the output dataframe. """ if quent_ir_execution_context is None: # pragma: no cover; return if result is not None: - output_rows = result.num_rows output_capacity_bytes = result._size_bytes() else: - output_rows = 0 output_capacity_bytes = 0 - quent_ir_execution_context.logger.emit( - quent_ir_execution_context.quent_operator.statistics( - statistics=Statistics( - input_bytes=sum(frame._size_bytes() for frame in frames), - output_bytes=output_capacity_bytes, - output_rows=output_rows, - ) - ) - ) if ir_type.is_io_node: quent_processor = quent_ir_execution_context.get_or_declare_processor( thread_ident=threading.get_ident(), diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py index c318affa4d2a..5decd5681653 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py @@ -21,6 +21,7 @@ from cudf_polars.utils.config import ConfigOptions +@dataclasses.dataclass(slots=True) class ActorTracer: """ Tracer for a single streaming actor (IR node). @@ -46,24 +47,15 @@ class ActorTracer: (e.g., after an allgather). Affects how rows are merged. """ - __slots__ = ( - "chunk_count", - "decision", - "duplicated", - "extra", - "ir_id", - "ir_type", - "row_count", - ) - - def __init__(self, ir_id: int | None = None, ir_type: str | None = None) -> None: - self.ir_id = ir_id - self.ir_type = ir_type - self.row_count: int | None = None - self.chunk_count: int = 0 - self.decision: str | None = None - self.duplicated: bool = False - self.extra: dict[str, Any] = {} + ir_id: int | None = None + ir_type: str | None = None + row_count: int | None = None + chunk_count: int = 0 + input_bytes: int = 0 + output_bytes: int = 0 + decision: str | None = None + duplicated: bool = False + extra: dict[str, Any] = dataclasses.field(default_factory=dict) def add_chunk(self, *, chunk: TableChunk | None = None) -> None: """ diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py index 99ad2ae07421..856320ca7586 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -271,9 +271,11 @@ async def shutdown_on_error( ir_type = type(trace_ir).__name__ tracer = ActorTracer(ir_id, ir_type) contextvars = {"actor_ir_id": ir_id, "actor_ir_type": ir_type} + import dataclasses if ir_context is not None: contextvars["cudf_polars_query_id"] = str(ir_context.query_id) + ir_context = dataclasses.replace(ir_context, tracer=tracer) with cudf_polars.dsl.tracing.bound_contextvars(**contextvars): start = time.monotonic_ns() @@ -349,8 +351,8 @@ async def shutdown_on_error( quent_ir_execution_context.quent_operator.statistics( statistics=cudf_polars.quent._types.Statistics( output_rows=output_rows, - input_bytes=0, - output_bytes=0, + input_bytes=tracer.input_bytes, + output_bytes=tracer.output_bytes, custom_attributes=custom_attributes, ) ) @@ -778,6 +780,8 @@ def _evaluate_chunk_sync( The IR execution context. br The buffer resource for lifetime tracking. + tracer + The actor tracer. Returns ------- @@ -831,7 +835,11 @@ async def evaluate_chunk( with opaque_memory_usage(extra): for single_ir in irs: chunk = await ir_context.to_thread( - _evaluate_chunk_sync, chunk, single_ir, ir_context, context.br() + _evaluate_chunk_sync, + chunk, + single_ir, + ir_context, + context.br(), ) return chunk diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index 187c965a9536..0a137e48dafd 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -1052,7 +1052,6 @@ def test_emit_task_events_computing_node() -> None: Filter, task, quent_ir_execution_context, - [], result, ) @@ -1088,7 +1087,6 @@ def test_emit_task_events_io_node(disk_to_device_channel: Channel) -> None: DataFrameScan, task, quent_ir_execution_context, - [], result, ) @@ -1101,26 +1099,3 @@ def test_emit_task_events_io_node(disk_to_device_channel: Channel) -> None: assert "Loading" in task_events[2]["data"]["Task"]["state"] assert "Computing" in task_events[3]["data"]["Task"]["state"] assert "Exit" in task_events[4]["data"]["Task"]["state"] - - -def test_emit_task_end_events_on_failure() -> None: - logger, quent_ir_execution_context = _make_quent_ir_execution_context() - task = Task(operator_id=quent_ir_execution_context.quent_operator.id) - input_frame = _make_dataframe(pl.DataFrame({"x": [10, 20]})) - - quent_ir_execution_context.context._emit_task_end_events( - Filter, - task, - quent_ir_execution_context, - [input_frame], - None, - ) - - events = _drained_events(logger) - operator_events = [event for event in events if "Operator" in event["data"]] - stats = operator_events[0]["data"]["Operator"]["Statistics"]["custom_attributes"] - assert stats == [ - {"key": "input_bytes", "value": {"U64": input_frame._size_bytes()}}, - {"key": "output_bytes", "value": {"U64": 0}}, - {"key": "output_rows", "value": {"U64": 0}}, - ] From e42355450e75b0c1b0937df95415fcf4f074ffe2 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 17 Jul 2026 10:47:07 -0700 Subject: [PATCH 29/60] Set tracer in IRExecutionContext --- python/cudf_polars/cudf_polars/dsl/tracing.py | 4 +++- .../cudf_polars/streaming/actor_graph/io.py | 13 +++++++++++-- .../cudf_polars/streaming/actor_graph/join.py | 5 +++++ .../cudf_polars/streaming/actor_graph/nodes.py | 3 +++ .../cudf_polars/streaming/actor_graph/utils.py | 16 ++++++++-------- 5 files changed, 30 insertions(+), 11 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index a2a941fe13d7..b317108d69ad 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -184,6 +184,7 @@ def wrapper( # And the kwonly 'context' argument has the IR execution context. ir_execution_context: IRExecutionContext = kwargs["context"] # type: ignore[assignment] + # assert ir_execution_context.tracer is not None if ir_execution_context.quent_ir_execution_context is not None: quent_task = Task.from_ir( cls, ir_execution_context.quent_ir_execution_context @@ -262,11 +263,12 @@ def wrapper( ) log.info("Execute IR", **record) - # Update the tracer with all the relevant information. if (tracer := ir_execution_context.tracer) is not None: # ActorTracer.send updates row_count and chunk_count tracer.input_bytes += sum(frame._size_bytes() for frame in frames) tracer.output_bytes += result._size_bytes() + # TODO: Emit some chunk-level statistics. Might need a custom Quent + # schema for that. return result diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py index 567c60036f2b..60064d291d58 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio +import dataclasses import functools import io import math @@ -281,7 +282,9 @@ async def _producer(producer_id: int, ch_out: Channel) -> None: await ch_out.drain(context) async with ( - shutdown_on_error(context, *lineariser.input_channels, trace_ir=ir), + shutdown_on_error( + context, *lineariser.input_channels, trace_ir=ir, ir_context=ir_context + ), ): await gather_in_task_group( lineariser.drain(), @@ -536,6 +539,9 @@ async def read_chunk( context, size=estimated_chunk_bytes, net_memory_delta=estimated_chunk_bytes ) ): + if ir_context.tracer is None: + ir_context = dataclasses.replace(ir_context, tracer=tracer) + assert ir_context.tracer is not None df = await ir_context.to_thread( scan.do_evaluate, *scan._non_child_args, @@ -585,6 +591,7 @@ async def scan_node( context, ch_out, trace_ir=ir, ir_context=ir_context ) as tracer: # Send basic metadata + ir_context = dataclasses.replace(ir_context, tracer=tracer) await send_metadata( ch_out, context, @@ -639,7 +646,9 @@ async def _producer(producer_id: int, ch_out: Channel) -> None: await ch_out.drain(context) async with ( - shutdown_on_error(context, *lineariser.input_channels, trace_ir=ir), + shutdown_on_error( + context, *lineariser.input_channels, trace_ir=ir, ir_context=ir_context + ), ): await gather_in_task_group( lineariser.drain(), diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py index 5773ffa1bce7..7844ed3975b4 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -282,6 +282,11 @@ async def _broadcast_join_large_chunk( join_results: list[DataFrame] = [] input_bytes = large_chunk_size + small_size + + import dataclasses + + ir_context = dataclasses.replace(ir_context, tracer=tracer) + with opaque_memory_usage( await reserve_memory(context, size=input_bytes, net_memory_delta=0) ): diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py index d1741482af73..b2c74e0a1006 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py @@ -87,6 +87,9 @@ async def default_node_single( ), duplicated=metadata_in.duplicated, ) + import dataclasses + + ir_context = dataclasses.replace(ir_context, tracer=tracer) # Process chunks (handle empty input for aggregation-like operations) await chunkwise_evaluate( diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py index 856320ca7586..09039d1d73c0 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -347,17 +347,17 @@ async def shutdown_on_error( output_rows = 0 else: output_rows = tracer.row_count - quent_ir_execution_context.logger.emit( - quent_ir_execution_context.quent_operator.statistics( - statistics=cudf_polars.quent._types.Statistics( - output_rows=output_rows, - input_bytes=tracer.input_bytes, - output_bytes=tracer.output_bytes, - custom_attributes=custom_attributes, - ) + stats = quent_ir_execution_context.quent_operator.statistics( + statistics=cudf_polars.quent._types.Statistics( + output_rows=output_rows, + input_bytes=tracer.input_bytes, + output_bytes=tracer.output_bytes, + custom_attributes=custom_attributes, ) ) + quent_ir_execution_context.logger.emit(stats) + def _update_ordering_indices( ordering: Ordering, new_indices: tuple[int, ...] From ba69bf0d9474265a45b39666bb4dd2478b3b0d56 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 17 Jul 2026 11:05:59 -0700 Subject: [PATCH 30/60] Ignore new format --- .gitignore | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index cb4b70da0503..21de3f8d9c74 100644 --- a/.gitignore +++ b/.gitignore @@ -189,4 +189,5 @@ rmm_log.txt python/cudf/cudf_pandas_tests/data/rmm_log.txt # Quent traces -logs/*.ndjson +logs/**/*.ndjson +logs/**/*.qmi From af692d899a9abfa28e7f0dd405e5699e1b3ba75c Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 17 Jul 2026 11:30:19 -0700 Subject: [PATCH 31/60] remove io_loading_at hck --- python/cudf_polars/cudf_polars/dsl/ir.py | 2 +- python/cudf_polars/cudf_polars/quent/_context.py | 4 ---- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index e370eeff3f22..c5bd73f6b68c 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -143,7 +143,7 @@ class IRExecutionContext: quent_ir_execution_context Optional Quent tracing context bound to a physical operator. tracer - The actor tracer. + The actor tracer. Used to propagate statistics. """ py_executor: concurrent.futures.ThreadPoolExecutor | None = field(default=None) diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index fd8ebca9d95a..4a4007783699 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -119,10 +119,6 @@ class QuentContext: def __post_init__(self) -> None: object.__setattr__(self, "_query_group_cache_", set()) - # Maps task id -> unix timestamp when an I/O task entered Allocating. - # Used to backdate the Loading transition so channel rate spans cover - # the full disk-to-device transfer duration. - object.__setattr__(self, "_io_loading_started_at_", {}) def serialize(self) -> bytes: """ From 908f785209a7defc4ad7924f973f3f98dfd475f1 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 17 Jul 2026 11:31:26 -0700 Subject: [PATCH 32/60] query_for docs --- python/cudf_polars/cudf_polars/quent/_context.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index 4a4007783699..54880a8c88d3 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -214,11 +214,15 @@ def query_for(self, query_id: uuid.UUID) -> Query: """ Build a per-collect Quent Query with a unique id. - The engine-scoped ``QuentContext`` is reused across many - ``.collect()`` calls, so each collect must derive its own - :class:`Query` (identified by the per-collect ``query_id``) rather - than reusing the shared ``self.query``. The ``instance_name`` from - the template ``self.query`` is preserved. + Parameters + ---------- + query_id: uuid.UUID + The unique ID for the query. + + Returns + ------- + A new Quent Query with the given ID and the same instance name as the + engine-scoped query. """ return Query(id=query_id, instance_name=self.query.instance_name) From 83d5f9dcf7f1f32c054bad7c680114b36d977cc1 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 17 Jul 2026 11:37:25 -0700 Subject: [PATCH 33/60] docstrings --- .../cudf_polars/cudf_polars/quent/_context.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index 54880a8c88d3..10f27311ec75 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -399,7 +399,7 @@ def _emit_task_begin_events( This emits the following events: - queueing - - allocating + - allocating (with the Quent Processor for the current thread) - loading (I/O nodes only) - computing (non-I/O nodes only) """ @@ -440,6 +440,9 @@ def _emit_task_end_events( """ Build and emit Quent events for the end of an IR node's evaluation. + The timestamp here represents when the IR node completed **host**-side + processing. Work work may be happening asynchronously on the GPU. + Parameters ---------- ir_type: type[IR] @@ -459,14 +462,10 @@ def _emit_task_end_events( Notes ----- - This method emits an ``Exit`` event for the Quent Task, whose timestamp represents - when the IR node completed host-side processing. - - The Quent Task FSM only permits a task to exit from the ``Computing`` - state. Non-I/O nodes already entered ``Computing`` in - :meth:`_emit_task_begin_events`. I/O nodes were left in ``Allocating`` - during the load; here they transition through ``Loading`` (with byte - counts), ``Computing``, and ``Exit``. + This emits the following events: + + - computing (I/O nodes only) + - exit """ if quent_ir_execution_context is None: # pragma: no cover; return From b098c603693430f2a3cd585ba4c74fb24b420d45 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 17 Jul 2026 11:47:09 -0700 Subject: [PATCH 34/60] Reuse get device memory --- python/cudf_polars/cudf_polars/quent/_context.py | 15 ++------------- 1 file changed, 2 insertions(+), 13 deletions(-) diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index 10f27311ec75..8c52b5aad2ed 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -27,6 +27,7 @@ QueryGroup, ThreadPool, ) +from cudf_polars.utils.config import get_total_device_memory if TYPE_CHECKING: from typing import Self @@ -488,18 +489,6 @@ def _emit_task_end_events( quent_ir_execution_context.logger.emit(quent_task.exit()) -def _device_memory_capacity_bytes() -> int | None: - """Return total device memory in bytes, or ``None`` if unavailable.""" - try: - import pynvml - - pynvml.nvmlInit() - handle = pynvml.nvmlDeviceGetHandleByIndex(0) - return int(pynvml.nvmlDeviceGetMemoryInfo(handle).total) - except Exception: # pragma: no cover - no GPU / pynvml unavailable - return None - - def declare_worker_resources( logger: QuentLogger, *, @@ -530,7 +519,7 @@ def declare_worker_resources( target=device_memory, ) thread_pool = ThreadPool(worker_id=worker_id) - device_memory_capacity = _device_memory_capacity_bytes() + device_memory_capacity = get_total_device_memory() or 0 logger.emit(device_memory.initializing()) logger.emit(device_memory.operating(device_memory_capacity)) logger.emit(filesystem.initializing()) From 42d0429fe01aef34e38b9d1bd7f1001e1921216c Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 20 Jul 2026 07:44:02 -0700 Subject: [PATCH 35/60] ir_context setting --- .../cudf_polars/cudf_polars/streaming/actor_graph/join.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py index 3dbe27c5c10f..df8811030fc0 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -4,7 +4,7 @@ from __future__ import annotations -from dataclasses import asdict, dataclass +from dataclasses import asdict, dataclass, replace from typing import TYPE_CHECKING, Any, Literal import pylibcudf as plc @@ -175,6 +175,7 @@ async def broadcast_join_actor( trace_ir=ir, ir_context=ir_context, ) as tracer: + ir_context = replace(ir_context, tracer=tracer) await _broadcast_join( context, comm, @@ -293,10 +294,6 @@ async def _broadcast_join_large_chunk( join_results: list[DataFrame] = [] input_bytes = large_chunk_size + small_size - import dataclasses - - ir_context = dataclasses.replace(ir_context, tracer=tracer) - with opaque_memory_usage( await reserve_memory(context, size=input_bytes, net_memory_delta=0) ): @@ -1364,6 +1361,7 @@ async def join_actor( trace_ir=ir, ir_context=ir_context, ) as tracer: + ir_context = replace(ir_context, tracer=tracer) left_metadata, right_metadata = await gather_in_task_group( recv_metadata(ch_left, context), recv_metadata(ch_right, context), From 90edf423c0b19dcdc75f87ed81bb08b7026746d5 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 20 Jul 2026 14:35:08 -0700 Subject: [PATCH 36/60] task/chunk level stats --- python/cudf_polars/cudf_polars/dsl/tracing.py | 2 ++ python/cudf_polars/cudf_polars/quent/_context.py | 16 ++++++++++++++++ python/cudf_polars/cudf_polars/quent/_types.py | 14 ++++++++++++++ python/cudf_polars/tests/quent/test_quent.py | 8 ++++++-- 4 files changed, 38 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index b317108d69ad..1b29be29e669 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -225,6 +225,7 @@ def wrapper( cls, quent_task, ir_execution_context.quent_ir_execution_context, + frames, None, ) raise @@ -238,6 +239,7 @@ def wrapper( cls, quent_task, ir_execution_context.quent_ir_execution_context, + frames, result, ) stop = time.monotonic_ns() diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index 8c52b5aad2ed..0a4b4ec51ffa 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -25,6 +25,7 @@ Processor, Query, QueryGroup, + Statistics, ThreadPool, ) from cudf_polars.utils.config import get_total_device_memory @@ -436,6 +437,7 @@ def _emit_task_end_events( ir_type: type[IR], quent_task: Task, quent_ir_execution_context: QuentIRExecutionContext, + frames: list[DataFrame], result: DataFrame | None, ) -> None: """ @@ -473,8 +475,10 @@ def _emit_task_end_events( if result is not None: output_capacity_bytes = result._size_bytes() + output_rows = result.num_rows else: output_capacity_bytes = 0 + output_rows = 0 if ir_type.is_io_node: quent_processor = quent_ir_execution_context.get_or_declare_processor( thread_ident=threading.get_ident(), @@ -488,6 +492,18 @@ def _emit_task_end_events( ) quent_ir_execution_context.logger.emit(quent_task.exit()) + # We emit a Statistics event here, despite knowing it'll be overwritten + # by the node-level Statistics later. + quent_ir_execution_context.logger.emit( + quent_task.statistics( + statistics=Statistics( + input_bytes=sum(frame._size_bytes() for frame in frames), + output_bytes=output_capacity_bytes, + output_rows=output_rows, + ), + ) + ) + def declare_worker_resources( logger: QuentLogger, diff --git a/python/cudf_polars/cudf_polars/quent/_types.py b/python/cudf_polars/cudf_polars/quent/_types.py index 144ee126ef6d..f7e4ef696d57 100644 --- a/python/cudf_polars/cudf_polars/quent/_types.py +++ b/python/cudf_polars/cudf_polars/quent/_types.py @@ -1017,3 +1017,17 @@ def exit(self, timestamp: int | None = None) -> Event: timestamp=timestamp if timestamp is not None else time.time_ns(), data={EventName.TASK.value: {"seq": next(self._seq), "state": "Exit"}}, ) + + def statistics(self, statistics: Statistics, timestamp: int | None = None) -> Event: + """Build a Quent Task Statistics event.""" + return Event( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={ + EventName.TASK.value: { + "seq": next(self._seq), + "state": "Statistics", + "data": statistics.to_dict(), + } + }, + ) diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index 0a137e48dafd..47d58790a216 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -1052,17 +1052,19 @@ def test_emit_task_events_computing_node() -> None: Filter, task, quent_ir_execution_context, + [], result, ) events = _drained_events(logger) task_events = [event for event in events if "Task" in event["data"]] # queueing -> allocating -> computing -> exit - assert [event["data"]["Task"]["seq"] for event in task_events] == [0, 1, 2, 3] + assert [event["data"]["Task"]["seq"] for event in task_events] == [0, 1, 2, 3, 4] assert "Queueing" in task_events[0]["data"]["Task"]["state"] assert "Allocating" in task_events[1]["data"]["Task"]["state"] assert "Computing" in task_events[2]["data"]["Task"]["state"] assert "Exit" in task_events[3]["data"]["Task"]["state"] + assert "Statistics" in task_events[4]["data"]["Task"]["state"] processor_events = [event for event in events if "Processor" in event["data"]] assert len(processor_events) == 2 @@ -1087,15 +1089,17 @@ def test_emit_task_events_io_node(disk_to_device_channel: Channel) -> None: DataFrameScan, task, quent_ir_execution_context, + [], result, ) events = _drained_events(logger) # queueing -> allocating -> loading -> computing -> exit task_events = [event for event in events if "Task" in event["data"]] - assert [event["data"]["Task"]["seq"] for event in task_events] == [0, 1, 2, 3, 4] + assert [event["data"]["Task"]["seq"] for event in task_events] == [0, 1, 2, 3, 4, 5] assert "Queueing" in task_events[0]["data"]["Task"]["state"] assert "Allocating" in task_events[1]["data"]["Task"]["state"] assert "Loading" in task_events[2]["data"]["Task"]["state"] assert "Computing" in task_events[3]["data"]["Task"]["state"] assert "Exit" in task_events[4]["data"]["Task"]["state"] + assert "Statistics" in task_events[5]["data"]["Task"]["state"] From 9e72712319b561be51c44a127f05d04fd7315d8c Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 21 Jul 2026 10:17:49 -0700 Subject: [PATCH 37/60] merge the conditions --- python/cudf_polars/cudf_polars/engine/ray.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index fbee5e04f271..c7c4c57e2fce 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -794,15 +794,11 @@ def __init__( quent_context: QuentContext | None = executor_options.get("quent_context") if quent_context is not None: self._quent_logger = cudf_polars.quent._logging.QuentLogger() - else: - self._quent_logger = None - - if quent_context is not None: executor_options.setdefault("quent_context", quent_context) - assert self._quent_logger is not None quent_context._emit_engine_init_events(self._quent_logger) engine = quent_context.engine else: + self._quent_logger = None engine = cudf_polars.quent.Engine(id=uuid.uuid4()) # This engine's store uid, used to key its partitions in each actor's process rank-local store. From f535f70785281a9c599c4859e7d971b3fa4b6756 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 22 Jul 2026 05:00:12 -0700 Subject: [PATCH 38/60] instance name --- python/cudf_polars/cudf_polars/quent/_plan.py | 1 - python/cudf_polars/cudf_polars/quent/_types.py | 3 +-- python/cudf_polars/tests/quent/test_quent.py | 1 - 3 files changed, 1 insertion(+), 4 deletions(-) diff --git a/python/cudf_polars/cudf_polars/quent/_plan.py b/python/cudf_polars/cudf_polars/quent/_plan.py index 542b525c2f0d..88f5db929b48 100644 --- a/python/cudf_polars/cudf_polars/quent/_plan.py +++ b/python/cudf_polars/cudf_polars/quent/_plan.py @@ -95,7 +95,6 @@ def build_plan( id=operator_id, plan=plan, parent_operators=parent_ops.get(node_id, []), - instance_name=serializable_node.type, type_name=serializable_node.type, custom_attributes=custom_attributes, ) diff --git a/python/cudf_polars/cudf_polars/quent/_types.py b/python/cudf_polars/cudf_polars/quent/_types.py index f7e4ef696d57..eefe0b3dfdfc 100644 --- a/python/cudf_polars/cudf_polars/quent/_types.py +++ b/python/cudf_polars/cudf_polars/quent/_types.py @@ -149,7 +149,6 @@ class Operator: id: uuid.UUID plan: Plan parent_operators: list[Operator] - instance_name: str type_name: str custom_attributes: list[Attribute] = dataclasses.field(default_factory=list) @@ -166,7 +165,7 @@ def to_dict(self) -> dict[str, Any]: "parent_operator_ids": [ str(operator.id) for operator in self.parent_operators ], - "instance_name": self.instance_name, + "instance_name": self.id.hex[:8], "type_name": self.type_name, "custom_attributes": [attr.serialize() for attr in self.custom_attributes], } diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index 47d58790a216..e1a3fe902abb 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -97,7 +97,6 @@ def _make_quent_ir_execution_context( id=operator_id, plan=plan, parent_operators=[], - instance_name="Filter-NodeIndex(0)", type_name="Filter", ) local_context = LocalQuentContext( From bc352122657ab5945d13d8fc3e505db94e638ef4 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Thu, 23 Jul 2026 14:57:36 -0700 Subject: [PATCH 39/60] fixups --- .../cudf_polars/cudf_polars/quent/_context.py | 17 +++-------------- python/cudf_polars/cudf_polars/quent/_types.py | 14 -------------- python/cudf_polars/tests/quent/test_quent.py | 6 ++---- 3 files changed, 5 insertions(+), 32 deletions(-) diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index 0a4b4ec51ffa..e7ccd3cb237a 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -25,7 +25,6 @@ Processor, Query, QueryGroup, - Statistics, ThreadPool, ) from cudf_polars.utils.config import get_total_device_memory @@ -475,10 +474,8 @@ def _emit_task_end_events( if result is not None: output_capacity_bytes = result._size_bytes() - output_rows = result.num_rows else: output_capacity_bytes = 0 - output_rows = 0 if ir_type.is_io_node: quent_processor = quent_ir_execution_context.get_or_declare_processor( thread_ident=threading.get_ident(), @@ -492,17 +489,9 @@ def _emit_task_end_events( ) quent_ir_execution_context.logger.emit(quent_task.exit()) - # We emit a Statistics event here, despite knowing it'll be overwritten - # by the node-level Statistics later. - quent_ir_execution_context.logger.emit( - quent_task.statistics( - statistics=Statistics( - input_bytes=sum(frame._size_bytes() for frame in frames), - output_bytes=output_capacity_bytes, - output_rows=output_rows, - ), - ) - ) + # TODO: Figure out how to emit some chunk/task-level statistics. + # We can't do it directly on the Task object, because that (seems to) + # break operator-level aggregation like duration_s. def declare_worker_resources( diff --git a/python/cudf_polars/cudf_polars/quent/_types.py b/python/cudf_polars/cudf_polars/quent/_types.py index eefe0b3dfdfc..3c337a41a62b 100644 --- a/python/cudf_polars/cudf_polars/quent/_types.py +++ b/python/cudf_polars/cudf_polars/quent/_types.py @@ -1016,17 +1016,3 @@ def exit(self, timestamp: int | None = None) -> Event: timestamp=timestamp if timestamp is not None else time.time_ns(), data={EventName.TASK.value: {"seq": next(self._seq), "state": "Exit"}}, ) - - def statistics(self, statistics: Statistics, timestamp: int | None = None) -> Event: - """Build a Quent Task Statistics event.""" - return Event( - id=self.id, - timestamp=timestamp if timestamp is not None else time.time_ns(), - data={ - EventName.TASK.value: { - "seq": next(self._seq), - "state": "Statistics", - "data": statistics.to_dict(), - } - }, - ) diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index 50e7302f2947..eb4d53f2d32b 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -1057,12 +1057,11 @@ def test_emit_task_events_computing_node() -> None: events = _drained_events(logger) task_events = [event for event in events if "Task" in event["data"]] # queueing -> allocating -> computing -> exit - assert [event["data"]["Task"]["seq"] for event in task_events] == [0, 1, 2, 3, 4] + assert [event["data"]["Task"]["seq"] for event in task_events] == [0, 1, 2, 3] assert "Queueing" in task_events[0]["data"]["Task"]["state"] assert "Allocating" in task_events[1]["data"]["Task"]["state"] assert "Computing" in task_events[2]["data"]["Task"]["state"] assert "Exit" in task_events[3]["data"]["Task"]["state"] - assert "Statistics" in task_events[4]["data"]["Task"]["state"] processor_events = [event for event in events if "Processor" in event["data"]] assert len(processor_events) == 2 @@ -1094,10 +1093,9 @@ def test_emit_task_events_io_node(disk_to_device_channel: Channel) -> None: events = _drained_events(logger) # queueing -> allocating -> loading -> computing -> exit task_events = [event for event in events if "Task" in event["data"]] - assert [event["data"]["Task"]["seq"] for event in task_events] == [0, 1, 2, 3, 4, 5] + assert [event["data"]["Task"]["seq"] for event in task_events] == [0, 1, 2, 3, 4] assert "Queueing" in task_events[0]["data"]["Task"]["state"] assert "Allocating" in task_events[1]["data"]["Task"]["state"] assert "Loading" in task_events[2]["data"]["Task"]["state"] assert "Computing" in task_events[3]["data"]["Task"]["state"] assert "Exit" in task_events[4]["data"]["Task"]["state"] - assert "Statistics" in task_events[5]["data"]["Task"]["state"] From f1f047e720b51d2a5db686bed6da0dfb94303cbc Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 24 Jul 2026 08:45:30 -0700 Subject: [PATCH 40/60] cleanpu --- python/cudf_polars/cudf_polars/dsl/tracing.py | 4 +--- python/cudf_polars/cudf_polars/engine/core.py | 1 - 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index 1b29be29e669..90888996c4d0 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -184,7 +184,7 @@ def wrapper( # And the kwonly 'context' argument has the IR execution context. ir_execution_context: IRExecutionContext = kwargs["context"] # type: ignore[assignment] - # assert ir_execution_context.tracer is not None + if ir_execution_context.quent_ir_execution_context is not None: quent_task = Task.from_ir( cls, ir_execution_context.quent_ir_execution_context @@ -269,8 +269,6 @@ def wrapper( # ActorTracer.send updates row_count and chunk_count tracer.input_bytes += sum(frame._size_bytes() for frame in frames) tracer.output_bytes += result._size_bytes() - # TODO: Emit some chunk-level statistics. Might need a custom Quent - # schema for that. return result diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 9a210a7f0a02..b26b400ef380 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -768,7 +768,6 @@ def evaluate_on_rank( if config_options.executor.quent_context is not None: assert local_quent_context is not None - # logical_plan_id = optimized.get_stable_plan_id() plan, ops, ports, logical_op_by_id = build_plan( optimized, config_options, From d933028036276c2c0bd1ccc4967df468cd640c99 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 24 Jul 2026 10:06:31 -0700 Subject: [PATCH 41/60] WorkerResources --- python/cudf_polars/cudf_polars/engine/dask.py | 78 ++----- python/cudf_polars/cudf_polars/engine/ray.py | 75 ++---- python/cudf_polars/cudf_polars/engine/spmd.py | 84 ++----- .../cudf_polars/cudf_polars/quent/_context.py | 218 +++++++++--------- .../cudf_polars/streaming/actor_graph/core.py | 12 +- .../cudf_polars/cudf_polars/utils/config.py | 27 +-- python/cudf_polars/tests/quent/test_quent.py | 156 ++++++------- 7 files changed, 238 insertions(+), 412 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 8d9c90775ad5..44646ad4e284 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -51,11 +51,7 @@ ) from cudf_polars.quent._context import ( LocalQuentContext, - ProcessorRegistry, - declare_network_channels, - declare_worker_resources, - finalize_network_channels, - finalize_worker_resources, + WorkerResources, ) from cudf_polars.unstable import unstable from cudf_polars.utils.config import DaskContext, MemoryResourceConfig @@ -139,14 +135,7 @@ class _WorkerContext: quent_worker: cudf_polars.quent._types.Worker statistics: Statistics mr: RmmResourceAdaptor | None = None - device_memory: cudf_polars.quent._types.Memory | None = None - disk_to_device_channel: cudf_polars.quent._types.Channel | None = None - thread_pool: cudf_polars.quent._types.ThreadPool | None = None - processor_registry: ProcessorRegistry | None = None - network: cudf_polars.quent._types.Network | None = None - link_channels: dict[int, cudf_polars.quent._types.Channel] = dataclasses.field( - default_factory=dict - ) + worker_resources: WorkerResources | None = None def _worker_evaluate_persisted( @@ -429,32 +418,18 @@ def _setup_worker( if quent_context is not None: quent_logger = cudf_polars.quent._logging.QuentLogger() - else: - quent_logger = None - - device_memory = None - disk_to_device_channel = None - thread_pool = None - processor_registry = None - network = None - link_channels: dict[int, cudf_polars.quent._types.Channel] = {} - if quent_logger is not None: - processor_registry = ProcessorRegistry() - device_memory, disk_to_device_channel, thread_pool = declare_worker_resources( - quent_logger, + worker_resources = WorkerResources.build( instance_suffix=f"rank-{comm.rank}", engine_id=engine_id, worker_id=worker_id, - ) - # Inter-rank network topology is engine-scoped: declare it once here - # alongside the other worker resources (a no-op for single-rank runs). - network, link_channels = declare_network_channels( - quent_logger, rank=comm.rank, nranks=comm.nranks, - engine_id=engine_id, - device_memory=device_memory, ) + quent_logger.emit(quent_worker._init()) + worker_resources.declare(quent_logger) + else: + quent_logger = None + worker_resources = None mp_ctx = _WorkerContext( comm=comm, @@ -464,17 +439,10 @@ def _setup_worker( mr=mr, quent_worker=quent_worker, quent_logger=quent_logger, - device_memory=device_memory, - disk_to_device_channel=disk_to_device_channel, - thread_pool=thread_pool, - processor_registry=processor_registry, - network=network, - link_channels=link_channels, + worker_resources=worker_resources, statistics=statistics, ) setattr(dask_worker, attr, mp_ctx) - if mp_ctx.quent_logger is not None: - mp_ctx.quent_logger.emit(quent_worker._init()) def _teardown_worker( @@ -498,20 +466,9 @@ def _teardown_worker( mp_ctx: _WorkerContext | None = getattr(dask_worker, attr, None) traces = [] if mp_ctx is not None: - if mp_ctx.quent_worker is not None and mp_ctx.quent_logger is not None: - if mp_ctx.processor_registry is not None: - mp_ctx.processor_registry._emit_processor_exit_events( - mp_ctx.quent_logger - ) - finalize_network_channels( - mp_ctx.quent_logger, link_channels=mp_ctx.link_channels - ) - if mp_ctx.device_memory is not None: - finalize_worker_resources( - mp_ctx.quent_logger, - device_memory=mp_ctx.device_memory, - disk_to_device_channel=mp_ctx.disk_to_device_channel, - ) + if mp_ctx.quent_logger is not None: + if mp_ctx.worker_resources is not None: + mp_ctx.worker_resources.finalize(mp_ctx.quent_logger) mp_ctx.quent_logger.emit(mp_ctx.quent_worker._exit()) traces = mp_ctx.quent_logger.drain() @@ -670,21 +627,14 @@ def _worker_evaluate( raise RuntimeError("_setup_worker must be called before _worker_evaluate") local_quent_context: LocalQuentContext | None = None if quent_context is not None: + assert mp_ctx.worker_resources is not None assert mp_ctx.quent_logger is not None - assert mp_ctx.device_memory is not None - assert mp_ctx.thread_pool is not None - assert mp_ctx.processor_registry is not None local_quent_context = LocalQuentContext( context=quent_context, query=quent_context.query_for(query_id), worker=mp_ctx.quent_worker, logger=mp_ctx.quent_logger, - thread_pool_id=mp_ctx.thread_pool.id, - processor_registry=mp_ctx.processor_registry, - device_memory=mp_ctx.device_memory, - disk_to_device_channel=mp_ctx.disk_to_device_channel, - network=mp_ctx.network, - link_channels=mp_ctx.link_channels, + worker_resources=mp_ctx.worker_resources, ) # evaluate_on_rank always collects metadata internally so we can read # metadata[-1].duplicated to decide whether to suppress this rank's output. diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index c7c4c57e2fce..eb9fe3bf3669 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -47,11 +47,7 @@ ) from cudf_polars.quent._context import ( LocalQuentContext, - ProcessorRegistry, - declare_network_channels, - declare_worker_resources, - finalize_network_channels, - finalize_worker_resources, + WorkerResources, ) from cudf_polars.quent._types import Worker from cudf_polars.unstable import unstable @@ -71,7 +67,6 @@ from cudf_polars.engine.options import StreamingOptions from cudf_polars.engine.persisted_result import PersistedQueryResult from cudf_polars.quent._context import QuentContext - from cudf_polars.quent._types import Channel from cudf_polars.streaming.parallel import ConfigOptions from cudf_polars.utils.config import StreamingExecutor @@ -282,32 +277,15 @@ def __init__( ) else: self._quent_logger = None + self._quent_engine = engine + self._worker_id = worker_id self._quent_worker = Worker( id=worker_id, engine=engine, instance_name=f"RankActor-{worker_id.hex[:8]}", ) - self._device_memory = None - self._disk_to_device_channel = None - self._quent_thread_pool = None - # Network topology is engine-scoped, but ``comm`` is not available until - # ``setup_worker``; declare it there and store the results here. - self._network: cudf_polars.quent._types.Network | None = None - self._link_channels: dict[int, Channel] = {} - self._processor_registry: ProcessorRegistry | None = None - if self._quent_logger is not None: - self._processor_registry = ProcessorRegistry() - self._quent_logger.emit(self._quent_worker._init()) - ( - self._device_memory, - self._disk_to_device_channel, - self._quent_thread_pool, - ) = declare_worker_resources( - self._quent_logger, - instance_suffix=f"RankActor-{worker_id.hex[:8]}", - engine_id=engine.id, - worker_id=worker_id, - ) + # Initialized later in setup_worker once ``comm`` is available. + self.worker_resources: WorkerResources | None = None def setup_root(self) -> bytes: """ @@ -358,14 +336,17 @@ def setup_worker(self, root_ucxx_address_as_bytes: bytes) -> None: barrier(self._comm) # Now that ``comm`` exists, declare the engine-scoped inter-rank network # topology once (a no-op for single-rank runs). - if self._quent_logger is not None and self._device_memory is not None: - self._network, self._link_channels = declare_network_channels( - self._quent_logger, + if self._quent_logger is not None: + self._quent_logger.emit(self._quent_worker._init()) + self.worker_resources = WorkerResources.build( + instance_suffix=f"RankActor-{self._quent_worker.id.hex[:8]}", + engine_id=self._quent_engine.id, + worker_id=self._worker_id, rank=self._comm.rank, - nranks=self._comm.nranks, - engine_id=self._quent_worker.engine.id, - device_memory=self._device_memory, + nranks=self._nranks, ) + self.worker_resources.declare(self._quent_logger) + assert self._base_mr is not None self._ctx = Context.from_options( self._comm.logger, @@ -426,20 +407,9 @@ def _exit(self) -> list[dict[str, Any]]: # Maybe generalize this to all application-level things, # followed by framework (ray) level things. if self._quent_worker is not None and self._quent_logger is not None: - if self._processor_registry is not None: - self._processor_registry._emit_processor_exit_events(self._quent_logger) - finalize_network_channels( - self._quent_logger, link_channels=self._link_channels - ) - if ( - self._device_memory is not None - and self._disk_to_device_channel is not None - ): - finalize_worker_resources( - self._quent_logger, - device_memory=self._device_memory, - disk_to_device_channel=self._disk_to_device_channel, - ) + if self.worker_resources is not None: + self.worker_resources.finalize(self._quent_logger) + self._quent_logger.emit(self._quent_worker._exit()) return self._drain_quent_events() return [] @@ -553,20 +523,13 @@ def evaluate_polars_ir( local_quent_context: LocalQuentContext | None = None if quent_context is not None: assert self._quent_logger is not None - assert self._device_memory is not None - assert self._quent_thread_pool is not None - assert self._processor_registry is not None + assert self.worker_resources is not None local_quent_context = LocalQuentContext( context=quent_context, query=quent_context.query_for(query_id), worker=self._quent_worker, logger=self._quent_logger, - thread_pool_id=self._quent_thread_pool.id, - processor_registry=self._processor_registry, - device_memory=self._device_memory, - disk_to_device_channel=self._disk_to_device_channel, - network=self._network, - link_channels=self._link_channels, + worker_resources=self.worker_resources, ) # evaluate_on_rank always collects metadata internally so we can read # metadata[-1].duplicated to decide whether to suppress this rank's diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index 809ab37a45c0..e894ca66e8d0 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -51,12 +51,8 @@ ) from cudf_polars.quent._context import ( LocalQuentContext, - ProcessorRegistry, QuentContext, - declare_network_channels, - declare_worker_resources, - finalize_network_channels, - finalize_worker_resources, + WorkerResources, ) from cudf_polars.quent._types import Worker from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id @@ -82,7 +78,6 @@ from cudf_polars.engine.core import T from cudf_polars.engine.options import StreamingOptions from cudf_polars.engine.persisted_result import PersistedQueryResult - from cudf_polars.quent._types import Channel, Memory, Network, ThreadPool from cudf_polars.streaming.parallel import ConfigOptions from cudf_polars.utils.config import StreamingExecutor @@ -138,13 +133,12 @@ def evaluate_pipeline_spmd_mode( if quent_context is not None: quent_logger = config_options.executor.spmd_context.quent_logger assert quent_logger is not None + assert spmd_context.worker_resources is not None + query = quent_context.query_for(query_id) quent_context._emit_query_group_events(quent_logger) quent_context._emit_query_events(quent_logger, query) worker_id = config_options.executor.spmd_context.worker_id - assert spmd_context.processor_registry is not None - assert spmd_context.thread_pool_id is not None - assert spmd_context.device_memory is not None local_quent_context = LocalQuentContext( context=quent_context, query=query, @@ -154,12 +148,7 @@ def evaluate_pipeline_spmd_mode( instance_name=f"rank-{comm.rank}", ), logger=quent_logger, - thread_pool_id=spmd_context.thread_pool_id, - processor_registry=spmd_context.processor_registry, - device_memory=spmd_context.device_memory, - disk_to_device_channel=spmd_context.disk_to_device_channel, - network=spmd_context.network, - link_channels=spmd_context.link_channels, + worker_resources=spmd_context.worker_resources, ) df, metadata = evaluate_on_rank( @@ -476,12 +465,6 @@ def __init__( self._comm: Communicator | None = comm self._ctx: Context | None = None self._py_executor: ThreadPoolExecutor | None = None - self._processor_registry: ProcessorRegistry | None = None - self._quent_thread_pool: ThreadPool | None = None - self._device_memory: Memory | None = None - self._disk_to_device_channel: Channel | None = None - self._network: Network | None = None - self._link_channels: dict[int, Channel] = {} self._store_uid = uuid.uuid4().hex exit_stack = contextlib.ExitStack() @@ -516,32 +499,21 @@ def __init__( instance_name=f"rank-{self.rank}", # relies on self.comm ) + worker_resources: WorkerResources | None = None if quent_context is not None: assert self._quent_logger is not None self._quent_logger.emit(self._quent_worker._init()) - self._processor_registry = ProcessorRegistry() - # Device memory, the disk->device channel, and the thread pool - # are all engine-scoped: they live for the lifetime of the - # engine's `py_executor`, not any single query. Declaring them - # once here (rather than per collect) means a single - # init/finalize pair per engine. - ( - self._device_memory, - self._disk_to_device_channel, - self._quent_thread_pool, - ) = declare_worker_resources( - self._quent_logger, + + worker_resources = WorkerResources.build( instance_suffix=f"rank-{self.rank}", engine_id=engine_id, worker_id=self._quent_worker.id, - ) - self._network, self._link_channels = declare_network_channels( - self._quent_logger, rank=comm.rank, nranks=comm.nranks, - engine_id=engine_id, - device_memory=self._device_memory, ) + worker_resources.declare(self._quent_logger) + + self._worker_resources = worker_resources # Register after `_cleanup_ctx` so on teardown (LIFO) the # executor shuts down first. `wait=True` is safe because @@ -568,16 +540,7 @@ def __init__( quent_logger=self._quent_logger, context=self._ctx, py_executor=self._py_executor, - processor_registry=self._processor_registry, - thread_pool_id=( - self._quent_thread_pool.id - if self._quent_thread_pool is not None - else None - ), - device_memory=self._device_memory, - disk_to_device_channel=self._disk_to_device_channel, - network=self._network, - link_channels=self._link_channels, + worker_resources=self._worker_resources, ), }, engine_options={ @@ -721,16 +684,7 @@ def _reset( engine_id=engine_id, worker_id=self._quent_worker.id, quent_logger=self._quent_logger, - processor_registry=self._processor_registry, - thread_pool_id=( - self._quent_thread_pool.id - if self._quent_thread_pool is not None - else None - ), - device_memory=self._device_memory, - disk_to_device_channel=self._disk_to_device_channel, - network=self._network, - link_channels=self._link_channels, + worker_resources=self._worker_resources, ), }, engine_options={ @@ -874,18 +828,8 @@ def shutdown(self) -> None: ) if self._quent_logger is not None: - if quent_context is not None: - assert self._processor_registry is not None - self._processor_registry._emit_processor_exit_events(self._quent_logger) - finalize_network_channels( - self._quent_logger, link_channels=self._link_channels - ) - if self._device_memory is not None: - finalize_worker_resources( - self._quent_logger, - device_memory=self._device_memory, - disk_to_device_channel=self._disk_to_device_channel, - ) + if self._worker_resources is not None: + self._worker_resources.finalize(self._quent_logger) self._quent_logger.emit(self._quent_worker._exit()) if quent_context is not None: assert self._quent_logger is not None diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index e7ccd3cb237a..e7ea0ee1f969 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -68,7 +68,6 @@ class ProcessorRegistry: def __init__(self) -> None: self._processors: dict[int, Processor] = {} self._lock = threading.Lock() - self._closed = False def get_or_declare_processor( self, logger: QuentLogger, thread_ident: int, pool_id: uuid.UUID @@ -415,9 +414,9 @@ def _emit_task_begin_events( quent_ir_execution_context.logger.emit( quent_task.loading( use_thread=quent_processor, - use_channel=quent_ir_execution_context.disk_to_device_channel, + use_channel=quent_ir_execution_context.worker_resources.disk_to_device_channel, channel_capacity_bytes=input_frames_bytes, - use_memory=quent_ir_execution_context.device_memory, + use_memory=quent_ir_execution_context.worker_resources.device_memory, memory_capacity_bytes=input_frames_bytes, ) ) @@ -425,7 +424,7 @@ def _emit_task_begin_events( quent_ir_execution_context.logger.emit( quent_task.computing( use_thread=quent_processor, - use_memory=quent_ir_execution_context.device_memory, + use_memory=quent_ir_execution_context.worker_resources.device_memory, input_bytes=input_frames_bytes, memory_capacity_bytes=input_frames_bytes, ) @@ -483,7 +482,7 @@ def _emit_task_end_events( quent_ir_execution_context.logger.emit( quent_task.computing( use_thread=quent_processor, - use_memory=quent_ir_execution_context.device_memory, + use_memory=quent_ir_execution_context.worker_resources.device_memory, memory_capacity_bytes=output_capacity_bytes, ) ) @@ -494,89 +493,116 @@ def _emit_task_end_events( # break operator-level aggregation like duration_s. -def declare_worker_resources( - logger: QuentLogger, - *, - instance_suffix: str, - engine_id: uuid.UUID, - worker_id: uuid.UUID, -) -> tuple[Memory, Channel, ThreadPool]: - """ - Declare per-worker Quent resources and emit their lifecycle events. +@dataclasses.dataclass(kw_only=True) +class WorkerResources: + """A simple container for per-worker Quent resources.""" - Returns device memory, disk-to-device channel, and thread pool handles. - """ - device_memory = Memory( - instance_name=f"{instance_suffix} device memory", - resource_type_name="memory", - parent_group_id=engine_id, - ) - filesystem = Memory( - instance_name=f"{instance_suffix} filesystem", - resource_type_name="filesystem", - parent_group_id=worker_id, - ) - disk_to_device_channel = Channel( - instance_name=f"{instance_suffix} disk -> device", - resource_type_name="DiskToDevice", - parent_group_id=worker_id, - source=filesystem, - target=device_memory, - ) - thread_pool = ThreadPool(worker_id=worker_id) - device_memory_capacity = get_total_device_memory() or 0 - logger.emit(device_memory.initializing()) - logger.emit(device_memory.operating(device_memory_capacity)) - logger.emit(filesystem.initializing()) - # Filesystem capacity is unknown; declare as unbounded. - logger.emit(filesystem.operating(None)) - logger.emit(disk_to_device_channel.initializing()) - # Channel capacity is a rate bound; unbounded when unknown. - logger.emit(disk_to_device_channel.operating(None)) - logger.emit(thread_pool.declare()) - return device_memory, disk_to_device_channel, thread_pool - - -def finalize_worker_resources( - logger: QuentLogger, - *, - device_memory: Memory, - disk_to_device_channel: Channel | None, -) -> None: - """Emit finalizing/exit events for per-worker Quent resources.""" - if disk_to_device_channel is not None: - logger.emit(disk_to_device_channel.finalizing()) - logger.emit(disk_to_device_channel.exit()) - logger.emit(disk_to_device_channel.source.finalizing()) - logger.emit(disk_to_device_channel.source.exit()) - logger.emit(device_memory.finalizing()) - logger.emit(device_memory.exit()) + thread_pool: ThreadPool + processor_registry: ProcessorRegistry + device_memory: Memory + filesystem: Memory + disk_to_device_channel: Channel + device_memory_capacity: int + network: Network + link_channels: dict[int, Channel] + + @classmethod + def build( + cls, + instance_suffix: str, + engine_id: uuid.UUID, + worker_id: uuid.UUID, + rank: int, + nranks: int, + ) -> Self: + processor_registry = ProcessorRegistry() + device_memory = Memory( + instance_name=f"{instance_suffix} device memory", + resource_type_name="memory", + parent_group_id=engine_id, + ) + filesystem = Memory( + instance_name=f"{instance_suffix} filesystem", + resource_type_name="filesystem", + parent_group_id=worker_id, + ) + disk_to_device_channel = Channel( + instance_name=f"{instance_suffix} disk -> device", + resource_type_name="DiskToDevice", + parent_group_id=worker_id, + source=filesystem, + target=device_memory, + ) + thread_pool = ThreadPool(worker_id=worker_id) + + network, link_channels = declare_network_channels( + rank=rank, + nranks=nranks, + engine_id=engine_id, + device_memory=device_memory, + ) + + return cls( + processor_registry=processor_registry, + device_memory=device_memory, + filesystem=filesystem, + disk_to_device_channel=disk_to_device_channel, + thread_pool=thread_pool, + device_memory_capacity=get_total_device_memory() or 0, + network=network, + link_channels=link_channels, + ) + + def declare(self, logger: QuentLogger) -> None: + logger.emit(self.device_memory.initializing()) + logger.emit(self.device_memory.operating(self.device_memory_capacity)) + logger.emit(self.filesystem.initializing()) + # Filesystem capacity is unknown; declare as unbounded. + logger.emit(self.filesystem.operating(None)) + logger.emit(self.disk_to_device_channel.initializing()) + # Channel capacity is a rate bound; unbounded when unknown. + logger.emit(self.disk_to_device_channel.operating(None)) + logger.emit(self.thread_pool.declare()) + + logger.emit(self.network.declare()) + for link in self.link_channels.values(): + logger.emit(link.initializing()) + logger.emit(link.operating(None)) + + def finalize(self, logger: QuentLogger) -> None: + self.processor_registry._emit_processor_exit_events(logger) + + logger.emit(self.disk_to_device_channel.finalizing()) + logger.emit(self.disk_to_device_channel.exit()) + logger.emit(self.disk_to_device_channel.source.finalizing()) + logger.emit(self.disk_to_device_channel.source.exit()) + logger.emit(self.device_memory.finalizing()) + logger.emit(self.device_memory.exit()) + + # network + for link in self.link_channels.values(): + logger.emit(link.finalizing()) + logger.emit(link.exit()) def declare_network_channels( - logger: QuentLogger, *, rank: int, nranks: int, engine_id: uuid.UUID, device_memory: Memory, -) -> tuple[Network | None, dict[int, Channel]]: +) -> tuple[Network, dict[int, Channel]]: """ - Declare engine-scoped network link channels for inter-rank communication. + Build engine-scoped network link channels for inter-rank communication. - Creates a Network resource group and one Channel per remote rank, emitting - their lifecycle events to the quent logger. This is engine/worker-scoped: - the inter-rank topology is fixed for the lifetime of the engine, so it is - declared once at worker setup rather than per query. + Creates a ``Network`` resource group and one ``Channel`` per remote + rank. Does not emit lifecycle events; the caller is responsible for + declaring them (see ``WorkerResources.declare``). - Returns ``(None, {})`` for single-rank runs, which have no inter-rank - communication. + For single-rank runs the returned ``Network`` has no link channels, since + there is no inter-rank communication. """ - if nranks <= 1: - return None, {} - network = Network(engine_id=engine_id) - logger.emit(network.declare()) link_channels: dict[int, Channel] = {} for target_rank in range(nranks): @@ -589,24 +615,11 @@ def declare_network_channels( source=device_memory, target=device_memory, ) - logger.emit(link.initializing()) - logger.emit(link.operating(None)) link_channels[target_rank] = link return network, link_channels -def finalize_network_channels( - logger: QuentLogger, - *, - link_channels: dict[int, Channel], -) -> None: - """Emit finalizing/exit events for engine-scoped network link channels.""" - for link in link_channels.values(): - logger.emit(link.finalizing()) - logger.emit(link.exit()) - - @dataclasses.dataclass(kw_only=True) class LocalQuentContext: """ @@ -615,42 +628,32 @@ class LocalQuentContext: This can contain non-serializable objects (like a ``QuentLogger``) and entities that are only valid on the local rank. - The ``processor_registry`` is engine/worker-scoped and outlives - individual queries. It is injected by the backend that owns the - ``ThreadPoolExecutor``. - The ``query`` is per-collect: each ``.collect()`` derives a fresh :class:`Query` from its unique ``query_id`` (see :meth:`QuentContext.query_for`), rather than reusing the shared ``context.query``. - The ``device_memory``, ``disk_to_device_channel``, ``network``, and - ``link_channels`` resources are all engine/worker-scoped: they are declared - once at worker setup (see :func:`declare_worker_resources` and - :func:`declare_network_channels`) and injected into each per-collect - context, rather than being declared per query. + The ``worker_resources`` (device memory, disk-to-device channel, network, + link channels, thread pool, and processor registry) are engine/worker-scoped: + they are declared once at worker setup via :class:`WorkerResources` and + injected into each per-collect context, rather than being declared per query. """ context: QuentContext query: Query worker: Worker logger: QuentLogger - thread_pool_id: uuid.UUID - processor_registry: ProcessorRegistry - device_memory: Memory - disk_to_device_channel: Channel | None = None - network: Network | None = None - link_channels: dict[int, Channel] = dataclasses.field(default_factory=dict) + worker_resources: WorkerResources def get_or_declare_processor( self, thread_ident: int, ) -> Processor: """Get (or declare a new) Quent Processor for a CPU thread.""" - return self.processor_registry.get_or_declare_processor( + return self.worker_resources.processor_registry.get_or_declare_processor( self.logger, thread_ident=thread_ident, - pool_id=self.thread_pool_id, + pool_id=self.worker_resources.thread_pool.id, ) @@ -671,10 +674,5 @@ def from_execution_context( query=execution_context.query, worker=execution_context.worker, logger=execution_context.logger, - thread_pool_id=execution_context.thread_pool_id, - processor_registry=execution_context.processor_registry, - device_memory=execution_context.device_memory, - disk_to_device_channel=execution_context.disk_to_device_channel, - network=execution_context.network, - link_channels=execution_context.link_channels, + worker_resources=execution_context.worker_resources, ) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py index ea9dff02995b..c91d3691e49a 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py @@ -96,17 +96,7 @@ def evaluate_logical_plan( engine_id=engine_id, worker_id=engine._quent_worker.id, quent_logger=engine._quent_logger, - # The engine owns these engine-scoped Quent resources; the - # per-collect config must forward them so query evaluation - # reuses them instead of re-declaring per collect. - processor_registry=engine._processor_registry, - thread_pool_id=( - engine._quent_thread_pool.id - if engine._quent_thread_pool is not None - else None - ), - device_memory=engine._device_memory, - disk_to_device_channel=engine._disk_to_device_channel, + worker_resources=engine._worker_resources, ), ), ) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 49a37c38d404..24b64376ec91 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -44,9 +44,8 @@ from rapidsmpf.streaming.core.context import Context from cudf_polars.engine.ray import RankActor - from cudf_polars.quent._context import ProcessorRegistry, QuentContext + from cudf_polars.quent._context import QuentContext, WorkerResources from cudf_polars.quent._logging import QuentLogger - from cudf_polars.quent._types import Channel, Memory, Network __all__ = [ @@ -626,19 +625,10 @@ class SPMDContext: The active RapidsMPF context. py_executor Thread-pool executor used to drive the actor network on each rank. - processor_registry - Engine/worker-scoped registry of dynamically declared Processors. - thread_pool_id - ID of the engine-scoped Quent ThreadPool resource. - device_memory - The engine-scoped Quent device Memory resource. - disk_to_device_channel - The engine-scoped Quent disk-to-device Channel resource. - network - The engine-scoped Quent Network resource group (``None`` for - single-rank runs). - link_channels - The engine-scoped Quent inter-rank Link channels, keyed by target rank. + worker_resources + Engine/worker-scoped Quent resources (device memory, channels, thread + pool, processor registry, network topology). ``None`` when Quent is + disabled. """ comm: Communicator @@ -647,12 +637,7 @@ class SPMDContext: engine_id: uuid.UUID worker_id: uuid.UUID quent_logger: QuentLogger | None - processor_registry: ProcessorRegistry | None = None - thread_pool_id: uuid.UUID | None = None - device_memory: Memory | None = None - disk_to_device_channel: Channel | None = None - network: Network | None = None - link_channels: dict[int, Channel] = dataclasses.field(default_factory=dict) + worker_resources: WorkerResources | None = None @dataclasses.dataclass(frozen=True) diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index eb4d53f2d32b..111980058d46 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -23,8 +23,7 @@ ProcessorRegistry, QuentContext, QuentIRExecutionContext, - declare_network_channels, - finalize_network_channels, + WorkerResources, ) from cudf_polars.quent._plan import build_plan, port_names_for_node from cudf_polars.quent._types import ( @@ -74,15 +73,17 @@ def _make_quent_ir_execution_context( context = QuentContext() engine_id = context.engine.id worker_id = uuid.uuid4() - pool_id = uuid.uuid4() + worker_resources = WorkerResources.build( + instance_suffix="test", + engine_id=engine_id, + worker_id=worker_id, + rank=0, + nranks=1, + ) if disk_to_device_channel is not None: - device_memory = disk_to_device_channel.target - else: - device_memory = Memory( - instance_name="device", - resource_type_name="memory", - parent_group_id=engine_id, - ) + worker_resources.disk_to_device_channel = disk_to_device_channel + worker_resources.device_memory = disk_to_device_channel.target + worker_resources.filesystem = disk_to_device_channel.source operator_id = operator_id or uuid.uuid4() query = context.query_for(uuid.uuid4()) plan = Plan( @@ -104,10 +105,7 @@ def _make_quent_ir_execution_context( query=query, worker=Worker(id=worker_id, engine=context.engine, instance_name="rank-0"), logger=logger, - thread_pool_id=pool_id, - processor_registry=ProcessorRegistry(), - device_memory=device_memory, - disk_to_device_channel=disk_to_device_channel, + worker_resources=worker_resources, ) quent_ir_execution_context = QuentIRExecutionContext.from_execution_context( local_context, operator @@ -728,37 +726,32 @@ def get_processor(_: int) -> Processor: def test_processor_registry_reused_across_quent_contexts() -> None: pytest.importorskip("structlog") logger = cudf_polars.quent._logging.QuentLogger() - registry = ProcessorRegistry() - pool_id = uuid.uuid4() thread_ident = 99 context_a = QuentContext() context_b = QuentContext() + # Share a single WorkerResources so both contexts reuse the same + # processor registry / thread pool. + worker_resources = WorkerResources.build( + instance_suffix="test", + engine_id=context_a.engine.id, + worker_id=uuid.uuid4(), + rank=0, + nranks=1, + ) local_a = LocalQuentContext( context=context_a, query=context_a.query_for(uuid.uuid4()), worker=Worker(id=uuid.uuid4(), engine=context_a.engine, instance_name="rank-0"), logger=logger, - thread_pool_id=pool_id, - processor_registry=registry, - device_memory=Memory( - instance_name="device", - resource_type_name="memory", - parent_group_id=context_a.engine.id, - ), + worker_resources=worker_resources, ) local_b = LocalQuentContext( context=context_b, query=context_b.query_for(uuid.uuid4()), worker=Worker(id=uuid.uuid4(), engine=context_b.engine, instance_name="rank-0"), logger=logger, - thread_pool_id=pool_id, - processor_registry=registry, - device_memory=Memory( - instance_name="device", - resource_type_name="memory", - parent_group_id=context_b.engine.id, - ), + worker_resources=worker_resources, ) processor_a = local_a.get_or_declare_processor(thread_ident=thread_ident) @@ -945,28 +938,43 @@ def test_network_declare_serialization() -> None: } -def test_declare_network_channels_single_rank(device_memory: Memory) -> None: +def test_declare_network_channels_single_rank() -> None: pytest.importorskip("structlog") logger = cudf_polars.quent._logging.QuentLogger() - - network, link_channels = declare_network_channels( - logger, + worker_resources = WorkerResources.build( + instance_suffix="test", + engine_id=uuid.uuid4(), + worker_id=uuid.uuid4(), rank=0, nranks=1, - engine_id=uuid.uuid4(), - device_memory=device_memory, + ) + worker_resources.declare(logger) + assert worker_resources.link_channels == {} + events = _drained_events(logger) + network_events = [event for event in events if "Network" in event["data"]] + assert len(network_events) == 1 + assert ( + network_events[0]["data"]["Network"]["Declaration"]["instance_name"] + == "Network" ) - assert network is None - assert link_channels == {} - assert _drained_events(logger) == [] + channel_events = [event for event in events if "Channel" in event["data"]] + network_channels = [ + event + for event in channel_events + if event["data"]["Channel"]["seq"] == 0 + and event["data"]["Channel"]["state"]["ChannelInitializing"][ + "resource_type_name" + ] + == "Link" + ] + assert len(network_channels) == 0 @pytest.mark.parametrize( "rank,nranks,expected_targets", [(0, 3, [1, 2]), (1, 3, [0, 2])] ) def test_declare_network_channels_multi_rank( - device_memory: Memory, rank: int, nranks: int, expected_targets: list[int], @@ -975,60 +983,48 @@ def test_declare_network_channels_multi_rank( logger = cudf_polars.quent._logging.QuentLogger() engine_id = uuid.uuid4() - network, link_channels = declare_network_channels( - logger, + worker_resources = WorkerResources.build( + instance_suffix="test", + engine_id=engine_id, + worker_id=uuid.uuid4(), rank=rank, nranks=nranks, - engine_id=engine_id, - device_memory=device_memory, ) + worker_resources.declare(logger) - assert network is not None - assert set(link_channels) == set(expected_targets) - for target_rank, link in link_channels.items(): + assert worker_resources.network is not None + assert set(worker_resources.link_channels) == set(expected_targets) + for target_rank, link in worker_resources.link_channels.items(): assert link.instance_name == f"rank-{rank} -> rank-{target_rank}" assert link.resource_type_name == "Link" - assert link.parent_group_id == network.id - assert link.source is device_memory - assert link.target is device_memory + assert link.parent_group_id == worker_resources.network.id + assert link.source is worker_resources.device_memory + assert link.target is worker_resources.device_memory + + worker_resources.finalize(logger) events = _drained_events(logger) network_events = [event for event in events if "Network" in event["data"]] channel_events = [event for event in events if "Channel" in event["data"]] + network_channel_ids = { + event["id"] + for event in channel_events + if event["data"]["Channel"]["seq"] == 0 + and event["data"]["Channel"]["state"]["ChannelInitializing"][ + "resource_type_name" + ] + == "Link" + } + + network_channel_events = [ + event for event in channel_events if event["id"] in network_channel_ids + ] assert len(network_events) == 1 assert network_events[0]["data"]["Network"]["Declaration"][ "parent_group_id" ] == str(engine_id) - assert len(channel_events) == len(expected_targets) * 2 - - -def test_finalize_network_channels(device_memory: Memory) -> None: - pytest.importorskip("structlog") - logger = cudf_polars.quent._logging.QuentLogger() - link_channels = { - target_rank: Channel( - instance_name=f"rank-0 -> rank-{target_rank}", - resource_type_name="Link", - parent_group_id=uuid.uuid4(), - source=device_memory, - target=device_memory, - ) - for target_rank in (1, 2) - } - - finalize_network_channels(logger, link_channels=link_channels) - - events = _drained_events(logger) - finalizing_events = [ - event - for event in events - if event["data"]["Channel"]["state"] == {"ChannelFinalizing": None} - ] - exit_events = [ - event for event in events if event["data"]["Channel"]["state"] == "Exit" - ] - assert len(finalizing_events) == 2 - assert len(exit_events) == 2 + # One event for Initializing, Operating, Finalizing, and Exit + assert len(network_channel_events) == len(expected_targets) * 4 def test_emit_task_events_computing_node() -> None: From da2484363ab9787c0fc40395f13dfd56409c5170 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 24 Jul 2026 13:43:54 -0700 Subject: [PATCH 42/60] remove unused parameter --- python/cudf_polars/cudf_polars/dsl/tracing.py | 2 -- python/cudf_polars/cudf_polars/quent/_context.py | 3 --- python/cudf_polars/tests/quent/test_quent.py | 2 -- 3 files changed, 7 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index 90888996c4d0..bbf3f08f0885 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -225,7 +225,6 @@ def wrapper( cls, quent_task, ir_execution_context.quent_ir_execution_context, - frames, None, ) raise @@ -239,7 +238,6 @@ def wrapper( cls, quent_task, ir_execution_context.quent_ir_execution_context, - frames, result, ) stop = time.monotonic_ns() diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index e7ea0ee1f969..6989f46b7664 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -435,7 +435,6 @@ def _emit_task_end_events( ir_type: type[IR], quent_task: Task, quent_ir_execution_context: QuentIRExecutionContext, - frames: list[DataFrame], result: DataFrame | None, ) -> None: """ @@ -452,8 +451,6 @@ def _emit_task_end_events( The Quent Task to emit events for. quent_ir_execution_context: QuentIRExecutionContext The Quent IR execution context. - frames - The input dataframes passed to the IR node. result The output dataframe returned from the IR node. This will be ``None`` if an exception was raised while evaluating the IR node. diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index 111980058d46..4eac3210890f 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -1046,7 +1046,6 @@ def test_emit_task_events_computing_node() -> None: Filter, task, quent_ir_execution_context, - [], result, ) @@ -1082,7 +1081,6 @@ def test_emit_task_events_io_node(disk_to_device_channel: Channel) -> None: DataFrameScan, task, quent_ir_execution_context, - [], result, ) From 9dc8d863bab7c3b494a1ba78a9ed1a07e8d5efc1 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 24 Jul 2026 13:47:39 -0700 Subject: [PATCH 43/60] remove unused parameter --- .../cudf_polars/cudf_polars/quent/_context.py | 57 +++++-------------- 1 file changed, 15 insertions(+), 42 deletions(-) diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index 6989f46b7664..f218f4e4fa8b 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -532,12 +532,20 @@ def build( ) thread_pool = ThreadPool(worker_id=worker_id) - network, link_channels = declare_network_channels( - rank=rank, - nranks=nranks, - engine_id=engine_id, - device_memory=device_memory, - ) + # Network / Link Channels + network = Network(engine_id=engine_id) + link_channels: dict[int, Channel] = {} + for target_rank in range(nranks): + if target_rank == rank: + continue + link = Channel( + instance_name=f"rank-{rank} -> rank-{target_rank}", + resource_type_name="Link", + parent_group_id=network.id, + source=device_memory, + target=device_memory, + ) + link_channels[target_rank] = link return cls( processor_registry=processor_registry, @@ -576,47 +584,12 @@ def finalize(self, logger: QuentLogger) -> None: logger.emit(self.device_memory.finalizing()) logger.emit(self.device_memory.exit()) - # network + # Network / Link Channels for link in self.link_channels.values(): logger.emit(link.finalizing()) logger.emit(link.exit()) -def declare_network_channels( - *, - rank: int, - nranks: int, - engine_id: uuid.UUID, - device_memory: Memory, -) -> tuple[Network, dict[int, Channel]]: - """ - Build engine-scoped network link channels for inter-rank communication. - - Creates a ``Network`` resource group and one ``Channel`` per remote - rank. Does not emit lifecycle events; the caller is responsible for - declaring them (see ``WorkerResources.declare``). - - For single-rank runs the returned ``Network`` has no link channels, since - there is no inter-rank communication. - """ - network = Network(engine_id=engine_id) - - link_channels: dict[int, Channel] = {} - for target_rank in range(nranks): - if target_rank == rank: - continue - link = Channel( - instance_name=f"rank-{rank} -> rank-{target_rank}", - resource_type_name="Link", - parent_group_id=network.id, - source=device_memory, - target=device_memory, - ) - link_channels[target_rank] = link - - return network, link_channels - - @dataclasses.dataclass(kw_only=True) class LocalQuentContext: """ From bd7b29e1429bc0e9af3a444d63016905dfab4512 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 24 Jul 2026 13:51:23 -0700 Subject: [PATCH 44/60] refactor --- python/cudf_polars/cudf_polars/engine/spmd.py | 9 +++++---- .../cudf_polars/streaming/actor_graph/join.py | 1 - .../cudf_polars/streaming/actor_graph/utils.py | 5 ++--- 3 files changed, 7 insertions(+), 8 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index e894ca66e8d0..c87e2877f236 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -823,14 +823,15 @@ def shutdown(self) -> None: # quent traces before that. # Clear the references only after shutdown completes. - quent_context: QuentContext | None = self.config["executor_options"].get( - "quent_context" - ) - if self._quent_logger is not None: if self._worker_resources is not None: self._worker_resources.finalize(self._quent_logger) self._quent_logger.emit(self._quent_worker._exit()) + + quent_context: QuentContext | None = self.config["executor_options"].get( + "quent_context" + ) + if quent_context is not None: assert self._quent_logger is not None quent_context._emit_engine_exit_events(self._quent_logger) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py index df8811030fc0..3bdaaeba9c20 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -293,7 +293,6 @@ async def _broadcast_join_large_chunk( join_results: list[DataFrame] = [] input_bytes = large_chunk_size + small_size - with opaque_memory_usage( await reserve_memory(context, size=input_bytes, net_memory_delta=0) ): diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py index ddc5acf65a22..2aa4ca524c20 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -12,7 +12,7 @@ import time from collections import defaultdict, deque from contextlib import asynccontextmanager -from dataclasses import dataclass +from dataclasses import dataclass, replace from functools import reduce from typing import TYPE_CHECKING, Any, Literal, TypeAlias, cast @@ -292,11 +292,10 @@ async def shutdown_on_error( ir_type = type(trace_ir).__name__ tracer = ActorTracer(ir_id, ir_type) contextvars = {"actor_ir_id": ir_id, "actor_ir_type": ir_type} - import dataclasses if ir_context is not None: contextvars["cudf_polars_query_id"] = str(ir_context.query_id) - ir_context = dataclasses.replace(ir_context, tracer=tracer) + ir_context = replace(ir_context, tracer=tracer) with cudf_polars.dsl.tracing.bound_contextvars(**contextvars): start = time.monotonic_ns() From 0ebe378812f752047bf29d3df29ec6aa8cab5051 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 27 Jul 2026 11:47:11 -0700 Subject: [PATCH 45/60] cvoerage --- python/cudf_polars/cudf_polars/quent/_context.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index f218f4e4fa8b..e0a0b1ce408c 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -470,7 +470,7 @@ def _emit_task_end_events( if result is not None: output_capacity_bytes = result._size_bytes() - else: + else: # pragma: no cover; output_capacity_bytes = 0 if ir_type.is_io_node: quent_processor = quent_ir_execution_context.get_or_declare_processor( From 716109e4da8a503d4705e9b6eee918fb7d99008a Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 27 Jul 2026 11:55:38 -0700 Subject: [PATCH 46/60] cleanup --- python/cudf_polars/cudf_polars/dsl/tracing.py | 10 +++------- python/cudf_polars/cudf_polars/quent/_types.py | 2 +- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index bbf3f08f0885..ea3036698f69 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -189,13 +189,6 @@ def wrapper( quent_task = Task.from_ir( cls, ir_execution_context.quent_ir_execution_context ) - else: - quent_task = None - - if ( - quent_task is not None - and ir_execution_context.quent_ir_execution_context is not None - ): ir_execution_context.quent_ir_execution_context.context._emit_task_begin_events( cls, quent_task, @@ -203,6 +196,9 @@ def wrapper( input_frames_bytes=sum(frame._size_bytes() for frame in frames), ) + else: + quent_task = None + before_start = time.monotonic_ns() before = make_snapshot( cls, frames, phase="input", device_handle=maybe_handle, pid=pid diff --git a/python/cudf_polars/cudf_polars/quent/_types.py b/python/cudf_polars/cudf_polars/quent/_types.py index 3c337a41a62b..7847a35b4e91 100644 --- a/python/cudf_polars/cudf_polars/quent/_types.py +++ b/python/cudf_polars/cudf_polars/quent/_types.py @@ -841,7 +841,7 @@ class Task: @classmethod def from_ir( cls, ir_type: type[IR], quent_ir_execution_context: QuentIRExecutionContext - ) -> Self | None: + ) -> Self: """ Build an operator-scoped Quent Task from an IR execution context. From 20483d427bdfc21bf6ad3f66208783dd9fb3ced5 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 27 Jul 2026 12:17:48 -0700 Subject: [PATCH 47/60] revert --- python/cudf_polars/tests/quent/conftest.py | 5 ++--- python/cudf_polars/tests/quent/test_quent.py | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/python/cudf_polars/tests/quent/conftest.py b/python/cudf_polars/tests/quent/conftest.py index ee94d385f66c..5335d8b33081 100644 --- a/python/cudf_polars/tests/quent/conftest.py +++ b/python/cudf_polars/tests/quent/conftest.py @@ -11,11 +11,10 @@ import pytest import cudf_polars.quent -import cudf_polars.quent._context from cudf_polars.quent._types import Channel, Memory, Processor if TYPE_CHECKING: - from cudf_polars.quent._context import QuentContext + from cudf_polars.quent import QuentContext @pytest.fixture @@ -51,7 +50,7 @@ def disk_to_device_channel(device_memory: Memory) -> Channel: @pytest.fixture def quent_context() -> QuentContext: """A Quent Context with a QueryGroup and Query set.""" - return cudf_polars.quent._context.QuentContext( + return cudf_polars.quent.QuentContext( query_group=cudf_polars.quent.QueryGroup(instance_name="test_query_group"), query=cudf_polars.quent.Query(instance_name="test_query"), ) diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index 4eac3210890f..c7faf8aefd82 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -18,10 +18,10 @@ from cudf_polars.containers import DataFrame from cudf_polars.dsl.ir import DataFrameScan, Filter from cudf_polars.dsl.translate import Translator +from cudf_polars.quent import QuentContext from cudf_polars.quent._context import ( LocalQuentContext, ProcessorRegistry, - QuentContext, QuentIRExecutionContext, WorkerResources, ) From 90c624c67cda3aa2ec0095f92ad44388a2e46bd8 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 28 Jul 2026 07:29:07 -0700 Subject: [PATCH 48/60] Update quent output - write a zipfile - include in ndjson output - add quent role --- .gitignore | 1 + .../cudf_polars/cudf_polars/quent/_export.py | 54 ++++---- .../cudf_polars/streaming/benchmarks/utils.py | 96 ++++++++++---- python/cudf_polars/tests/quent/test_export.py | 118 ++++++++++-------- 4 files changed, 169 insertions(+), 100 deletions(-) diff --git a/.gitignore b/.gitignore index 21de3f8d9c74..ac3cc314ab23 100644 --- a/.gitignore +++ b/.gitignore @@ -191,3 +191,4 @@ python/cudf/cudf_pandas_tests/data/rmm_log.txt # Quent traces logs/**/*.ndjson logs/**/*.qmi +logs/*.zip diff --git a/python/cudf_polars/cudf_polars/quent/_export.py b/python/cudf_polars/cudf_polars/quent/_export.py index d652a4d36c64..478edce0af5d 100644 --- a/python/cudf_polars/cudf_polars/quent/_export.py +++ b/python/cudf_polars/cudf_polars/quent/_export.py @@ -1,11 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Export Quent events to the filesystem directory layout.""" +"""Export Quent events to an archive.""" from __future__ import annotations import json +import zipfile from typing import TYPE_CHECKING, Any from cudf_polars.quent._types import EventName, new_quent_id @@ -58,7 +59,7 @@ def unwrap_event_data(data: dict[str, Any]) -> tuple[str, Any]: """ Extract the entity name and unwrapped payload from a buffered event. - Buffered events wrap payloads as ``{"Engine": {...}}``; directory export + Buffered events wrap payloads as ``{"Engine": {...}}``; archive export stores the payload directly because the entity type is implied by the subdirectory name. """ @@ -76,7 +77,7 @@ def unwrap_event_data(data: dict[str, Any]) -> tuple[str, Any]: def to_export_line(event: dict[str, Any]) -> tuple[str, dict[str, Any]]: - """Convert a buffered event envelope to directory export line format.""" + """Convert a buffered event envelope to archive export line format.""" entity_name, payload = unwrap_event_data(event["data"]) directory = ENTITY_DIRECTORIES[entity_name] export_line = { @@ -87,14 +88,6 @@ def to_export_line(event: dict[str, Any]) -> tuple[str, dict[str, Any]]: return directory, export_line -def write_sidecar(context_dir: Path, sidecar: dict[str, Any]) -> None: - """Atomically write the ``model.qmi`` provenance sidecar.""" - tmp_path = context_dir / f".{SIDECAR_FILE_NAME}.tmp" - final_path = context_dir / SIDECAR_FILE_NAME - tmp_path.write_text(json.dumps(sidecar, indent=2) + "\n", encoding="utf-8") - tmp_path.replace(final_path) - - def write_quent_export( events: list[dict[str, Any]], export_root: Path, @@ -103,14 +96,14 @@ def write_quent_export( sidecar: dict[str, Any] | None = None, ) -> Path: """ - Write Quent events to the filesystem export layout. + Write Quent events to a ZIP archive. Parameters ---------- events Buffered Quent event envelopes from ``engine._quent_events``. export_root - Root directory for exported contexts (e.g. ``logs``). + Directory for exported archives (e.g. ``logs``). context_id Context UUID, typically the engine/run id. sidecar @@ -120,25 +113,30 @@ def write_quent_export( Returns ------- Path - The context directory ``export_root//``. + The archive path ``export_root/.zip``. The archive contains + the Quent export layout under a top-level ``/`` directory. """ - context_dir = export_root / str(context_id) - context_dir.mkdir(parents=True, exist_ok=True) - - write_sidecar(context_dir, sidecar or MODEL_QMI) - grouped: dict[str, list[dict[str, Any]]] = {} for event in events: directory, export_line = to_export_line(event) grouped.setdefault(directory, []).append(export_line) - for directory, lines in grouped.items(): - entity_dir = context_dir / directory - entity_dir.mkdir(parents=True, exist_ok=True) - stream_path = entity_dir / f"{new_quent_id()}.{EXTENSION}" - with stream_path.open("w", encoding="utf-8") as stream_file: - for line in lines: - stream_file.write(json.dumps(line, separators=(",", ":"))) - stream_file.write("\n") + export_root.mkdir(parents=True, exist_ok=True) + archive_path = export_root / f"{context_id}.zip" + tmp_path = export_root / f".{context_id}.zip.tmp" + context_dir = str(context_id) + + with zipfile.ZipFile( + tmp_path, mode="w", compression=zipfile.ZIP_DEFLATED + ) as archive: + archive.writestr( + f"{context_dir}/{SIDECAR_FILE_NAME}", + json.dumps(MODEL_QMI if sidecar is None else sidecar, indent=2) + "\n", + ) + for directory, lines in grouped.items(): + stream_path = f"{context_dir}/{directory}/{new_quent_id()}.{EXTENSION}" + contents = [json.dumps(line, separators=(",", ":")) for line in lines] + archive.writestr(stream_path, "\n".join(contents) + "\n") - return context_dir + tmp_path.replace(archive_path) + return archive_path diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index 84c85ab70d0e..d609d6a63a1c 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -141,7 +141,13 @@ class NsysRole: type: Literal["nsys"] = dataclasses.field(default="nsys", init=False) -Role = NightlyRole | NsysRole +@dataclasses.dataclass +class QuentRole: + type: Literal["quent"] = dataclasses.field(default="quent", init=False) + filename: str + + +Role = NightlyRole | NsysRole | QuentRole @dataclasses.dataclass @@ -648,9 +654,27 @@ def from_args(cls, args: argparse.Namespace) -> RunConfig: roles=roles, ) - def serialize(self, engine: StreamingEngine | None) -> dict: - """Serialize the run config to a dictionary.""" + def serialize( + self, engine: StreamingEngine | None, quent_archive: Path | None + ) -> dict: + """ + Serialize the run config to a dictionary. + + Parameters + ---------- + engine + The engine that was used to run the benchmark. + quent_archive + The path to the Quent archive that was written during the benchmark, if any. + This path will be inserted in ``extra_info.quent-archive``. + """ opts = self.streaming_options + extra_info = dict(self.extra_info) + if quent_archive is not None: + extra_info["quent-archive"] = str(quent_archive.absolute()) + roles = list(self.roles) + if quent_archive is not None: + roles.append(QuentRole(filename=quent_archive.name)) result: dict[str, Any] = { "engine_name": self.engine_name, "queries": self.queries, @@ -666,7 +690,7 @@ def serialize(self, engine: StreamingEngine | None) -> dict: "native_parquet": self.native_parquet, "max_io_threads": self.max_io_threads, "n_workers": self.n_workers, - "extra_info": self.extra_info, + "extra_info": extra_info, "run_id": str(self.run_id), "timestamp": self.timestamp, "command_line": self.command_line, @@ -684,7 +708,7 @@ def serialize(self, engine: StreamingEngine | None) -> dict: "validation_method": dataclasses.asdict(self.validation_method) if self.validation_method else None, - "roles": [dataclasses.asdict(r) for r in self.roles], + "roles": [dataclasses.asdict(r) for r in roles], } if engine is not None: config_options = ConfigOptions.from_polars_engine(engine) @@ -1218,6 +1242,7 @@ def _finalize_benchmark_run( validation_failures: list[int], query_failures: list[tuple[int, int]], engine: StreamingEngine | None, + quent_archive: Path | None, ) -> None: """Summarize, serialize, and exit after a benchmark run.""" if args.summarize: @@ -1235,7 +1260,9 @@ def _finalize_benchmark_run( ) else: print("✅ All validated queries passed.") - args.output.write(json.dumps(run_config.serialize(engine=engine))) + args.output.write( + json.dumps(run_config.serialize(engine=engine, quent_archive=quent_archive)) + ) args.output.write("\n") sys.exit(1 if (query_failures or validation_failures) else 0) @@ -1258,7 +1285,12 @@ def run_polars_cpu( ) run_config = dataclasses.replace(run_config, records=dict(records), plans=plans) _finalize_benchmark_run( - args, run_config, validation_failures, query_failures, engine=None + args, + run_config, + validation_failures, + query_failures, + engine=None, + quent_archive=None, ) @@ -1291,7 +1323,12 @@ def run_polars_in_memory( run_config = dataclasses.replace(run_config, records=dict(records), plans=plans) run_config = _consolidate_logs(run_config, engine=None) _finalize_benchmark_run( - args, run_config, validation_failures, query_failures, engine=None + args, + run_config, + validation_failures, + query_failures, + engine=None, + quent_archive=None, ) @@ -1352,13 +1389,20 @@ def _allgather_result(df: pl.DataFrame) -> pl.DataFrame: ) if is_rank_0: - _write_quent_traces( + quent_archive = _write_quent_traces( engine=engine, run_id=run_config.run_id, collect_traces=run_config.collect_traces, ) + else: + quent_archive = None _finalize_benchmark_run( - args, run_config, validation_failures, query_failures, engine=engine + args, + run_config, + validation_failures, + query_failures, + engine=engine, + quent_archive=quent_archive, ) @@ -1405,13 +1449,18 @@ def run_polars_ray( run_config = dataclasses.replace(run_config, records=dict(records), plans=plans) run_config = _consolidate_logs(run_config, engine=engine) - _write_quent_traces( + quent_archive = _write_quent_traces( engine=engine, run_id=run_config.run_id, collect_traces=run_config.collect_traces, ) _finalize_benchmark_run( - args, run_config, validation_failures, query_failures, engine=engine + args, + run_config, + validation_failures, + query_failures, + engine=engine, + quent_archive=quent_archive, ) @@ -1464,7 +1513,7 @@ def run_polars_dask( ) run_config = _consolidate_logs(run_config, engine) - _write_quent_traces( + quent_archive = _write_quent_traces( engine=engine, run_id=run_config.run_id, collect_traces=run_config.collect_traces, @@ -1473,7 +1522,12 @@ def run_polars_dask( if dask_client is not None: dask_client.close() _finalize_benchmark_run( - args, run_config, validation_failures, query_failures, engine=engine + args, + run_config, + validation_failures, + query_failures, + engine=engine, + quent_archive=quent_archive, ) @@ -1552,10 +1606,10 @@ def inject( def _write_quent_traces( engine: StreamingEngine, run_id: uuid.UUID, *, collect_traces: bool -) -> None: - """Write collected Quent events to logs// directory export layout.""" +) -> Path | None: + """Write collected Quent events to a ``logs/.zip`` archive.""" if not (_HAS_STRUCTLOG or collect_traces): - return + return None from cudf_polars.quent._export import write_quent_export @@ -1574,10 +1628,8 @@ def _write_quent_traces( logs_dir = Path("logs") output_path = write_quent_export(quent_logs, logs_dir, run_id) - print( - f"Wrote {len(quent_logs)} Quent trace events to {output_path} " - f"(directory export layout)" - ) + print(f"Wrote {len(quent_logs)} Quent trace events to {output_path}") + return output_path def _consolidate_logs( @@ -1832,7 +1884,7 @@ def run_duckdb(duckdb_queries_cls: Any, args: argparse.Namespace) -> None: if args.summarize: run_config.summarize() - args.output.write(json.dumps(run_config.serialize(engine=None))) + args.output.write(json.dumps(run_config.serialize(engine=None, quent_archive=None))) args.output.write("\n") diff --git a/python/cudf_polars/tests/quent/test_export.py b/python/cudf_polars/tests/quent/test_export.py index ed5a599c305c..83f22cc92458 100644 --- a/python/cudf_polars/tests/quent/test_export.py +++ b/python/cudf_polars/tests/quent/test_export.py @@ -7,6 +7,7 @@ import json import uuid +import zipfile from typing import TYPE_CHECKING, Any import pytest @@ -45,8 +46,8 @@ def _buffered_events() -> list[dict[str, Any]]: ] -def _read_ndjson_lines(path: Path) -> list[dict[str, Any]]: - return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines()] +def _read_ndjson_lines(archive: zipfile.ZipFile, path: str) -> list[dict[str, Any]]: + return [json.loads(line) for line in archive.read(path).decode().splitlines()] def test_unwrap_event_data() -> None: @@ -82,65 +83,80 @@ def test_to_export_line_unwraps_payload() -> None: } -def test_write_quent_export_creates_context_layout(tmp_path: Path) -> None: +def test_write_quent_export_creates_archive(tmp_path: Path) -> None: context_id = uuid.UUID("019dd571-105a-7c53-a15b-713cbdd7666b") events = _buffered_events() - context_dir = write_quent_export(events, tmp_path, context_id) + archive_path = write_quent_export(events, tmp_path, context_id) - assert context_dir == tmp_path / str(context_id) - assert (context_dir / SIDECAR_FILE_NAME).exists() - assert ( - json.loads((context_dir / SIDECAR_FILE_NAME).read_text(encoding="utf-8")) - == MODEL_QMI - ) + assert archive_path == tmp_path / f"{context_id}.zip" + assert zipfile.is_zipfile(archive_path) - expected_dirs = {"engine", "query_group", "query", "network"} - created_dirs = {path.name for path in context_dir.iterdir() if path.is_dir()} - assert created_dirs == expected_dirs + context_dir = str(context_id) + with zipfile.ZipFile(archive_path) as archive: + assert ( + json.loads(archive.read(f"{context_dir}/{SIDECAR_FILE_NAME}")) == MODEL_QMI + ) - for entity_dir in expected_dirs: - stream_files = list((context_dir / entity_dir).glob(f"*.{EXTENSION}")) - assert len(stream_files) == 1 - lines = _read_ndjson_lines(stream_files[0]) - assert lines - for line in lines: - assert "id" in line - assert "timestamp" in line - assert isinstance(line["data"], dict) - assert len(line["data"]) == 1 or "seq" in line["data"] + expected_dirs = {"engine", "query_group", "query", "network"} + names = archive.namelist() + for entity_dir in expected_dirs: + stream_files = [ + name + for name in names + if name.startswith(f"{context_dir}/{entity_dir}/") + and name.endswith(f".{EXTENSION}") + ] + assert len(stream_files) == 1 + lines = _read_ndjson_lines(archive, stream_files[0]) + assert lines + for line in lines: + assert "id" in line + assert "timestamp" in line + assert isinstance(line["data"], dict) + assert len(line["data"]) == 1 or "seq" in line["data"] def test_write_quent_export_unwraps_buffered_envelopes(tmp_path: Path) -> None: context_id = uuid.UUID("019dd571-105a-7c53-a15b-713cbdd7666b") events = _buffered_events() - write_quent_export(events, tmp_path, context_id) - - engine_stream = next((tmp_path / str(context_id) / "engine").glob(f"*.{EXTENSION}")) - engine_lines = _read_ndjson_lines(engine_stream) - assert engine_lines[0]["data"] == { - "Init": { - "implementation": { - "name": "cudf-polars", - "version": engine_lines[0]["data"]["Init"]["implementation"]["version"], - "custom_attributes": [], - }, - "instance_name": "cudf-polars-019dd571", + archive_path = write_quent_export(events, tmp_path, context_id) + + with zipfile.ZipFile(archive_path) as archive: + context_dir = str(context_id) + engine_stream = next( + name + for name in archive.namelist() + if name.startswith(f"{context_dir}/engine/") + ) + engine_lines = _read_ndjson_lines(archive, engine_stream) + assert engine_lines[0]["data"] == { + "Init": { + "implementation": { + "name": "cudf-polars", + "version": engine_lines[0]["data"]["Init"]["implementation"][ + "version" + ], + "custom_attributes": [], + }, + "instance_name": "cudf-polars-019dd571", + } } - } - assert engine_lines[1]["data"] == {"Exit": None} + assert engine_lines[1]["data"] == {"Exit": None} - network_stream = next( - (tmp_path / str(context_id) / "network").glob(f"*.{EXTENSION}") - ) - network_lines = _read_ndjson_lines(network_stream) - assert network_lines[0]["data"] == { - "Declaration": { - "instance_name": "Network", - "parent_group_id": str(context_id), + network_stream = next( + name + for name in archive.namelist() + if name.startswith(f"{context_dir}/network/") + ) + network_lines = _read_ndjson_lines(archive, network_stream) + assert network_lines[0]["data"] == { + "Declaration": { + "instance_name": "Network", + "parent_group_id": str(context_id), + } } - } def test_write_quent_export_rejects_malformed_event(tmp_path: Path) -> None: @@ -171,7 +187,9 @@ class FakeEngine: collect_traces=True, ) - context_dir = tmp_path / "logs" / str(run_id) - assert (context_dir / SIDECAR_FILE_NAME).exists() - assert (context_dir / "engine").is_dir() - assert (context_dir / "query").is_dir() + archive_path = tmp_path / "logs" / f"{run_id}.zip" + with zipfile.ZipFile(archive_path) as archive: + names = archive.namelist() + assert f"{run_id}/{SIDECAR_FILE_NAME}" in names + assert any(name.startswith(f"{run_id}/engine/") for name in names) + assert any(name.startswith(f"{run_id}/query/") for name in names) From 3b8be3209d8fd7e87780aff3e516e16d21fb2c5a Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 28 Jul 2026 14:44:30 -0700 Subject: [PATCH 49/60] Include type name in physical operator instance name --- python/cudf_polars/cudf_polars/quent/_types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/quent/_types.py b/python/cudf_polars/cudf_polars/quent/_types.py index 7847a35b4e91..d47e3c9a0404 100644 --- a/python/cudf_polars/cudf_polars/quent/_types.py +++ b/python/cudf_polars/cudf_polars/quent/_types.py @@ -165,7 +165,7 @@ def to_dict(self) -> dict[str, Any]: "parent_operator_ids": [ str(operator.id) for operator in self.parent_operators ], - "instance_name": self.id.hex[:8], + "instance_name": f"{self.type_name}-{self.id.hex[:8]}", "type_name": self.type_name, "custom_attributes": [attr.serialize() for attr in self.custom_attributes], } From b883dac22beb8e0030504407369d725fa997a934 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 28 Jul 2026 14:45:03 -0700 Subject: [PATCH 50/60] require tracing with --collect-traces --- python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index d609d6a63a1c..de6a2914dc50 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -2273,6 +2273,11 @@ def run_polars(benchmark: Any, args: argparse.Namespace) -> None: "(in-memory, dask, ray, spmd)." ) + if run_config.collect_traces and not cudf_polars.dsl.tracing.LOG_TRACES: + raise ValueError( + "--collect-traces is not supported when CUDF_POLARS_LOG_TRACES is not enabled. Set CUDF_POLARS_LOG_TRACES=1 and rerun." + ) + if run_config.validation_method is not None: validate_against = run_config.validation_method.expected_source if validate_against == run_config.frontend: From e2fa3a1ea386ba2dbadfc643b89bbfc1c58bb615 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Thu, 30 Jul 2026 05:23:18 -0700 Subject: [PATCH 51/60] Serialize all properties --- python/cudf_polars/cudf_polars/quent/_plan.py | 12 +- .../cudf_polars/cudf_polars/quent/_types.py | 134 ++++++++++++++---- python/cudf_polars/tests/quent/test_quent.py | 108 +++++++++++++- 3 files changed, 217 insertions(+), 37 deletions(-) diff --git a/python/cudf_polars/cudf_polars/quent/_plan.py b/python/cudf_polars/cudf_polars/quent/_plan.py index 88f5db929b48..c020bbcb8917 100644 --- a/python/cudf_polars/cudf_polars/quent/_plan.py +++ b/python/cudf_polars/cudf_polars/quent/_plan.py @@ -6,7 +6,7 @@ from __future__ import annotations import functools -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast from cudf_polars.dsl.traversal import traversal from cudf_polars.quent._types import ( @@ -23,7 +23,7 @@ import uuid from cudf_polars.dsl.ir import IR - from cudf_polars.quent._types import Query, Worker + from cudf_polars.quent._types import Query, Value, Worker from cudf_polars.utils.config import ConfigOptions, StreamingExecutor _JOIN_TYPES = frozenset({"Join", "ConditionalJoin"}) @@ -86,10 +86,14 @@ def build_plan( serializable_node = serializable_plan.nodes[node_id] operator_id = new_quent_id() - # TODO: Include serializable_node.properties as custom attributes - # We need to handle serialization of lists and dicts properly. custom_attributes = [ Attribute(name="node_id", value=node_id), + *( + # SerializablePlan properties are JSON-shaped values that map + # onto Quent Attribute Value (scalars, homogeneous lists, structs). + Attribute(name=key, value=cast("Value | None", value)) + for key, value in serializable_node.properties.items() + ), ] operator = Operator( id=operator_id, diff --git a/python/cudf_polars/cudf_polars/quent/_types.py b/python/cudf_polars/cudf_polars/quent/_types.py index d47e3c9a0404..d1f336c9e591 100644 --- a/python/cudf_polars/cudf_polars/quent/_types.py +++ b/python/cudf_polars/cudf_polars/quent/_types.py @@ -481,8 +481,23 @@ def _declare(self, engine: Engine, timestamp: int | None = None) -> Event: ScalarValue = int | float | str | bool -HomogeneousListValue = list[int] | list[float] | list[str] | list[bool] -Value: TypeAlias = ScalarValue | HomogeneousListValue | dict[str, "Value"] +StructValue: TypeAlias = dict[str, "Value | None"] +HomogeneousListValue = ( + list[int] | list[float] | list[str] | list[bool] | list[StructValue] +) +Value: TypeAlias = ScalarValue | HomogeneousListValue | StructValue + +_INT_VARIANTS: tuple[tuple[str, int, int], ...] = ( + ("U8", 0, 2**8 - 1), + ("U16", 0, 2**16 - 1), + ("U32", 0, 2**32 - 1), + ("U64", 0, 2**64 - 1), + ("I8", -(2**7), 2**7 - 1), + ("I16", -(2**15), 2**15 - 1), + ("I32", -(2**31), 2**31 - 1), + ("I64", -(2**63), 2**63 - 1), +) +_INT_VARIANT_NAMES = frozenset(name for name, _, _ in _INT_VARIANTS) @dataclasses.dataclass(frozen=True, slots=True) @@ -500,6 +515,58 @@ def deserialize(cls, payload: dict[str, Any]) -> Attribute: return cls(name=payload["key"], value=_deserialize_value(payload["value"])) +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(f"Integer list values {values} do not fit any Quent integer type.") + + +def _serialize_struct(value: StructValue) -> list[dict[str, Any]]: + """Serialize a dict as a Quent ``Struct`` (list of attributes).""" + return [ + {"key": key, "value": _serialize_value(item)} for key, item in value.items() + ] + + +def _serialize_list(values: list[Any]) -> dict[str, Any]: + """ + Serialize a homogeneous list as a Quent ``List`` payload. + + Empty lists default to ``String`` because the element type cannot be + inferred. Nested lists are not supported by Quent + (see https://github.com/rapidsai/quent/issues/79). + """ + if not values: + return {"String": []} + # bool is a subclass of int; check it before int. + if all(isinstance(item, bool) for item in values): + return {"U8": [int(item) for item in values]} + if all(isinstance(item, int) for item in values): + return {_common_integer_variant(values): values} + if all(isinstance(item, float) for item in values): + return {"F64": values} + if all(isinstance(item, str) for item in values): + return {"String": values} + if all(isinstance(item, dict) for item in values): + return {"Struct": [_serialize_struct(item) for item in values]} + if any(isinstance(item, list) for item in values): + raise NotImplementedError("Nested list attributes are not supported by Quent.") + raise TypeError( + "Quent list attributes must be homogeneous " + f"(int, float, str, bool, or dict); got {[type(v).__name__ for v in values]}" + ) + + def _serialize_value(value: Value | None) -> dict[str, Any] | None: match value: case None: @@ -508,40 +575,49 @@ def _serialize_value(value: Value | None) -> dict[str, Any] | None: # Bool is not a native Quent Value variant. return {"U8": int(value)} case int(): - if value >= 0: - if value <= 2**8 - 1: - return {"U8": value} - if value <= 2**16 - 1: - return {"U16": value} - if value <= 2**32 - 1: - return {"U32": value} - if value <= 2**64 - 1: - return {"U64": value} - else: - if -(2**7) <= value <= 2**7 - 1: - return {"I8": value} - if -(2**15) <= value <= 2**15 - 1: - return {"I16": value} - if -(2**31) <= value <= 2**31 - 1: - return {"I32": value} - if -(2**63) <= value <= 2**63 - 1: - return {"I64": value} - raise ValueError( - f"Integer value {value} does not fit any Quent integer type." - ) + return {_integer_variant(value): value} case float(): return {"F64": value} case str(): return {"String": value} - case list() | dict(): - raise NotImplementedError("List and dict attributes are not supported yet.") + case list(): + return {"List": _serialize_list(value)} + case dict(): + return {"Struct": _serialize_struct(value)} case _: # pragma: no cover; should be exhaustive raise TypeError(f"Unsupported Quent custom attribute type: {type(value)}") -def _deserialize_value(value: dict[str, Any] | None) -> Value | None: +def _deserialize_struct(payload: list[dict[str, Any]]) -> StructValue: + return {item["key"]: _deserialize_value(item["value"]) for item in payload} + + +def _deserialize_list(payload: dict[str, Any]) -> HomogeneousListValue: + n = len(payload) + if n != 1: + raise ValueError( + f"Expected Quent List envelope with exactly one variant, got '{n}' instead." + ) + variant, deserialized = next(iter(payload.items())) + if variant in _INT_VARIANT_NAMES: + return [int(item) for item in deserialized] + if variant == "F64": + return [float(item) for item in deserialized] + if variant == "String": + return [str(item) for item in deserialized] + if variant == "Struct": + return [_deserialize_struct(item) for item in deserialized] + raise ValueError(f"Unsupported Quent List variant: '{variant}'") + + +def _deserialize_value(value: dict[str, Any] | list[Any] | None) -> Value | None: if value is None: return None + if not isinstance(value, dict): + raise TypeError( + "Expected Quent attribute value envelope as a single-variant object, " + f"got {type(value).__name__}." + ) n = len(value) if n != 1: raise ValueError( @@ -549,12 +625,16 @@ def _deserialize_value(value: dict[str, Any] | None) -> Value | None: ) variant, deserialized = next(iter(value.items())) - if variant in {"U8", "U16", "U32", "U64", "I8", "I16", "I32", "I64"}: + if variant in _INT_VARIANT_NAMES: return int(deserialized) if variant == "F64": return float(deserialized) if variant == "String": return str(deserialized) + if variant == "Struct": + return _deserialize_struct(deserialized) + if variant == "List": + return _deserialize_list(deserialized) raise ValueError(f"Unsupported Quent custom attribute variant: '{variant}'") diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index c7faf8aefd82..3801680a0ccc 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -798,14 +798,78 @@ def _drained_events( return [x["event"] for x in logger.drain()] -def test_serialize_list_raises(): - with pytest.raises(NotImplementedError, match="not supported yet"): - Attribute("list", [1, 2]).serialize() +def test_serialize_list() -> None: + assert Attribute("keys", ["a", "b"]).serialize() == { + "key": "keys", + "value": {"List": {"String": ["a", "b"]}}, + } + assert Attribute("counts", [1, 2, 300]).serialize() == { + "key": "counts", + "value": {"List": {"U16": [1, 2, 300]}}, + } + assert Attribute("flags", [True, False]).serialize() == { + "key": "flags", + "value": {"List": {"U8": [1, 0]}}, + } + assert Attribute("empty", []).serialize() == { + "key": "empty", + "value": {"List": {"String": []}}, + } + assert Attribute( + "events", + [{"bytes": 1024, "kind": "disk"}], # type: ignore[arg-type] + ).serialize() == { + "key": "events", + "value": { + "List": { + "Struct": [ + [ + {"key": "bytes", "value": {"U16": 1024}}, + {"key": "kind", "value": {"String": "disk"}}, + ] + ] + } + }, + } + + +def test_serialize_nested_list_raises() -> None: + with pytest.raises(NotImplementedError, match="Nested list"): + Attribute("nested", [[1, 2], [3, 4]]).serialize() # type: ignore[arg-type] -def test_serialize_dict_raises(): - with pytest.raises(NotImplementedError, match="not supported yet"): - Attribute("dict", {"a": 1, "b": 2}).serialize() +def test_serialize_heterogeneous_list_raises() -> None: + with pytest.raises(TypeError, match="homogeneous"): + Attribute("mixed", [1, "a"]).serialize() # type: ignore[arg-type] + + +def test_serialize_dict() -> None: + assert Attribute("expr", {"type": "Col", "name": "x"}).serialize() == { + "key": "expr", + "value": { + "Struct": [ + {"key": "type", "value": {"String": "Col"}}, + {"key": "name", "value": {"String": "x"}}, + ] + }, + } + assert Attribute("nullable", {"predicate": None}).serialize() == { + "key": "nullable", + "value": {"Struct": [{"key": "predicate", "value": None}]}, + } + + +def test_attribute_list_and_dict_roundtrip() -> None: + cases = [ + Attribute("keys", ["a", "b"]), + Attribute("counts", [1, 40000]), + Attribute("ratios", [1.5, 2.5]), + Attribute("expr", {"type": "Col", "name": "x", "child": None}), + Attribute("events", [{"bytes": 1024, "kind": "disk"}]), # type: ignore[arg-type] + Attribute("empty", []), + ] + for attr in cases: + assert Attribute.deserialize(attr.serialize()) == attr def test_quent_serialize_none(): @@ -815,6 +879,38 @@ def test_quent_serialize_none(): } +def test_build_plan_includes_node_properties( + ir_and_config: tuple[IR, ConfigOptions[StreamingExecutor]], +) -> None: + ir, config_options = ir_and_config + _, operators, _, _ = build_plan( + ir, config_options, Query(), uuid.uuid4(), _make_worker() + ) + filter_op = next(op for op in operators if op.type_name == "Filter") + attrs = {attr.name: attr.value for attr in filter_op.custom_attributes} + + assert "node_id" in attrs + # Filter properties come from _serialize_properties / _serialize_expr. + assert attrs["op"] == "GREATER" + assert attrs["left"] == {"type": "Col", "name": "x"} + assert attrs["right"] == { + "type": "Literal", + "value": {"type": "int", "value": 1}, + } + assert attrs["predicate"] == "x" + + # Ensure the nested properties serialize into Quent's List/Struct envelopes. + serialized = { + attr.name: attr.serialize()["value"] for attr in filter_op.custom_attributes + } + assert serialized["left"] == { + "Struct": [ + {"key": "type", "value": {"String": "Col"}}, + {"key": "name", "value": {"String": "x"}}, + ] + } + + def test_task_from_ir() -> None: operator_id = uuid.uuid4() _logger, quent_ir_execution_context = _make_quent_ir_execution_context( From 4d406df4a592fca4b8e0b025d9cd8de2e67d43fa Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Mon, 3 Aug 2026 12:09:24 -0700 Subject: [PATCH 52/60] coverage --- .../cudf_polars/cudf_polars/quent/_types.py | 2 +- python/cudf_polars/tests/quent/test_quent.py | 35 +++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/python/cudf_polars/cudf_polars/quent/_types.py b/python/cudf_polars/cudf_polars/quent/_types.py index d1f336c9e591..46fc3bd87a91 100644 --- a/python/cudf_polars/cudf_polars/quent/_types.py +++ b/python/cudf_polars/cudf_polars/quent/_types.py @@ -528,7 +528,7 @@ def _common_integer_variant(values: list[int]) -> str: for variant, lo, hi in _INT_VARIANTS: if all(lo <= value <= hi for value in values): return variant - raise ValueError(f"Integer list values {values} do not fit any Quent integer type.") + raise ValueError("Integer list values do not fit any Quent integer type.") def _serialize_struct(value: StructValue) -> list[dict[str, Any]]: diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index 3801680a0ccc..f7670d44d571 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -173,6 +173,14 @@ def test_deserialize_value_requires_single_variant() -> None: _deserialize_value({"U8": 1, "I8": -1}) +def test_deserialize_value_requires_dict_envelope() -> None: + with pytest.raises( + TypeError, + match=r"Expected Quent attribute value envelope as a single-variant object, got list\.", + ): + _deserialize_value([{"U8": 1}]) + + def test_deserialize_value_raises_on_unknown_variant() -> None: with pytest.raises( ValueError, @@ -838,11 +846,38 @@ def test_serialize_nested_list_raises() -> None: Attribute("nested", [[1, 2], [3, 4]]).serialize() # type: ignore[arg-type] +def test_serialize_list_integer_overflow_raises() -> None: + with pytest.raises( + ValueError, + match="Integer list values", + ): + Attribute("x", [2**64]).serialize() + + def test_serialize_heterogeneous_list_raises() -> None: with pytest.raises(TypeError, match="homogeneous"): Attribute("mixed", [1, "a"]).serialize() # type: ignore[arg-type] +def test_deserialize_invalid_heterogeneous_list_raises() -> None: + with pytest.raises( + ValueError, + match="Expected Quent List envelope with exactly one variant, got '2' instead", + ): + Attribute.deserialize( + {"key": "mixed", "value": {"List": {"String": ["a", "b"], "U8": [1, 2]}}} + ) + + +def test_deserialize_unsupported_attribute_type_raises() -> None: + with pytest.raises( + ValueError, match="Unsupported Quent List variant: 'Unsupported'" + ): + Attribute.deserialize( + {"key": "unsupported", "value": {"List": {"Unsupported": ["a", "b"]}}} + ) + + def test_serialize_dict() -> None: assert Attribute("expr", {"type": "Col", "name": "x"}).serialize() == { "key": "expr", From 8b146541760e5a975ae1cc88196d6b17ce6fbf21 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 4 Aug 2026 05:09:42 -0700 Subject: [PATCH 53/60] Bump quent version this gives us the timeline animation. --- python/cudf_polars/cudf_polars/quent/_export.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/quent/_export.py b/python/cudf_polars/cudf_polars/quent/_export.py index 478edce0af5d..49308e9b1b88 100644 --- a/python/cudf_polars/cudf_polars/quent/_export.py +++ b/python/cudf_polars/cudf_polars/quent/_export.py @@ -22,7 +22,7 @@ MODEL_QMI: dict[str, Any] = { "quent": { "version": "0.1.0", - "commit": "9c2924510756d9eca288819e125330805053051f", + "commit": "153d422ae3392c24dfaf6ac5743a8682f783f864", "remote": "https://github.com/rapidsai/quent", }, "model": { @@ -31,7 +31,7 @@ "type_path": "quent_simulator_instrumentation::SimulatorEvent", "source": { "version": "0.1.0", - "commit": "9c2924510756d9eca288819e125330805053051f", + "commit": "153d422ae3392c24dfaf6ac5743a8682f783f864", "remote": "https://github.com/rapidsai/quent", }, "analyzer_package": "quent-simulator-analyzer", From 90af2f984019337ad6d834f12f66910e400c2555 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 4 Aug 2026 11:30:59 -0700 Subject: [PATCH 54/60] Comm --- python/cudf_polars/cudf_polars/engine/ray.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index eb9fe3bf3669..0ce861d7f0ee 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -334,8 +334,7 @@ def setup_worker(self, root_ucxx_address_as_bytes: bytes) -> None: progress_thread=ProgressThread(self._rapidsmpf_statistics), ) barrier(self._comm) - # Now that ``comm`` exists, declare the engine-scoped inter-rank network - # topology once (a no-op for single-rank runs). + # Now we can declare the Quent worker resources, which depends on self._comm if self._quent_logger is not None: self._quent_logger.emit(self._quent_worker._init()) self.worker_resources = WorkerResources.build( From e3193eb2e1bd4933f5c46464a6232e9e0f296667 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 4 Aug 2026 11:53:02 -0700 Subject: [PATCH 55/60] Reuse via closure --- python/cudf_polars/cudf_polars/dsl/tracing.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index ea3036698f69..ce7fc777a09d 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -162,6 +162,10 @@ def log_do_evaluate( if not LOG_TRACES: return func else: # pragma: no cover; requires CUDF_POLARS_LOG_TRACES=1 + # do this just once + pynvml.nvmlInit() + maybe_handle = get_device_handle() + pid = _getpid() @functools.wraps(func) def wrapper( @@ -171,9 +175,6 @@ def wrapper( ) -> cudf_polars.containers.DataFrame: from cudf_polars.quent._types import Task - pynvml.nvmlInit() - maybe_handle = get_device_handle() - pid = _getpid() log = structlog.get_logger() # By convention, all non-dataframe arguments (non-child) come first. From 7e60dc933ac1bfbd6479a7bc1a97eb526c88288e Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 4 Aug 2026 11:55:46 -0700 Subject: [PATCH 56/60] Cleanup exception handling --- python/cudf_polars/cudf_polars/dsl/tracing.py | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/tracing.py b/python/cudf_polars/cudf_polars/dsl/tracing.py index ce7fc777a09d..e7c405d4988b 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -214,18 +214,9 @@ def wrapper( try: result = func(cls, *args, **kwargs) except Exception: # pragma: no cover; - if ( - quent_task is not None - and ir_execution_context.quent_ir_execution_context is not None - ): - ir_execution_context.quent_ir_execution_context.context._emit_task_end_events( - cls, - quent_task, - ir_execution_context.quent_ir_execution_context, - None, - ) + result = None raise - else: + finally: if ( quent_task is not None and ir_execution_context.quent_ir_execution_context is not None @@ -239,6 +230,8 @@ def wrapper( ) stop = time.monotonic_ns() + assert result is not None + after_start = time.monotonic_ns() after = make_snapshot( cls, From ec427500aba1dde7c3092fc7fed4d25ac7bc7960 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 4 Aug 2026 12:03:37 -0700 Subject: [PATCH 57/60] Updated size_bytes test --- python/cudf_polars/tests/containers/test_dataframe.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/tests/containers/test_dataframe.py b/python/cudf_polars/tests/containers/test_dataframe.py index 941dfd28eab8..8ae919d61a5a 100644 --- a/python/cudf_polars/tests/containers/test_dataframe.py +++ b/python/cudf_polars/tests/containers/test_dataframe.py @@ -228,6 +228,10 @@ def test_serialization_roundtrip(polars_tbl): def test_size_bytes(): stream = get_cuda_stream() - df = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]}) + df = pl.DataFrame( + { + "a": pl.Series([1, 2, 3], dtype=pl.Int64()), + } + ) df = DataFrame.from_polars(df, stream=stream) - assert df._size_bytes() == df.to_polars().estimated_size() + assert df._size_bytes() == 24 From 15153e72e39941a9810f8a7eb5adb0d665d44905 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 4 Aug 2026 12:19:51 -0700 Subject: [PATCH 58/60] simplify int sizes --- .../cudf_polars/cudf_polars/quent/_types.py | 34 ++++++++----------- python/cudf_polars/tests/quent/test_quent.py | 21 ++++-------- 2 files changed, 21 insertions(+), 34 deletions(-) diff --git a/python/cudf_polars/cudf_polars/quent/_types.py b/python/cudf_polars/cudf_polars/quent/_types.py index 46fc3bd87a91..df13ed45f6d8 100644 --- a/python/cudf_polars/cudf_polars/quent/_types.py +++ b/python/cudf_polars/cudf_polars/quent/_types.py @@ -487,17 +487,11 @@ def _declare(self, engine: Engine, timestamp: int | None = None) -> Event: ) Value: TypeAlias = ScalarValue | HomogeneousListValue | StructValue -_INT_VARIANTS: tuple[tuple[str, int, int], ...] = ( - ("U8", 0, 2**8 - 1), - ("U16", 0, 2**16 - 1), - ("U32", 0, 2**32 - 1), - ("U64", 0, 2**64 - 1), - ("I8", -(2**7), 2**7 - 1), - ("I16", -(2**15), 2**15 - 1), - ("I32", -(2**31), 2**31 - 1), - ("I64", -(2**63), 2**63 - 1), -) -_INT_VARIANT_NAMES = frozenset(name for name, _, _ in _INT_VARIANTS) +_I64_MIN = -(2**63) +_I64_MAX = 2**63 - 1 +_U64_MAX = 2**64 - 1 +# Accepted on deserialize for compatibility with other Quent producers. +_INT_VARIANT_NAMES = frozenset({"U8", "U16", "U32", "U64", "I8", "I16", "I32", "I64"}) @dataclasses.dataclass(frozen=True, slots=True) @@ -516,18 +510,20 @@ def deserialize(cls, payload: dict[str, Any]) -> Attribute: 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 + """Return ``I64`` if ``value`` fits, otherwise ``U64`` up to ``2**64 - 1``.""" + if _I64_MIN <= value <= _I64_MAX: + return "I64" + if _I64_MAX < value <= _U64_MAX: + return "U64" 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 + """Return ``I64`` if all ``values`` fit, otherwise ``U64`` when all are unsigned.""" + if all(_I64_MIN <= value <= _I64_MAX for value in values): + return "I64" + if all(0 <= value <= _U64_MAX for value in values): + return "U64" raise ValueError("Integer list values do not fit any Quent integer type.") diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index f7670d44d571..fae4c4511d45 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -116,20 +116,11 @@ def _make_quent_ir_execution_context( @pytest.mark.parametrize( "value,expected_variant", [ - (0, "U8"), - (2**8 - 1, "U8"), - (2**8, "U16"), - (2**16 - 1, "U16"), - (2**16, "U32"), - (2**32 - 1, "U32"), - (2**32, "U64"), + (0, "I64"), + (2**63 - 1, "I64"), + (2**63, "U64"), (2**64 - 1, "U64"), - (-(2**7), "I8"), - (2**7 - 1, "U8"), - (-(2**15), "I16"), - (-(2**15) - 1, "I32"), - (-(2**31), "I32"), - (-(2**31) - 1, "I64"), + (-1, "I64"), (-(2**63), "I64"), ], ) @@ -813,7 +804,7 @@ def test_serialize_list() -> None: } assert Attribute("counts", [1, 2, 300]).serialize() == { "key": "counts", - "value": {"List": {"U16": [1, 2, 300]}}, + "value": {"List": {"I64": [1, 2, 300]}}, } assert Attribute("flags", [True, False]).serialize() == { "key": "flags", @@ -832,7 +823,7 @@ def test_serialize_list() -> None: "List": { "Struct": [ [ - {"key": "bytes", "value": {"U16": 1024}}, + {"key": "bytes", "value": {"I64": 1024}}, {"key": "kind", "value": {"String": "disk"}}, ] ] From 49eb2322d471f621beb6b7e61a76869ee5348153 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 4 Aug 2026 12:21:58 -0700 Subject: [PATCH 59/60] Revert "simplify int sizes" This reverts commit 15153e72e39941a9810f8a7eb5adb0d665d44905. --- .../cudf_polars/cudf_polars/quent/_types.py | 34 +++++++++++-------- python/cudf_polars/tests/quent/test_quent.py | 21 ++++++++---- 2 files changed, 34 insertions(+), 21 deletions(-) diff --git a/python/cudf_polars/cudf_polars/quent/_types.py b/python/cudf_polars/cudf_polars/quent/_types.py index df13ed45f6d8..46fc3bd87a91 100644 --- a/python/cudf_polars/cudf_polars/quent/_types.py +++ b/python/cudf_polars/cudf_polars/quent/_types.py @@ -487,11 +487,17 @@ def _declare(self, engine: Engine, timestamp: int | None = None) -> Event: ) Value: TypeAlias = ScalarValue | HomogeneousListValue | StructValue -_I64_MIN = -(2**63) -_I64_MAX = 2**63 - 1 -_U64_MAX = 2**64 - 1 -# Accepted on deserialize for compatibility with other Quent producers. -_INT_VARIANT_NAMES = frozenset({"U8", "U16", "U32", "U64", "I8", "I16", "I32", "I64"}) +_INT_VARIANTS: tuple[tuple[str, int, int], ...] = ( + ("U8", 0, 2**8 - 1), + ("U16", 0, 2**16 - 1), + ("U32", 0, 2**32 - 1), + ("U64", 0, 2**64 - 1), + ("I8", -(2**7), 2**7 - 1), + ("I16", -(2**15), 2**15 - 1), + ("I32", -(2**31), 2**31 - 1), + ("I64", -(2**63), 2**63 - 1), +) +_INT_VARIANT_NAMES = frozenset(name for name, _, _ in _INT_VARIANTS) @dataclasses.dataclass(frozen=True, slots=True) @@ -510,20 +516,18 @@ def deserialize(cls, payload: dict[str, Any]) -> Attribute: def _integer_variant(value: int) -> str: - """Return ``I64`` if ``value`` fits, otherwise ``U64`` up to ``2**64 - 1``.""" - if _I64_MIN <= value <= _I64_MAX: - return "I64" - if _I64_MAX < value <= _U64_MAX: - return "U64" + """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 ``I64`` if all ``values`` fit, otherwise ``U64`` when all are unsigned.""" - if all(_I64_MIN <= value <= _I64_MAX for value in values): - return "I64" - if all(0 <= value <= _U64_MAX for value in values): - return "U64" + """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.") diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index fae4c4511d45..f7670d44d571 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -116,11 +116,20 @@ def _make_quent_ir_execution_context( @pytest.mark.parametrize( "value,expected_variant", [ - (0, "I64"), - (2**63 - 1, "I64"), - (2**63, "U64"), + (0, "U8"), + (2**8 - 1, "U8"), + (2**8, "U16"), + (2**16 - 1, "U16"), + (2**16, "U32"), + (2**32 - 1, "U32"), + (2**32, "U64"), (2**64 - 1, "U64"), - (-1, "I64"), + (-(2**7), "I8"), + (2**7 - 1, "U8"), + (-(2**15), "I16"), + (-(2**15) - 1, "I32"), + (-(2**31), "I32"), + (-(2**31) - 1, "I64"), (-(2**63), "I64"), ], ) @@ -804,7 +813,7 @@ def test_serialize_list() -> None: } assert Attribute("counts", [1, 2, 300]).serialize() == { "key": "counts", - "value": {"List": {"I64": [1, 2, 300]}}, + "value": {"List": {"U16": [1, 2, 300]}}, } assert Attribute("flags", [True, False]).serialize() == { "key": "flags", @@ -823,7 +832,7 @@ def test_serialize_list() -> None: "List": { "Struct": [ [ - {"key": "bytes", "value": {"I64": 1024}}, + {"key": "bytes", "value": {"U16": 1024}}, {"key": "kind", "value": {"String": "disk"}}, ] ] From 6090c4793434cc5cc689624e5d833ea0e33445ea Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 5 Aug 2026 07:37:02 -0700 Subject: [PATCH 60/60] Fixed engine config serialization --- .../cudf_polars/streaming/benchmarks/utils.py | 37 ++++++++++++------- 1 file changed, 24 insertions(+), 13 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index 7e0f737ed8d5..c5722f532c54 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -655,7 +655,7 @@ def from_args(cls, args: argparse.Namespace) -> RunConfig: ) def serialize( - self, engine: StreamingEngine | None, quent_archive: Path | None + self, engine: pl.GPUEngine | None, quent_archive: Path | None ) -> dict: """ Serialize the run config to a dictionary. @@ -713,11 +713,14 @@ def serialize( if engine is not None: config_options = ConfigOptions.from_polars_engine(engine) config_options = config_options.drop_unserializable() - rapidsmpf_options = engine.rapidsmpf_options.get_strings() - result["config_options"] = { + extra = { "config_options": dataclasses.asdict(config_options), - "rapidsmpf_options": rapidsmpf_options, } + + if isinstance(engine, StreamingEngine): + extra["rapidsmpf_options"] = engine.rapidsmpf_options.get_strings() + + result["config_options"] = extra # discard unserializable / unnecessary UUIDs result["config_options"]["config_options"]["executor"].pop( "quent_context", None @@ -1182,7 +1185,9 @@ def _run_query_loop( for q_id in run_config.queries: if engine is not None: - quent_context = engine.config["executor_options"].get("quent_context") + quent_context = engine.config.get("executor_options", {}).get( + "quent_context" + ) if quent_context is not None: engine.config["executor_options"]["quent_context"] = ( dataclasses.replace( @@ -1267,7 +1272,7 @@ def _finalize_benchmark_run( def run_polars_cpu( benchmark: Any, args: argparse.Namespace, - run_config: Any, + run_config: RunConfig, numeric_type: str, date_type: str, ) -> None: @@ -1286,14 +1291,16 @@ def run_polars_cpu( run_config, validation_failures, query_failures, - serializable_engine_config=run_config.serialize(engine=None), + serializable_engine_config=run_config.serialize( + engine=None, quent_archive=None + ), ) def run_polars_in_memory( benchmark: Any, args: argparse.Namespace, - run_config: Any, + run_config: RunConfig, parquet_options: dict[str, Any], numeric_type: str, date_type: str, @@ -1323,14 +1330,16 @@ def run_polars_in_memory( run_config, validation_failures, query_failures, - serializable_engine_config=run_config.serialize(engine=engine), + serializable_engine_config=run_config.serialize( + engine=engine, quent_archive=None + ), ) def run_polars_spmd( benchmark: Any, args: argparse.Namespace, - run_config: Any, + run_config: RunConfig, parquet_options: dict[str, Any], numeric_type: str, date_type: str, @@ -1407,7 +1416,7 @@ def _allgather_result(df: pl.DataFrame) -> pl.DataFrame: def run_polars_ray( benchmark: Any, args: argparse.Namespace, - run_config: Any, + run_config: RunConfig, parquet_options: dict[str, Any], numeric_type: str, date_type: str, @@ -1447,7 +1456,9 @@ def run_polars_ray( run_config = dataclasses.replace(run_config, records=dict(records), plans=plans) run_config = _consolidate_logs(run_config, engine=engine) # We need to create this before StreamingEngine.shutdown(), which clears engine.config - serializable_engine_config = run_config.serialize(engine=engine) + serializable_engine_config = run_config.serialize( + engine=engine, quent_archive=None + ) quent_archive = Path("logs") / f"{run_config.run_id}.zip" _write_quent_traces( @@ -1468,7 +1479,7 @@ def run_polars_ray( def run_polars_dask( benchmark: Any, args: argparse.Namespace, - run_config: Any, + run_config: RunConfig, parquet_options: dict[str, Any], numeric_type: str, date_type: str,