nemo_rl.models.generation.fleet_health

Liveness bookkeeping for the vLLM generation fleet.

Both SingleController rollout paths pick a generation shard by static round-robin with no idea whether the shard is alive: the native path in VllmGeneration._async_generate_base, and the NeMo-Gym path inside Gym’s own _resolve_client. This module owns the missing half — which shards are eligible to serve, and why.

Deliberately a pure state machine. Probing, restarting and pushing membership are I/O and belong to the caller, which keeps every transition here testable without Ray, a network, or a GPU — and keeps one description of “eligible” that both routing adapters read.

The transition that carries the most weight is the one that does not exist: a shard cannot go from DEAD back to HEALTHY on its own. A restarted engine holds whatever weights it loaded at init, so re-admitting it because it started answering probes again would feed training rollouts generated from a checkpoint hundreds of steps stale – invisible, and far worse than the outage that caused it. Recovery must pass through STALE and a completed refit.

Module Contents

Classes

ShardState	Lifecycle of one generation data-parallel shard.
ShardHealth	Everything known about one shard.
FleetHealthPolicy	Thresholds resolved from async_rl.generation_fleet_health.
GenerationFleetHealth	Tracks per-shard health and exposes the current serving set.
HealthyShardSelector	Picks a serving shard, preferring the one with the fewest requests in flight.

Data

_SERVING_STATES

API

class nemo_rl.models.generation.fleet_health.ShardState

Bases: str, enum.Enum

Lifecycle of one generation data-parallel shard.

Initialization

Initialize self. See help(type(self)) for accurate signature.

HEALTHY

‘healthy’

SUSPECT

‘suspect’

DEAD

‘dead’

RESTARTING

‘restarting’

STALE

‘stale’

RETIRED

‘retired’

nemo_rl.models.generation.fleet_health._SERVING_STATES

‘frozenset(…)’

class nemo_rl.models.generation.fleet_health.ShardHealth

Everything known about one shard.

dp_shard_idx: int

None

base_url: Optional[str]

None

state: nemo_rl.models.generation.fleet_health.ShardState

None

consecutive_probe_failures: int

0

consecutive_probe_successes: int

0

consecutive_reported_failures: int

0

restart_attempts: int

0

weight_version: int

0

last_ok_at: float

0.0

last_error: str = <Multiline-String>

property is_serving: bool

exception nemo_rl.models.generation.fleet_health.GenerationFleetExhausted

Bases: RuntimeError

Too few shards remain for the run to be worth continuing.

Initialization

Initialize self. See help(type(self)) for accurate signature.

class nemo_rl.models.generation.fleet_health.FleetHealthPolicy

Thresholds resolved from async_rl.generation_fleet_health.

Mirrors :class:nemo_rl.experience.rollout_manager.RolloutTimeouts in shape: an internal dataclass so the BaseModel stays the single home for user-facing defaults.

unhealthy_threshold: int

3

healthy_threshold: int

2

max_restart_attempts_per_shard: int

5

min_healthy_shards: int

1

__post_init__() → None

class nemo_rl.models.generation.fleet_health.GenerationFleetHealth(

*,

shard_count: int,

policy: nemo_rl.models.generation.fleet_health.FleetHealthPolicy,

base_urls: Optional[list[Optional[str]]] = None,

clock: Optional[Callable[[], float]] = None,

)

Tracks per-shard health and exposes the current serving set.

Parameters:

• shard_count – Number of generation data-parallel shards.

• policy – Thresholds governing the transitions.

• base_urls – Per-shard OpenAI base URLs, when the deployment has them.

• clock – Monotonic time source; injectable so tests need not sleep.

Initialization

property membership_epoch: int

property shard_count: int

snapshot() → list[nemo_rl.models.generation.fleet_health.ShardHealth]

Per-shard health, ordered by shard index.

serving_shards() → list[int]

Shard indices currently eligible to be handed traffic.

state_of(

shard_idx: int,

) → nemo_rl.models.generation.fleet_health.ShardState

serving_base_urls() → list[str]

Base URLs of the serving shards, for pushing to the NeMo-Gym router.

counts_by_state() → dict[nemo_rl.models.generation.fleet_health.ShardState, int]

as_metrics() → dict[str, float]

Flatten into a metric dict for the SingleController logger.

record_probe(shard_idx: int, *, ok: bool, error: str = '') → None

Fold one probe result into a shard’s state.

Probes never resurrect a shard that has left the serving set: DEAD, RESTARTING, STALE and RETIRED all ignore them, because getting an answer says nothing about whether the weights are current.

report_failure(shard_idx: int, error: BaseException) → None

Record a failure observed by a routing adapter rather than by a probe.

The adapters are the only components actually issuing generation requests, so they see failures a liveness probe cannot – a shard that answers /health and still errors on every generation.

Counted on its own streak rather than folded into :meth:record_probe. That used to be the implementation, and it could not condemn the very shard it exists for: a wedged engine answers is_alive, so an ok-probe every probe_interval_s reset the shared counter, and reported failures only accumulated if unhealthy_threshold of them landed inside one probe window. Under load they do; under a trickle the shard oscillated HEALTHY<->SUSPECT forever. A streak that only a successful generation clears says what is actually meant: this shard has failed N requests in a row.

report_success(shard_idx: int) → None

Record a generation that completed on this shard.

The reset half of :meth:report_failure. Without it the reported streak is monotonic and every shard eventually reaches unhealthy_threshold given a long enough run, however healthy it is.

shard_for_base_url(url: str) → Optional[int]

Reverse the shard -> base URL mapping, for failures reported by URL.

The NeMo-Gym router knows its backends only as URLs; the ledger keys everything by shard index.

mark_restarting(shard_idx: int) → None

A replacement engine is being brought up for a dead shard.

mark_loaded(shard_idx: int) → None

The replacement finished loading. It holds stale weights until refit.

report_refit(shard_idx: int, *, weight_version: int) → None

A completed refit is the only way back into the serving set.

retire(shard_idx: int, *, reason: str) → None

Remove a shard permanently. Training continues on what is left.

raise_if_exhausted() → None

Raise once too few shards remain for the run to be worth continuing.

_transition(

shard: nemo_rl.models.generation.fleet_health.ShardHealth,

new_state: nemo_rl.models.generation.fleet_health.ShardState,

) → None

_refresh_membership() → None

class nemo_rl.models.generation.fleet_health.HealthyShardSelector

Picks a serving shard, preferring the one with the fewest requests in flight.

Least-outstanding rather than round-robin because it is strictly better for LLM serving: it steers away from a shard that is merely slow or wedged without needing that to be diagnosed first.

monitor: nemo_rl.models.generation.fleet_health.GenerationFleetHealth

None

_inflight: dict[int, int]

‘field(…)’

next_shard() → int

Return the shard index to serve the next request.

Raises:

NoHealthyShards – No shard is currently eligible.

acquire(shard_idx: int) → None

release(shard_idx: int) → None

inflight(shard_idx: int) → int