`nemo_rl.data_plane.adapters.mooncake_ssd_offload`#

mooncake_cpu.ssd_offload: Mooncake’s embedded offload in segment owners.

Mooncake’s master offloads an object only to the SSD of the process owning the segment that holds it, so every process that owns Mooncake memory runs offload in its own directory. TQ’s own MooncakeStore.offload is not used: its single offload client on the driver owns none of the data.

Module Contents#

Functions#

`master_args`	mooncake_master flags for `ssd_offload`; empty when it is off.
`install`	Set up this process’s Mooncake client for SSD offload.
`wait_for_memory_on_put`	Make TQ’s puts wait for offload to free memory instead of failing.
`upsert_waiting_for_memory`	Upsert, putting keys that find every segment full again until they fit.
`_pick`
`require_offloading_store`	Raise if this process owns memory but TQ never built the offloading store.
`_hold_dir_lock`
`_remove_stale_dirs`	Delete offload dirs whose owning process is gone (crash, kill, restart).
`remove_dir`	Delete this process’s offload directory, if it has one.
`_file_key`	The default tenant’s key in a per-key offload file, or None.
`sweep_cleared_files`	Delete offload files under `root` whose key the master no longer has.
`_stamp`
`_start_sweep`
`disk_evictions`	The master’s SSD-eviction count from its Prometheus `/metrics` text.
`watch_disk_evictions`	Poll `read` and stop the run at the first SSD eviction.
`watch_master_disk_evictions`	Watch the master’s `/metrics` at `metrics_url` from this process.

Data#

`LOGGER`
`REQUIRED_ENV`
`_dir`
`_lock_fd`
`_offloading_store_set_up`
`NO_AVAILABLE_HANDLE`
`PUT_WAIT_S`
`_FIRST_BACKOFF_S`
`PUT_BACKOFF_MAX_S`
`EVICTION_COUNTER`
`_COUNTER_LINE`
`_watch_stop`
`SWEEP_INTERVAL_S`
`_SWEEP_BATCH`
`_TENANT_PREFIX`
`_sweep_stop`

API#

nemo_rl.data_plane.adapters.mooncake_ssd_offload.LOGGER#: ‘getLogger(…)’

nemo_rl.data_plane.adapters.mooncake_ssd_offload.REQUIRED_ENV#: None

nemo_rl.data_plane.adapters.mooncake_ssd_offload._dir: str | None#: None

nemo_rl.data_plane.adapters.mooncake_ssd_offload._lock_fd: int | None#: None

nemo_rl.data_plane.adapters.mooncake_ssd_offload._offloading_store_set_up#: False

nemo_rl.data_plane.adapters.mooncake_ssd_offload.NO_AVAILABLE_HANDLE#: None

nemo_rl.data_plane.adapters.mooncake_ssd_offload.PUT_WAIT_S#: 600.0

nemo_rl.data_plane.adapters.mooncake_ssd_offload._FIRST_BACKOFF_S#: 0.05

nemo_rl.data_plane.adapters.mooncake_ssd_offload.PUT_BACKOFF_MAX_S#: 1.0

nemo_rl.data_plane.adapters.mooncake_ssd_offload.EVICTION_COUNTER#: ‘master_evict_disk_replica_requests_total’

nemo_rl.data_plane.adapters.mooncake_ssd_offload._COUNTER_LINE#: ‘compile(…)’

nemo_rl.data_plane.adapters.mooncake_ssd_offload._watch_stop: threading.Event | None#: None

nemo_rl.data_plane.adapters.mooncake_ssd_offload.SWEEP_INTERVAL_S#: 30.0

nemo_rl.data_plane.adapters.mooncake_ssd_offload._SWEEP_BATCH#: 1000

nemo_rl.data_plane.adapters.mooncake_ssd_offload._TENANT_PREFIX#: ‘default\x00’

nemo_rl.data_plane.adapters.mooncake_ssd_offload._sweep_stop: threading.Event | None#: None

nemo_rl.data_plane.adapters.mooncake_ssd_offload.master_args( cfg: nemo_rl.data_plane.interfaces.MooncakeCpuConfig, ) → tuple[str, ...]#

mooncake_master flags for ssd_offload; empty when it is off.

In Mooncake, offload happens on eviction: past the high watermark the master moves objects to SSD and frees eviction_ratio of the memory. These replace the no-eviction flags TQ passes (gflags keeps the last occurrence). offload_force_evict=false: an object whose SSD write failed stays in memory rather than being dropped. A fresh object cannot be evicted until its lease expires: TQ’s 999999 ms would pin ~17 min of writes, and even 10 s let a 512 MiB segment fill at 64 MiB/s before anything was evictable (job 1321643). The 1 s lease must cover one whole batch_get_into call (up to staging_buffer_size per call), not one key: a slower call fails with LEASE_EXPIRED.

nemo_rl.data_plane.adapters.mooncake_ssd_offload.install( cfg: nemo_rl.data_plane.interfaces.MooncakeCpuConfig, segment_size: int | None, ) → None#

Set up this process’s Mooncake client for SSD offload.

No-op unless offload is on. Every client’s puts then wait for memory (:func:wait_for_memory_on_put). A process that owns memory also offloads to its own SSD directory: TQ’s client calls setup with no offload arguments, so this swaps the store class it builds for one that adds them. Each owner gets a fresh, owner-only <path>/nemo_rl_mooncake_offload/<ray job>-<pid>-<random>, locked for its lifetime: no live run can pick up, or wipe, another’s files, and dirs whose owner died are removed here. Must run before the store is set up.

nemo_rl.data_plane.adapters.mooncake_ssd_offload.wait_for_memory_on_put(client_cls: Any) → None#

Make TQ’s puts wait for offload to free memory instead of failing.

TQ’s _batch_upsert_with_retry gives up after 3 x 1 s. Raising TQ’s MAX_RETRIES instead would also retry permanent errors, logging each at ERROR. Patches client_cls once.

nemo_rl.data_plane.adapters.mooncake_ssd_offload.upsert_waiting_for_memory( upsert: collections.abc.Callable[[list[str], list[int], list[int]], Any], keys: list[str], ptrs: list[int], sizes: list[int], *, on_error: collections.abc.Callable[[list[str], list[int], list[int], list[int]], None], ) → None#

Upsert, putting keys that find every segment full again until they fit.

Mooncake admits a put only into free segment memory, while offload frees memory later, in the background. upsert returns one status per key; keys answered NO_AVAILABLE_HANDLE are put again with capped backoff for up to PUT_WAIT_S, with one warning. on_error(keys, ptrs, sizes, statuses) gets the keys that failed any other way.

nemo_rl.data_plane.adapters.mooncake_ssd_offload._pick( columns: tuple[list[Any], ...], indexes: list[int], ) → list[list[Any]]#

nemo_rl.data_plane.adapters.mooncake_ssd_offload.require_offloading_store() → None#

Raise if this process owns memory but TQ never built the offloading store.

The swap in :func:install works only while TQ looks the store class up by name; call this after connecting, when TQ has built its store.

nemo_rl.data_plane.adapters.mooncake_ssd_offload._hold_dir_lock(path: str) → None#

nemo_rl.data_plane.adapters.mooncake_ssd_offload._remove_stale_dirs(parent: str) → None#: Delete offload dirs whose owning process is gone (crash, kill, restart).

nemo_rl.data_plane.adapters.mooncake_ssd_offload.remove_dir() → None#

Delete this process’s offload directory, if it has one.

Call only once nothing will read this process’s data again: the files may be the only copy of offloaded objects.

nemo_rl.data_plane.adapters.mooncake_ssd_offload._file_key(path: str) → str | None#

The default tenant’s key in a per-key offload file, or None.

The file holds a protobuf KVEntry whose field 1 is the tenant, a NUL, then the key; the file name replaces the NUL and filesystem-unsafe characters with _, so it cannot be mapped back. Other tenants’ files are skipped.

nemo_rl.data_plane.adapters.mooncake_ssd_offload.sweep_cleared_files( root: str, exists: collections.abc.Callable[[list[str]], list[int]], ) → int#

Delete offload files under root whose key the master no longer has.

exists is batch_is_exist: 1 present, 0 absent, negative on error. A file is deleted only on an explicit 0 and only if it is unchanged since it was listed (before the check): a key re-put and offloaded again after the check rewrites its file, which is then kept. Returns the count deleted.

nemo_rl.data_plane.adapters.mooncake_ssd_offload._stamp(path: str) → tuple[int, int, int] | None#

nemo_rl.data_plane.adapters.mooncake_ssd_offload._start_sweep(store: Any, root: str) → None#

nemo_rl.data_plane.adapters.mooncake_ssd_offload.disk_evictions(metrics_text: str) → int#: The master’s SSD-eviction count from its Prometheus /metrics text.

nemo_rl.data_plane.adapters.mooncake_ssd_offload.watch_disk_evictions( read: collections.abc.Callable[[], int], path: str, *, interval_s: float = 30.0, on_eviction: collections.abc.Callable[[], None] = _thread.interrupt_main, ) → threading.Event#

Poll read and stop the run at the first SSD eviction.

Mooncake’s per-key backend deletes the oldest offloaded files once the disk is nearly full and cannot be told not to (0.3.13), and a deleted file may be the only copy of an object. The default on_eviction interrupts this process’s main thread, ending the run. Returns an event that stops the watcher when set; the watcher also sets it once it has fired.

nemo_rl.data_plane.adapters.mooncake_ssd_offload.watch_master_disk_evictions(metrics_url: str, path: str) → None#: Watch the master’s /metrics at metrics_url from this process.

nemo_rl.data_plane.adapters.mooncake_ssd_offload#

Module Contents#

Functions#

Data#

API#

`nemo_rl.data_plane.adapters.mooncake_ssd_offload`#