unilab.ipc.memory_budget

Memory budget estimation for async RL training buffers.

Pure functions that estimate memory usage and warn if the system is likely to OOM before allocating large shared buffers.

Functions

estimate_appo_bytes(num_envs, steps_per_env, ...)

Estimate memory for APPO rollout ring buffer.

estimate_offpolicy_bytes(num_envs, ...)

Estimate memory for off-policy replay buffer + double-buffer slots.

get_available_memory_bytes()

Best-effort available memory detection.

warn_if_over_budget(estimated, label[, ...])

Print a warning if estimated memory exceeds threshold of available.
unilab.ipc.memory_budget.estimate_offpolicy_bytes(num_envs, replay_buffer_n, obs_dim, action_dim, critic_dim, batch_size, updates_per_step)[source]

Estimate memory for off-policy replay buffer + double-buffer slots.

Parameters:

  • num_envs (int)

  • replay_buffer_n (int)

  • obs_dim (int)

  • action_dim (int)

  • critic_dim (int)

  • batch_size (int)

  • updates_per_step (int)

Return type:

dict[str, int | str]

unilab.ipc.memory_budget.estimate_appo_bytes(num_envs, steps_per_env, obs_dim, action_dim, critic_dim, num_slots=4)[source]

Estimate memory for APPO rollout ring buffer.

Parameters:

  • num_envs (int)

  • steps_per_env (int)

  • obs_dim (int)

  • action_dim (int)

  • critic_dim (int)

  • num_slots (int)

Return type:

dict[str, int | str]

unilab.ipc.memory_budget.get_available_memory_bytes()[source]

Best-effort available memory detection.

Return type:

int | None

unilab.ipc.memory_budget.warn_if_over_budget(estimated, label, threshold=0.8)[source]

Print a warning if estimated memory exceeds threshold of available.

Parameters:
Return type:

None