nemo_gym.base_environment_server

Shared lifecycle for environment servers.

Module Contents

Classes

Name	Description
`BaseEnvironmentServer`	Expose a typed episode protocol with shared limits and cleanup.
`BaseEnvironmentServerConfig`	Configure protocol-neutral episode limits.
`CleanupContext`	Hold bounded process-local cleanup callbacks for one episode.
`CleanupHandle`	Close one registered participant at a protocol boundary.
`HandledEpisodeError`	Carry a failure that belongs in the episode response.
`_CleanupEntry`	-

Data

API

class nemo_gym.base_environment_server.BaseEnvironmentServer()

Bases: SimpleServer, Generic[EpisodeRequestT, EpisodeResponseT]

Expose a typed episode protocol with shared limits and cleanup.

_admission

Semaphore | None = None

_checkpoint

EnvironmentParticipant | None = None

checkpoint_boundaries

bool = False

config

BaseEnvironmentServerConfig

request_model

type[EpisodeRequestT]

response_model

type[EpisodeResponseT]

nemo_gym.base_environment_server.BaseEnvironmentServer._unhandled_failure_response(
    request: nemo_gym.base_environment_server.EpisodeRequestT,
    error: Exception
) -> nemo_gym.base_environment_server.EpisodeResponseT

nemo_gym.base_environment_server.BaseEnvironmentServer.aggregate_metrics(
    body: nemo_gym.config_types.AggregateMetricsRequest = Body()
) -> nemo_gym.config_types.AggregateMetrics

asyncabstract

Aggregate per-rollout scores into task-level metrics.

nemo_gym.base_environment_server.BaseEnvironmentServer.checkpoint_boundary(
    request: nemo_gym.base_environment_server.EpisodeRequestT,
    state: dict[str, pydantic.JsonValue] | collections.abc.Callable[[], dict[str, pydantic.JsonValue]]
) -> None

async

Record that a protocol step completed; state names the next step and what it needs.

While a checkpoint is open the episode parks here until resume. A replacement attempt that continues a checkpoint receives the latest recorded state from checkpoint_continuation. Pass a function when building state costs something: it is called only with checkpointing on.

nemo_gym.base_environment_server.BaseEnvironmentServer.checkpoint_continuation(
    request: nemo_gym.base_environment_server.EpisodeRequestT
) -> dict[str, pydantic.JsonValue] | None

Return the protocol state this attempt continues from a checkpoint, once, or None.

nemo_gym.base_environment_server.BaseEnvironmentServer.checkpoint_restart(
    request: nemo_gym.base_environment_server.EpisodeRequestT
) -> None

async

Mark this episode as a restart: a server it uses cannot capture its part, such as a restart-only agent.

Call it when a seed reply says so (nemo_gym._checkpoint.steps.seed_restarts). The episode then never holds up a checkpoint and is never exported; after a crash, the controller starts it over from its input.

nemo_gym.base_environment_server.BaseEnvironmentServer.checkpoint_step(
    request: nemo_gym.base_environment_server.EpisodeRequestT,
    mode: nemo_gym._checkpoint.steps.StepMode
) -> contextlib.AbstractAsyncContextManager[None]

Run one protocol step in wait or replay mode (see nemo_gym._checkpoint.steps).

Use replay only for a step that is safe to run again after a crash and does not change checkpointed state; everything else waits.

nemo_gym.base_environment_server.BaseEnvironmentServer.failure_response(
    request: nemo_gym.base_environment_server.EpisodeRequestT,
    failure: nemo_gym.episode_types.EpisodeFailure
) -> nemo_gym.base_environment_server.EpisodeResponseT

nemo_gym.base_environment_server.BaseEnvironmentServer.model_post_init(
    context: typing.Any
) -> None

nemo_gym.base_environment_server.BaseEnvironmentServer.run(
    request: nemo_gym.base_environment_server.EpisodeRequestT,
    cleanup: nemo_gym.base_environment_server.CleanupContext
) -> nemo_gym.base_environment_server.EpisodeResponseT

asyncabstract

Run one concrete environment protocol.

nemo_gym.base_environment_server.BaseEnvironmentServer.run_request(
    request: nemo_gym.base_environment_server.EpisodeRequestT
) -> nemo_gym.base_environment_server.EpisodeResponseT

async

nemo_gym.base_environment_server.BaseEnvironmentServer.setup_environment_checkpoint(
    app: fastapi.FastAPI
) -> None

Take part in partial-rollout checkpoints when they are enabled.

Every episode is tracked. A protocol that records boundaries (checkpoint_boundaries) continues after a restore; the episodes of a protocol that records none are restarts, which never hold up a checkpoint and start over from their input after a crash. Call this from any setup_webserver that builds its own app.

nemo_gym.base_environment_server.BaseEnvironmentServer.setup_webserver() -> fastapi.FastAPI

nemo_gym.base_environment_server.BaseEnvironmentServer.validate_response_identity(
    request: nemo_gym.base_environment_server.EpisodeRequestT,
    response: nemo_gym.base_environment_server.EpisodeResponseT
) -> None

staticmethod

class nemo_gym.base_environment_server.BaseEnvironmentServerConfig()

Bases: BaseRunServerInstanceConfig

Configure protocol-neutral episode limits.

cleanup_timeout_seconds

PositiveFloat

default_episode_timeout_seconds

PositiveFloat | None = None

max_concurrent_episodes

PositiveInt | None = None

model_config

= ConfigDict(extra='forbid')

queue_timeout_seconds

PositiveFloat | None = None

nemo_gym.base_environment_server.BaseEnvironmentServerConfig.validate_queue_timeout() -> typing_extensions.Self

class nemo_gym.base_environment_server.CleanupContext(
    episode_id: nemo_gym.episode_types.EpisodeId,
    cleanup_timeout_seconds: float,
    _cleanups: list[nemo_gym.base_environment_server._CleanupEntry] = list()
)

Dataclass

Hold bounded process-local cleanup callbacks for one episode.

Callbacks:

are process-local Python objects, although they may issue remote close requests;
run sequentially in LIFO order, each with its own cleanup timeout;
must be idempotent because a timed-out remote request may have succeeded;
are lost on process or host failure, so remote owners need expiry or reaping.

Callback failures or timeouts do not prevent later owners from tearing down their resources. Total cleanup is bounded by the number of registered callbacks times the per-callback timeout. Cleanup has no durable retry after this context is discarded.

_cleanups

list[_CleanupEntry] = field(default_factory=list)

cleanup_timeout_seconds

float

episode_id

EpisodeId

nemo_gym.base_environment_server.CleanupContext.aclose() -> None

async

nemo_gym.base_environment_server.CleanupContext.register_cleanup(
    name: str,
    callback: nemo_gym.base_environment_server.CleanupCallback
) -> nemo_gym.base_environment_server.CleanupHandle

class nemo_gym.base_environment_server.CleanupHandle(
    entry: nemo_gym.base_environment_server._CleanupEntry,
    timeout_seconds: float
)

Dataclass

Close one registered participant at a protocol boundary.

entry

_CleanupEntry

timeout_seconds

float

nemo_gym.base_environment_server.CleanupHandle.close() -> None

async

Run this idempotent callback once.

The episode timeout bounds calls made during the protocol. The context’s cleanup timeout bounds callbacks left for final unwinding.

class nemo_gym.base_environment_server.HandledEpisodeError(
    failure: nemo_gym.episode_types.EpisodeFailure
)

Exception

Bases: Exception

Carry a failure that belongs in the episode response.

class nemo_gym.base_environment_server._CleanupEntry(
    name: str,
    callback: nemo_gym.base_environment_server.CleanupCallback,
    active: bool = True,
    lock: asyncio.Lock = asyncio.Lock()
)

Dataclass

active

bool = True

callback

CleanupCallback

lock

Lock = field(default_factory=(asyncio.Lock))

name

str

nemo_gym.base_environment_server._CleanupEntry.close() -> None

async

nemo_gym.base_environment_server.CleanupCallback = Callable[[], Awaitable[None]]

nemo_gym.base_environment_server.EpisodeRequestT = TypeVar('EpisodeRequestT', bound=(BaseEpisodeRequest[Any]))

nemo_gym.base_environment_server.EpisodeResponseT = TypeVar('EpisodeResponseT', bound=(BaseEpisodeResponse[Any]))

nemo_gym.base_environment_server.LOGGER = logging.getLogger(__name__)