diff --git a/python/cudf_polars/cudf_polars/quent/_context.py b/python/cudf_polars/cudf_polars/quent/_context.py index 4bbf1da907ce..ea5ad30b26a9 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,62 @@ 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). By using + the current thread, we're assuming that the same thread that requested + the memory reservation also emitted the memory reservation events. + - exit + """ + 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..f2e88b06ea53 100644 --- a/python/cudf_polars/cudf_polars/quent/_types.py +++ b/python/cudf_polars/cudf_polars/quent/_types.py @@ -881,6 +881,98 @@ 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: + """ + 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. 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`. + """ + + purpose: MemoryReservationPurpose + 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.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.value, + "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 +1013,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 +1072,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..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 @@ -9,7 +9,11 @@ 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 ( + MemoryReservationPurpose, + reserve_memory_traced, +) if TYPE_CHECKING: import pylibcudf as plc @@ -117,10 +121,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=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 1c57765dd4d7..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 @@ -22,10 +22,13 @@ 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 ( + MemoryReservationPurpose, + reserve_memory_traced, +) from cudf_polars.streaming.actor_graph.utils import ( ChunkStore, concat_batch, @@ -324,6 +327,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 +337,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=MemoryReservationPurpose.ORDERING_UNPACK_REMOTE, + sequence_number=partition_id, ) return TableChunk.from_pylibcudf_table( unpack_and_concat(partitions, stream=stream, br=br, reservation=reservation), @@ -693,7 +701,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..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 @@ -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,10 @@ generate_ir_sub_network, ir_context_for_node, ) +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, @@ -87,6 +90,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 +116,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=MemoryReservationPurpose.SHUFFLE_INSERT_HASH, ) self._manager.shuffler.insert( py_partition_and_pack( @@ -140,12 +148,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=MemoryReservationPurpose.SHUFFLE_INSERT_HASH_KEYS, ) with opaque_memory_usage(reservation.split(chunk_nbytes)): key_table = _evaluate_key_table(chunk, keys, schema) @@ -170,12 +180,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=MemoryReservationPurpose.SHUFFLE_INSERT_SPLIT, ) self._manager.shuffler.insert( py_split_and_pack( @@ -205,12 +217,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=MemoryReservationPurpose.SHUFFLE_INSERT_INDEX, ) reorder_nbytes = py_split_and_pack_cost( chunk.table_view(), chunk.stream, br @@ -254,11 +268,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 +307,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=MemoryReservationPurpose.SHUFFLE_EXTRACT, + sequence_number=partition_id, ) return py_unpack_and_concat( partitions=partitions, @@ -346,6 +365,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 +373,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=MemoryReservationPurpose.REPARTITION_EXTRACT, + sequence_number=partition_id, ) table = py_unpack_and_concat( pieces, stream=stream, br=self._br, reservation=reservation @@ -545,7 +568,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..d5ad971f3b29 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,10 @@ generate_ir_sub_network, ir_context_for_node, ) +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, @@ -394,8 +397,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=MemoryReservationPurpose.PYTHON_SCAN, + sequence_number=seq_num, ) ): df = await ir_context.to_thread(process) @@ -561,10 +569,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=MemoryReservationPurpose.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..f31edeea8118 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,10 @@ generate_ir_sub_network, ir_context_for_node, ) +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, @@ -314,7 +315,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=MemoryReservationPurpose.BROADCAST_JOIN, + sequence_number=seq_num, + ) ): for sdf in dfs_to_join: result = await ir_context.to_thread( @@ -615,7 +623,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=MemoryReservationPurpose.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..c26a8be628fa --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/memory.py @@ -0,0 +1,123 @@ +# 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 + +from cudf_polars.quent._types import ( + MemoryReservationPurpose, + MemoryReservationRequest, + Task, +) + +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__ = ["MemoryReservationPurpose", "reserve_memory_traced"] + + +async def reserve_memory_traced( + context: Context, + size: int, + *, + net_memory_delta: int, + ir_context: IRExecutionContext | None, + purpose: MemoryReservationPurpose, + 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. 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 = 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( + 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..16f592e2715b 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -32,6 +32,8 @@ Engine, Implementation, Memory, + MemoryReservationPurpose, + MemoryReservationRequest, Network, Operator, Plan, @@ -1176,3 +1178,85 @@ def test_emit_task_events_io_node(disk_to_device_channel: Channel) -> None: assert "Loading" in task_events[2]["data"]["Task"]["state"] assert "Computing" in task_events[3]["data"]["Task"]["state"] assert "Exit" in task_events[4]["data"]["Task"]["state"] + + +def test_memory_reservation_request_serialization() -> None: + request = MemoryReservationRequest( + purpose=MemoryReservationPurpose.SCAN, + size_bytes=2 * 1024**2, + mem_type="DEVICE", + net_memory_delta=1024**2, + sequence_number=3, + ) + assert request.label == "scan-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=MemoryReservationPurpose.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=MemoryReservationPurpose.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-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: