RAPIDS Accelerator for Apache Spark SQL Plugin 22.08.0.0-eep-900 API

case class BlockRange[T <: BlockWithSize](block: T, rangeStart: Long, rangeEnd: Long) extends Product with Serializable

Specifies a start and end range of bytes for a block.

T: - the specific type of BlockWithSize
block: - a BlockWithSize instance
rangeStart: - byte offset for the start of the range (inclusive)
rangeEnd: - byte offset for the end of the range (exclusive)

trait BlockWithSize extends AnyRef

abstract class BounceBuffer extends AutoCloseable

Class to hold a bounce buffer reference in buffer.

It is AutoCloseable, where a call to close puts the bounce buffer back into the corresponding source BounceBufferManager

class BounceBufferManager[T <: MemoryBuffer] extends AutoCloseable with Logging

This classes manages a set of bounce buffers, that are instances of MemoryBuffer.

This classes manages a set of bounce buffers, that are instances of MemoryBuffer. The size/quantity of buffers is configurable, and so is the allocator.

T: the specific type of MemoryBuffer i.e. DeviceMemoryBuffer, HostMemoryBuffer, etc.

class BufferReceiveState extends AutoCloseable with Logging with Arm

A helper case class to maintain the state associated with a transfer request to a peer.

On getBufferWhenReady(finalizeCb, size), a bounce buffer is either made immediately available to finalizeCb or it will be made available later, via toFinalize.

By convention, the id is used as the header for receives for this BufferReceiveState

consumeWindow is called when data has arrived at the bounce buffer, copying onto the ranges the bytes received.

It also is AutoCloseable. close() should be called to free bounce buffers.

class BufferSendState extends AutoCloseable with Logging with Arm

A helper case class to maintain the server side state in response to a transfer request initiated by a peer.

The class implements the Iterator interface.

On next(), a set of RapidsBuffer are copied onto a bounce buffer, and a MemoryBuffer slice of the bounce buffer is returned. Buffers are copied to the bounce buffer in TransferRequest order. The receiver has the same conventions.

Synchronization with serverStream is outside of this class. It is assumed that the caller has several BufferSendState iterators and on calling next on this collection, it will perform a sync on serverStream before handing it off to the transport.

It also is AutoCloseable. close() should be called to free bounce buffers.

In terms of the lifecycle of this object, it begins with the client asking for transfers to start, it lasts through all buffers being transmitted, and ultimately finishes when a TransferResponse is sent back to the client.

trait ClientConnection extends AnyRef

Trait used by the clients to interact with the transport.

case class ConsumedBatchFromBounceBuffer(contigBuffer: DeviceMemoryBuffer, meta: TableMeta, handler: RapidsShuffleFetchHandler) extends Product with Serializable

class CudfTransportBuffer extends TransportBuffer

CudfTransportBuffer encapsulates a cuDF MemoryBuffer used for shuffle block transfers via the transport

class DirectByteBufferPool extends Logging

A pool of direct byte buffers, sized to be bufferSize.

A pool of direct byte buffers, sized to be bufferSize. This is a controlled leak at the moment, there is no reclaiming of buffers.

NOTE: this is used for metadata messages.

trait MemoryRegistrationCallback extends AnyRef

class MetadataTransportBuffer extends TransportBuffer

MetadataTransportBuffer encapsulates a direct byte buffer used exclusively for metadata transfers via the transport.

case class PendingTransferRequest(client: RapidsShuffleClient, tableMeta: TableMeta, handler: RapidsShuffleFetchHandler) extends Product with Serializable

A helper case class that describes a pending table.

A helper case class that describes a pending table. It allows the transport to schedule/throttle these requests to fit within the maximum bytes in flight.

client: client used to issue the requests
tableMeta: shuffle metadata describing the table
handler: a specific handler that is waiting for this batch

class RapidsShuffleClient extends Logging with Arm with AutoCloseable

The client makes requests via a ClientConnection obtained from the RapidsShuffleTransport.

This class handles fetch requests from RapidsShuffleIterator, turning them into ShuffleMetadata messages, and shuffle TransferRequests.

Its counterpart is the RapidsShuffleServer on a specific peer executor, specified by connection.

trait RapidsShuffleFetchHandler extends AnyRef

trait used by client consumers (RapidsShuffleIterator) to gather what the expected number of batches (Tables) is, and the ids for each table as they are received.

class RapidsShuffleIterator extends Iterator[ColumnarBatch] with Logging with Arm

An Iterator over columnar batches that fetches blocks using RapidsShuffleClients.

A transport instance is used to make RapidsShuffleClients that are able to fetch blocks.

trait RapidsShuffleRequestHandler extends AnyRef

Trait used for the server to get buffer metadata (for metadata requests), and also to acquire a buffer (for transfer requests)

class RapidsShuffleServer extends AutoCloseable with Logging with Arm

A server that replies to shuffle metadata messages, and issues device/host memory sends.

A single command thread is used to orchestrate sends/receives and to remove from transport's progress thread.

trait RapidsShuffleTransport extends AutoCloseable

This defines what a "transport" should support.

This defines what a "transport" should support. The intention is to allow for various transport implementations to exist, for different communication frameworks.

It is an AutoCloseable and so the caller should close when the transport is no longer needed.

class RefCountedDirectByteBuffer extends AutoCloseable

RefCountedDirectByteBuffer is a simple wrapper on top of a ByteBuffer that has been allocated in direct mode.

The pool is used to return the ByteBuffer to be reused, but not all of these buffers are pooled (hence the argument is optional)

The user should always close a RefCountedDirectByteBuffer. The close could hard destroy the buffer, or return the object to the pool

case class SendBounceBuffers(deviceBounceBuffer: BounceBuffer, hostBounceBuffer: Option[BounceBuffer]) extends AutoCloseable with Product with Serializable

This class can hold 1 or 2 BounceBuffers and is only used in the send case.

Ideally, the device buffer is used if most of the data to be sent is on the device. The host buffer is used in the opposite case.

deviceBounceBuffer: - device buffer to use for sends
hostBounceBuffer: - optional host buffer to use for sends

trait ServerConnection extends AnyRef

A server-side interface to the transport.

The RapidsShuffleServer uses a ServerConnection to start the management port in the transport (in order to allow for incoming connections)

trait Transaction extends AutoCloseable

This trait represents a shuffle "transaction", and it is specific to a transfer (or set of transfers).

It is useful in that it groups a set of sends and receives requires in order to carry an action against a peer. It can be used to find statistics about the transfer (bytes send/received, throughput), and it also can be waited on, for blocking clients.

NOTE: a Transaction is thread safe w.r.t. a connection's callback. Calling methods on the transaction outside of waitForCompletion produces undefined behavior.

trait TransactionCallback extends AnyRef

case class TransactionStats(txTimeMs: Double, sendSize: Long, receiveSize: Long, sendThroughput: Double, recvThroughput: Double) extends Product with Serializable

Case class representing stats for the a transaction

txTimeMs: amount of time this Transaction took
sendSize: amount of bytes sent
receiveSize: amount of bytes received
sendThroughput: send throughput in GB/sec
recvThroughput: receive throughput in GB/sec

trait TransportBuffer extends AutoCloseable

TransportBuffer represents a buffer with an address and length.

There are two implementations:

MetadataTransportBuffer: This is host memory backing a flat buffer, used for block metadata and transfer requests.
CudfTransportBuffer: This is data backing a cuDF MemoryBuffer, and can be on the host or device (sender side) and exclusively on the device (client side)

class WindowedBlockIterator[T <: BlockWithSize] extends Iterator[Seq[BlockRange[T]]]

Given a set of blocks, this iterator returns BlockRanges of such blocks that fit windowSize.

Given a set of blocks, this iterator returns BlockRanges of such blocks that fit windowSize. The ranges are just logical chunks of the blocks, so this class performs no memory management or copying.

If a block is too large for the window, the block will be returned in next() until the full block can be covered.

For example, given a block that is 4 window-sizes in length: block = [sb1, sb2, sb3, sb4]

The window will return on next() four "sub-blocks", governed by windowSize: window.next() // sb1 window.next() // sb2 window.next() // sb3 window.next() // sb4

If blocks are smaller than the windowSize, they will be packed: block1 = [b1] block2 = [b2] window.next() // [b1, b2]

A mix of both scenarios above is possible: block1 = [sb11, sb12, sb13] // where sb13 is smaller than window length block2 = [b2]

window.next() // sb11 window.next() // sb12 window.next() // [sb13, b2]

T: - the specific type of BlockWithSize

Note: this class does not own transferBlocks
,
this class is not thread safe

Packages

shuffle

package shuffle

Type Members

Value Members

Ungrouped