metrale_model_engine/model/
impl_lora_swap.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Runtime LoRA slot swap and promotion, from an RDMA peer or from disk.
4//!
5//! Owner: model-engine.
6//! Invariants:
7//! - A swap refused for a busy slot, a disarmed rotation or a missing pool has not
8//!   touched the slot's bytes or identity.
9
10use anyhow::{Context, Result};
11
12use metrale_gpu_runtime::gpu::DevicePtr;
13
14use super::types::TransformerModel;
15use metrale_model_layers::layers::ops;
16
17impl TransformerModel {
18    /// 2026-09-25: Stage the adapter `adapter_name` (held by the peer at
19    /// `peer_addr` as `adapter_id`) into pool `slot` in place and make it that
20    /// slot's resident adapter. Call it only at a scheduler-quiescent point.
21    ///
22    /// Refuses when no pool is loaded, rotation is not armed
23    /// (`lora_rotatable`), `slot >= max_loras`, or the slot has in-flight
24    /// sequences. Otherwise the slot region is zeroed, the tensors land over
25    /// RDMA, the slot's pairs are rebuilt with the new r/scale, its tables are
26    /// refreshed and its generation bumped. If the slot is active, its pairs
27    /// are re-installed and the LoRA graph caches destroyed.
28    #[cfg(feature = "cuda")]
29    // 2026-09-25: Peer staging lands the tensors over RDMA, so this is
30    // unix-only. The `_from_disk` functions below are plain file I/O.
31    #[cfg(unix)]
32    pub fn swap_lora_slot_from_peer(
33        &mut self,
34        peer_addr: &str,
35        adapter_id: &str,
36        adapter_name: &str,
37        slot: usize,
38        peft: metrale_config::PeftAdapterConfig,
39    ) -> Result<()> {
40        use metrale_model_layers::lora::rdma_stage;
41
42        let (pool, max_rank, max_loras) = {
43            let lw = self
44                .lora
45                .as_ref()
46                .ok_or_else(|| anyhow::anyhow!("LoRA RDMA swap: no adapter pool loaded"))?;
47            (lw.pool, lw.max_rank, lw.max_loras)
48        };
49        if !self.lora_rotatable {
50            anyhow::bail!(
51                "LoRA RDMA swap needs rotation armed (set $METRALE_LORA_PEER or \
52                 METRALE_LORA_ROTATE=1 so decode runs eager)"
53            );
54        }
55        if slot >= max_loras {
56            anyhow::bail!("LoRA RDMA swap: slot {slot} >= max_loras {max_loras}");
57        }
58        // 2026-09-25: Refuse a busy slot before the memset below, so a refused
59        // swap leaves the slot's bytes and identity untouched.
60        {
61            let busy = self.lora.as_ref().unwrap().slot_ref_count(slot);
62            if busy > 0 {
63                anyhow::bail!(
64                    "LoRA RDMA swap REFUSED: slot {slot} has {busy} in-flight \
65                     sequence(s) (ref_count>0); cannot replace an adapter mid-decode"
66                );
67            }
68        }
69
70        let manifest = rdma_stage::fetch_adapter_manifest(peer_addr, adapter_id)?;
71        let targets =
72            rdma_stage::build_land_targets(&manifest, &self.config, pool, slot, max_rank)?;
73
74        // 2026-09-25: Zero the slot region before landing: a reused slot still
75        // holds the previous adapter's bytes.
76        let slot_bytes = rdma_stage::slot_bytes(&self.config, max_rank);
77        let slot_base = DevicePtr(pool.0 + (slot * slot_bytes) as u64);
78        self.gpu.memset(slot_base, 0, slot_bytes)?;
79        let loader = metrale_model_weights::weight_lora_rdma::RdmaLoraLoader::new(
80            peer_addr.to_string(),
81            adapter_id.to_string(),
82        );
83        loader.stage_into_slot(self.gpu.as_ref(), &targets)?;
84
85        let layers =
86            rdma_stage::rebuild_slot_layers(&targets, &self.config, &peft, pool, slot, max_rank)?;
87        // 2026-09-25: Refresh this slot's a/b pointer tables and scale table from
88        // the new adapter's module coverage, so the slot keeps no route entry
89        // from the adapter it replaced. `pack_store_into_slot` does the same for
90        // the disk swap.
91        self.lora.as_ref().unwrap().refresh_slot_tables(
92            slot,
93            &layers,
94            peft.scaling(),
95            self.gpu.as_ref(),
96        )?;
97        {
98            let lw = self.lora.as_mut().unwrap();
99            let s = lw
100                .slots
101                .get_mut(slot)
102                .ok_or_else(|| anyhow::anyhow!("LoRA RDMA swap: slot {slot} not resident"))?;
103            s.name = adapter_name.to_string();
104            s.adapter_config = peft;
105            s.layers = layers;
106            // 2026-09-25: Bump the generation so the slot yields a new adapter
107            // id, and a same-name request cannot hit prefix-cache entries keyed
108            // to the old contents.
109            s.generation = s.generation.wrapping_add(1);
110        }
111
112        let active = self.lora.as_ref().unwrap().active;
113        if active == slot {
114            let installed_layers = self.lora.as_ref().unwrap().slots[slot].layers.clone();
115            let tables = self.lora.as_ref().unwrap().tables.clone();
116            let scale_table = self.lora.as_ref().unwrap().scale_table;
117            let kernels = ops::lora_delta::LoraKernels::new(self.gpu.as_ref())?;
118            self.install_lora_layers(&installed_layers, kernels, &tables, scale_table)?;
119            self.lora.as_mut().unwrap().name = adapter_name.to_string();
120            self.destroy_lora_decode_graphs();
121        }
122        tracing::info!(
123            "LoRA RDMA swap: '{adapter_name}' landed in slot {slot} \
124             ({} targets, active_slot={active})",
125            targets.len()
126        );
127        Ok(())
128    }
129
130    /// 2026-09-25: Promote the adapter `adapter_name` from the peer into a
131    /// cache-region pool slot and make it active; returns `(slot,
132    /// evicted_name)`. Call it only at a scheduler-quiescent point.
133    ///
134    /// The victim (`select_victim_slot`) is a never-filled placeholder first,
135    /// else the least recently used idle (`ref_count == 0`) cache slot. With
136    /// every cache slot busy it fails with `POOL_FULL`; a busy slot is never
137    /// evicted. [`Self::swap_lora_slot_from_peer`] re-checks the ref count and
138    /// bumps the slot generation.
139    #[cfg(feature = "cuda")]
140    // 2026-09-25: Peer staging lands the tensors over RDMA, so this is
141    // unix-only. The `_from_disk` functions below are plain file I/O.
142    #[cfg(unix)]
143    pub fn promote_lora_slot_from_peer(
144        &mut self,
145        peer_addr: &str,
146        adapter_id: &str,
147        adapter_name: &str,
148        peft: metrale_config::PeftAdapterConfig,
149    ) -> Result<(usize, Option<String>)> {
150        let (slot, evicted) = {
151            let lw = self
152                .lora
153                .as_ref()
154                .ok_or_else(|| anyhow::anyhow!("LoRA promote: no adapter pool loaded"))?;
155            let views = lw.cache_slot_views();
156            let slot =
157                metrale_model_layers::lora::select_victim_slot(&views).map_err(|e| match e {
158                    metrale_model_layers::lora::VictimError::PoolFull => anyhow::anyhow!(
159                        "POOL_FULL: all {} cache slot(s) are busy (ref_count>0); retry",
160                        views.len()
161                    ),
162                })?;
163            // 2026-09-25: The name the victim held, if any, returned as `evicted_name`.
164            let evicted = lw
165                .slots
166                .get(slot)
167                .map(|s| s.name.clone())
168                .filter(|n| !n.is_empty());
169            (slot, evicted)
170        };
171
172        self.swap_lora_slot_from_peer(peer_addr, adapter_id, adapter_name, slot, peft)?;
173
174        // 2026-09-25: Make the promoted slot active. If it already was,
175        // `swap_lora_slot_from_peer` has re-installed it; otherwise install its
176        // pairs here.
177        let already_active = self.lora.as_ref().unwrap().active == slot;
178        if !already_active {
179            let (layers, tables, scale_table) = {
180                let lw = self.lora.as_mut().unwrap();
181                lw.active = slot;
182                lw.name = lw.slots[slot].name.clone();
183                lw.adapter_config = lw.slots[slot].adapter_config.clone();
184                (
185                    lw.slots[slot].layers.clone(),
186                    lw.tables.clone(),
187                    lw.scale_table,
188                )
189            };
190            let kernels = ops::lora_delta::LoraKernels::new(self.gpu.as_ref())?;
191            self.install_lora_layers(&layers, kernels, &tables, scale_table)?;
192            self.destroy_lora_decode_graphs();
193        }
194        // 2026-09-25: Mark the promoted slot most recently used, so a
195        // back-to-back promote of a different adapter picks an older victim
196        // before the request that triggered this one has taken its ref.
197        self.lora.as_ref().unwrap().touch_slot(slot);
198        tracing::info!(
199            "LoRA promote: '{adapter_name}' hot in cache slot {slot} \
200             (evicted={:?}), now active",
201            evicted
202        );
203        Ok((slot, evicted))
204    }
205
206    /// 2026-09-25: [`Self::promote_lora_slot_from_peer`] with the adapter
207    /// loaded from `adapter_dir` (named `name`) instead of a peer: same victim
208    /// choice, same make-active step, returns `(slot, evicted_name)`.
209    /// [`Self::swap_lora_slot_from_disk`] reads the directory's
210    /// `adapter_config.json`, refuses when rotation is not armed, re-checks
211    /// the ref count and bumps the slot generation. Call it only at a
212    /// scheduler-quiescent point.
213    pub fn promote_lora_slot_from_disk(
214        &mut self,
215        adapter_dir: &std::path::Path,
216        name: &str,
217    ) -> Result<(usize, Option<String>)> {
218        let (slot, evicted) = {
219            let lw = self
220                .lora
221                .as_ref()
222                .ok_or_else(|| anyhow::anyhow!("LoRA disk promote: no adapter pool loaded"))?;
223            let views = lw.cache_slot_views();
224            let slot =
225                metrale_model_layers::lora::select_victim_slot(&views).map_err(|e| match e {
226                    metrale_model_layers::lora::VictimError::PoolFull => anyhow::anyhow!(
227                        "POOL_FULL: all {} cache slot(s) are busy (ref_count>0); retry",
228                        views.len()
229                    ),
230                })?;
231            // 2026-09-25: The name the victim held, if any, returned as `evicted_name`.
232            let evicted = lw
233                .slots
234                .get(slot)
235                .map(|s| s.name.clone())
236                .filter(|n| !n.is_empty());
237            (slot, evicted)
238        };
239
240        self.swap_lora_slot_from_disk(adapter_dir, name, slot)?;
241
242        // 2026-09-25: Make the promoted slot active. If it already was,
243        // `swap_lora_slot_from_disk` has re-installed it; otherwise install its
244        // pairs here.
245        let already_active = self.lora.as_ref().unwrap().active == slot;
246        if !already_active {
247            let (layers, tables, scale_table) = {
248                let lw = self.lora.as_mut().unwrap();
249                lw.active = slot;
250                lw.name = lw.slots[slot].name.clone();
251                lw.adapter_config = lw.slots[slot].adapter_config.clone();
252                (
253                    lw.slots[slot].layers.clone(),
254                    lw.tables.clone(),
255                    lw.scale_table,
256                )
257            };
258            let kernels = ops::lora_delta::LoraKernels::new(self.gpu.as_ref())?;
259            self.install_lora_layers(&layers, kernels, &tables, scale_table)?;
260            self.destroy_lora_decode_graphs();
261        }
262        // 2026-09-25: Mark the promoted slot most recently used, so a
263        // back-to-back promote of a different adapter picks an older victim
264        // before the request that triggered this one has taken its ref.
265        self.lora.as_ref().unwrap().touch_slot(slot);
266        tracing::info!(
267            "LoRA disk-promote: '{name}' hot in cache slot {slot} \
268             (evicted={:?}), now active",
269            evicted
270        );
271        Ok((slot, evicted))
272    }
273
274    /// 2026-09-25: Load the adapter at `adapter_dir` into pool `slot` in place
275    /// and make it that slot's resident adapter, using the directory's own
276    /// `adapter_config.json`. If the slot is active, its pairs are re-installed
277    /// and the LoRA graph caches destroyed. Call it only at a
278    /// scheduler-quiescent point.
279    ///
280    /// Refuses when rotation is not armed (`lora_rotatable`) or the slot has
281    /// in-flight sequences. `pack_store_into_slot` also refuses adapters with
282    /// router/expert deltas or token-overlay tensors.
283    pub fn swap_lora_slot_from_disk(
284        &mut self,
285        adapter_dir: &std::path::Path,
286        name: &str,
287        slot: usize,
288    ) -> Result<()> {
289        if !self.lora_rotatable {
290            anyhow::bail!(
291                "LoRA disk swap needs rotation armed (set METRALE_LORA_ROTATE=1 so \
292                 decode runs eager); a single startup adapter with no rotation env \
293                 is baked into the decode graph and a re-point would replay stale"
294            );
295        }
296        // 2026-09-25: Refuse a busy slot before the disk load.
297        // `pack_store_into_slot` checks again before its memset.
298        if let Some(lw) = self.lora.as_ref() {
299            let busy = lw.slot_ref_count(slot);
300            if busy > 0 {
301                anyhow::bail!(
302                    "LoRA disk swap REFUSED: slot {slot} has {busy} in-flight \
303                     sequence(s) (ref_count>0); cannot replace an adapter mid-decode"
304                );
305            }
306        }
307        let cfg_path = adapter_dir.join("adapter_config.json");
308        let raw = std::fs::read_to_string(&cfg_path)
309            .with_context(|| format!("read {}", cfg_path.display()))?;
310        let peft = metrale_config::parse_peft_adapter_config(&raw)
311            .with_context(|| format!("parse {}", cfg_path.display()))?;
312        let store = metrale_model_weights::weights::adapter::load_adapter_safetensors(
313            adapter_dir,
314            self.gpu.as_ref(),
315            0,
316        )
317        .context("load LoRA adapter weights for disk swap")?;
318        let layers = {
319            let lw = self
320                .lora
321                .as_mut()
322                .ok_or_else(|| anyhow::anyhow!("LoRA disk swap: no adapter pool loaded"))?;
323            metrale_model_layers::lora::pack_store_into_slot(
324                lw,
325                slot,
326                name,
327                &store,
328                &peft,
329                &self.config,
330                self.gpu.as_ref(),
331            )?
332        };
333        let active = self.lora.as_ref().unwrap().active;
334        if active == slot {
335            let tables = self.lora.as_ref().unwrap().tables.clone();
336            let scale_table = self.lora.as_ref().unwrap().scale_table;
337            let kernels = ops::lora_delta::LoraKernels::new(self.gpu.as_ref())?;
338            self.install_lora_layers(&layers, kernels, &tables, scale_table)?;
339            self.lora.as_mut().unwrap().name = name.to_string();
340            self.destroy_lora_decode_graphs();
341        }
342        tracing::info!(
343            "LoRA disk swap: '{name}' packed into slot {slot} (r={}, active_slot={active})",
344            peft.r
345        );
346        Ok(())
347    }
348}
metrale_model_engine/model/impl_lora_swap.rs

metrale_model_engine/model/
impl_lora_swap.rs
