From f716bbbe7d77ecbe1e6add7c1e0fc8e6082b384f Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Fri, 4 Sep 2026 13:07:26 -0700 Subject: [PATCH 1/3] FEA: Track memory reservations --- .../cudf_polars/cudf_polars/quent/_context.py | 61 +++++++ .../cudf_polars/cudf_polars/quent/_types.py | 156 +++++++++++++++++- .../actor_graph/collectives/allgather.py | 7 +- .../actor_graph/collectives/ordering.py | 13 +- .../actor_graph/collectives/shuffle.py | 38 ++++- .../streaming/actor_graph/collectives/sort.py | 1 + .../streaming/actor_graph/groupby.py | 1 + .../cudf_polars/streaming/actor_graph/io.py | 16 +- .../cudf_polars/streaming/actor_graph/join.py | 24 ++- .../streaming/actor_graph/memory.py | 121 ++++++++++++++ .../cudf_polars/streaming/actor_graph/over.py | 7 +- python/cudf_polars/docs/overview.md | 15 ++ python/cudf_polars/tests/quent/test_quent.py | 95 +++++++++++ .../tests/quent/test_quent_integration.py | 39 +++++ 14 files changed, 561 insertions(+), 33 deletions(-) create mode 100644 python/cudf_polars/cudf_polars/streaming/actor_graph/memory.py diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index 4bbf1da907ce..9f7fae0ed8c0 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -35,6 +35,7 @@ from cudf_polars.dsl.ir import IR from cudf_polars.quent._logging import QuentLogger from cudf_polars.quent._types import ( + MemoryReservationRequest, Operator, Plan, Port, @@ -481,6 +482,66 @@ def _emit_task_end_events( # We can't do it directly on the Task object, because that (seems to) # break operator-level aggregation like duration_s. + def _emit_memory_reservation_events( + self, + quent_task: Task, + quent_ir_execution_context: QuentIRExecutionContext, + request: MemoryReservationRequest, + *, + requested_at: int, + satisfied_at: int, + ) -> None: + """ + Emit Quent events describing a single memory reservation. + + The task enters ``Allocating`` when the reservation is requested and + exits once rapidsmpf satisfies it (or fails to), so the duration of + the ``Allocating`` state is how long the operator waited for memory. + + Parameters + ---------- + quent_task: Task + The reservation's Quent Task, from + :meth:`~cudf_polars.quent._types.Task.for_memory_reservation`. + quent_ir_execution_context: QuentIRExecutionContext + The Quent IR execution context, which binds the reservation to the + operator that made it. + request: MemoryReservationRequest + The size, memory tier and purpose of the reservation. + requested_at: int + Timestamp (unix nanoseconds) at which the reservation was requested. + satisfied_at: int + Timestamp (unix nanoseconds) at which the reservation was satisfied. + + Notes + ----- + This emits the following events: + + - queueing + - allocating (with the Quent Processor for the current thread) + - exit + + They're all built after the fact from recorded timestamps, so emitting + them doesn't inflate the wait we're trying to measure. + + A reservation that induces spilling ought to pass through the Task's + ``Spilling`` state, but rapidsmpf doesn't report that back to us yet. + """ + quent_processor = quent_ir_execution_context.get_or_declare_processor( + thread_ident=threading.get_ident(), + ) + quent_ir_execution_context.logger.emit( + quent_task.queueing(timestamp=requested_at) + ) + quent_ir_execution_context.logger.emit( + quent_task.allocating( + resource_id=quent_processor.id, + timestamp=requested_at, + reservation=request, + ) + ) + quent_ir_execution_context.logger.emit(quent_task.exit(timestamp=satisfied_at)) + @dataclasses.dataclass(kw_only=True) class WorkerResources: diff --git a/python/cudf_polars/cudf_polars/quent/_types.py b/python/cudf_polars/cudf_polars/quent/_types.py index 3b1304d63e38..c6d6ee53d2be 100644 --- a/python/cudf_polars/cudf_polars/quent/_types.py +++ b/python/cudf_polars/cudf_polars/quent/_types.py @@ -881,6 +881,92 @@ def exit(self, timestamp: int | None = None) -> Event: ) +def _format_bytes(nbytes: int) -> str: + """Format a byte count compactly (e.g. ``256.0MiB``) for display.""" + value = float(nbytes) + for unit in ("B", "KiB", "MiB", "GiB", "TiB"): + if abs(value) < 1024.0: + return f"{value:.0f}B" if unit == "B" else f"{value:.1f}{unit}" + value /= 1024.0 + return f"{value:.1f}PiB" + + +@dataclasses.dataclass(frozen=True, slots=True, kw_only=True) +class MemoryReservationRequest: + """ + A request to reserve memory, recorded on a Quent Task. + + cudf-polars reserves (but does not necessarily allocate) device or host + memory before operations that grow their memory footprint. A reservation + is not always satisfiable right away, which gives the runtime a chance to + apply backpressure or spill rather than run out of memory. + + Parameters + ---------- + purpose + What the memory is reserved for (e.g. ``"scan"``). Distinguishes + reservations made by a single operator. + size_bytes + The number of bytes requested. + mem_type + The memory tier reserved from (e.g. ``"DEVICE"``). + net_memory_delta + The expected lasting change in memory usage, which is smaller than + ``size_bytes`` for operations whose peak usage is transient. + allow_overbooking + Whether the runtime may hand out a reservation it cannot back, + or ``None`` to use the rapidsmpf default. + sequence_number + The sequence number of the chunk this reservation is for, if any. + granted + Whether the reservation was satisfied. A failed reservation still + took time, so it is worth recording. + + Notes + ----- + Reservations are recorded as the ``Allocating`` state of a Quent + :class:`Task`, so the time spent in that state is the time it took to + satisfy the request. See + :meth:`~cudf_polars.quent._context.QuentContext._emit_memory_reservation_events`. + + Quent's data processing domain doesn't declare attributes on the + ``Allocating`` state yet, so an analyzer built against the current model + reads the timing but ignores the attributes below. They're still written + to the event stream, and will show up once the model declares them. + """ + + purpose: str + size_bytes: int + mem_type: str + net_memory_delta: int | None = None + allow_overbooking: bool | None = None + sequence_number: int | None = None + granted: bool = True + + @property + def label(self) -> str: + """A compact description, e.g. ``scan-256.0MiB-device``.""" + return ( + f"{self.purpose}-{_format_bytes(self.size_bytes)}-{self.mem_type.lower()}" + ) + + def to_dict(self) -> dict[str, Any]: + """Serialize to the flat attribute layout used by Quent FSM states.""" + attributes: dict[str, Any] = { + "purpose": self.purpose, + "size_bytes": self.size_bytes, + "mem_type": self.mem_type, + "granted": self.granted, + } + if self.net_memory_delta is not None: + attributes["net_memory_delta"] = self.net_memory_delta + if self.allow_overbooking is not None: + attributes["allow_overbooking"] = self.allow_overbooking + if self.sequence_number is not None: + attributes["sequence_number"] = self.sequence_number + return attributes + + @dataclasses.dataclass(frozen=True, slots=True, kw_only=True) class Task: """A Quent Task representing a unit of work on an operator.""" @@ -921,6 +1007,43 @@ def from_ir( operator_id=quent_ir_execution_context.quent_operator.id, ) + @classmethod + def for_memory_reservation( + cls, + request: MemoryReservationRequest, + quent_ir_execution_context: QuentIRExecutionContext, + ) -> Self: + """ + Build an operator-scoped Quent Task recording a memory reservation. + + Parameters + ---------- + request + The reservation being made. Its :attr:`MemoryReservationRequest.label` + goes into the task's instance name, so the size and memory tier are + legible without reading the task's attributes. + quent_ir_execution_context + The Quent IR execution context, which is used to get the operator + the reservation is made on behalf of. + + Returns + ------- + The operator-scoped Quent Task. + """ + operator = quent_ir_execution_context.quent_operator + # Reservations for a single chunk are already distinguished by their + # sequence number. Fall back to a token for the reservations that + # aren't per-chunk (e.g. gathering a whole collective's output). + suffix = ( + uuid.uuid4().hex[:8] + if request.sequence_number is None + else request.sequence_number + ) + return cls( + instance_name=f"reserve-{request.label}-{operator.id.hex[:8]}-{suffix}", + operator_id=operator.id, + ) + def queueing(self, timestamp: int | None = None) -> Event: """Build a Quent Task Queueing event.""" return Event( @@ -943,22 +1066,37 @@ def allocating( self, resource_id: uuid.UUID, timestamp: int | None = None, + *, + reservation: MemoryReservationRequest | None = None, ) -> Event: - """Build a Quent Task Allocating event.""" + """ + Build a Quent Task Allocating event. + + Parameters + ---------- + resource_id + The Quent Processor (thread) doing the allocation. + timestamp + The event timestamp, defaulting to now. + reservation + The memory reservation being waited on, when this transition + represents a call into rapidsmpf's memory admission control. + """ + allocating_data: dict[str, Any] = { + "use_thread": { + "resource_id": str(resource_id), + "capacity": None, + } + } + if reservation is not None: + allocating_data.update(reservation.to_dict()) 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, - } - } - }, + "state": {"Allocating": allocating_data}, } }, ) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/allgather.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/allgather.py index 672b8cb150f1..7e0dacf0ca1e 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/allgather.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/allgather.py @@ -9,7 +9,8 @@ from cudf_streaming.partition_utils import unpack_and_concat, unpack_and_concat_cost from cudf_streaming.table_chunk import make_table_chunks_available_or_wait from rapidsmpf.streaming.coll.allgather import AllGather -from rapidsmpf.streaming.core.memory_reserve_or_wait import reserve_memory + +from cudf_polars.streaming.actor_graph.memory import reserve_memory_traced if TYPE_CHECKING: import pylibcudf as plc @@ -117,10 +118,12 @@ async def extract_concatenated( # host-resident partitions to device. The packed inputs stay live # until the concat finishes and are released after, so the net # change is about zero. - reservation = await reserve_memory( + reservation = await reserve_memory_traced( self.context, unpack_and_concat_cost(partitions), net_memory_delta=0, + ir_context=ir_context, + purpose="allgather-extract", ) return await ir_context.to_thread( unpack_and_concat, diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py index 1c57765dd4d7..c3c5fe0cc5b2 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py @@ -22,10 +22,10 @@ from pylibcudf.contiguous_split import pack from rapidsmpf.memory.memory_reservation import opaque_memory_usage from rapidsmpf.streaming.coll.sparse_alltoall import SparseAlltoall -from rapidsmpf.streaming.core.memory_reserve_or_wait import reserve_memory from rapidsmpf.streaming.core.message import Message from cudf_polars.containers import DataFrame, DataType +from cudf_polars.streaming.actor_graph.memory import reserve_memory_traced from cudf_polars.streaming.actor_graph.utils import ( ChunkStore, concat_batch, @@ -324,6 +324,8 @@ async def _unpack_remote_partition( context: Context, packed: PackedData, stream: Stream, + ir_context: IRExecutionContext, + partition_id: int, ) -> TableChunk: """Unpack one remote output-partition payload.""" br = context.br() @@ -332,10 +334,13 @@ async def _unpack_remote_partition( # host-resident partitions to device. The packed inputs stay live # until the concat finishes and are released after, so the net # change is about zero. - reservation = await reserve_memory( + reservation = await reserve_memory_traced( context, unpack_and_concat_cost(partitions), net_memory_delta=0, + ir_context=ir_context, + purpose="ordering-unpack-remote", + sequence_number=partition_id, ) return TableChunk.from_pylibcudf_table( unpack_and_concat(partitions, stream=stream, br=br, reservation=reservation), @@ -693,7 +698,9 @@ async def _adjust_ordering_impl( exchange.extract(source_rank), strict=True, ): - chunk = await _unpack_remote_partition(context, packed, stream) + chunk = await _unpack_remote_partition( + context, packed, stream, ir_context, pid + ) if chunk.table_view().num_rows() > 0: _store_chunk(context, remote_pieces, pid, chunk) pieces_by_source[source_rank] = remote_pieces 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 d7bf5d9cf7f2..17644965ba44 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,7 +32,6 @@ from rapidsmpf.streaming.coll.shuffler import ShufflerAsync from rapidsmpf.streaming.core.actor import define_actor from rapidsmpf.streaming.core.context import Context -from rapidsmpf.streaming.core.memory_reserve_or_wait import reserve_memory from rapidsmpf.streaming.core.message import Message from cudf_polars.containers import DataFrame @@ -42,6 +41,7 @@ generate_ir_sub_network, ir_context_for_node, ) +from cudf_polars.streaming.actor_graph.memory import reserve_memory_traced from cudf_polars.streaming.actor_graph.nodes import shutdown_on_error from cudf_polars.streaming.actor_graph.tracing import ( trace_channel, @@ -87,6 +87,9 @@ class ShuffleManager: How to assign partition IDs to ranks: ROUND_ROBIN (default) or CONTIGUOUS. Use CONTIGUOUS for sort so each rank gets adjacent partition IDs and concatenation order matches global order. + ir_context, optional + The execution context for the IR node this shuffle is performed for. + Attributes the manager's memory reservations to that node's operator. """ class Inserter: @@ -110,12 +113,14 @@ async def insert_hash( ) -> None: """Partition chunk by hash and insert into the shuffler.""" br = self._manager.context.br() - reservation = await reserve_memory( + reservation = await reserve_memory_traced( self._manager.context, py_partition_and_pack_cost(chunk.table_view(), chunk.stream, br), # The chunk's data moves into shuffler-owned packed buffers, # nothing lasting is added. net_memory_delta=0, + ir_context=self._manager.ir_context, + purpose="shuffle-insert-hash", ) self._manager.shuffler.insert( py_partition_and_pack( @@ -140,12 +145,14 @@ async def insert_hash_keys( # bound. A key expression that expands its input evaluates to more # than the chunk's packed size and under-reserves. chunk_nbytes = py_split_and_pack_cost(chunk.table_view(), chunk.stream, br) - reservation = await reserve_memory( + reservation = await reserve_memory_traced( self._manager.context, 3 * chunk_nbytes, # The chunk's data moves into shuffler-owned packed buffers, # nothing lasting is added. net_memory_delta=0, + ir_context=self._manager.ir_context, + purpose="shuffle-insert-hash-keys", ) with opaque_memory_usage(reservation.split(chunk_nbytes)): key_table = _evaluate_key_table(chunk, keys, schema) @@ -170,12 +177,14 @@ async def insert_hash_keys( async def insert_split(self, chunk: TableChunk, splits: list[int]) -> None: """Split chunk at the given indices and insert into the shuffler.""" br = self._manager.context.br() - reservation = await reserve_memory( + reservation = await reserve_memory_traced( self._manager.context, py_split_and_pack_cost(chunk.table_view(), chunk.stream, br), # The chunk's data moves into shuffler-owned packed buffers, # nothing lasting is added. net_memory_delta=0, + ir_context=self._manager.ir_context, + purpose="shuffle-insert-split", ) self._manager.shuffler.insert( py_split_and_pack( @@ -205,12 +214,14 @@ async def insert_index( """ br = self._manager.context.br() # As in `insert_hash_keys`, this covers the reorder plus the pack. - reservation = await reserve_memory( + reservation = await reserve_memory_traced( self._manager.context, py_partition_and_pack_cost(chunk.table_view(), chunk.stream, br), # The chunk's data moves into shuffler-owned packed buffers, # nothing lasting is added. net_memory_delta=0, + ir_context=self._manager.ir_context, + purpose="shuffle-insert-index", ) reorder_nbytes = py_split_and_pack_cost( chunk.table_view(), chunk.stream, br @@ -254,11 +265,13 @@ def __init__( collective_id: int, *, partition_assignment: PartitionAssignment = PartitionAssignment.ROUND_ROBIN, + ir_context: IRExecutionContext | None = None, ): self.context = context self.comm = comm self.num_partitions = num_partitions self.collective_id = collective_id + self.ir_context = ir_context self.shuffler = ShufflerAsync( context, comm, @@ -291,12 +304,15 @@ async def extract_chunk(self, partition_id: int, stream: Stream) -> plc.Table: The extracted table. """ partitions = self.shuffler.extract(partition_id) - reservation = await reserve_memory( + reservation = await reserve_memory_traced( self.context, py_unpack_and_concat_cost(partitions), # Representation change: the packed input is consumed as the # unpacked table is produced, at roughly the same size. net_memory_delta=0, + ir_context=self.ir_context, + purpose="shuffle-extract", + sequence_number=partition_id, ) return py_unpack_and_concat( partitions=partitions, @@ -346,6 +362,7 @@ def __init__(self, shuffle: ShuffleManager, local_count: int) -> None: local_comm, local_count, shuffle.collective_id, + ir_context=shuffle.ir_context, ) async def _iter_chunks(self, stream: Stream) -> AsyncGenerator[plc.Table, None]: @@ -353,12 +370,15 @@ async def _iter_chunks(self, stream: Stream) -> AsyncGenerator[plc.Table, None]: for piece in self._global_shuffle.extract_pieces(partition_id): # TODO: batch pieces up to target_partition_size before unpacking pieces = [piece] - reservation = await reserve_memory( + reservation = await reserve_memory_traced( self._global_shuffle.context, py_unpack_and_concat_cost(pieces), # Representation change: the packed input is consumed as the # unpacked table is produced, at roughly the same size. net_memory_delta=0, + ir_context=self._global_shuffle.ir_context, + purpose="repartition-extract", + sequence_number=partition_id, ) table = py_unpack_and_concat( pieces, stream=stream, br=self._br, reservation=reservation @@ -545,7 +565,9 @@ async def _global_shuffle( # Other ranks still participate in the shuffle protocol. skip_insert = metadata_in.duplicated and comm.rank != 0 - shuffle = ShuffleManager(context, comm, num_partitions, collective_id) + shuffle = ShuffleManager( + context, comm, num_partitions, collective_id, ir_context=ir_context + ) async with shuffle.inserting() as inserter: while (msg := await ch_in.recv(context)) is not None: if not skip_insert: 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 be19e12e372b..7f6082ea52c9 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 @@ -583,6 +583,7 @@ async def _insert_chunks_into_shuffle( num_partitions, collective_ids.pop(), partition_assignment=PartitionAssignment.CONTIGUOUS, + ir_context=ir_context, ) async with shuffle.inserting() as inserter: while (msg := await ch_in.recv(context)) is not None: 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 3dcada6efeed..26be0781df9a 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py @@ -412,6 +412,7 @@ async def _shuffle_reduce( shuffle_comm, modulus, collective_id, + ir_context=ir_context, ) async with shuffle.inserting() as inserter: await inserter.insert_hash( 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 24925730b8d2..b1519ed4a336 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -21,7 +21,6 @@ make_table_chunks_available_or_wait, ) from rapidsmpf.memory.memory_reservation import opaque_memory_usage -from rapidsmpf.streaming.core.memory_reserve_or_wait import reserve_memory from rapidsmpf.streaming.core.message import Message from cudf_polars.containers import DataFrame @@ -31,6 +30,7 @@ generate_ir_sub_network, ir_context_for_node, ) +from cudf_polars.streaming.actor_graph.memory import reserve_memory_traced from cudf_polars.streaming.actor_graph.nodes import define_actor, shutdown_on_error from cudf_polars.streaming.actor_graph.tracing import ( send_chunk, @@ -394,8 +394,13 @@ async def _process_and_send_chunk( net_memory_delta = input_bytes reservation = input_bytes * (1 + (ir.predicate is not None)) with opaque_memory_usage( - await reserve_memory( - context, size=reservation, net_memory_delta=net_memory_delta + await reserve_memory_traced( + context, + size=reservation, + net_memory_delta=net_memory_delta, + ir_context=ir_context, + purpose="python-scan", + sequence_number=seq_num, ) ): df = await ir_context.to_thread(process) @@ -561,10 +566,13 @@ async def read_chunk( else 2 * estimated_chunk_bytes ) start = time.monotonic_ns() - reservation = await reserve_memory( + reservation = await reserve_memory_traced( context, size=reservation_bytes, net_memory_delta=estimated_chunk_bytes, + ir_context=ir_context, + purpose="scan", + sequence_number=seq_num, ) admitted = time.monotonic_ns() with opaque_memory_usage(reservation): 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 71833df27785..25eb00216eac 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -20,10 +20,7 @@ ) from rapidsmpf.memory.memory_reservation import opaque_memory_usage from rapidsmpf.streaming.core.actor import define_actor -from rapidsmpf.streaming.core.memory_reserve_or_wait import ( - missing_net_memory_delta, - reserve_memory, -) +from rapidsmpf.streaming.core.memory_reserve_or_wait import missing_net_memory_delta from cudf_polars.containers import DataFrame from cudf_polars.dsl.ir import IR, Join @@ -43,6 +40,7 @@ generate_ir_sub_network, ir_context_for_node, ) +from cudf_polars.streaming.actor_graph.memory import reserve_memory_traced from cudf_polars.streaming.actor_graph.nodes import default_node_multi from cudf_polars.streaming.actor_graph.tracing import ( send_chunk, @@ -314,7 +312,14 @@ async def _broadcast_join_large_chunk( join_results: list[DataFrame] = [] input_bytes = large_chunk_size + small_size with opaque_memory_usage( - await reserve_memory(context, size=input_bytes, net_memory_delta=0) + await reserve_memory_traced( + context, + size=input_bytes, + net_memory_delta=0, + ir_context=ir_context, + purpose="broadcast-join", + sequence_number=seq_num, + ) ): for sdf in dfs_to_join: result = await ir_context.to_thread( @@ -615,7 +620,14 @@ async def _join_chunks( ) ) with opaque_memory_usage( - await reserve_memory(context, size=input_bytes, net_memory_delta=0) + await reserve_memory_traced( + context, + size=input_bytes, + net_memory_delta=0, + ir_context=ir_context, + purpose="join", + sequence_number=left_msg.sequence_number, + ) ): df = await ir_context.to_thread( ir.do_evaluate, diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/memory.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/memory.py new file mode 100644 index 000000000000..280e152aedc2 --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/memory.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Memory reservations for the RapidsMPF streaming runtime.""" + +from __future__ import annotations + +import time +from typing import TYPE_CHECKING + +from rapidsmpf.memory.buffer import MemoryType +from rapidsmpf.streaming.core.memory_reserve_or_wait import reserve_memory + +import cudf_polars.quent._types + +if TYPE_CHECKING: + from rapidsmpf.memory.memory_reservation import MemoryReservation + from rapidsmpf.streaming.core.context import Context + + from cudf_polars.dsl.ir import IRExecutionContext + +__all__ = ["reserve_memory_traced"] + + +async def reserve_memory_traced( + context: Context, + size: int, + *, + net_memory_delta: int, + ir_context: IRExecutionContext | None, + purpose: str, + sequence_number: int | None = None, + mem_type: MemoryType = MemoryType.DEVICE, + allow_overbooking: bool | None = None, +) -> MemoryReservation: + """ + Reserve memory, recording the wait as a Quent event. + + This is a drop-in replacement for + :func:`rapidsmpf.streaming.core.memory_reserve_or_wait.reserve_memory` for + the reservations made on behalf of an IR node. The reservation is recorded + as a Quent Task bound to that node's operator, so a trace shows which + operators are waiting on memory admission and for how much. + + Parameters + ---------- + context + The rapidsmpf context. + size + The number of bytes to reserve. + net_memory_delta + The expected lasting change in memory usage. This is smaller than + ``size`` for operations whose peak usage is transient. + ir_context + The execution context for the IR node making the reservation, which + supplies the Quent operator the reservation is attributed to. ``None`` + for a reservation that can't be attributed to an IR node, in which + case nothing is recorded. + purpose + What the memory is reserved for (e.g. ``"scan"``). Distinguishes + reservations made by a single operator. + sequence_number + The sequence number of the chunk this reservation is for, if any. + mem_type + The memory tier to reserve from. + allow_overbooking + Whether the runtime may hand out a reservation it cannot back, or + ``None`` to use the rapidsmpf default. + + Returns + ------- + The satisfied memory reservation. + + Notes + ----- + Nothing is recorded unless the query is being traced, in which case this + behaves exactly like ``reserve_memory``. + """ + quent_ir_execution_context = ( + None if ir_context is None else ir_context.quent_ir_execution_context + ) + if quent_ir_execution_context is None: + return await reserve_memory( + context, + size, + net_memory_delta=net_memory_delta, + mem_type=mem_type, + allow_overbooking=allow_overbooking, + ) + + granted = False + requested_at = time.time_ns() + try: + reservation = await reserve_memory( + context, + size, + net_memory_delta=net_memory_delta, + mem_type=mem_type, + allow_overbooking=allow_overbooking, + ) + granted = True + finally: + # A reservation that failed still spent time waiting, so record it too. + request = cudf_polars.quent._types.MemoryReservationRequest( + purpose=purpose, + size_bytes=size, + mem_type=mem_type.name, + net_memory_delta=net_memory_delta, + allow_overbooking=allow_overbooking, + sequence_number=sequence_number, + granted=granted, + ) + quent_ir_execution_context.context._emit_memory_reservation_events( + cudf_polars.quent._types.Task.for_memory_reservation( + request, quent_ir_execution_context + ), + quent_ir_execution_context, + request, + requested_at=requested_at, + satisfied_at=time.time_ns(), + ) + return reservation 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 514b4e67d7f8..c3e042726d9f 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py @@ -666,7 +666,11 @@ async def _shuffle_and_reassemble( ) forward_shuffle = ShuffleManager( - context, comm, forward_modulus, forward_shuffle_collective_id + context, + comm, + forward_modulus, + forward_shuffle_collective_id, + ir_context=ir_context, ) return_shuffle = ShuffleManager( context, @@ -674,6 +678,7 @@ async def _shuffle_and_reassemble( comm.nranks, return_shuffle_collective_id, partition_assignment=PartitionAssignment.CONTIGUOUS, + ir_context=ir_context, ) ch_replay = context.create_channel() diff --git a/python/cudf_polars/docs/overview.md b/python/cudf_polars/docs/overview.md index f2a11f91bc71..fa2dcb52a111 100644 --- a/python/cudf_polars/docs/overview.md +++ b/python/cudf_polars/docs/overview.md @@ -777,6 +777,21 @@ the engine start and exit events. Upon `StreamingEngine.shutdown`, all events are gathered from the workers and persisted on the (now closed) engine at `StreamingEngine._quent_events`. +### Memory Reservations + +Before an operation that grows its memory footprint, cudf-polars reserves the +device or host memory necessary to complete that operation (or our best estimate +of it). A reservation isn't always satisfiable right away, which gives the +runtime a chance to apply backpressure or spill rather than run out of memory. + +Reservations go through +`cudf_polars.streaming.actor_graph.memory.reserve_memory_traced` instead of +rapidsmpf's `reserve_memory`. Each reservation request is recorded as a Quent +`Task` bound its parent `Operator`, moving through `Queueing`, `Allocating` and +`Exit`. The duration of the `Allocating` state is how long the operator waited +for the `reserve_memory` to complete. The size, memory tier, purpose and net +memory delta ride along as attributes on the `Allocating` transition. + ### Concepts cudf-polars and Quent's Query Engine domain have somewhat overlapping names for diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index 5171da3a2af9..85de28a9d689 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -32,6 +32,7 @@ Engine, Implementation, Memory, + MemoryReservationRequest, Network, Operator, Plan, @@ -40,6 +41,7 @@ Statistics, Task, Worker, + _format_bytes, ) from cudf_polars.utils.config import ConfigOptions from cudf_polars.utils.cuda_stream import get_cuda_stream @@ -1176,3 +1178,96 @@ def test_emit_task_events_io_node(disk_to_device_channel: Channel) -> None: assert "Loading" in task_events[2]["data"]["Task"]["state"] assert "Computing" in task_events[3]["data"]["Task"]["state"] assert "Exit" in task_events[4]["data"]["Task"]["state"] + + +@pytest.mark.parametrize( + "nbytes,expected", + [ + (0, "0B"), + (512, "512B"), + (1024, "1.0KiB"), + (256 * 1024**2, "256.0MiB"), + (3 * 1024**5, "3.0PiB"), + ], +) +def test_format_bytes(nbytes: int, expected: str) -> None: + assert _format_bytes(nbytes) == expected + + +def test_memory_reservation_request_serialization() -> None: + request = MemoryReservationRequest( + purpose="scan", + size_bytes=2 * 1024**2, + mem_type="DEVICE", + net_memory_delta=1024**2, + sequence_number=3, + ) + assert request.label == "scan-2.0MiB-device" + assert request.to_dict() == { + "purpose": "scan", + "size_bytes": 2 * 1024**2, + "mem_type": "DEVICE", + "granted": True, + "net_memory_delta": 1024**2, + "sequence_number": 3, + } + # The optional attributes are dropped rather than serialized as null. + minimal = MemoryReservationRequest( + purpose="join", size_bytes=0, mem_type="HOST", granted=False + ) + assert minimal.to_dict() == { + "purpose": "join", + "size_bytes": 0, + "mem_type": "HOST", + "granted": False, + } + + +def test_emit_memory_reservation_events() -> None: + operator_id = uuid.uuid4() + logger, quent_ir_execution_context = _make_quent_ir_execution_context( + operator_id=operator_id + ) + request = MemoryReservationRequest( + purpose="scan", + size_bytes=1024**2, + mem_type="DEVICE", + net_memory_delta=1024**2, + sequence_number=2, + ) + task = Task.for_memory_reservation(request, quent_ir_execution_context) + # The size and memory tier are legible from the instance name, since Quent + # doesn't declare attributes on the Allocating state yet. + assert task.instance_name is not None + assert task.instance_name.startswith("reserve-scan-1.0MiB-device-") + assert task.instance_name.endswith(f"-{operator_id.hex[:8]}-2") + + quent_ir_execution_context.context._emit_memory_reservation_events( + task, + quent_ir_execution_context, + request, + requested_at=10, + satisfied_at=30, + ) + + events = _drained_events(logger) + task_events = [event for event in events if "Task" in event["data"]] + # queueing -> allocating -> exit + assert [event["data"]["Task"]["seq"] for event in task_events] == [0, 1, 2] + queueing, allocating, exit_ = task_events + assert all(event["id"] == str(task.id) for event in task_events) + assert queueing["data"]["Task"]["state"]["Queueing"]["operator_id"] == str( + operator_id + ) + # The Allocating state spans the wait: it is entered when the reservation + # is requested and left when it is satisfied. + assert queueing["timestamp"] == allocating["timestamp"] == 10 + assert exit_["timestamp"] == 30 + + allocating_state = allocating["data"]["Task"]["state"]["Allocating"] + processor_events = [event for event in events if "Processor" in event["data"]] + assert allocating_state["use_thread"] == { + "resource_id": processor_events[0]["id"], + "capacity": None, + } + assert allocating_state.items() >= request.to_dict().items() diff --git a/python/cudf_polars/tests/quent/test_quent_integration.py b/python/cudf_polars/tests/quent/test_quent_integration.py index ba9ca4d25dc3..6807e16bd3e7 100644 --- a/python/cudf_polars/tests/quent/test_quent_integration.py +++ b/python/cudf_polars/tests/quent/test_quent_integration.py @@ -178,11 +178,50 @@ def check_quent_events(engine: StreamingEngine, quent_context: QuentContext) -> if LOG_TRACES: assert len(task_events) > 0 + check_memory_reservations(quent_events, task_events) + # 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 check_memory_reservations( + quent_events: list[dict], task_events: list[dict] +) -> None: + """Check the reservations the scan makes before reading each chunk.""" + operator_ids = { + x["id"] + for x in quent_events + if "Operator" in x["data"] and "Declaration" in x["data"]["Operator"] + } + task_states = [ + (x, x["data"]["Task"]["state"]) + for x in task_events + if isinstance(x["data"]["Task"]["state"], dict) + ] + reservation_ids = set() + for event, state in task_states: + if "Queueing" in state and state["Queueing"]["instance_name"].startswith( + "reserve-" + ): + # Every reservation is attributed to the operator that made it. + assert state["Queueing"]["operator_id"] in operator_ids + reservation_ids.add(event["id"]) + assert len(reservation_ids) > 0 + + allocating = [ + state["Allocating"] + for event, state in task_states + if event["id"] in reservation_ids and "Allocating" in state + ] + assert len(allocating) == len(reservation_ids) + for state in allocating: + assert state["purpose"] == "scan" + assert state["mem_type"] == "DEVICE" + assert state["size_bytes"] > 0 + assert state["granted"] is True + + def test_quent_events_multiple_collects( engine_with_quent_context: StreamingEngine, quent_context: QuentContext ) -> None: From 43361a6a17401b80e7632ce3800ab750f3822b68 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 9 Sep 2026 09:04:09 -0700 Subject: [PATCH 2/3] fixup --- .../cudf_polars/cudf_polars/quent/_context.py | 10 +++------- .../cudf_polars/cudf_polars/quent/_types.py | 19 +------------------ python/cudf_polars/tests/quent/test_quent.py | 19 ++----------------- 3 files changed, 6 insertions(+), 42 deletions(-) diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index 9f7fae0ed8c0..ea5ad30b26a9 100644 --- a/python/cudf_polars/cudf_polars/quent/_context.py +++ b/python/cudf_polars/cudf_polars/quent/_context.py @@ -518,14 +518,10 @@ def _emit_memory_reservation_events( This emits the following events: - queueing - - allocating (with the Quent Processor for the current thread) + - allocating (with the Quent Processor for the current thread). By using + the current thread, we're assuming that the same thread that requested + the memory reservation also emitted the memory reservation events. - exit - - They're all built after the fact from recorded timestamps, so emitting - them doesn't inflate the wait we're trying to measure. - - A reservation that induces spilling ought to pass through the Task's - ``Spilling`` state, but rapidsmpf doesn't report that back to us yet. """ quent_processor = quent_ir_execution_context.get_or_declare_processor( thread_ident=threading.get_ident(), diff --git a/python/cudf_polars/cudf_polars/quent/_types.py b/python/cudf_polars/cudf_polars/quent/_types.py index c6d6ee53d2be..c83f2898bdcd 100644 --- a/python/cudf_polars/cudf_polars/quent/_types.py +++ b/python/cudf_polars/cudf_polars/quent/_types.py @@ -881,16 +881,6 @@ def exit(self, timestamp: int | None = None) -> Event: ) -def _format_bytes(nbytes: int) -> str: - """Format a byte count compactly (e.g. ``256.0MiB``) for display.""" - value = float(nbytes) - for unit in ("B", "KiB", "MiB", "GiB", "TiB"): - if abs(value) < 1024.0: - return f"{value:.0f}B" if unit == "B" else f"{value:.1f}{unit}" - value /= 1024.0 - return f"{value:.1f}PiB" - - @dataclasses.dataclass(frozen=True, slots=True, kw_only=True) class MemoryReservationRequest: """ @@ -928,11 +918,6 @@ class MemoryReservationRequest: :class:`Task`, so the time spent in that state is the time it took to satisfy the request. See :meth:`~cudf_polars.quent._context.QuentContext._emit_memory_reservation_events`. - - Quent's data processing domain doesn't declare attributes on the - ``Allocating`` state yet, so an analyzer built against the current model - reads the timing but ignores the attributes below. They're still written - to the event stream, and will show up once the model declares them. """ purpose: str @@ -946,9 +931,7 @@ class MemoryReservationRequest: @property def label(self) -> str: """A compact description, e.g. ``scan-256.0MiB-device``.""" - return ( - f"{self.purpose}-{_format_bytes(self.size_bytes)}-{self.mem_type.lower()}" - ) + return f"{self.purpose}-{self.mem_type.lower()}" def to_dict(self) -> dict[str, Any]: """Serialize to the flat attribute layout used by Quent FSM states.""" diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index 85de28a9d689..eb4a00e754f3 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -41,7 +41,6 @@ Statistics, Task, Worker, - _format_bytes, ) from cudf_polars.utils.config import ConfigOptions from cudf_polars.utils.cuda_stream import get_cuda_stream @@ -1180,20 +1179,6 @@ def test_emit_task_events_io_node(disk_to_device_channel: Channel) -> None: assert "Exit" in task_events[4]["data"]["Task"]["state"] -@pytest.mark.parametrize( - "nbytes,expected", - [ - (0, "0B"), - (512, "512B"), - (1024, "1.0KiB"), - (256 * 1024**2, "256.0MiB"), - (3 * 1024**5, "3.0PiB"), - ], -) -def test_format_bytes(nbytes: int, expected: str) -> None: - assert _format_bytes(nbytes) == expected - - def test_memory_reservation_request_serialization() -> None: request = MemoryReservationRequest( purpose="scan", @@ -1202,7 +1187,7 @@ def test_memory_reservation_request_serialization() -> None: net_memory_delta=1024**2, sequence_number=3, ) - assert request.label == "scan-2.0MiB-device" + assert request.label == "scan-device" assert request.to_dict() == { "purpose": "scan", "size_bytes": 2 * 1024**2, @@ -1239,7 +1224,7 @@ def test_emit_memory_reservation_events() -> None: # The size and memory tier are legible from the instance name, since Quent # doesn't declare attributes on the Allocating state yet. assert task.instance_name is not None - assert task.instance_name.startswith("reserve-scan-1.0MiB-device-") + assert task.instance_name.startswith("reserve-scan-device-") assert task.instance_name.endswith(f"-{operator_id.hex[:8]}-2") quent_ir_execution_context.context._emit_memory_reservation_events( From 5b61b3d7ccc32f1416cef59741d3568d6febee59 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 9 Sep 2026 09:11:58 -0700 Subject: [PATCH 3/3] enum --- .../cudf_polars/cudf_polars/quent/_types.py | 33 ++++++++++++++++--- .../actor_graph/collectives/allgather.py | 7 ++-- .../actor_graph/collectives/ordering.py | 7 ++-- .../actor_graph/collectives/shuffle.py | 17 ++++++---- .../cudf_polars/streaming/actor_graph/io.py | 9 +++-- .../cudf_polars/streaming/actor_graph/join.py | 9 +++-- .../streaming/actor_graph/memory.py | 20 ++++++----- python/cudf_polars/tests/quent/test_quent.py | 10 ++++-- 8 files changed, 78 insertions(+), 34 deletions(-) diff --git a/python/cudf_polars/cudf_polars/quent/_types.py b/python/cudf_polars/cudf_polars/quent/_types.py index c83f2898bdcd..f2e88b06ea53 100644 --- a/python/cudf_polars/cudf_polars/quent/_types.py +++ b/python/cudf_polars/cudf_polars/quent/_types.py @@ -881,6 +881,29 @@ def exit(self, timestamp: int | None = None) -> Event: ) +class MemoryReservationPurpose(enum.StrEnum): + """ + Why an operator reserved memory. + + Operators that make more than one reservation use distinct values so a + trace can tell the reservations apart. Values are the strings written into + Quent ``Allocating`` attributes. + """ + + PYTHON_SCAN = "python-scan" + SCAN = "scan" + BROADCAST_JOIN = "broadcast-join" + JOIN = "join" + ALLGATHER_EXTRACT = "allgather-extract" + ORDERING_UNPACK_REMOTE = "ordering-unpack-remote" + SHUFFLE_INSERT_HASH = "shuffle-insert-hash" + SHUFFLE_INSERT_HASH_KEYS = "shuffle-insert-hash-keys" + SHUFFLE_INSERT_SPLIT = "shuffle-insert-split" + SHUFFLE_INSERT_INDEX = "shuffle-insert-index" + SHUFFLE_EXTRACT = "shuffle-extract" + REPARTITION_EXTRACT = "repartition-extract" + + @dataclasses.dataclass(frozen=True, slots=True, kw_only=True) class MemoryReservationRequest: """ @@ -894,8 +917,8 @@ class MemoryReservationRequest: Parameters ---------- purpose - What the memory is reserved for (e.g. ``"scan"``). Distinguishes - reservations made by a single operator. + What the memory is reserved for. Distinguishes reservations made by a + single operator. size_bytes The number of bytes requested. mem_type @@ -920,7 +943,7 @@ class MemoryReservationRequest: :meth:`~cudf_polars.quent._context.QuentContext._emit_memory_reservation_events`. """ - purpose: str + purpose: MemoryReservationPurpose size_bytes: int mem_type: str net_memory_delta: int | None = None @@ -931,12 +954,12 @@ class MemoryReservationRequest: @property def label(self) -> str: """A compact description, e.g. ``scan-256.0MiB-device``.""" - return f"{self.purpose}-{self.mem_type.lower()}" + return f"{self.purpose.value}-{self.mem_type.lower()}" def to_dict(self) -> dict[str, Any]: """Serialize to the flat attribute layout used by Quent FSM states.""" attributes: dict[str, Any] = { - "purpose": self.purpose, + "purpose": self.purpose.value, "size_bytes": self.size_bytes, "mem_type": self.mem_type, "granted": self.granted, diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/allgather.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/allgather.py index 7e0dacf0ca1e..9b76ce6f1f12 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/allgather.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/allgather.py @@ -10,7 +10,10 @@ from cudf_streaming.table_chunk import make_table_chunks_available_or_wait from rapidsmpf.streaming.coll.allgather import AllGather -from cudf_polars.streaming.actor_graph.memory import reserve_memory_traced +from cudf_polars.streaming.actor_graph.memory import ( + MemoryReservationPurpose, + reserve_memory_traced, +) if TYPE_CHECKING: import pylibcudf as plc @@ -123,7 +126,7 @@ async def extract_concatenated( unpack_and_concat_cost(partitions), net_memory_delta=0, ir_context=ir_context, - purpose="allgather-extract", + purpose=MemoryReservationPurpose.ALLGATHER_EXTRACT, ) return await ir_context.to_thread( unpack_and_concat, diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py index c3c5fe0cc5b2..fcca22e7a6e8 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py @@ -25,7 +25,10 @@ from rapidsmpf.streaming.core.message import Message from cudf_polars.containers import DataFrame, DataType -from cudf_polars.streaming.actor_graph.memory import reserve_memory_traced +from cudf_polars.streaming.actor_graph.memory import ( + MemoryReservationPurpose, + reserve_memory_traced, +) from cudf_polars.streaming.actor_graph.utils import ( ChunkStore, concat_batch, @@ -339,7 +342,7 @@ async def _unpack_remote_partition( unpack_and_concat_cost(partitions), net_memory_delta=0, ir_context=ir_context, - purpose="ordering-unpack-remote", + purpose=MemoryReservationPurpose.ORDERING_UNPACK_REMOTE, sequence_number=partition_id, ) return TableChunk.from_pylibcudf_table( 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 17644965ba44..71206f507829 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 @@ -41,7 +41,10 @@ generate_ir_sub_network, ir_context_for_node, ) -from cudf_polars.streaming.actor_graph.memory import reserve_memory_traced +from cudf_polars.streaming.actor_graph.memory import ( + MemoryReservationPurpose, + reserve_memory_traced, +) from cudf_polars.streaming.actor_graph.nodes import shutdown_on_error from cudf_polars.streaming.actor_graph.tracing import ( trace_channel, @@ -120,7 +123,7 @@ async def insert_hash( # nothing lasting is added. net_memory_delta=0, ir_context=self._manager.ir_context, - purpose="shuffle-insert-hash", + purpose=MemoryReservationPurpose.SHUFFLE_INSERT_HASH, ) self._manager.shuffler.insert( py_partition_and_pack( @@ -152,7 +155,7 @@ async def insert_hash_keys( # nothing lasting is added. net_memory_delta=0, ir_context=self._manager.ir_context, - purpose="shuffle-insert-hash-keys", + purpose=MemoryReservationPurpose.SHUFFLE_INSERT_HASH_KEYS, ) with opaque_memory_usage(reservation.split(chunk_nbytes)): key_table = _evaluate_key_table(chunk, keys, schema) @@ -184,7 +187,7 @@ async def insert_split(self, chunk: TableChunk, splits: list[int]) -> None: # nothing lasting is added. net_memory_delta=0, ir_context=self._manager.ir_context, - purpose="shuffle-insert-split", + purpose=MemoryReservationPurpose.SHUFFLE_INSERT_SPLIT, ) self._manager.shuffler.insert( py_split_and_pack( @@ -221,7 +224,7 @@ async def insert_index( # nothing lasting is added. net_memory_delta=0, ir_context=self._manager.ir_context, - purpose="shuffle-insert-index", + purpose=MemoryReservationPurpose.SHUFFLE_INSERT_INDEX, ) reorder_nbytes = py_split_and_pack_cost( chunk.table_view(), chunk.stream, br @@ -311,7 +314,7 @@ async def extract_chunk(self, partition_id: int, stream: Stream) -> plc.Table: # unpacked table is produced, at roughly the same size. net_memory_delta=0, ir_context=self.ir_context, - purpose="shuffle-extract", + purpose=MemoryReservationPurpose.SHUFFLE_EXTRACT, sequence_number=partition_id, ) return py_unpack_and_concat( @@ -377,7 +380,7 @@ async def _iter_chunks(self, stream: Stream) -> AsyncGenerator[plc.Table, None]: # unpacked table is produced, at roughly the same size. net_memory_delta=0, ir_context=self._global_shuffle.ir_context, - purpose="repartition-extract", + purpose=MemoryReservationPurpose.REPARTITION_EXTRACT, sequence_number=partition_id, ) table = py_unpack_and_concat( 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 b1519ed4a336..d5ad971f3b29 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -30,7 +30,10 @@ generate_ir_sub_network, ir_context_for_node, ) -from cudf_polars.streaming.actor_graph.memory import reserve_memory_traced +from cudf_polars.streaming.actor_graph.memory import ( + MemoryReservationPurpose, + reserve_memory_traced, +) from cudf_polars.streaming.actor_graph.nodes import define_actor, shutdown_on_error from cudf_polars.streaming.actor_graph.tracing import ( send_chunk, @@ -399,7 +402,7 @@ async def _process_and_send_chunk( size=reservation, net_memory_delta=net_memory_delta, ir_context=ir_context, - purpose="python-scan", + purpose=MemoryReservationPurpose.PYTHON_SCAN, sequence_number=seq_num, ) ): @@ -571,7 +574,7 @@ async def read_chunk( size=reservation_bytes, net_memory_delta=estimated_chunk_bytes, ir_context=ir_context, - purpose="scan", + purpose=MemoryReservationPurpose.SCAN, sequence_number=seq_num, ) admitted = time.monotonic_ns() 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 25eb00216eac..f31edeea8118 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -40,7 +40,10 @@ generate_ir_sub_network, ir_context_for_node, ) -from cudf_polars.streaming.actor_graph.memory import reserve_memory_traced +from cudf_polars.streaming.actor_graph.memory import ( + MemoryReservationPurpose, + reserve_memory_traced, +) from cudf_polars.streaming.actor_graph.nodes import default_node_multi from cudf_polars.streaming.actor_graph.tracing import ( send_chunk, @@ -317,7 +320,7 @@ async def _broadcast_join_large_chunk( size=input_bytes, net_memory_delta=0, ir_context=ir_context, - purpose="broadcast-join", + purpose=MemoryReservationPurpose.BROADCAST_JOIN, sequence_number=seq_num, ) ): @@ -625,7 +628,7 @@ async def _join_chunks( size=input_bytes, net_memory_delta=0, ir_context=ir_context, - purpose="join", + purpose=MemoryReservationPurpose.JOIN, sequence_number=left_msg.sequence_number, ) ): diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/memory.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/memory.py index 280e152aedc2..c26a8be628fa 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/memory.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/memory.py @@ -10,7 +10,11 @@ from rapidsmpf.memory.buffer import MemoryType from rapidsmpf.streaming.core.memory_reserve_or_wait import reserve_memory -import cudf_polars.quent._types +from cudf_polars.quent._types import ( + MemoryReservationPurpose, + MemoryReservationRequest, + Task, +) if TYPE_CHECKING: from rapidsmpf.memory.memory_reservation import MemoryReservation @@ -18,7 +22,7 @@ from cudf_polars.dsl.ir import IRExecutionContext -__all__ = ["reserve_memory_traced"] +__all__ = ["MemoryReservationPurpose", "reserve_memory_traced"] async def reserve_memory_traced( @@ -27,7 +31,7 @@ async def reserve_memory_traced( *, net_memory_delta: int, ir_context: IRExecutionContext | None, - purpose: str, + purpose: MemoryReservationPurpose, sequence_number: int | None = None, mem_type: MemoryType = MemoryType.DEVICE, allow_overbooking: bool | None = None, @@ -56,8 +60,8 @@ async def reserve_memory_traced( for a reservation that can't be attributed to an IR node, in which case nothing is recorded. purpose - What the memory is reserved for (e.g. ``"scan"``). Distinguishes - reservations made by a single operator. + What the memory is reserved for. Distinguishes reservations made by a + single operator. sequence_number The sequence number of the chunk this reservation is for, if any. mem_type @@ -100,7 +104,7 @@ async def reserve_memory_traced( granted = True finally: # A reservation that failed still spent time waiting, so record it too. - request = cudf_polars.quent._types.MemoryReservationRequest( + request = MemoryReservationRequest( purpose=purpose, size_bytes=size, mem_type=mem_type.name, @@ -110,9 +114,7 @@ async def reserve_memory_traced( granted=granted, ) quent_ir_execution_context.context._emit_memory_reservation_events( - cudf_polars.quent._types.Task.for_memory_reservation( - request, quent_ir_execution_context - ), + Task.for_memory_reservation(request, quent_ir_execution_context), quent_ir_execution_context, request, requested_at=requested_at, diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index eb4a00e754f3..16f592e2715b 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -32,6 +32,7 @@ Engine, Implementation, Memory, + MemoryReservationPurpose, MemoryReservationRequest, Network, Operator, @@ -1181,7 +1182,7 @@ def test_emit_task_events_io_node(disk_to_device_channel: Channel) -> None: def test_memory_reservation_request_serialization() -> None: request = MemoryReservationRequest( - purpose="scan", + purpose=MemoryReservationPurpose.SCAN, size_bytes=2 * 1024**2, mem_type="DEVICE", net_memory_delta=1024**2, @@ -1198,7 +1199,10 @@ def test_memory_reservation_request_serialization() -> None: } # The optional attributes are dropped rather than serialized as null. minimal = MemoryReservationRequest( - purpose="join", size_bytes=0, mem_type="HOST", granted=False + purpose=MemoryReservationPurpose.JOIN, + size_bytes=0, + mem_type="HOST", + granted=False, ) assert minimal.to_dict() == { "purpose": "join", @@ -1214,7 +1218,7 @@ def test_emit_memory_reservation_events() -> None: operator_id=operator_id ) request = MemoryReservationRequest( - purpose="scan", + purpose=MemoryReservationPurpose.SCAN, size_bytes=1024**2, mem_type="DEVICE", net_memory_delta=1024**2,