Skip to main content

rlx_distributed/cluster/
config.rs

1// RLX — versatile ML compiler + runtime.
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! **Declarative cluster config** (TOML) — the DX surface. A coordinator loads
5//! one file describing the model, the placement policy, and per-node device /
6//! precision / RNG / KV-cache / RAM settings, then probes + plans + runs.
7//!
8//! ```toml
9//! model = "mlx-community/DeepSeek-V4-Flash-2bit-DQ"   # HF id or local dir
10//! seq = 6
11//! rng_seed = 42            # global; a node may override
12//! reserve_ram_gb = 6.0     # headroom left free on every node
13//!
14//! [placement]
15//! policy = "auto"          # auto | ram_balanced | throughput | manual
16//!
17//! # Nodes can be discovered instead of listed:
18//! # [discover]
19//! # ssh_hosts = ["node-a", "node-b"]
20//! # ckpt_dir  = "/data/models/GLM-5.3-Flash"
21//!
22//! [[node]]
23//! addr = "127.0.0.1:9100"
24//! ssh = "node-a"                     # ~/.ssh/config alias — for probe / launch / weight sync
25//! ckpt_dir = "/path/to/DeepSeek-V4-Flash-2bit-DQ"
26//! device = "cpu"                     # cpu | metal | cuda | ane | vulkan | "metal+cpu"
27//! precision = "bf16"                 # f32 | f16 | bf16 | mixed
28//! kv_cache = "host"                  # none | host | device
29//! max_ram_gb = 44                    # cap (e.g. leave room for other apps)
30//! # layers = "0:12"                  # manual override (policy = "manual")
31//! ```
32
33use anyhow::{Context, Result};
34use rlx_runtime::{Device, parse_device_list};
35use serde::{Deserialize, Serialize};
36use std::ops::Range;
37use std::path::Path;
38
39/// Where a node keeps the attention KV cache across decode steps.
40#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
41#[serde(rename_all = "lowercase")]
42pub enum KvPolicy {
43    /// No persistent cache (prefill-only / recompute).
44    #[default]
45    None,
46    /// Keep KV in host RAM (works for any device; PCIe copy on GPU).
47    Host,
48    /// Keep KV resident on the compute device (fastest; device-mem permitting).
49    Device,
50}
51
52/// How the coordinator assigns layers to nodes.
53#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
54#[serde(rename_all = "snake_case")]
55pub enum PlacementPolicy {
56    /// Balance so each node's stage fills a similar fraction of its RAM budget
57    /// (fits first; the default — a model too big for any one node still runs).
58    #[default]
59    RamBalanced,
60    /// Balance per-stage compute time by node throughput (GFLOP/s) — minimizes
61    /// the pipeline's critical path when everything already fits.
62    Throughput,
63    /// **Everything the probe knows.** Capacity is the better of "experts
64    /// resident in RAM/VRAM" and "dense weights resident, expert banks on
65    /// disk"; the share each node gets is proportional to how fast it can run
66    /// one layer, counting compute *and* the disk time to stream in the experts
67    /// that layer touches.
68    ///
69    /// This is the policy to use for a fine-grained MoE, where RAM alone is the
70    /// wrong question: the banks are most of the model, only `top_k` of them are
71    /// read per token, and whether a node is a good host depends on its NVMe as
72    /// much as its GPU.
73    Auto,
74    /// Use each node's explicit `layers` range verbatim.
75    Manual,
76}
77
78/// What a node is doing right now.
79///
80/// A cluster that runs for days needs machines to come and go without a
81/// restart, so a node's *presence* and its *participation* are separate things.
82#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
83#[serde(rename_all = "lowercase")]
84pub enum NodeRole {
85    /// Carries a stage and serves traffic.
86    #[default]
87    Active,
88    /// Loaded with the **same layers** as the node it backs, kept warm, serving
89    /// nothing. Promoting it is a pointer swap rather than a cold weight load —
90    /// which for a 90 GB stage is the difference between seconds and minutes.
91    Standby,
92    /// Being retired. Keeps its weights until traffic has moved off, then leaves
93    /// the plan. `Cluster::retire` puts a node here.
94    Draining,
95}
96
97/// Per-node settings.
98#[derive(Debug, Clone, Serialize, Deserialize)]
99pub struct NodeConfig {
100    /// `host:port` the coordinator connects to (and the worker binds).
101    pub addr: String,
102    /// SSH alias/host for probing, launching, and syncing weights (None = local).
103    #[serde(default)]
104    pub ssh: Option<String>,
105    /// Local checkpoint directory ON THAT NODE.
106    pub ckpt_dir: String,
107    /// Device string (`cpu`, `metal`, `cuda`, `ane`, `vulkan`, or `a+b` combo —
108    /// first is primary, rest are CPU/host spill for weights that don't fit).
109    #[serde(default)]
110    pub device: DeviceList,
111    /// Numeric precision string. Float widths: `f32` | `f16` | `bf16` | `mixed`.
112    /// Scaled-GEMM shorthands: `fp8` (e4m3) | `fp8e5m2` | `mxfp8` | `nvfp4` |
113    /// `mxfp4`. Or ANY minifloat by name — `f8e4m3`, `f6e3m2`, `f4e2m1`,
114    /// `f4e3m0`, … (the `fNeXmY` family). The model crate maps it to compile
115    /// flags; exotic formats change numerics and want backend/tensor-core support.
116    #[serde(default)]
117    pub precision: Precision,
118    #[serde(default)]
119    pub kv_cache: KvPolicy,
120    /// Per-node RNG seed (overrides the global). Keeps sampling reproducible and
121    /// *distinct* per stage when desired.
122    #[serde(default)]
123    pub rng_seed: Option<u64>,
124    /// Hard cap on this node's stage RAM (GB). The planner never exceeds it.
125    #[serde(default)]
126    pub max_ram_gb: Option<f64>,
127    /// Manual `"a:b"` layer range (used when policy = manual, or to pin a node).
128    #[serde(default)]
129    pub layers: Option<String>,
130    /// Active / standby / draining. See [`NodeRole`].
131    #[serde(default)]
132    pub role: NodeRole,
133    /// For a standby: the `addr` of the active node it mirrors. When unset, a
134    /// standby backs whichever active stage is largest — the one whose loss
135    /// would cost the most to re-plan.
136    #[serde(default)]
137    pub standby_for: Option<String>,
138}
139
140impl NodeConfig {
141    /// Device preference list (primary first). CPU spill devices follow.
142    ///
143    /// Infallible: the string was parsed and validated when the config loaded.
144    pub fn devices(&self) -> &[Device] {
145        self.device.as_slice()
146    }
147    /// Primary compute device (first in the list).
148    pub fn primary_device(&self) -> Device {
149        self.device.primary()
150    }
151    /// True if the node lists a CPU/host spill target after its primary.
152    pub fn cpu_offload(&self) -> bool {
153        self.device.has_cpu_spill()
154    }
155    /// Manual layer range if given.
156    pub fn manual_range(&self) -> Option<Range<usize>> {
157        let s = self.layers.as_ref()?;
158        let (a, b) = s.split_once(':')?;
159        Some(a.trim().parse().ok()?..b.trim().parse().ok()?)
160    }
161}
162
163/// Whole-cluster config.
164#[derive(Debug, Clone, Serialize, Deserialize)]
165pub struct ClusterConfig {
166    /// HF repo id or a local directory holding the checkpoint.
167    pub model: String,
168    /// Prompt length the pipeline graphs are built for.
169    #[serde(default = "default_seq")]
170    pub seq: usize,
171    /// Global RNG seed (a node's `rng_seed` overrides it).
172    #[serde(default)]
173    pub rng_seed: Option<u64>,
174    /// RAM (GB) to leave free on every node beyond its stage.
175    #[serde(default = "default_reserve")]
176    pub reserve_ram_gb: f64,
177    #[serde(default)]
178    pub placement: PlacementSection,
179    /// Find machines instead of listing them. See
180    /// [`super::discover::DiscoverSection`]; discovered nodes are appended to
181    /// any written by hand.
182    #[serde(default)]
183    pub discover: super::discover::DiscoverSection,
184    /// The nodes, in pipeline order. May be empty when `[discover]` is set.
185    #[serde(rename = "node", default)]
186    pub nodes: Vec<NodeConfig>,
187}
188
189#[derive(Debug, Clone, Default, Serialize, Deserialize)]
190pub struct PlacementSection {
191    #[serde(default)]
192    pub policy: PlacementPolicy,
193    #[serde(default)]
194    pub objective: Objective,
195}
196
197/// What the planner is trying to make small.
198///
199/// Distinct from [`PlacementPolicy`], which says *how* to split. Two clusters
200/// can both be RAM-balanced and want opposite plans: a chat endpoint wants the
201/// fewest hops it can get away with, a batch job wants every stage the same
202/// length so the pipeline never stalls.
203#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
204#[serde(rename_all = "lowercase")]
205pub enum Optimize {
206    /// **Time to first token.** Every extra stage is a network hop on the
207    /// critical path, so this prefers to use as FEW nodes as will hold the
208    /// model, packing each one full before reaching for the next.
209    Latency,
210    /// **Tokens per second.** With the pipeline full, throughput is set by the
211    /// slowest stage, so spread work until stage times match — extra hops cost
212    /// nothing once they overlap.
213    #[default]
214    Throughput,
215    /// **Headroom.** Spread as thinly as possible across every node, leaving the
216    /// most free RAM per machine. For sharing a cluster with other work.
217    Memory,
218}
219
220/// Objective plus the constraints that shape it.
221#[derive(Debug, Clone, Serialize, Deserialize)]
222pub struct Objective {
223    #[serde(default)]
224    pub optimize: Optimize,
225    /// Context length the KV cache must hold. The planner reserves KV bytes per
226    /// node *before* fitting weights — a plan that fits at seq 8 and OOMs at 128 k
227    /// is not a plan.
228    #[serde(default = "default_context")]
229    pub context: usize,
230    /// Seconds a network hop adds to the critical path. Only [`Optimize::Latency`]
231    /// uses it; the default is a LAN round trip for a hidden-state relay.
232    #[serde(default = "default_hop_secs")]
233    pub hop_secs: f64,
234    /// Cache bytes per layer per token, when you know it and the checkpoint
235    /// cannot be read for it. Overrides inference.
236    #[serde(default)]
237    pub kv_bytes_per_layer_token: Option<u64>,
238    /// What to do when the cache size cannot be determined.
239    #[serde(default)]
240    pub on_unknown_kv: UnknownKv,
241    /// Cache element dtype at plan time (`f32` default, `f16`/`bf16`, `f8`).
242    ///
243    /// Mirrors `CostOptions::kv_dtype`, which applies when the profile is
244    /// *inferred*; this one applies when the planner has to synthesize a bound
245    /// itself. Defaulting to `f32` keeps that bound the conservative one.
246    #[serde(default)]
247    pub kv_dtype: Option<String>,
248    /// Precisions to try, widest first. The planner walks down until the model
249    /// fits and reports what it settled on, instead of failing with "too big"
250    /// and leaving you to guess the next step.
251    #[serde(default)]
252    pub precision_ladder: Vec<Precision>,
253}
254
255impl Default for Objective {
256    fn default() -> Self {
257        Self {
258            optimize: Optimize::default(),
259            context: default_context(),
260            hop_secs: default_hop_secs(),
261            kv_bytes_per_layer_token: None,
262            on_unknown_kv: UnknownKv::default(),
263            kv_dtype: None,
264            precision_ladder: Vec::new(),
265        }
266    }
267}
268
269/// What the planner does when the KV cache size cannot be worked out.
270///
271/// The default is to stop. Reserving nothing is the tempting choice and the
272/// wrong one: the plan looks fine, every stage fits, and the node dies when the
273/// cache fills — at which point the failure is a long way from its cause.
274#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
275#[serde(rename_all = "snake_case")]
276pub enum UnknownKv {
277    /// Refuse to plan, and say how to resolve it.
278    #[default]
279    Fail,
280    /// Reserve a deliberately pessimistic bound: the MHA worst case of
281    /// `2 × hidden × elem` bytes per layer per token, as if every layer cached
282    /// full keys and values at model width. Real caches are smaller — GQA by the
283    /// head ratio, MLA by ~64× — so a plan that fits under this fits in practice.
284    ///
285    /// Requires the model's width. Without it there is no bound to be
286    /// conservative *about*, so this errors rather than inventing one — see
287    /// [`super::placement::resolve_kv`].
288    Assume,
289    /// Reserve nothing and carry on. Only when you have sized the cache
290    /// yourself, or there is no cache at all.
291    Ignore,
292}
293
294fn default_context() -> usize {
295    4096
296}
297fn default_hop_secs() -> f64 {
298    0.002
299}
300
301/// A numeric precision, validated at config load.
302///
303/// Kept as a newtype over the name rather than an enum because the set is open
304/// by design: any `fNeXmY` minifloat is legal, and the model crate maps the name
305/// to compile flags. What a type buys here is that a typo fails when the config
306/// is *read*, naming the file, instead of surfacing as a puzzling compile flag
307/// on node three — and that [`Self::bits`] lives with the thing it describes.
308///
309/// Serialized as the plain string, so existing configs are unchanged.
310#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
311#[serde(try_from = "String", into = "String")]
312pub struct Precision(String);
313
314impl Precision {
315    pub fn as_str(&self) -> &str {
316        &self.0
317    }
318
319    /// Bits per weight, for scaling a cost model between precisions.
320    ///
321    /// Understands the `fNeXmY` family as well as the named shorthands, so a
322    /// precision ladder can step onto `f6e3m2` and have it actually cost
323    /// something — the previous lookup table silently ignored anything it did
324    /// not list by name.
325    pub fn bits(&self) -> Option<f64> {
326        Some(match self.0.as_str() {
327            "f32" => 32.0,
328            "f16" | "bf16" | "mixed" => 16.0,
329            "fp8" | "fp8e5m2" | "mxfp8" => 8.0,
330            // Int8 with per-block scales costs slightly more than 8 bits.
331            "int8" | "q8_0" => 8.5,
332            // 4-bit blocks likewise carry a shared exponent per group.
333            "nvfp4" | "mxfp4" => 4.25,
334            other => return minifloat_bits(other),
335        })
336    }
337}
338
339/// Width of an `fNeXmY` minifloat name (`f8e4m3` → 8), or a bare `fN`.
340fn minifloat_bits(name: &str) -> Option<f64> {
341    let rest = name.strip_prefix('f')?;
342    let width: String = rest.chars().take_while(char::is_ascii_digit).collect();
343    if width.is_empty() {
344        return None;
345    }
346    let bits: u32 = width.parse().ok()?;
347    let tail = &rest[width.len()..];
348    if tail.is_empty() {
349        // A bare `fN` is only a float width for the widths that exist.
350        return matches!(bits, 8 | 16 | 32 | 64).then_some(bits as f64);
351    }
352    // `eXmY`: exponent and mantissa must both be present and numeric.
353    let em = tail.strip_prefix('e')?;
354    let (e, m) = em.split_once('m')?;
355    if e.is_empty() || m.is_empty() || !e.chars().all(|c| c.is_ascii_digit()) {
356        return None;
357    }
358    m.chars().all(|c| c.is_ascii_digit()).then_some(bits as f64)
359}
360
361impl std::str::FromStr for Precision {
362    type Err = anyhow::Error;
363    fn from_str(s: &str) -> Result<Self> {
364        let p = Self(s.trim().to_string());
365        if p.bits().is_none() {
366            anyhow::bail!(
367                "unknown precision `{s}`; expected f32 | f16 | bf16 | mixed | fp8 | \
368                 mxfp8 | mxfp4 | nvfp4 | int8, or an fNeXmY minifloat such as f8e4m3"
369            );
370        }
371        Ok(p)
372    }
373}
374
375impl TryFrom<String> for Precision {
376    type Error = anyhow::Error;
377    fn try_from(s: String) -> Result<Self> {
378        s.parse()
379    }
380}
381impl From<Precision> for String {
382    fn from(p: Precision) -> String {
383        p.0
384    }
385}
386impl std::fmt::Display for Precision {
387    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
388        f.write_str(&self.0)
389    }
390}
391impl Default for Precision {
392    fn default() -> Self {
393        Self("bf16".into())
394    }
395}
396
397/// A node's device preference, primary first, with CPU/host spill after.
398///
399/// Parsed once at config load, so `"cuda+cpu"` is validated there rather than
400/// at every use — which also lets [`NodeConfig::primary_device`] be infallible
401/// instead of silently defaulting to CPU when the string was malformed.
402///
403/// Serialized as the original string, so existing configs round-trip verbatim.
404#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
405#[serde(try_from = "String", into = "String")]
406pub struct DeviceList {
407    raw: String,
408    devices: Vec<Device>,
409}
410
411impl DeviceList {
412    /// Primary compute device (first listed).
413    pub fn primary(&self) -> Device {
414        self.devices.first().copied().unwrap_or(Device::Cpu)
415    }
416    pub fn as_slice(&self) -> &[Device] {
417        &self.devices
418    }
419    pub fn as_str(&self) -> &str {
420        &self.raw
421    }
422    /// True if a CPU/host spill target follows the primary.
423    pub fn has_cpu_spill(&self) -> bool {
424        self.devices.len() > 1 && self.devices.contains(&Device::Cpu)
425    }
426}
427
428impl std::str::FromStr for DeviceList {
429    type Err = anyhow::Error;
430    fn from_str(s: &str) -> Result<Self> {
431        let raw = s.trim().to_string();
432        let devices = parse_device_list(&raw.replace('+', ","))
433            .map_err(|e| anyhow::anyhow!("bad device `{raw}`: {e:?}"))?;
434        if devices.is_empty() {
435            anyhow::bail!("device list `{raw}` is empty");
436        }
437        Ok(Self { raw, devices })
438    }
439}
440
441impl TryFrom<String> for DeviceList {
442    type Error = anyhow::Error;
443    fn try_from(s: String) -> Result<Self> {
444        s.parse()
445    }
446}
447impl From<DeviceList> for String {
448    fn from(d: DeviceList) -> String {
449        d.raw
450    }
451}
452impl std::fmt::Display for DeviceList {
453    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
454        f.write_str(&self.raw)
455    }
456}
457impl Default for DeviceList {
458    fn default() -> Self {
459        Self {
460            raw: "cpu".into(),
461            devices: vec![Device::Cpu],
462        }
463    }
464}
465
466fn default_seq() -> usize {
467    8
468}
469fn default_reserve() -> f64 {
470    6.0
471}
472
473impl ClusterConfig {
474    pub fn from_toml_str(s: &str) -> Result<Self> {
475        // Inline the cause rather than `.context(..)`: a validation failure from
476        // `DeviceList` / `Precision` is the whole reason those types exist, and
477        // `to_string()` on a context-wrapped anyhow error shows only the outer
478        // layer — leaving the caller with "parse cluster TOML" and no idea which
479        // field was wrong.
480        toml::from_str(s).map_err(|e| anyhow::anyhow!("parse cluster TOML: {e}"))
481    }
482    pub fn from_path(p: impl AsRef<Path>) -> Result<Self> {
483        let s = std::fs::read_to_string(p.as_ref())
484            .with_context(|| format!("read {}", p.as_ref().display()))?;
485        Self::from_toml_str(&s)
486    }
487    /// Effective RNG seed for node `i` (per-node override, else global, else 0).
488    pub fn seed_for(&self, i: usize) -> u64 {
489        self.nodes
490            .get(i)
491            .and_then(|n| n.rng_seed)
492            .or(self.rng_seed)
493            .unwrap_or(0)
494    }
495}
496
497#[cfg(test)]
498mod tests {
499    use super::*;
500
501    /// Hand-listed and discovered nodes must coexist: pin the machines you care
502    /// about, let the rest be found. Hand-written entries come first, so they
503    /// keep the head of the pipeline (the embedding stage).
504    #[test]
505    fn discover_and_hand_listed_nodes_coexist() {
506        let toml = r#"
507model = "some/model"
508
509[placement]
510policy = "auto"
511
512[discover]
513ssh_hosts = ["node-b", "node-c"]
514ckpt_dir  = "/data/model"
515port      = 9100
516
517[[node]]
518addr = "10.0.0.5:9100"
519ckpt_dir = "/pinned/model"
520device = "cuda"
521max_ram_gb = 40.0
522"#;
523        let cfg = ClusterConfig::from_toml_str(toml).unwrap();
524        assert_eq!(cfg.placement.policy, PlacementPolicy::Auto);
525        // The hand-written node survives verbatim...
526        assert_eq!(cfg.nodes.len(), 1);
527        assert_eq!(cfg.nodes[0].addr, "10.0.0.5:9100");
528        assert_eq!(cfg.nodes[0].max_ram_gb, Some(40.0));
529        // ...and discovery is configured alongside it.
530        assert!(!cfg.discover.is_empty());
531        assert_eq!(cfg.discover.ssh_hosts, ["node-b", "node-c"]);
532        assert_eq!(cfg.discover.ckpt_dir.as_deref(), Some("/data/model"));
533        assert_eq!(cfg.discover.port, 9100);
534    }
535
536    /// A config with ONLY `[discover]` is legal — `nodes` defaults to empty.
537    #[test]
538    fn discover_only_config_needs_no_node_blocks() {
539        let cfg = ClusterConfig::from_toml_str(
540            r#"
541model = "m"
542[discover]
543ssh_hosts = ["a"]
544ckpt_dir = "/m"
545"#,
546        )
547        .unwrap();
548        assert!(cfg.nodes.is_empty());
549        assert!(!cfg.discover.is_empty());
550    }
551
552    /// And an all-manual config still parses with no `[discover]` at all, so
553    /// existing cluster files keep working unchanged.
554    #[test]
555    fn hand_listed_only_config_is_unchanged() {
556        let cfg = ClusterConfig::from_toml_str(
557            r#"
558model = "m"
559[[node]]
560addr = "a:1"
561ckpt_dir = "/m"
562"#,
563        )
564        .unwrap();
565        assert_eq!(cfg.nodes.len(), 1);
566        assert!(cfg.discover.is_empty());
567        assert_eq!(cfg.placement.policy, PlacementPolicy::RamBalanced);
568    }
569
570    /// Existing configs must be byte-identical on the wire — the types are a
571    /// load-time check, not a format change.
572    #[test]
573    fn device_and_precision_still_serialize_as_plain_strings() {
574        let cfg = ClusterConfig::from_toml_str(
575            r#"
576model = "m"
577[[node]]
578addr = "a:1"
579ckpt_dir = "/m"
580device = "cuda+cpu"
581precision = "f8e4m3"
582"#,
583        )
584        .unwrap();
585        assert_eq!(cfg.nodes[0].device.as_str(), "cuda+cpu");
586        assert_eq!(cfg.nodes[0].precision.as_str(), "f8e4m3");
587        // Round-trips verbatim, `+` separator and all.
588        let back = toml::to_string(&cfg).unwrap();
589        assert!(back.contains(r#"device = "cuda+cpu""#), "{back}");
590        assert!(back.contains(r#"precision = "f8e4m3""#), "{back}");
591    }
592
593    /// The point of parsing at load: a typo names the config, not a compile flag
594    /// on node three.
595    #[test]
596    fn bad_device_and_precision_fail_at_load() {
597        let bad_dev = ClusterConfig::from_toml_str(
598            r#"model="m"
599[[node]]
600addr="a:1"
601ckpt_dir="/m"
602device="quantum"
603"#,
604        )
605        .unwrap_err()
606        .to_string();
607        assert!(bad_dev.contains("quantum"), "{bad_dev}");
608
609        let bad_prec = ClusterConfig::from_toml_str(
610            r#"model="m"
611[[node]]
612addr="a:1"
613ckpt_dir="/m"
614precision="f7ish"
615"#,
616        )
617        .unwrap_err()
618        .to_string();
619        assert!(bad_prec.contains("unknown precision"), "{bad_prec}");
620    }
621
622    /// Device access is infallible after load, and `+` means spill.
623    #[test]
624    fn device_list_exposes_primary_and_spill() {
625        let d: DeviceList = "cuda+cpu".parse().unwrap();
626        assert_eq!(d.primary(), Device::Cuda);
627        assert!(d.has_cpu_spill());
628        let solo: DeviceList = "metal".parse().unwrap();
629        assert_eq!(solo.primary(), Device::Metal);
630        assert!(!solo.has_cpu_spill());
631        assert_eq!(DeviceList::default().primary(), Device::Cpu);
632    }
633
634    /// The old lookup table silently ignored anything it did not list, so a
635    /// ladder step onto a minifloat did nothing. The family is now parsed.
636    #[test]
637    fn minifloat_widths_are_parsed_not_looked_up() {
638        for (name, bits) in [
639            ("f8e4m3", 8.0),
640            ("f8e5m2", 8.0),
641            ("f6e3m2", 6.0),
642            ("f4e2m1", 4.0),
643            ("f4e3m0", 4.0),
644        ] {
645            let p: Precision = name.parse().unwrap();
646            assert_eq!(p.bits(), Some(bits), "{name}");
647        }
648        // Named shorthands still work, block formats keep their overhead.
649        assert_eq!(Precision::default().bits(), Some(16.0));
650        assert_eq!("f32".parse::<Precision>().unwrap().bits(), Some(32.0));
651        assert_eq!("mxfp4".parse::<Precision>().unwrap().bits(), Some(4.25));
652        assert_eq!("int8".parse::<Precision>().unwrap().bits(), Some(8.5));
653    }
654
655    /// Near-misses must be rejected rather than read as a width.
656    #[test]
657    fn malformed_minifloats_are_rejected() {
658        for bad in ["f8e4", "f8m3", "fe4m3", "f8e4mx", "f9", "bf", "f8e4m"] {
659            assert!(
660                bad.parse::<Precision>().is_err(),
661                "`{bad}` should not parse"
662            );
663        }
664    }
665
666    #[test]
667    fn policy_names_round_trip() {
668        for (name, want) in [
669            ("auto", PlacementPolicy::Auto),
670            ("ram_balanced", PlacementPolicy::RamBalanced),
671            ("throughput", PlacementPolicy::Throughput),
672            ("manual", PlacementPolicy::Manual),
673        ] {
674            let cfg = ClusterConfig::from_toml_str(&format!(
675                "model = \"m\"\n[placement]\npolicy = \"{name}\"\n[[node]]\naddr=\"a:1\"\nckpt_dir=\"/m\"\n"
676            ))
677            .unwrap();
678            assert_eq!(cfg.placement.policy, want, "{name}");
679        }
680    }
681}