clawft_kernel/model_manifest/
locality.rs

1//! Locality-aware fetch-versus-relocate decision (ADR-101 section 3,
2//! ADR-099 section 6).
3//!
4//! Weights of tens of gigabytes are never re-shipped when a node already
5//! holds them. [`decide`] is a pure function over what nodes advertise:
6//!
7//! 1. an eligible node that holds every shard wins (relocate the workload to
8//!    the bytes);
9//! 2. otherwise, if the [`TransferPolicy`] allows it, the manifest is
10//!    redistributable ([`Sharing`]) and a complete holder exists, the
11//!    eligible node with the fewest missing bytes that has the space fetches
12//!    the rest (through the artifact exchange, chunked and resumable);
13//! 3. otherwise the workload is unplaceable, with the reason named.
14//!
15//! The plan carries `--explain` notes. Placement-engine integration is the
16//! [`locality_preference`] (soft, prefers the holder) and
17//! [`model_present_requirement`] (hard, place only where the bytes are).
18
19use std::collections::BTreeSet;
20
21use clawft_types::placement::engine::Preference;
22use clawft_types::placement::{
23    AttrPredicate, AttrValue, Capability, CapabilityId, CapabilityState, Requirement,
24};
25
26use super::body::{MODEL_MARKER_PREFIX, ModelPackageBody};
27
28/// Operator policy for moving weights between nodes.
29#[derive(Debug, Clone, PartialEq, Eq)]
30pub struct TransferPolicy {
31    /// Weights may be fetched from another node at all. Default false:
32    /// place where the bytes are, or report `Unplaceable`.
33    pub allow_weight_transfer: bool,
34    /// Largest transfer allowed. `None` means no ceiling (v1 default).
35    pub max_bytes: Option<u64>,
36    /// Free space that must remain on the target after the transfer.
37    pub free_headroom_bytes: u64,
38}
39
40impl Default for TransferPolicy {
41    fn default() -> Self {
42        Self {
43            allow_weight_transfer: false,
44            max_bytes: None,
45            free_headroom_bytes: 1024 * 1024 * 1024,
46        }
47    }
48}
49
50/// Whether the weights may be handed to other nodes. Fail closed: there is
51/// no default, and nothing but an explicit opt-in is `Allowed`.
52#[derive(Debug, Clone, Copy, PartialEq, Eq)]
53pub enum Sharing {
54    /// The signer opted in and the redistribution policy allows it.
55    Allowed,
56    /// Not opted in (or vetoed): weights stay where they are.
57    Refused,
58}
59
60impl Sharing {
61    /// From the manifest's own `redistributable` flag alone.
62    pub fn from_manifest(body: &ModelPackageBody) -> Self {
63        if body.redistributable { Self::Allowed } else { Self::Refused }
64    }
65}
66
67/// What one node holds of one model, and whether the placer could use it.
68#[derive(Debug, Clone, PartialEq, Eq)]
69pub struct NodeHolding {
70    /// Node id.
71    pub node_id: String,
72    /// The node passes every other placement constraint for the workload.
73    pub eligible: bool,
74    /// Usable shard hashes held (excludes an offline, detached holding).
75    pub shards: BTreeSet<String>,
76    /// Holds the whole attested model and offers it.
77    pub complete: bool,
78    /// Free bytes where a fetched model would land; `None` when unknown.
79    pub free_bytes: Option<u64>,
80}
81
82impl NodeHolding {
83    /// Read a node's holding of `body` (identified by `package_id`) from its
84    /// advertised capabilities. `free_bytes` is the largest `store.tier.*`
85    /// `free` attribute (fetches land on the tier with the most room).
86    pub fn from_capabilities(
87        node_id: &str,
88        package_id: &str,
89        body: &ModelPackageBody,
90        caps: &[Capability],
91        eligible: bool,
92    ) -> Self {
93        let marker = format!("{MODEL_MARKER_PREFIX}{package_id}");
94        let wanted: BTreeSet<&str> = body.shard_hashes().collect();
95        let mut shards = BTreeSet::new();
96        let mut complete = false;
97        for c in caps.iter().filter(|c| c.id.as_str() == "model.present") {
98            let list = match c.attrs.get("shards") {
99                Some(AttrValue::List(l)) => l,
100                _ => continue,
101            };
102            let items: Vec<&str> = list
103                .iter()
104                .filter_map(|v| match v {
105                    AttrValue::Str(s) => Some(s.as_str()),
106                    _ => None,
107                })
108                .collect();
109            let has_marker = items.contains(&marker.as_str());
110            match (c.state, has_marker) {
111                (CapabilityState::Available | CapabilityState::Busy, true) => {
112                    complete = true;
113                    shards.extend(wanted.iter().map(|s| s.to_string()));
114                }
115                // Offline (marker, degraded): nothing usable.
116                (_, true) => {}
117                (_, false) => {
118                    shards.extend(items.into_iter().filter(|s| wanted.contains(s)).map(String::from));
119                }
120            }
121        }
122        let free_bytes = caps
123            .iter()
124            .filter(|c| c.id.as_str().starts_with("store.tier."))
125            .filter(|c| c.state == CapabilityState::Available)
126            .filter_map(|c| c.attrs.get("free").and_then(|v| v.as_f64()))
127            .filter(|f| *f >= 0.0)
128            .map(|f| f as u64)
129            .max();
130        Self {
131            node_id: node_id.to_string(),
132            eligible,
133            shards,
134            complete,
135            free_bytes,
136        }
137    }
138
139    fn missing_bytes(&self, body: &ModelPackageBody) -> u64 {
140        body.shards
141            .iter()
142            .filter(|s| !self.shards.contains(&s.blake3))
143            .map(|s| s.size)
144            .fold(0, u64::saturating_add)
145    }
146}
147
148/// Why a model cannot be placed.
149#[derive(Debug, Clone, PartialEq, Eq)]
150pub enum Unplaceable {
151    /// No node offers these weights, so there is nothing to fetch from.
152    NoHolder,
153    /// No node passes the other placement constraints.
154    NoEligibleNode,
155    /// The bytes are elsewhere and weight transfer is not allowed.
156    TransferDisabled,
157    /// The bytes are elsewhere and the manifest is not redistributable.
158    NotRedistributable,
159    /// The transfer exceeds the policy ceiling.
160    OverCeiling {
161        /// Bytes that would move.
162        bytes: u64,
163        /// The ceiling.
164        max: u64,
165    },
166    /// The best target lacks room (or its free space is unknown).
167    NoSpace {
168        /// Bytes that would move plus headroom.
169        needed: u64,
170        /// Free bytes, when known.
171        free: Option<u64>,
172    },
173}
174
175impl std::fmt::Display for Unplaceable {
176    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
177        match self {
178            Self::NoHolder => write!(f, "no node holds these weights; adopt them on a node first"),
179            Self::NoEligibleNode => write!(f, "no node passes the other placement constraints"),
180            Self::TransferDisabled => {
181                write!(f, "weights are held elsewhere and allow_weight_transfer is off")
182            }
183            Self::NotRedistributable => {
184                write!(f, "weights are held elsewhere and the manifest is not marked redistributable")
185            }
186            Self::OverCeiling { bytes, max } => {
187                write!(f, "transfer of {bytes} bytes exceeds the {max} byte ceiling")
188            }
189            Self::NoSpace { needed, free: Some(free) } => {
190                write!(f, "target needs {needed} bytes free, has {free}")
191            }
192            Self::NoSpace { needed, free: None } => {
193                write!(f, "target needs {needed} bytes free and its free space is unknown")
194            }
195        }
196    }
197}
198
199/// The decision.
200#[derive(Debug, Clone, PartialEq, Eq)]
201pub enum LocalityDecision {
202    /// Run on a node that already holds every shard; nothing moves.
203    PlaceAtHolder {
204        /// The chosen node.
205        node_id: String,
206    },
207    /// Fetch the missing shards to `target`, then run there.
208    FetchThenPlace {
209        /// The node that receives the shards.
210        target: String,
211        /// Complete holders that can serve them, by node id.
212        sources: Vec<String>,
213        /// Bytes that will move.
214        bytes: u64,
215        /// Shards that will move.
216        missing_shards: usize,
217    },
218    /// Not placeable, with the reason.
219    Unplaceable(Unplaceable),
220}
221
222/// A decision plus the notes behind it (for `--explain`).
223#[derive(Debug, Clone, PartialEq, Eq)]
224pub struct LocalityPlan {
225    /// The decision.
226    pub decision: LocalityDecision,
227    /// One line per consideration, in order.
228    pub notes: Vec<String>,
229}
230
231/// Choose between relocating to the bytes and fetching them.
232pub fn decide(
233    body: &ModelPackageBody,
234    nodes: &[NodeHolding],
235    policy: &TransferPolicy,
236    sharing: Sharing,
237) -> LocalityPlan {
238    let mut notes = vec![format!(
239        "model {} ({} shards, {} bytes)",
240        body.name,
241        body.shards.len(),
242        body.total_bytes()
243    )];
244    let mut sorted: Vec<&NodeHolding> = nodes.iter().collect();
245    sorted.sort_by(|a, b| a.node_id.cmp(&b.node_id));
246    for n in &sorted {
247        notes.push(format!(
248            "node {}: eligible={} complete={} holds {}/{} shards",
249            n.node_id,
250            n.eligible,
251            n.complete,
252            n.shards.len(),
253            body.shards.len()
254        ));
255    }
256    let done = |decision, mut notes: Vec<String>, line: String| {
257        notes.push(line);
258        LocalityPlan { decision, notes }
259    };
260    if let Some(n) = sorted.iter().find(|n| n.eligible && n.complete) {
261        return done(
262            LocalityDecision::PlaceAtHolder { node_id: n.node_id.clone() },
263            notes,
264            format!("relocate: {} already holds every shard, nothing is copied", n.node_id),
265        );
266    }
267    let sources: Vec<String> = sorted
268        .iter()
269        .filter(|n| n.complete)
270        .map(|n| n.node_id.clone())
271        .collect();
272    let fail = |u: Unplaceable, notes: Vec<String>| {
273        let line = format!("unplaceable: {u}");
274        done(LocalityDecision::Unplaceable(u), notes, line)
275    };
276    let eligible: Vec<&&NodeHolding> = sorted.iter().filter(|n| n.eligible).collect();
277    if eligible.is_empty() {
278        return fail(Unplaceable::NoEligibleNode, notes);
279    }
280    if sources.is_empty() {
281        return fail(Unplaceable::NoHolder, notes);
282    }
283    if !policy.allow_weight_transfer {
284        return fail(Unplaceable::TransferDisabled, notes);
285    }
286    if sharing != Sharing::Allowed {
287        return fail(Unplaceable::NotRedistributable, notes);
288    }
289    // Best target first: fewest missing bytes, then node id.
290    let mut targets: Vec<(&NodeHolding, u64)> =
291        eligible.iter().map(|n| (**n, n.missing_bytes(body))).collect();
292    targets.sort_by(|a, b| a.1.cmp(&b.1).then_with(|| a.0.node_id.cmp(&b.0.node_id)));
293    let mut first_failure = None;
294    for (n, bytes) in targets {
295        let verdict = (|| {
296            if let Some(max) = policy.max_bytes
297                && bytes > max
298            {
299                return Err(Unplaceable::OverCeiling { bytes, max });
300            }
301            let needed = bytes.saturating_add(policy.free_headroom_bytes);
302            match n.free_bytes {
303                Some(free) if free >= needed => Ok(()),
304                free => Err(Unplaceable::NoSpace { needed, free }),
305            }
306        })();
307        match verdict {
308            Ok(()) => {
309                let missing_shards = body
310                    .shards
311                    .iter()
312                    .filter(|s| !n.shards.contains(&s.blake3))
313                    .count();
314                let line = format!(
315                    "fetch: {} gets {missing_shards} shard(s), {bytes} bytes, from {}",
316                    n.node_id,
317                    sources.join(",")
318                );
319                return done(
320                    LocalityDecision::FetchThenPlace {
321                        target: n.node_id.clone(),
322                        sources,
323                        bytes,
324                        missing_shards,
325                    },
326                    notes,
327                    line,
328                );
329            }
330            Err(u) => {
331                notes.push(format!("target {} rejected: {u}", n.node_id));
332                first_failure.get_or_insert(u);
333            }
334        }
335    }
336    fail(first_failure.unwrap_or(Unplaceable::NoEligibleNode), notes)
337}
338
339/// Requirement: the node advertises this exact attested model as present
340/// and available (`model.present` with the `model:<id>` marker).
341pub fn model_present_requirement(package_id: &str) -> Requirement {
342    Requirement::exact(CapabilityId::new("model.present").expect("valid id"))
343        .with_where(AttrPredicate::has(
344            "shards",
345            AttrValue::Str(format!("{MODEL_MARKER_PREFIX}{package_id}")),
346        ))
347}
348
349/// Soft preference of weight `weight` for nodes that already hold the model,
350/// to push into a workload spec's `policy.preferences` so the placer prefers
351/// the node holding the shards. Name carries the model for `--explain`.
352pub fn locality_preference(body: &ModelPackageBody, package_id: &str, weight: f64) -> Preference {
353    Preference {
354        name: format!("model-local:{}", body.name),
355        requirement: model_present_requirement(package_id),
356        weight,
357    }
358}
clawft_kernel/model_manifest/locality.rs

clawft_kernel/model_manifest/
locality.rs
