`nemo_rl.environments.utils`#

Module Contents#

Classes#

EnvRegistryEntry

Functions#

`chunk_list_to_workers`	Chunk a list into a list of lists, where each sublist is assigned to a worker. Keeps ordering of elements.
`create_env`
`register_env`
`shutdown_environments`	Gracefully shut down every distinct environment actor in the given maps.

Data#

`DEFAULT_ENV_SHUTDOWN_TIMEOUT_SECONDS`
`ENV_REGISTRY`

API#

nemo_rl.environments.utils.DEFAULT_ENV_SHUTDOWN_TIMEOUT_SECONDS#: 10.0

class nemo_rl.environments.utils.EnvRegistryEntry#

Bases: typing.TypedDict

actor_class_fqn: str#: None

default_processor: NotRequired[str]#: None

nemo_rl.environments.utils.ENV_REGISTRY: Dict[str, nemo_rl.environments.utils.EnvRegistryEntry]#: None

nemo_rl.environments.utils.chunk_list_to_workers( to_chunk: list[Any], num_workers: int, ) → list[list[Any]]#

Chunk a list into a list of lists, where each sublist is assigned to a worker. Keeps ordering of elements.

If the list is not divisible by the number of workers, the last worker may have fewer elements. If there are more workers than elements, the first len(list) workers will have a single element each, and the remaining workers will have empty lists.

Parameters:

list – The list to be chunked.
num_workers – The number of workers to distribute the list to.

Returns:

A list of lists, where each sublist contains elements assigned to a worker.

Examples:

>>> from nemo_rl.environments.utils import chunk_list_to_workers
>>> chunk_list_to_workers([1, 2, 3, 4, 5], 3)
[[1, 2], [3, 4], [5]]

nemo_rl.environments.utils.create_env( env_name: str, env_config: dict, ) → nemo_rl.environments.interfaces.EnvironmentInterface#

nemo_rl.environments.utils.register_env(env_name: str, actor_class_fqn: str) → None#

nemo_rl.environments.utils.shutdown_environments( *env_maps: Optional[Dict[str, nemo_rl.environments.interfaces.EnvironmentInterface]], timeout: float | None = DEFAULT_ENV_SHUTDOWN_TIMEOUT_SECONDS, ) → None#

Gracefully shut down every distinct environment actor in the given maps.

Runners commonly bind the same actor — often the very same mapping — to both training and validation, so each handle is asked to stop once no matter how many maps contain it.

Environments must be torn down before generation workers: they may have in-flight HTTP requests to the vLLM endpoints, and killing generation first leaves them retrying dead connections.

An entry is either a Ray actor handle or a local object that owns actors of its own (a NeMo-Gym shard set); the latter is asked to shut itself down rather than killed, since only it knows what it holds.