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}