rlightning.policy.utils.router

class rlightning.policy.utils.router.AsyncRouter

Bases: ABC

Abstract base class for request routing strategies.

abstractmethod assign(current_loads: List[int], num_tasks: int, env_ids: Sequence[str] | None = None) → List[int]

Assign num_tasks to policy indices.

Parameters:

• current_loads – List of current loads for each policy.

• num_tasks – Number of tasks to assign.

• env_ids – Optional sequence of environment IDs for affinity-based routing.

Returns:

List of policy indices corresponding to the chosen target for each task.

class rlightning.policy.utils.router.NodeAffinityRouter(component_distribution: Dict[str, Dict[str, Dict[str, Any]]], policy_node_ids: List[str] | None = None)

Bases: AsyncRouter

Route tasks to policies on the same node as env workers.

assign(current_loads: List[int], num_tasks: int, env_ids: Sequence[str] | None = None) → List[int]

Assign num_tasks to policy indices.

Parameters:

• current_loads – List of current loads for each policy.

• num_tasks – Number of tasks to assign.

• env_ids – Optional sequence of environment IDs for affinity-based routing.

Returns:

List of policy indices corresponding to the chosen target for each task.

class rlightning.policy.utils.router.SimpleRouter

Bases: AsyncRouter

Load-balancing router that assigns tasks to policies with minimum load.

assign(current_loads: List[int], num_tasks: int) → List[int]

Assign num_tasks to indices based on current_loads using a min-heap. Returns a list of indices corresponding to the chosen target for each task in order.

class rlightning.policy.utils.router.SyncNodeAffinityRouter(eval_policies: Sequence[Any], component_distribution: Dict[str, Dict[str, Dict[str, Any]]])

Bases: SyncRouter

Sync router that prefers policies on the same node as env workers.

This router supports concurrent scheduling across different nodes while maintaining sequential scheduling within each node through per-node locks.

property num_nodes: int

Return the number of nodes with eval policies.

select_policy(env_id: str) → Any

Select policy using round-robin within the target node.

With max_pending_tasks > 1 on the eval TaskSubmitter, backpressure is handled by TaskSubmitter itself. The router only needs to distribute tasks evenly across eval policies on the correct node — no idle-check RPCs needed.

Parameters:

env_id – Environment ID to determine target node.

Returns:

Selected policy instance.

class rlightning.policy.utils.router.SyncRouter

Bases: ABC

Abstract base class for sync rollout routing strategies.

abstractmethod select_policy(env_id: str | None = None) → Any

Select a policy for a single env_id in sync mode.

class rlightning.policy.utils.router.SyncSimpleRouter(eval_policies: Sequence[Any])

Bases: SyncRouter

Sync router using a single idle deque.

select_policy(env_id: str | None = None) → Any

Select policy with internal locking (simple router uses single global lock).

rlightning.policy.utils.router.create_router(rollout_mode: str, router_type: str, eval_policies: Sequence[Any], global_resource_manager=None) → SyncRouter | AsyncRouter

Factory for sync routers.

Parameters:

• rollout_mode – Rollout mode (“sync” or “async”).

• router_type – Type of router (“simple” or “node_affinity”).

• eval_policies – List of eval policy instances.

• global_resource_manager – Global resource manager. Required for “node_affinity” router.

Returns:

A SyncRouter or AsyncRouter instance.