Channel API#

使用这些接口创建队列，并定制入队转换和消费者路由行为。执行模型和最小注册 hook 案例参见 使用 Channel 进行通信。

Channel#

class rlinf.scheduler.Channel

A FIFO queue-like channel for inter-worker communication.

Creation: Channel can be created both inside and outside of worker contexts. The recommended practice is to create channels outside of worker contexts using Channel.create(), and then pass them into workers as needed. You can also create channels inside worker contexts or connect to existing channels, using self.create_channel() or self.connect_channel().

Interface: Similar as the asyncio.Queue, the Channel provides interfaces like put, get, put_no_wait, and get_no_wait, as well as query interfaces like qsize, empty, and full. The semantics of these interfaces are identical to those of asyncio.Queue.

Features:

  1. Async operation: Channel supports both synchronous and asynchronous put and get operations, similar to Worker's send and recv APIs. Both operations accept arbitrary data item as long as it's serializable. The default behavior is synchronous, and async operations can be enabled by setting the async_op flag. This async can be used not only in asyncio context with await channel.get(async_op=True).async_wait(), but also in non-asyncio contexts by generating a communication handle that can be waited later, like async torch.distributed.send().

  2. Key-based routing: Channel allows specifying a key for each data item, which can be used to identify and route messages. For example, if you wish a specific data to be get and processed by a specific worker, you can assign a unique key to that data item when putting it into the channel. The target worker can then use this key to retrieve the specific data item.This is useful in multi-turn scenarios in agent and embodied RL, where a data is processed by a fixed set of workers.

  3. Weight and batch processing: Channel also supports assigning weights to individual data items, allowing for more fine-grained control over how messages are processed. A get_batch method can be used to retrieve a batch of messages which respects the assigned weights.

  4. Debugging: Channel allows you to print a Channel's internal data by directly print the Channel object.

Example:

>>> import sys
>>> import os
>>> import asyncio
>>> import torch
>>> from rlinf.scheduler import (
...     Worker,
...     Cluster,
...     PackedPlacementStrategy,
... )
>>>
>>> class Producer(Worker):
...     def __init__(self):
...         super().__init__()
...
...     def produce(self, channel: Channel):
...         # Synchronous put of common object
...         channel.put("Hello from Producer")
...
...         # Synchronous put of tensor
...         tensor = torch.ones(1, device=torch.cuda.current_device())
...         channel.put(tensor)
...
...         # Asynchronous put of common object
...         async_work = channel.put(
...             "Hello from Producer asynchronously", async_op=True
...         )
...         async_work.wait()
...
...         # Asynchronous put using asyncio
...         async_work = channel.put(tensor, async_op=True)
...
...         async def wait_async():
...             await async_work.async_wait()
...
...         asyncio.run(wait_async())
...
...         # Put object with weight
...         channel.put("Hello with weight", weight=1)
...         channel.put(tensor, weight=2)
>>>
>>> class Consumer(Worker):
...     def __init__(self):
...         super().__init__()
...
...     def consume(self, channel: Channel):
...         tensor = channel.get()
...
...         async_work = channel.get(async_op=True)
...         async_result = async_work.wait()
...
...         async_work = channel.get(async_op=True)
...
...         async def wait_async():
...             result = await async_work.async_wait()
...
...         asyncio.run(wait_async())
...
...         # Get batch of objects based on weight
...         batch = channel.get_batch(target_weight=3)
>>>
>>> cluster = Cluster(num_nodes=1)
>>> channel = Channel.create(name="channel")
>>> placement = PackedPlacementStrategy(
...     start_hardware_rank=0, end_hardware_rank=0
... )
>>> producer = Producer.create_group().launch(
...     cluster, name="test", placement_strategy=placement
... )
>>> consumer = Consumer.create_group().launch(
...     cluster, name="test2", placement_strategy=placement
... )
>>> r1 = producer.produce(channel)
>>> r2 = consumer.consume(channel)
>>> res = r1.wait()
>>> res = r2.wait()
classmethod create(name, maxsize=0, distributed=False, node_rank=None, local=False, disable_distributed_log=True, collector=None, dispatcher=None, cfg=None, producers=None, consumers=None, hook_options=None)

Create a new channel with the specified name, node ID, and accelerator ID.

参数:

  • name (str) -- The name of the channel.

  • maxsize (int) -- The maximum size of the channel queue. Defaults to 0 (unbounded).

  • distributed (bool) -- Whether the channel should be distributed. A distributed channel creates distributed workers on each node, and routes communications to the channel worker on the same node as the current worker, benefiting from the locality of the data. The routing is based on the key of the put/get APIs. So if you expect the key to be randomly distributed, you should set this to False to avoid unnecessary routing overhead.

  • node_rank (int) -- The node to place the channel on. Only valid when distributed is False. Defaults to None, which places the channel on the node of its first producer, falling back to its first consumer and then to the creating worker, so channel traffic stays node-local instead of crossing the network to node 0.

  • local (bool) -- Create the channel for intra-process communication. A local channel cannot be connected by other workers, and its data cannot be shared among different processes.

  • disable_distributed_log (bool) -- Whether to disable distributed log for the channel.

  • collector (Any) -- Transforms items on the way into the channel. A registered name, a "module:ClassName" import path, a Collector subclass, or an instance. Defaults to passing items through unchanged.

  • dispatcher (Any) -- Chooses which consumer receives each item, in the same forms. Defaults to one shared queue per key, which every consumer competes for.

  • cfg (Any) -- Run configuration handed to the hooks' setup.

  • producers (WorkerGroupSpec) -- Worker groups expected to put into this channel. Accepts worker groups, group names, or an iterable mixing the two.

  • consumers (WorkerGroupSpec) -- Worker groups expected to get from it, in the same forms. Each group is expanded into one consumer per rank, with the group's size read from the worker manager, so a dealing dispatcher can split work evenly across them.

  • hook_options (dict) -- Free-form options handed to the hooks' setup.

返回:

A new instance of the Channel class.

返回类型:

Channel

classmethod connect(name, current_worker)

Connect to an existing channel.

参数:

  • name (str) -- The name of the channel to connect to.

  • current_worker (Worker) -- The current worker that is connecting to the channel.

返回:

An instance of the Channel class connected to the specified channel.

返回类型:

Channel

property is_local

Check if the channel is a local channel.

qsize(key='default_queue', consumer=None)

Get the size of the channel queue.

参数:

  • key (Any) -- check the queue associated with the key.

  • consumer (str) -- With a dealing dispatcher, report only this consumer's share instead of the total outstanding across all consumers.

返回:

The number of items in the channel queue.

返回类型:

int

empty(key='default_queue')

Check if the channel queue is empty.

参数:

key (Any) -- The key to check the queue emptiness for.

返回:

True if the channel queue is empty, False otherwise.

返回类型:

bool

full(key='default_queue')

Check if the channel queue is full.

参数:

key (Any) -- The key to check the queue fullness for.

返回:

True if the channel queue is full, False otherwise.

返回类型:

bool

put(item, weight=0, key='default_queue', async_op=False)

Put an item into the channel queue.

参数:

  • item (Any) -- The item to put into the channel queue.

  • weight (int) -- The priority weight of the item. Defaults to 0.

  • key (Any) -- The key to get the item from. A unique identifier for a specific set of items.

  • given (When a key is)

  • key. (the channel will put the item in the queue associated with that)

  • exist (If the queue associated with the key does not)

  • created. (it will be)

  • async_op (bool) -- Whether to perform the operation asynchronously.

返回类型:

AsyncWork | None

put_nowait(item, weight=0, key='default_queue')

Put an item into the channel queue without waiting. Raises asyncio.QueueFull if the queue is full.

参数:

  • item (Any) -- The item to put into the channel queue.

  • weight (int) -- The priority weight of the item. Defaults to 0.

  • key (Any) -- The key to get the item from. A unique identifier for a specific set of items.

  • given (When a key is)

  • key. (the channel will put the item in the queue associated with that)

  • exist (If the queue associated with the key does not)

  • created. (it will be)

抛出:

asyncio.QueueFull -- If the queue is full.

get(key='default_queue', async_op=False)

Get an item from the channel queue.

参数:

  • key (Any) -- The key to get the item from. A unique identifier for a specific set of items.

  • given (When a key is)

  • key. (the channel will look for the item in the queue associated with that)

  • async_op (bool) -- Whether to perform the operation asynchronously.

返回:

The item retrieved from the channel queue.

返回类型:

Any

get_nowait(key='default_queue')

Get an item from the channel queue without waiting. Raises asyncio.QueueEmpty if the queue is empty.

参数:

  • key (Any) -- The key to get the item from. A unique identifier for a specific set of items.

  • given (When a key is)

  • key. (the channel will look for the item in the queue associated with that)

返回:

The item retrieved from the channel queue.

返回类型:

Any

抛出:

asyncio.QueueEmpty -- If the queue is empty.

get_batch(target_weight=0, key='default_queue', async_op=False)

Get a batch of items from the channel queue based on the set batch weight.

It will try to get items until the total weight of the items is about to (i.e., the next item will) exceed the set batch weight.

参数:

  • target_weight (int) -- The target weight for the batch.

  • key (Any) -- The key to get the item from. A unique identifier for a specific set of items.

  • given (When a key is)

  • key. (the channel will look for the item in the queue associated with that)

  • async_op (bool) -- Whether to perform the operation asynchronously.

返回:

A list of items retrieved from the channel queue.

返回类型:

List[Any]

Hook Context#

class rlinf.scheduler.ChannelContext#

Read-only description of the channel a hook is attached to.

Passed to Collector.setup() and Dispatcher.setup() once, on the channel worker, before any traffic flows.

name#

The channel's name.

Type:

str

cfg#

The run configuration, if one was supplied at channel creation.

Type:

Any

producers#

Worker group names expected to put into this channel.

Type:

tuple[str, ...]

consumers#

Consumer ids expected to get from it, each "<group>:<rank>".

Type:

tuple[str, ...]

options#

The hook's own free-form options, supplied at channel creation.

Type:

dict[str, Any]

__init__(name, cfg=None, producers=(), consumers=(), options=<factory>)#
参数:

  • name (str)

  • cfg (Any)

  • producers (tuple[str, ...])

  • consumers (tuple[str, ...])

  • options (dict[str, Any])

返回类型:

None

Collector#

class rlinf.scheduler.Collector#

Transform items on the way into a channel.

Subclass and override collect() to reshape, split, merge, filter, or re-key items before they are queued. The base implementation passes items through unchanged, which is the channel's default behavior.

setup(ctx)#

Prepare the collector. Called once before any item is collected.

参数:

ctx (ChannelContext) -- Description of the channel this collector is attached to.

返回类型:

None

collect(item, key)#

Turn one incoming item into the items that should be queued.

参数:

  • item (Any) -- The item as it was put into the channel.

  • key (str) -- The key it was put under.

生成器:

(key, item) pairs to enqueue. May yield nothing to drop the input, or several pairs to fan it out.

返回类型:

Iterable[tuple[str, Any]]

rlinf.scheduler.register_collector(name)#

Register a collector class under a name.

参数:

name (str) -- The name channels may select this collector by.

返回:

A class decorator that registers and returns the class unchanged.

返回类型:

Callable[[type], type]

Dispatcher#

class rlinf.scheduler.Dispatcher#

Decide which consumer receives each item.

The base implementation returns None for everything, meaning all consumers of a key share one queue -- the channel's default behavior. Override route() to hand each item to a specific consumer instead, which gives every consumer a private queue and makes non-blocking polling fair by construction.

setup(ctx)#

Prepare the dispatcher. Called once before any item is routed.

参数:

ctx (ChannelContext) -- Description of the channel this dispatcher is attached to.

返回类型:

None

property deals: bool#

Whether this dispatcher gives each consumer a private queue.

True as soon as route() is overridden. Override this property to force it either way.

route(item, key)#

Choose the consumer that should receive an item.

参数:

  • item (Any) -- The item about to be queued.

  • key (str) -- The key it will be queued under.

返回:

A consumer id ("<group>:<rank>"), or None to leave the item in the key's shared queue for whichever consumer asks first.

返回类型:

str | None

rebalance(key, starved, depths)#

Choose a consumer to take work from when one has run dry.

Only consulted for a blocking get on an empty private queue.

参数:

  • key (str) -- The key being read.

  • starved (str) -- The consumer id with nothing left to read.

  • depths (dict[str, int]) -- Current queue depth of every consumer on this key.

返回:

The consumer id to take one item from, or None to wait instead.

返回类型:

str | None

rlinf.scheduler.register_dispatcher(name)#

Register a dispatcher class under a name.

参数:

name (str) -- The name channels may select this dispatcher by.

返回:

A class decorator that registers and returns the class unchanged.

返回类型:

Callable[[type], type]