spark_model/model/impl_lora_swap.rs
1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Runtime LoRA slot swap: RDMA peer swap/promote (Tasks #26/#27) and the
4//! disk swap sibling. Split from `impl_lora.rs` (500-LoC cap); the install
5//! walk / rotation live there.
6
7use anyhow::{Context, Result};
8
9use spark_runtime::gpu::DevicePtr;
10
11use super::types::TransformerModel;
12use crate::layers::ops;
13
14impl TransformerModel {
15 /// RDMA-swap the adapter named `adapter_name` (staged on `$ATLAS_LORA_PEER`
16 /// at `adapter_id`) INTO pool `slot`, in place, then make it that slot's
17 /// resident adapter. Byte-identical to a disk pack (the loader does the same
18 /// F16/F32→BF16 convert + B row-repack). MUST be called at a scheduler
19 /// QUIESCENT point (no in-flight decode reading `slot`). Re-zeroes the slot
20 /// sub-region first (a reused slot may hold the prior adapter's bytes), then
21 /// rebuilds the slot's `LoraLayerWeights` with the NEW adapter's r/scale —
22 /// re-installing if the swapped slot is currently active. Requires rotation
23 /// armed (`ATLAS_LORA_ROTATE`/`$ATLAS_LORA_PEER`) so decode is eager.
24 #[cfg(feature = "cuda")]
25 // Peer staging pulls adapter tensors over RDMA (rdma-core), so this half
26 // is unix-only. The `_from_disk` twins below are plain file I/O and are
27 // portable.
28 #[cfg(unix)]
29 pub fn swap_lora_slot_from_peer(
30 &mut self,
31 peer_addr: &str,
32 adapter_id: &str,
33 adapter_name: &str,
34 slot: usize,
35 peft: atlas_core::config::PeftAdapterConfig,
36 ) -> Result<()> {
37 use crate::lora::rdma_stage;
38
39 let (pool, max_rank, max_loras) = {
40 let lw = self
41 .lora
42 .as_ref()
43 .ok_or_else(|| anyhow::anyhow!("LoRA RDMA swap: no adapter pool loaded"))?;
44 (lw.pool, lw.max_rank, lw.max_loras)
45 };
46 if !self.lora_rotatable {
47 anyhow::bail!(
48 "LoRA RDMA swap needs rotation armed (set $ATLAS_LORA_PEER or \
49 ATLAS_LORA_ROTATE=1 so decode runs eager)"
50 );
51 }
52 if slot >= max_loras {
53 anyhow::bail!("LoRA RDMA swap: slot {slot} >= max_loras {max_loras}");
54 }
55 // Task #25 busy-slot refusal: bail BEFORE the destructive memset/stage
56 // below so a refused swap leaves the slot's bytes + identity untouched.
57 // Replacing an adapter while sequences are mid-decode on it would corrupt
58 // their KV and replay a captured graph over swapped pool bytes.
59 {
60 let busy = self.lora.as_ref().unwrap().slot_ref_count(slot);
61 if busy > 0 {
62 anyhow::bail!(
63 "LoRA RDMA swap REFUSED: slot {slot} has {busy} in-flight \
64 sequence(s) (ref_count>0); cannot replace an adapter mid-decode"
65 );
66 }
67 }
68
69 // 1) Fetch manifest + build landing targets (classify + slot offsets).
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 // 2) Re-zero the slot sub-region (in-place reload of a dirty slot),
75 // then RDMA-land the adapter's A/B into it.
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 =
80 spark_storage::RdmaLoraLoader::new(peer_addr.to_string(), adapter_id.to_string());
81 loader.stage_into_slot(self.gpu.as_ref(), &targets)?;
82
83 // 3) Rebuild the slot's per-layer pairs (new r/scale), stamp the slot.
84 let layers =
85 rdma_stage::rebuild_slot_layers(&targets, &self.config, &peft, pool, slot, max_rank)?;
86 // Task #26: refresh this slot's a/b pointer tables + scale table from the
87 // freshly-staged adapter's actual coverage BEFORE `peft`/`layers` are moved
88 // into the slot stamp — a promoted adapter with different module coverage
89 // than the evicted one would otherwise keep a stale bgmv route entry for the
90 // reused cache slot (missed / wrong-scaled delta). Same fix as 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 // Task #25: contents changed → bump generation so this re-staged slot
107 // yields a FRESH adapter_id (a later same-name request misses the
108 // stale prior KV). Pure rotate does NOT reach here.
109 s.generation = s.generation.wrapping_add(1);
110 }
111
112 // 4) If the swapped slot is active, re-install onto the layer structs.
113 let active = self.lora.as_ref().unwrap().active;
114 if active == slot {
115 let installed_layers = self.lora.as_ref().unwrap().slots[slot].layers.clone();
116 let tables = self.lora.as_ref().unwrap().tables.clone();
117 let scale_table = self.lora.as_ref().unwrap().scale_table;
118 let kernels = ops::lora_delta::LoraKernels::new(self.gpu.as_ref())?;
119 self.install_lora_layers(&installed_layers, kernels, &tables, scale_table)?;
120 self.lora.as_mut().unwrap().name = adapter_name.to_string();
121 self.destroy_lora_decode_graphs();
122 }
123 tracing::info!(
124 "LoRA RDMA swap: '{adapter_name}' landed in slot {slot} \
125 ({} targets, active_slot={active})",
126 targets.len()
127 );
128 Ok(())
129 }
130
131 /// Task #27 (demand-driven promotion): promote the adapter `adapter_name`
132 /// (staged on `peer_addr` at `adapter_id`) from the peer into a CACHE-region
133 /// pool slot and make it ACTIVE, returning `(slot, evicted_name)`. Runs on
134 /// the scheduler thread at a QUIESCENT point (the only place per-slot
135 /// `ref_count` is authoritative). Victim policy (pure `select_victim_slot`):
136 /// a never-filled placeholder first, else the LRU idle (`ref_count == 0`)
137 /// cache slot, else `POOL_FULL` (retryable — a busy slot is NEVER evicted).
138 /// The underlying [`Self::swap_lora_slot_from_peer`] re-checks `ref_count>0`
139 /// and bails as a backstop, and bumps the slot generation so #24 KV stays
140 /// correct. Making the promoted slot active mirrors the rotate/load control
141 /// plane so the delta actually applies under batch-1 (the per-slot bgmv route
142 /// tables are still dormant — compute reads the installed active adapter).
143 #[cfg(feature = "cuda")]
144 // Peer staging pulls adapter tensors over RDMA (rdma-core), so this half
145 // is unix-only. The `_from_disk` twins below are plain file I/O and are
146 // portable.
147 #[cfg(unix)]
148 pub fn promote_lora_slot_from_peer(
149 &mut self,
150 peer_addr: &str,
151 adapter_id: &str,
152 adapter_name: &str,
153 peft: atlas_core::config::PeftAdapterConfig,
154 ) -> Result<(usize, Option<String>)> {
155 // 1) Snapshot the cache region + pick a victim (pure policy).
156 let (slot, evicted) = {
157 let lw = self
158 .lora
159 .as_ref()
160 .ok_or_else(|| anyhow::anyhow!("LoRA promote: no adapter pool loaded"))?;
161 let views = lw.cache_slot_views();
162 let slot = crate::lora::select_victim_slot(&views).map_err(|e| match e {
163 crate::lora::VictimError::PoolFull => anyhow::anyhow!(
164 "POOL_FULL: all {} cache slot(s) are busy (ref_count>0); retry",
165 views.len()
166 ),
167 })?;
168 // The name being replaced (if the victim already held an adapter) so
169 // the caller can drop the stale name->slot overlay entry.
170 let evicted = lw
171 .slots
172 .get(slot)
173 .map(|s| s.name.clone())
174 .filter(|n| !n.is_empty());
175 (slot, evicted)
176 };
177
178 // 2) RDMA-stage into the victim slot (re-checks ref_count>0, bumps gen).
179 self.swap_lora_slot_from_peer(peer_addr, adapter_id, adapter_name, slot, peft)?;
180
181 // 3) Make the promoted slot ACTIVE so its delta applies (batch-1 honest).
182 // `swap_lora_slot_from_peer` already re-installed if the victim WAS the
183 // active slot; otherwise re-point the installed pairs onto it here.
184 let already_active = self.lora.as_ref().unwrap().active == slot;
185 if !already_active {
186 let (layers, tables, scale_table) = {
187 let lw = self.lora.as_mut().unwrap();
188 lw.active = slot;
189 lw.name = lw.slots[slot].name.clone();
190 lw.adapter_config = lw.slots[slot].adapter_config.clone();
191 (
192 lw.slots[slot].layers.clone(),
193 lw.tables.clone(),
194 lw.scale_table,
195 )
196 };
197 let kernels = ops::lora_delta::LoraKernels::new(self.gpu.as_ref())?;
198 self.install_lora_layers(&layers, kernels, &tables, scale_table)?;
199 self.destroy_lora_decode_graphs();
200 }
201 // Stamp the freshly-promoted slot as most-recently-used so a back-to-back
202 // promote of a DIFFERENT cold adapter picks an older victim, not this one,
203 // before the request that triggered this promote has acquired its ref.
204 self.lora.as_ref().unwrap().touch_slot(slot);
205 tracing::info!(
206 "LoRA promote: '{adapter_name}' hot in cache slot {slot} \
207 (evicted={:?}), now active",
208 evicted
209 );
210 Ok((slot, evicted))
211 }
212
213 /// Demand-driven DISK promotion (no RDMA/peer): load the adapter at
214 /// `adapter_dir` (named `name`) into a CACHE-region pool slot (LRU victim)
215 /// and make it ACTIVE, returning `(slot, evicted_name)`. Local-disk analog
216 /// of [`Self::promote_lora_slot_from_peer`] — same victim policy (pure
217 /// `select_victim_slot`: never-filled placeholder first, else the LRU idle
218 /// (`ref_count == 0`) cache slot, else `POOL_FULL`, retryable — a busy slot
219 /// is NEVER evicted) and the same make-active control plane, but the inner
220 /// swap reads the adapter from disk instead of the peer.
221 /// [`Self::swap_lora_slot_from_disk`] re-parses the dir's
222 /// `adapter_config.json` (so no `peft` arg), re-checks `ref_count>0` as a
223 /// backstop, and bumps the slot generation so #24 KV stays correct. Runs on
224 /// the scheduler thread at a QUIESCENT point; requires rotation armed
225 /// (`ATLAS_LORA_ROTATE=1`) — the inner swap enforces it.
226 pub fn promote_lora_slot_from_disk(
227 &mut self,
228 adapter_dir: &std::path::Path,
229 name: &str,
230 ) -> Result<(usize, Option<String>)> {
231 // 1) Snapshot the cache region + pick a victim (pure policy).
232 let (slot, evicted) = {
233 let lw = self
234 .lora
235 .as_ref()
236 .ok_or_else(|| anyhow::anyhow!("LoRA disk promote: no adapter pool loaded"))?;
237 let views = lw.cache_slot_views();
238 let slot = crate::lora::select_victim_slot(&views).map_err(|e| match e {
239 crate::lora::VictimError::PoolFull => anyhow::anyhow!(
240 "POOL_FULL: all {} cache slot(s) are busy (ref_count>0); retry",
241 views.len()
242 ),
243 })?;
244 // The name being replaced (if the victim already held an adapter) so
245 // the caller can drop the stale name->slot overlay entry.
246 let evicted = lw
247 .slots
248 .get(slot)
249 .map(|s| s.name.clone())
250 .filter(|n| !n.is_empty());
251 (slot, evicted)
252 };
253
254 // 2) Disk-load into the victim slot (re-checks ref_count>0, bumps gen).
255 self.swap_lora_slot_from_disk(adapter_dir, name, slot)?;
256
257 // 3) Make the promoted slot ACTIVE so its delta applies (batch-1 honest).
258 // `swap_lora_slot_from_disk` already re-installed if the victim WAS the
259 // active slot; otherwise re-point the installed pairs onto it here.
260 let already_active = self.lora.as_ref().unwrap().active == slot;
261 if !already_active {
262 let (layers, tables, scale_table) = {
263 let lw = self.lora.as_mut().unwrap();
264 lw.active = slot;
265 lw.name = lw.slots[slot].name.clone();
266 lw.adapter_config = lw.slots[slot].adapter_config.clone();
267 (
268 lw.slots[slot].layers.clone(),
269 lw.tables.clone(),
270 lw.scale_table,
271 )
272 };
273 let kernels = ops::lora_delta::LoraKernels::new(self.gpu.as_ref())?;
274 self.install_lora_layers(&layers, kernels, &tables, scale_table)?;
275 self.destroy_lora_decode_graphs();
276 }
277 // Stamp the freshly-promoted slot as most-recently-used so a back-to-back
278 // promote of a DIFFERENT cold adapter picks an older victim, not this one,
279 // before the request that triggered this promote has acquired its ref.
280 self.lora.as_ref().unwrap().touch_slot(slot);
281 tracing::info!(
282 "LoRA disk-promote: '{name}' hot in cache slot {slot} \
283 (evicted={:?}), now active",
284 evicted
285 );
286 Ok((slot, evicted))
287 }
288
289 /// Disk-swap the adapter at `adapter_dir` INTO pool `slot`, in place, then
290 /// make it that slot's resident adapter (re-installing onto the layer structs
291 /// if the slot is currently active). The local-disk analog of
292 /// [`Self::swap_lora_slot_from_peer`] — same audit + pack + re-point, no RDMA.
293 /// This is the pool-size-1 dynamic-load path: load a DIFFERENT adapter into
294 /// the single slot at runtime (per-request weight change). MUST be called at
295 /// a scheduler QUIESCENT point (no in-flight decode reading `slot`) and needs
296 /// rotation armed (`ATLAS_LORA_ROTATE=1`/`$ATLAS_LORA_PEER`) so decode is
297 /// eager and no captured graph replays the swapped slot's stale pointers.
298 pub fn swap_lora_slot_from_disk(
299 &mut self,
300 adapter_dir: &std::path::Path,
301 name: &str,
302 slot: usize,
303 ) -> Result<()> {
304 if !self.lora_rotatable {
305 anyhow::bail!(
306 "LoRA disk swap needs rotation armed (set ATLAS_LORA_ROTATE=1 so \
307 decode runs eager); a single startup adapter with no rotation env \
308 is baked into the decode graph and a re-point would replay stale"
309 );
310 }
311 // Task #25 busy-slot refusal: fail fast (before the disk load + pack)
312 // when the target slot has in-flight sequences. `pack_store_into_slot`
313 // re-checks under `&mut lw` right before the destructive memset (the
314 // authoritative gate); this early check just avoids the wasted load.
315 if let Some(lw) = self.lora.as_ref() {
316 let busy = lw.slot_ref_count(slot);
317 if busy > 0 {
318 anyhow::bail!(
319 "LoRA disk swap REFUSED: slot {slot} has {busy} in-flight \
320 sequence(s) (ref_count>0); cannot replace an adapter mid-decode"
321 );
322 }
323 }
324 // Parse the adapter's own PEFT config (scaling read per adapter, never
325 // defaulted) — the same hard-fail parser the startup path uses.
326 let cfg_path = adapter_dir.join("adapter_config.json");
327 let raw = std::fs::read_to_string(&cfg_path)
328 .with_context(|| format!("read {}", cfg_path.display()))?;
329 let peft = atlas_core::config::parse_peft_adapter_config(&raw)
330 .with_context(|| format!("parse {}", cfg_path.display()))?;
331 // Load the adapter's A/B into a device WeightStore (host F16/F32→BF16),
332 // then pack it into the slot (same layout as a startup pack).
333 let store = spark_runtime::weights::adapter::load_adapter_safetensors(
334 adapter_dir,
335 self.gpu.as_ref(),
336 0,
337 )
338 .context("load LoRA adapter weights for disk swap")?;
339 let layers = {
340 let lw = self
341 .lora
342 .as_mut()
343 .ok_or_else(|| anyhow::anyhow!("LoRA disk swap: no adapter pool loaded"))?;
344 crate::lora::pack_store_into_slot(
345 lw,
346 slot,
347 name,
348 &store,
349 &peft,
350 &self.config,
351 self.gpu.as_ref(),
352 )?
353 };
354 // If the swapped slot is the active one, re-install onto the layer structs
355 // so subsequent requests apply the new adapter's delta.
356 let active = self.lora.as_ref().unwrap().active;
357 if active == slot {
358 let tables = self.lora.as_ref().unwrap().tables.clone();
359 let scale_table = self.lora.as_ref().unwrap().scale_table;
360 let kernels = ops::lora_delta::LoraKernels::new(self.gpu.as_ref())?;
361 self.install_lora_layers(&layers, kernels, &tables, scale_table)?;
362 self.lora.as_mut().unwrap().name = name.to_string();
363 self.destroy_lora_decode_graphs();
364 }
365 tracing::info!(
366 "LoRA disk swap: '{name}' packed into slot {slot} (r={}, active_slot={active})",
367 peft.r
368 );
369 Ok(())
370 }
371}