diff --git a/.gitignore b/.gitignore index cb4b70da0503..ac3cc314ab23 100644 --- a/.gitignore +++ b/.gitignore @@ -189,4 +189,6 @@ rmm_log.txt python/cudf/cudf_pandas_tests/data/rmm_log.txt # Quent traces -logs/*.ndjson +logs/**/*.ndjson +logs/**/*.qmi +logs/*.zip diff --git a/python/cudf_polars/cudf_polars/containers/dataframe.py b/python/cudf_polars/cudf_polars/containers/dataframe.py index fc671bb0e2b0..01261b035649 100644 --- a/python/cudf_polars/cudf_polars/containers/dataframe.py +++ b/python/cudf_polars/cudf_polars/containers/dataframe.py @@ -106,6 +106,10 @@ def __init__( self.table = plc.Table([c.obj for c in self.columns], num_rows=num_rows) self.stream = stream + 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/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 559777f43c7d..c5bd73f6b68c 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -84,6 +84,8 @@ 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 @@ -138,11 +140,17 @@ 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. + tracer + The actor tracer. Used to propagate statistics. """ 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 @@ -244,6 +252,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. @@ -697,6 +708,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, @@ -1637,6 +1650,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 2c1dd1f26f0c..e7c405d4988b 100644 --- a/python/cudf_polars/cudf_polars/dsl/tracing.py +++ b/python/cudf_polars/cudf_polars/dsl/tracing.py @@ -49,6 +49,7 @@ import cudf_polars.containers from cudf_polars.dsl import ir + from cudf_polars.dsl.ir import IRExecutionContext class Scope(enum.StrEnum): @@ -161,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( @@ -168,10 +173,8 @@ def wrapper( *args: P.args, **kwargs: P.kwargs, ) -> cudf_polars.containers.DataFrame: - # do this just once - pynvml.nvmlInit() - maybe_handle = get_device_handle() - pid = _getpid() + from cudf_polars.quent._types import Task + log = structlog.get_logger() # By convention, all non-dataframe arguments (non-child) come first. @@ -180,6 +183,23 @@ def wrapper( 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] + + if ir_execution_context.quent_ir_execution_context is not None: + quent_task = Task.from_ir( + cls, ir_execution_context.quent_ir_execution_context + ) + ir_execution_context.quent_ir_execution_context.context._emit_task_begin_events( + cls, + quent_task, + ir_execution_context.quent_ir_execution_context, + 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 @@ -191,9 +211,27 @@ 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: # pragma: no cover; + result = None + raise + finally: + if ( + 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, + result, + ) stop = time.monotonic_ns() + assert result is not None + after_start = time.monotonic_ns() after = make_snapshot( cls, @@ -215,6 +253,11 @@ def wrapper( ) log.info("Execute IR", **record) + 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/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index b382f12099a4..b26b400ef380 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -32,7 +32,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 @@ -54,8 +54,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.dsl.translate import Translator from cudf_polars.quent._context import LocalQuentContext @@ -461,6 +461,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._types.Operator] | None = None, + local_quent_context: LocalQuentContext | None = None, ) -> tuple[DataFrame, list[ChannelMetadata]]: """ Execute a Polars IR query on a single rank's GPU. @@ -487,6 +490,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 ------- @@ -507,6 +516,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: @@ -733,6 +744,19 @@ def evaluate_on_rank( Collected channel metadata. """ stats = allgather_stats(comm, ctx.br(), ir, config_options, py_executor) + # ``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())) + + 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 lowering, node_map = lower_ir_graph_with_node_map( ir, config_options, stats, rank=comm.rank, nranks=comm.nranks @@ -740,13 +764,14 @@ def evaluate_on_rank( optimized = lowering.optimized ir = lowering.lowered partition_info = lowering.partition_info + # TODO: figure out if we emit anything about optimized. + 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, - query=local_quent_context.context.query, + query=local_quent_context.query, plan_id=logical_plan_id, worker=local_quent_context.worker, instance_name="logical", @@ -764,7 +789,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, @@ -774,6 +799,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 ) @@ -796,6 +822,8 @@ def evaluate_on_rank( config_options, stats, collective_id_map, + quent_operator_map=quent_operator_map, + local_quent_context=local_quent_context, ) diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 20b0139650b0..44646ad4e284 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -49,7 +49,10 @@ PersistedBackend, execute_persisted_query, ) -from cudf_polars.quent._context import LocalQuentContext +from cudf_polars.quent._context import ( + LocalQuentContext, + WorkerResources, +) from cudf_polars.unstable import unstable from cudf_polars.utils.config import DaskContext, MemoryResourceConfig @@ -64,6 +67,7 @@ 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._context import QuentContext from cudf_polars.streaming.parallel import ConfigOptions from cudf_polars.utils.config import StreamingExecutor @@ -130,7 +134,8 @@ class _WorkerContext: quent_logger: cudf_polars.quent._logging.QuentLogger | None quent_worker: cudf_polars.quent._types.Worker statistics: Statistics - mr: RmmResourceAdaptor | None = None # set after `Context` is built (below). + mr: RmmResourceAdaptor | None = None + worker_resources: WorkerResources | None = None def _worker_evaluate_persisted( @@ -236,7 +241,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. @@ -326,7 +331,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: """ @@ -412,11 +417,19 @@ 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() + worker_resources = WorkerResources.build( + instance_suffix=f"rank-{comm.rank}", + engine_id=engine_id, + worker_id=worker_id, + rank=comm.rank, + nranks=comm.nranks, ) + quent_logger.emit(quent_worker._init()) + worker_resources.declare(quent_logger) else: quent_logger = None + worker_resources = None mp_ctx = _WorkerContext( comm=comm, @@ -426,11 +439,10 @@ def _setup_worker( mr=mr, quent_worker=quent_worker, quent_logger=quent_logger, + 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( @@ -454,7 +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.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() @@ -569,7 +583,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[int, pl.DataFrame, list[ChannelMetadata] | None]: """ @@ -613,11 +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 local_quent_context = LocalQuentContext( context=quent_context, + query=quent_context.query_for(query_id), worker=mp_ctx.quent_worker, logger=mp_ctx.quent_logger, + 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. @@ -702,8 +719,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) worker_config = config_options.drop_unserializable() result_map = dask_context.client.run( @@ -725,7 +743,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) ranked.sort(key=lambda p: p[0]) dfs = [df for _, df in ranked] @@ -832,9 +850,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: @@ -1118,9 +1134,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 295a081307fa..beda8f643400 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, JoinFilterPushdownOptions, diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index b8c7067119dd..0ce861d7f0ee 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -45,7 +45,10 @@ PersistedBackend, execute_persisted_query, ) -from cudf_polars.quent._context import LocalQuentContext +from cudf_polars.quent._context import ( + LocalQuentContext, + WorkerResources, +) from cudf_polars.quent._types import Worker from cudf_polars.unstable import unstable from cudf_polars.utils.config import MemoryResourceConfig, RayContext @@ -63,6 +66,7 @@ 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._context import QuentContext from cudf_polars.streaming.parallel import ConfigOptions from cudf_polars.utils.config import StreamingExecutor @@ -162,8 +166,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. @@ -192,7 +197,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 @@ -272,13 +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]}", ) - if self._quent_logger is not None: - self._quent_logger.emit(self._quent_worker._init()) + # Initialized later in setup_worker once ``comm`` is available. + self.worker_resources: WorkerResources | None = None def setup_root(self) -> bytes: """ @@ -327,6 +334,18 @@ def setup_worker(self, root_ucxx_address_as_bytes: bytes) -> None: progress_thread=ProgressThread(self._rapidsmpf_statistics), ) barrier(self._comm) + # 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( + 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._nranks, + ) + self.worker_resources.declare(self._quent_logger) + assert self._base_mr is not None self._ctx = Context.from_options( self._comm.logger, @@ -387,6 +406,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.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 [] @@ -454,7 +476,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]: """ @@ -500,10 +522,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.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, + 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 @@ -728,20 +753,14 @@ 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: - 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. @@ -985,9 +1004,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 bc1fc56fe8e7..c87e2877f236 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -49,7 +49,11 @@ PersistedBackend, execute_persisted_query, ) -from cudf_polars.quent._context import LocalQuentContext +from cudf_polars.quent._context import ( + LocalQuentContext, + QuentContext, + WorkerResources, +) 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 @@ -122,22 +126,29 @@ 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 + 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) + quent_context._emit_query_events(quent_logger, query) + worker_id = config_options.executor.spmd_context.worker_id local_quent_context = LocalQuentContext( context=quent_context, + query=query, 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, + worker_resources=spmd_context.worker_resources, ) df, metadata = evaluate_on_rank( @@ -151,8 +162,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 + # 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 @@ -229,7 +244,7 @@ def synchronize_quent_context( *, comm: Communicator, context: Context, -) -> cudf_polars.quent.QuentContext: +) -> QuentContext: """ Ensure all ranks use the same Quent engine ID. @@ -237,19 +252,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): @@ -408,9 +423,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: @@ -452,6 +465,7 @@ def __init__( self._comm: Communicator | None = comm self._ctx: Context | None = None self._py_executor: ThreadPoolExecutor | None = None + self._store_uid = uuid.uuid4().hex exit_stack = contextlib.ExitStack() @@ -485,6 +499,22 @@ 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()) + + worker_resources = WorkerResources.build( + instance_suffix=f"rank-{self.rank}", + engine_id=engine_id, + worker_id=self._quent_worker.id, + rank=comm.rank, + nranks=comm.nranks, + ) + 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 # rapidsmpf's `run_actor_network` awaits its only submitted @@ -510,6 +540,7 @@ def __init__( quent_logger=self._quent_logger, context=self._ctx, py_executor=self._py_executor, + worker_resources=self._worker_resources, ), }, engine_options={ @@ -518,9 +549,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 @@ -600,9 +628,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. @@ -658,6 +684,7 @@ def _reset( engine_id=engine_id, worker_id=self._quent_worker.id, quent_logger=self._quent_logger, + worker_resources=self._worker_resources, ), }, engine_options={ @@ -797,10 +824,14 @@ def shutdown(self) -> None: # Clear the references only after shutdown completes. 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: cudf_polars.quent.QuentContext | None = self.config[ - "executor_options" - ].get("quent_context") + + 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/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index 627d864cfa69..e0a0b1ce408c 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 @@ -16,30 +17,82 @@ ) from cudf_polars.quent._types import ( Attribute, + Channel, Engine, Implementation, + Memory, + Network, + Processor, Query, QueryGroup, + ThreadPool, ) +from cudf_polars.utils.config import get_total_device_memory if TYPE_CHECKING: from typing import Self + 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 __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 (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: + self._processors: dict[int, Processor] = {} + self._lock = threading.Lock() + + 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 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: + 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: """ @@ -157,19 +210,35 @@ 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. + + 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) + + 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()) + 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, @@ -304,16 +373,276 @@ 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, + input_frames_bytes: int, + ) -> None: + """ + Emit begin events for a Quent Task. -@dataclasses.dataclass(frozen=True, kw_only=True) + 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. + input_frames_bytes: int + The total size of the input dataframes in bytes. + + Notes + ----- + This emits the following events: + + - queueing + - allocating (with the Quent Processor for the current thread) + - 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()) + 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.loading( + use_thread=quent_processor, + use_channel=quent_ir_execution_context.worker_resources.disk_to_device_channel, + channel_capacity_bytes=input_frames_bytes, + use_memory=quent_ir_execution_context.worker_resources.device_memory, + memory_capacity_bytes=input_frames_bytes, + ) + ) + else: + quent_ir_execution_context.logger.emit( + quent_task.computing( + use_thread=quent_processor, + use_memory=quent_ir_execution_context.worker_resources.device_memory, + input_bytes=input_frames_bytes, + memory_capacity_bytes=input_frames_bytes, + ) + ) + + def _emit_task_end_events( + self, + ir_type: type[IR], + quent_task: Task, + quent_ir_execution_context: QuentIRExecutionContext, + result: DataFrame | None, + ) -> None: + """ + 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] + 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. + 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 emits the following events: + + - computing (I/O nodes only) + - exit + """ + if quent_ir_execution_context is None: # pragma: no cover; + return + + if result is not None: + output_capacity_bytes = result._size_bytes() + else: # pragma: no cover; + output_capacity_bytes = 0 + 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.worker_resources.device_memory, + memory_capacity_bytes=output_capacity_bytes, + ) + ) + quent_ir_execution_context.logger.emit(quent_task.exit()) + + # 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. + + +@dataclasses.dataclass(kw_only=True) +class WorkerResources: + """A simple container for per-worker Quent resources.""" + + 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 + 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, + 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 / Link Channels + for link in self.link_channels.values(): + logger.emit(link.finalizing()) + logger.emit(link.exit()) + + +@dataclasses.dataclass(kw_only=True) class LocalQuentContext: """ A Quent Context that is only ever used on the local worker rank. This can contain non-serializable objects (like a ``QuentLogger``) and entities that are only valid on the local rank. + + 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 ``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 + 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.worker_resources.processor_registry.get_or_declare_processor( + self.logger, + thread_ident=thread_ident, + pool_id=self.worker_resources.thread_pool.id, + ) + + +@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, + query=execution_context.query, + worker=execution_context.worker, + logger=execution_context.logger, + worker_resources=execution_context.worker_resources, + ) 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..7307fb949115 --- /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 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 + +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": "153d422ae3392c24dfaf6ac5743a8682f783f864", + "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": "153d422ae3392c24dfaf6ac5743a8682f783f864", + "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": {...}}``; archive 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 archive 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_quent_export( + events: list[dict[str, Any]], + export_root: Path, + context_id: uuid.UUID, + quent_archive: Path, + *, + sidecar: dict[str, Any] | None = None, +) -> Path: + """ + Write Quent events to a ZIP archive. + + Parameters + ---------- + events + Buffered Quent event envelopes from ``engine._quent_events``. + export_root + Directory for exported archives (e.g. ``logs``). + context_id + Context UUID, typically the engine/run id. + quent_archive + Quent archive path. The archive will be written to this path. + sidecar + Optional provenance payload for ``model.qmi``. Defaults to + :data:`MODEL_QMI`. + + Returns + ------- + Path + The archive path ``export_root/.zip``. The archive contains + the Quent export layout under a top-level ``/`` directory. + """ + grouped: dict[str, list[dict[str, Any]]] = {} + for event in events: + directory, export_line = to_export_line(event) + grouped.setdefault(directory, []).append(export_line) + + export_root.mkdir(parents=True, exist_ok=True) + 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") + + tmp_path.replace(quent_archive) + return quent_archive diff --git a/python/cudf_polars/cudf_polars/quent/_plan.py b/python/cudf_polars/cudf_polars/quent/_plan.py index edbb41d1efb9..c020bbcb8917 100644 --- a/python/cudf_polars/cudf_polars/quent/_plan.py +++ b/python/cudf_polars/cudf_polars/quent/_plan.py @@ -6,8 +6,9 @@ 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 ( Attribute, Edge, @@ -22,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"}) @@ -85,16 +86,19 @@ 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, plan=plan, parent_operators=parent_ops.get(node_id, []), - instance_name=serializable_node.type, type_name=serializable_node.type, custom_attributes=custom_attributes, ) @@ -177,3 +181,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..46fc3bd87a91 100644 --- a/python/cudf_polars/cudf_polars/quent/_types.py +++ b/python/cudf_polars/cudf_polars/quent/_types.py @@ -9,13 +9,20 @@ import dataclasses import enum +import itertools import sys import time import uuid -from typing import Any, TypeAlias +from typing import TYPE_CHECKING, Any, Literal, Self, TypeAlias 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 + QUENT_SCOPE = "QUENT" @@ -30,6 +37,11 @@ class EventName(enum.Enum): OPERATOR = "Operator" PORT = "Port" TASK = "Task" + MEMORY = "Memory" + CHANNEL = "Channel" + THREAD_POOL = "ThreadPool" + PROCESSOR = "Processor" + NETWORK = "Network" if sys.version_info >= (3, 14): # pragma: no cover; requires Python 3.14+ @@ -72,6 +84,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: """ @@ -93,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) @@ -110,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": f"{self.type_name}-{self.id.hex[:8]}", "type_name": self.type_name, "custom_attributes": [attr.serialize() for attr in self.custom_attributes], } @@ -123,6 +178,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: @@ -418,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) @@ -437,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("Integer list 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: @@ -445,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( @@ -486,10 +625,474 @@ 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}'") + + +# 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 + + +@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 | None = None, 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={ + EventName.NETWORK.value: { + "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 + _seq: Iterator[int] = dataclasses.field( + default_factory=itertools.count, compare=False, repr=False + ) + + @classmethod + def from_ir( + cls, ir_type: type[IR], quent_ir_execution_context: QuentIRExecutionContext + ) -> Self: + """ + Build an operator-scoped Quent Task from an IR execution context. + + Parameters + ---------- + ir_type + The IR type of the operator. + quent_ir_execution_context + The Quent IR execution context, which is used to get the operator ID. + + Returns + ------- + Task | None + The operator-scoped Quent Task, or ``None`` if the IR execution context + is not bound to a Quent operator. + """ + 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( + id=self.id, + timestamp=timestamp if timestamp is not None else time.time_ns(), + data={ + EventName.TASK.value: { + "seq": next(self._seq), + "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, + 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": next(self._seq), + "state": { + "Allocating": { + "use_thread": { + "resource_id": str(resource_id), + "capacity": None, + } + } + }, + } + }, + ) + + 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": rate_usage_capacity_bytes(channel_capacity_bytes), + } + if use_memory is not None: + loading_data["use_memory"] = { + "resource_id": str(use_memory.id), + "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": next(self._seq), + "state": {"Loading": loading_data}, + } + }, + ) + + 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, 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), + "capacity": None, + } + if use_memory is not None: + computing_data["use_memory"] = { + "resource_id": str(use_memory.id), + "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": next(self._seq), + "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": 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": next(self._seq), + "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": next(self._seq), "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 020224917f17..27c6042ac36f 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 @@ -32,6 +32,7 @@ from cudf_polars.dsl.traversal import traversal 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 ( @@ -571,6 +572,7 @@ def _( # Create output ChannelManager channels[ir] = ChannelManager(rec.state["context"]) + ir_context = ir_context_for_node(rec, ir) # Complete shuffle node nodes[ir] = [ @@ -578,7 +580,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(), keys_to_hash=ir.keys, 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 a1e78f43d42c..c91d3691e49a 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,8 @@ from rapidsmpf.streaming.core.leaf_actor import pull_from_channel import cudf_polars.dsl.tracing +import cudf_polars.quent._context +import cudf_polars.quent._types from cudf_polars.dsl.ir import ( DataFrameScan, Join, @@ -94,6 +96,7 @@ def evaluate_logical_plan( engine_id=engine_id, worker_id=engine._quent_worker.id, quent_logger=engine._quent_logger, + worker_resources=engine._worker_resources, ), ), ) @@ -219,6 +222,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._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. @@ -245,6 +250,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 ------- @@ -277,6 +288,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..e5c7e8a3507d 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,8 @@ from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.streaming.core.context import Context + 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 ( @@ -58,6 +61,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 +76,42 @@ class GenState(TypedDict): max_io_threads: int stats: StatsCollector collective_id_map: dict[IR, list[int]] + 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. + + 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"] + 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._context.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 f0c4979fef6f..3302cc959a33 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 ( @@ -776,6 +777,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)}" ) @@ -784,7 +786,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 0a6a568f0e1e..b9196cfb55f4 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 @@ -32,6 +33,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, @@ -280,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(), @@ -302,7 +306,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: [ @@ -498,7 +502,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: [ @@ -549,6 +553,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, @@ -598,6 +605,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, @@ -652,7 +660,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(), @@ -780,6 +790,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" @@ -825,7 +836,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: @@ -833,7 +844,7 @@ def _( scan_node( rec.state["context"], ir, - rec.state["ir_context"], + ir_context, ch_out, num_producers=num_producers, estimated_chunk_bytes=( @@ -965,12 +976,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 d944847e621b..1fa28241fbfb 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 dataclass +from dataclasses import dataclass, replace from typing import TYPE_CHECKING, Any, Literal from cudf_streaming.channel_metadata import ( @@ -34,6 +34,7 @@ ) 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 @@ -145,6 +146,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, @@ -986,6 +988,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), @@ -1115,6 +1118,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) @@ -1123,7 +1127,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(), @@ -1155,7 +1159,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(), @@ -1178,7 +1182,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 e3670a0e9a2d..b2c74e0a1006 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 ( @@ -86,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( @@ -525,6 +529,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 @@ -532,7 +537,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(), ) @@ -543,7 +548,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), ) @@ -600,7 +605,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())] @@ -635,6 +640,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( @@ -642,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, ) else: # "bounded" fanout_node = fanout_node_bounded( @@ -650,7 +656,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 8be6f71bdeda..9dfaf1dc0fbc 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, @@ -824,12 +827,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/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/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 1b7153954f2f..0a4320693d02 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 @@ -37,9 +37,17 @@ from rapidsmpf.streaming.core.message import Message import cudf_polars.dsl.tracing +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 Filter, GroupBy, HStack, Join, Projection, Select +from cudf_polars.dsl.ir import ( + Filter, + GroupBy, + HStack, + Join, + Projection, + Select, +) from cudf_polars.dsl.tracing import Scope from cudf_polars.dsl.utils.column_domain import column_domain_bindings from cudf_polars.dsl.utils.naming import names_to_indices @@ -293,6 +301,7 @@ async def shutdown_on_error( if ir_context is not None: contextvars["cudf_polars_query_id"] = str(ir_context.query_id) + ir_context = replace(ir_context, tracer=tracer) with cudf_polars.dsl.tracing.bound_contextvars(**contextvars): start = time.monotonic_ns() @@ -321,6 +330,54 @@ async def shutdown_on_error( "Streaming Actor", start=start, stop=stop, **record ) + 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( + 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 + 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, ...] @@ -742,6 +799,8 @@ def _evaluate_chunk_sync( The IR execution context. br The buffer resource for lifetime tracking. + tracer + The actor tracer. Returns ------- @@ -795,7 +854,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/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index ce4f7c0d4b81..c5722f532c54 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 @@ -140,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 @@ -647,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: pl.GPUEngine | 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, @@ -665,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, @@ -683,16 +708,19 @@ 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) 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 @@ -745,7 +773,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) ) @@ -1157,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( @@ -1242,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: @@ -1261,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, @@ -1298,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, @@ -1358,13 +1392,17 @@ def _allgather_result(df: pl.DataFrame) -> pl.DataFrame: run_config, engine=engine, gather_client_logs=False ) # We need to create this before StreamingEngine.shutdown(), which clears engine.config - serializable_engine_config = run_config.serialize(engine=engine) + quent_archive = Path("logs") / f"{run_config.run_id}.zip" + serializable_engine_config = run_config.serialize( + engine=engine, quent_archive=quent_archive + ) if is_rank_0: _write_quent_traces( engine=engine, run_id=run_config.run_id, collect_traces=run_config.collect_traces, + quent_archive=quent_archive, ) _finalize_benchmark_run( args, @@ -1378,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, @@ -1418,12 +1456,16 @@ 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( engine=engine, run_id=run_config.run_id, collect_traces=run_config.collect_traces, + quent_archive=quent_archive, ) _finalize_benchmark_run( args, @@ -1437,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, @@ -1483,12 +1525,16 @@ def run_polars_dask( ) run_config = _consolidate_logs(run_config, engine) # We need to create this before StreamingEngine.shutdown(), which clears engine.config - serializable_engine_config = run_config.serialize(engine=engine) + quent_archive = Path("logs") / f"{run_config.run_id}.zip" + serializable_engine_config = run_config.serialize( + engine=engine, quent_archive=quent_archive + ) _write_quent_traces( engine=engine, run_id=run_config.run_id, collect_traces=run_config.collect_traces, + quent_archive=quent_archive, ) finally: if dask_client is not None: @@ -1576,15 +1622,21 @@ 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.""" + engine: StreamingEngine, + run_id: uuid.UUID, + *, + collect_traces: bool, + quent_archive: Path, +) -> 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 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 @@ -1596,13 +1648,9 @@ 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") + output_path = write_quent_export(quent_logs, logs_dir, run_id, quent_archive) print(f"Wrote {len(quent_logs)} Quent trace events to {output_path}") + return output_path def _consolidate_logs( @@ -1857,7 +1905,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") @@ -2246,6 +2294,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: diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 5464b586f661..8e65653f5593 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", diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 63fc734a0ca7..cc4b18fc74c3 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -44,7 +44,7 @@ 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 QuentContext, WorkerResources from cudf_polars.quent._logging import QuentLogger @@ -564,6 +564,10 @@ class SPMDContext: The active RapidsMPF context. py_executor Thread-pool executor used to drive the actor network on each rank. + worker_resources + Engine/worker-scoped Quent resources (device memory, channels, thread + pool, processor registry, network topology). ``None`` when Quent is + disabled. """ comm: Communicator @@ -572,6 +576,7 @@ class SPMDContext: engine_id: uuid.UUID worker_id: uuid.UUID quent_logger: QuentLogger | None + worker_resources: WorkerResources | None = None @dataclasses.dataclass(frozen=True) diff --git a/python/cudf_polars/tests/containers/test_dataframe.py b/python/cudf_polars/tests/containers/test_dataframe.py index 9b5d7ba18d87..8ae919d61a5a 100644 --- a/python/cudf_polars/tests/containers/test_dataframe.py +++ b/python/cudf_polars/tests/containers/test_dataframe.py @@ -224,3 +224,14 @@ 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": pl.Series([1, 2, 3], dtype=pl.Int64()), + } + ) + df = DataFrame.from_polars(df, stream=stream) + assert df._size_bytes() == 24 diff --git a/python/cudf_polars/tests/quent/conftest.py b/python/cudf_polars/tests/quent/conftest.py index 8d90ced657fa..5335d8b33081 100644 --- a/python/cudf_polars/tests/quent/conftest.py +++ b/python/cudf_polars/tests/quent/conftest.py @@ -5,16 +5,48 @@ from __future__ import annotations +import uuid from typing import TYPE_CHECKING import pytest import cudf_polars.quent +from cudf_polars.quent._types import Channel, Memory, Processor if TYPE_CHECKING: from cudf_polars.quent 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_export.py b/python/cudf_polars/tests/quent/test_export.py new file mode 100644 index 000000000000..8b6386d425a9 --- /dev/null +++ b/python/cudf_polars/tests/quent/test_export.py @@ -0,0 +1,199 @@ +# 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 +import zipfile +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(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: + 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_archive(tmp_path: Path) -> None: + context_id = uuid.UUID("019dd571-105a-7c53-a15b-713cbdd7666b") + events = _buffered_events() + archive_path = tmp_path / f"{context_id}.zip" + + write_quent_export(events, tmp_path, context_id, archive_path) + + assert archive_path == tmp_path / f"{context_id}.zip" + assert zipfile.is_zipfile(archive_path) + + 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 + ) + + 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() + archive_path = tmp_path / f"{context_id}.zip" + write_quent_export(events, tmp_path, context_id, archive_path) + + 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} + + 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: + archive_path = tmp_path / f"{uuid.uuid4()}.zip" + with pytest.raises(ValueError, match="exactly one entity wrapper"): + write_quent_export( + [{"id": "x", "timestamp": 1, "data": {"Engine": {}, "Query": {}}}], + tmp_path, + uuid.uuid4(), + archive_path, + ) + + +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) + archive_path = tmp_path / "logs" / f"{run_id}.zip" + benchmark_utils._write_quent_traces( + FakeEngine(), # type: ignore[arg-type] + run_id, + collect_traces=True, + quent_archive=archive_path, + ) + + 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) diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index 2b7217d0689c..f7670d44d571 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -15,24 +15,39 @@ 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 import QuentContext +from cudf_polars.quent._context import ( + LocalQuentContext, + ProcessorRegistry, + QuentIRExecutionContext, + WorkerResources, +) 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, Query, + Statistics, + Task, Worker, _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 import QuentContext + from cudf_polars.quent._types import Processor from cudf_polars.utils.config import StreamingExecutor @@ -44,6 +59,60 @@ 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() + 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: + 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( + id=uuid.uuid4(), + query=query, + parent_plan=None, + instance_name="logical", + edges=[], + worker=None, + ) + operator = Operator( + id=operator_id, + plan=plan, + parent_operators=[], + type_name="Filter", + ) + local_context = LocalQuentContext( + context=context, + query=query, + worker=Worker(id=worker_id, engine=context.engine, instance_name="rank-0"), + logger=logger, + worker_resources=worker_resources, + ) + quent_ir_execution_context = QuentIRExecutionContext.from_execution_context( + local_context, operator + ) + return logger, quent_ir_execution_context + + @pytest.mark.parametrize( "value,expected_variant", [ @@ -104,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, @@ -237,6 +314,52 @@ 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 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"] == 2 + + def test_port_declare_serialization( ir_and_config: tuple[IR, ConfigOptions[StreamingExecutor]], ) -> None: @@ -507,20 +630,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 @@ -537,14 +660,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 @@ -557,14 +680,231 @@ def test_emit_query_group_events_idempotent(quent_context: QuentContext): assert len(logger._buffer) == 1 -def test_serialize_list_raises(): - with pytest.raises(NotImplementedError, match="not supported yet"): - Attribute("list", [1, 2]).serialize() +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() + 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, + 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, + worker_resources=worker_resources, + ) + + 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() -> 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) + + 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() -> 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_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_serialize_dict_raises(): - with pytest.raises(NotImplementedError, match="not supported yet"): - Attribute("dict", {"a": 1, "b": 2}).serialize() + +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", + "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(): @@ -572,3 +912,315 @@ def test_quent_serialize_none(): "key": "none", "value": 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( + 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( + processor: Processor, + device_memory: Memory, + disk_to_device_channel: Channel, +) -> None: + operator_id = uuid.uuid4() + task = Task(operator_id=operator_id, instance_name="scan-task") + + event = task.loading( + use_thread=processor, + use_channel=disk_to_device_channel, + channel_capacity_bytes=4096, + use_memory=device_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(disk_to_device_channel.id), + "capacity": {"capacity_bytes": 4096}, + } + assert loading["use_memory"] == { + "resource_id": str(device_memory.id), + "capacity": {"capacity_bytes": 8192}, + } + + +def test_task_computing_serialization( + processor: Processor, + device_memory: Memory, +) -> None: + operator_id = uuid.uuid4() + task = Task(operator_id=operator_id, instance_name="filter-task") + + event = task.computing( + use_thread=processor, + use_memory=device_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(device_memory.id), + "capacity": {"capacity_bytes": 16384}, + } + + +def test_task_sending_serialization( + processor: Processor, + device_memory: Memory, +) -> None: + operator_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=device_memory, + target=device_memory, + ) + + 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() + worker_resources = WorkerResources.build( + instance_suffix="test", + engine_id=uuid.uuid4(), + worker_id=uuid.uuid4(), + rank=0, + nranks=1, + ) + 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" + ) + + 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( + rank: int, + nranks: int, + expected_targets: list[int], +) -> None: + pytest.importorskip("structlog") + logger = cudf_polars.quent._logging.QuentLogger() + engine_id = uuid.uuid4() + + worker_resources = WorkerResources.build( + instance_suffix="test", + engine_id=engine_id, + worker_id=uuid.uuid4(), + rank=rank, + nranks=nranks, + ) + worker_resources.declare(logger) + + 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 == 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) + # 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: + 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, + 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"]] + # 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_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 + ) + 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, + 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( + 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 "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"] diff --git a/python/cudf_polars/tests/quent/test_quent_integration.py b/python/cudf_polars/tests/quent/test_quent_integration.py index 3c4c53710e93..ba9ca4d25dc3 100644 --- a/python/cudf_polars/tests/quent/test_quent_integration.py +++ b/python/cudf_polars/tests/quent/test_quent_integration.py @@ -11,11 +11,13 @@ import polars as pl +from cudf_polars.dsl.tracing import LOG_TRACES + if TYPE_CHECKING: 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. @@ -152,15 +154,152 @@ 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 + + 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 + + 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) + + +def test_quent_events_multiple_collects( + engine_with_quent_context: StreamingEngine, quent_context: QuentContext +) -> None: + # 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: + q.collect(engine=engine_with_quent_context) + q.collect(engine=engine_with_quent_context) + + quent_events = engine_with_quent_context._quent_events + + # 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 + 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} + ] + 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) + + # 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 + assert len(set(query_init_ids)) == 2 + assert str(quent_context.query.id) not in query_init_ids + + # 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 + 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) + + 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 diff --git a/python/cudf_polars/tests/streaming/test_spmd.py b/python/cudf_polars/tests/streaming/test_spmd.py index 326d39421265..c12fa306ca41 100644 --- a/python/cudf_polars/tests/streaming/test_spmd.py +++ b/python/cudf_polars/tests/streaming/test_spmd.py @@ -20,6 +20,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 @@ -461,7 +462,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 aa01207e7e51..48a3178f335c 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, @@ -423,7 +423,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)