nemo_rl.utils.compile_cache#

Persistent JIT compile caches for NeMo RL’s vLLM and Megatron workers (opt-in).

Set NRL_COMPILE_CACHE_SHARED_DIR (kept between jobs, on a shared filesystem) and NRL_COMPILE_CACHE_LOCAL_DIR (node-local). Each cache type (vllm, megatron) lives in <shared>/<type>, or in NRL_COMPILE_CACHE_<TYPE>_DIR if set; a node copies it once to <local>/<type> (cache_sync.sh seed-type) and points that type’s compilers there. Writing new entries back is the launcher’s job (cache_sync.sh sidecar-types). See docs/guides/persistent-jit-caches.md.

Module Contents#

Functions#

local_dir

The node-local cache root, or None if persistent caches are off.

script

The cache_sync.sh to run: the one next to this module unless overridden.

env_overrides

Point the compilers at <local>/<cache_type>; {} if persistent caches are off.

seed

Copy cache_type’s persistent cache onto this node once (no-op if off).

vllm_engine_dir

This engine’s own VLLM_CACHE_ROOT under the local root, or None if off.

_vllm_patches

_prepare_vllm_engine_on_this_node

prepare_vllm_engine

Before an engine compiles: seed every node it runs on and fill its folder.

Data#

SHARED_VAR

LOCAL_VAR

SCRIPT_VAR

COMPILER_DIRS

API#

nemo_rl.utils.compile_cache.SHARED_VAR#

‘NRL_COMPILE_CACHE_SHARED_DIR’

nemo_rl.utils.compile_cache.LOCAL_VAR#

‘NRL_COMPILE_CACHE_LOCAL_DIR’

nemo_rl.utils.compile_cache.SCRIPT_VAR#

‘NRL_COMPILE_CACHE_SYNC_SCRIPT’

nemo_rl.utils.compile_cache.COMPILER_DIRS#

None

nemo_rl.utils.compile_cache.local_dir() → str | None[source]#

The node-local cache root, or None if persistent caches are off.

nemo_rl.utils.compile_cache.script() → str[source]#

The cache_sync.sh to run: the one next to this module unless overridden.

nemo_rl.utils.compile_cache.env_overrides(cache_type: str) → dict[str, str][source]#

Point the compilers at <local>/<cache_type>; {} if persistent caches are off.

nemo_rl.utils.compile_cache.seed(cache_type: str) → None[source]#

Copy cache_type’s persistent cache onto this node once (no-op if off).

Returns when the copy is done, also one another process on this node started (e.g. the launcher, to overlap Ray start-up). A failed copy only warns (in cache_sync.sh); a missing script or variable raises.

nemo_rl.utils.compile_cache.vllm_engine_dir(seed: int) → str | None[source]#

This engine’s own VLLM_CACHE_ROOT under the local root, or None if off.

Engines on one node must not share a folder (vLLM writes its cache index without locking), so each gets its own, filled from <local>/vllm/vllm.

nemo_rl.utils.compile_cache._vllm_patches() → dict[source]#
nemo_rl.utils.compile_cache._prepare_vllm_engine_on_this_node(engine_dir: str) → None[source]#
nemo_rl.utils.compile_cache.prepare_vllm_engine(bundle_indices: list[int] | None) → None[source]#

Before an engine compiles: seed every node it runs on and fill its folder.

Runs one task per node of the engine’s placement-group bundles (the vLLM process’s own node included) and waits for all of them.