ruccl/rank/
protocol.rs

1//! Versioned wire contract shared by host-staged, TCP, peer-memory, and RDMA
2//! collective transports.
3
4use std::collections::BTreeMap;
5use std::error::Error;
6use std::fmt::{Display, Formatter};
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::time::{SystemTime, UNIX_EPOCH};
9
10pub const PROTOCOL_MAGIC: [u8; 4] = *b"GXCL";
11pub const PROTOCOL_VERSION: u16 = 3;
12pub const FRAME_HEADER_BYTES: usize = 80;
13pub const MAX_FRAME_PAYLOAD_BYTES: usize = 1 << 30;
14pub const ANY_RANK: u32 = u32::MAX;
15pub const FLAG_PAYLOAD: u16 = 1;
16pub const FLAG_COUNTS_PREFIX: u16 = 1 << 1;
17pub const FLAG_P2P_CHANNEL: u16 = 1 << 2;
18
19static NEXT_UNIQUE_ID: AtomicU64 = AtomicU64::new(1);
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
22pub struct UniqueId([u8; 16]);
23
24impl UniqueId {
25    pub fn new() -> Self {
26        let counter = NEXT_UNIQUE_ID.fetch_add(1, Ordering::Relaxed);
27        let time = SystemTime::now()
28            .duration_since(UNIX_EPOCH)
29            .map_or(0, |duration| duration.as_nanos());
30        let process = u128::from(std::process::id());
31        let mixed = time ^ (u128::from(counter) << 64) ^ (process << 32);
32        Self(mixed.to_le_bytes())
33    }
34
35    pub const fn from_bytes(bytes: [u8; 16]) -> Self {
36        Self(bytes)
37    }
38
39    pub const fn as_bytes(&self) -> &[u8; 16] {
40        &self.0
41    }
42}
43
44impl Default for UniqueId {
45    fn default() -> Self {
46        Self::new()
47    }
48}
49
50#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
51#[repr(u8)]
52pub enum Opcode {
53    Join = 1,
54    Ready = 2,
55    Broadcast = 3,
56    AllGather = 4,
57    Gather = 5,
58    Scatter = 6,
59    Reduce = 7,
60    AllReduce = 8,
61    ReduceScatter = 9,
62    AllToAll = 10,
63    AllToAllV = 11,
64    Barrier = 12,
65    Send = 13,
66    Receive = 14,
67    Abort = 15,
68    Heartbeat = 16,
69    Leave = 17,
70    SetTimeout = 18,
71    PeerEndpoint = 19,
72}
73
74impl TryFrom<u8> for Opcode {
75    type Error = ProtocolError;
76
77    fn try_from(value: u8) -> Result<Self, Self::Error> {
78        match value {
79            1 => Ok(Self::Join),
80            2 => Ok(Self::Ready),
81            3 => Ok(Self::Broadcast),
82            4 => Ok(Self::AllGather),
83            5 => Ok(Self::Gather),
84            6 => Ok(Self::Scatter),
85            7 => Ok(Self::Reduce),
86            8 => Ok(Self::AllReduce),
87            9 => Ok(Self::ReduceScatter),
88            10 => Ok(Self::AllToAll),
89            11 => Ok(Self::AllToAllV),
90            12 => Ok(Self::Barrier),
91            13 => Ok(Self::Send),
92            14 => Ok(Self::Receive),
93            15 => Ok(Self::Abort),
94            16 => Ok(Self::Heartbeat),
95            17 => Ok(Self::Leave),
96            18 => Ok(Self::SetTimeout),
97            19 => Ok(Self::PeerEndpoint),
98            _ => Err(ProtocolError::UnknownOpcode(value)),
99        }
100    }
101}
102
103#[derive(Debug, Clone, Copy, PartialEq, Eq)]
104#[repr(u8)]
105pub enum ElementType {
106    None = 0,
107    U8 = 1,
108    U32 = 2,
109    I32 = 3,
110    F32 = 4,
111    F16 = 5,
112    BF16 = 6,
113    Bool = 7,
114    I8 = 8,
115    I16 = 9,
116    I64 = 10,
117    F64 = 11,
118    U16 = 12,
119    U64 = 13,
120    Complex64 = 14,
121    Complex128 = 15,
122    F8E4M3Fn = 16,
123    F8E5M2 = 17,
124    F8E4M3Fnuz = 18,
125    F8E5M2Fnuz = 19,
126    F8E8M0Fnu = 20,
127    F4E2M1FnX2 = 21,
128}
129
130impl ElementType {
131    pub const fn byte_width(self) -> usize {
132        match self {
133            Self::None => 0,
134            Self::U8
135            | Self::Bool
136            | Self::I8
137            | Self::F8E4M3Fn
138            | Self::F8E5M2
139            | Self::F8E4M3Fnuz
140            | Self::F8E5M2Fnuz
141            | Self::F8E8M0Fnu
142            | Self::F4E2M1FnX2 => 1,
143            Self::F16 | Self::BF16 | Self::I16 | Self::U16 => 2,
144            Self::U32 | Self::I32 | Self::F32 => 4,
145            Self::I64 | Self::F64 | Self::U64 | Self::Complex64 => 8,
146            Self::Complex128 => 16,
147        }
148    }
149
150    pub const fn is_low_precision_storage(self) -> bool {
151        matches!(
152            self,
153            Self::F8E4M3Fn
154                | Self::F8E5M2
155                | Self::F8E4M3Fnuz
156                | Self::F8E5M2Fnuz
157                | Self::F8E8M0Fnu
158                | Self::F4E2M1FnX2
159        )
160    }
161}
162
163impl TryFrom<u8> for ElementType {
164    type Error = ProtocolError;
165
166    fn try_from(value: u8) -> Result<Self, Self::Error> {
167        match value {
168            0 => Ok(Self::None),
169            1 => Ok(Self::U8),
170            2 => Ok(Self::U32),
171            3 => Ok(Self::I32),
172            4 => Ok(Self::F32),
173            5 => Ok(Self::F16),
174            6 => Ok(Self::BF16),
175            7 => Ok(Self::Bool),
176            8 => Ok(Self::I8),
177            9 => Ok(Self::I16),
178            10 => Ok(Self::I64),
179            11 => Ok(Self::F64),
180            12 => Ok(Self::U16),
181            13 => Ok(Self::U64),
182            14 => Ok(Self::Complex64),
183            15 => Ok(Self::Complex128),
184            16 => Ok(Self::F8E4M3Fn),
185            17 => Ok(Self::F8E5M2),
186            18 => Ok(Self::F8E4M3Fnuz),
187            19 => Ok(Self::F8E5M2Fnuz),
188            20 => Ok(Self::F8E8M0Fnu),
189            21 => Ok(Self::F4E2M1FnX2),
190            _ => Err(ProtocolError::UnknownElementType(value)),
191        }
192    }
193}
194
195#[derive(Debug, Clone, PartialEq, Eq)]
196pub struct FrameHeader {
197    pub opcode: Opcode,
198    pub element_type: ElementType,
199    pub flags: u16,
200    pub source_rank: u32,
201    pub destination_rank: u32,
202    pub root_rank: u32,
203    pub world_size: u32,
204    pub sequence: u64,
205    pub tag: u64,
206    pub element_count: u64,
207    pub payload_bytes: u64,
208    pub unique_id: UniqueId,
209    pub payload_checksum: u32,
210}
211
212impl FrameHeader {
213    #[allow(clippy::too_many_arguments)]
214    pub fn collective(
215        unique_id: UniqueId,
216        opcode: Opcode,
217        element_type: ElementType,
218        source_rank: u32,
219        root_rank: u32,
220        world_size: u32,
221        sequence: u64,
222        element_count: u64,
223    ) -> Self {
224        Self {
225            opcode,
226            element_type,
227            flags: 0,
228            source_rank,
229            destination_rank: ANY_RANK,
230            root_rank,
231            world_size,
232            sequence,
233            tag: 0,
234            element_count,
235            payload_bytes: 0,
236            unique_id,
237            payload_checksum: checksum(&[]),
238        }
239    }
240
241    pub fn with_payload(mut self, payload: &[u8]) -> Result<Self, ProtocolError> {
242        if payload.len() > MAX_FRAME_PAYLOAD_BYTES {
243            return Err(ProtocolError::PayloadTooLarge(payload.len()));
244        }
245        self.flags |= FLAG_PAYLOAD;
246        self.payload_bytes = payload.len() as u64;
247        self.payload_checksum = checksum(payload);
248        // The checksum above was computed from this exact slice. Validate the
249        // resulting wire contract without scanning the payload a second time.
250        self.validate(None)?;
251        Ok(self)
252    }
253
254    pub fn encode(&self) -> Result<[u8; FRAME_HEADER_BYTES], ProtocolError> {
255        self.validate(None)?;
256        let mut bytes = [0_u8; FRAME_HEADER_BYTES];
257        bytes[0..4].copy_from_slice(&PROTOCOL_MAGIC);
258        put_u16(&mut bytes, 4, PROTOCOL_VERSION);
259        put_u16(&mut bytes, 6, FRAME_HEADER_BYTES as u16);
260        bytes[8] = self.opcode as u8;
261        bytes[9] = self.element_type as u8;
262        put_u16(&mut bytes, 10, self.flags);
263        put_u32(&mut bytes, 12, self.source_rank);
264        put_u32(&mut bytes, 16, self.destination_rank);
265        put_u32(&mut bytes, 20, self.root_rank);
266        put_u32(&mut bytes, 24, self.world_size);
267        put_u64(&mut bytes, 28, self.sequence);
268        put_u64(&mut bytes, 36, self.tag);
269        put_u64(&mut bytes, 44, self.element_count);
270        put_u64(&mut bytes, 52, self.payload_bytes);
271        bytes[60..76].copy_from_slice(self.unique_id.as_bytes());
272        put_u32(&mut bytes, 76, self.payload_checksum);
273        Ok(bytes)
274    }
275
276    pub fn decode(bytes: &[u8]) -> Result<Self, ProtocolError> {
277        if bytes.len() < FRAME_HEADER_BYTES {
278            return Err(ProtocolError::TruncatedHeader(bytes.len()));
279        }
280        if bytes[0..4] != PROTOCOL_MAGIC {
281            return Err(ProtocolError::BadMagic(bytes[0..4].try_into().unwrap()));
282        }
283        let version = get_u16(bytes, 4);
284        if version != PROTOCOL_VERSION {
285            return Err(ProtocolError::UnsupportedVersion(version));
286        }
287        let header_bytes = get_u16(bytes, 6) as usize;
288        if header_bytes != FRAME_HEADER_BYTES {
289            return Err(ProtocolError::UnsupportedHeaderLength(header_bytes));
290        }
291        let header = Self {
292            opcode: Opcode::try_from(bytes[8])?,
293            element_type: ElementType::try_from(bytes[9])?,
294            flags: get_u16(bytes, 10),
295            source_rank: get_u32(bytes, 12),
296            destination_rank: get_u32(bytes, 16),
297            root_rank: get_u32(bytes, 20),
298            world_size: get_u32(bytes, 24),
299            sequence: get_u64(bytes, 28),
300            tag: get_u64(bytes, 36),
301            element_count: get_u64(bytes, 44),
302            payload_bytes: get_u64(bytes, 52),
303            unique_id: UniqueId::from_bytes(bytes[60..76].try_into().unwrap()),
304            payload_checksum: get_u32(bytes, 76),
305        };
306        header.validate(None)?;
307        Ok(header)
308    }
309
310    pub fn validate(&self, payload: Option<&[u8]>) -> Result<(), ProtocolError> {
311        let unknown_flags = self.flags & !(FLAG_PAYLOAD | FLAG_COUNTS_PREFIX | FLAG_P2P_CHANNEL);
312        if unknown_flags != 0 {
313            return Err(ProtocolError::UnknownFlags(unknown_flags));
314        }
315        let has_counts_prefix = self.flags & FLAG_COUNTS_PREFIX != 0;
316        if has_counts_prefix && self.opcode != Opcode::AllToAllV {
317            return Err(ProtocolError::UnexpectedCountsPrefix(self.opcode));
318        }
319        if self.flags & FLAG_P2P_CHANNEL != 0
320            && !matches!(self.opcode, Opcode::Join | Opcode::Ready)
321        {
322            return Err(ProtocolError::UnexpectedP2pChannelFlag(self.opcode));
323        }
324        if self.world_size == 0 {
325            return Err(ProtocolError::EmptyWorld);
326        }
327        validate_rank("source", self.source_rank, self.world_size, false)?;
328        validate_rank("destination", self.destination_rank, self.world_size, true)?;
329        validate_rank("root", self.root_rank, self.world_size, true)?;
330        let payload_bytes = usize::try_from(self.payload_bytes)
331            .map_err(|_| ProtocolError::PayloadTooLarge(usize::MAX))?;
332        if payload_bytes > MAX_FRAME_PAYLOAD_BYTES {
333            return Err(ProtocolError::PayloadTooLarge(payload_bytes));
334        }
335        let has_payload = self.flags & FLAG_PAYLOAD != 0;
336        if has_payload != (self.payload_bytes != 0) {
337            return Err(ProtocolError::PayloadFlagMismatch);
338        }
339        if self.element_type == ElementType::None {
340            if self.element_count != 0 || self.payload_bytes != 0 {
341                return Err(ProtocolError::ElementLengthMismatch);
342            }
343        } else if has_payload {
344            let element_bytes = self
345                .element_count
346                .checked_mul(self.element_type.byte_width() as u64)
347                .ok_or(ProtocolError::ElementLengthMismatch)?;
348            let prefix_bytes = if has_counts_prefix {
349                u64::from(self.world_size)
350                    .checked_mul(8)
351                    .ok_or(ProtocolError::ElementLengthMismatch)?
352            } else {
353                0
354            };
355            let expected = element_bytes
356                .checked_add(prefix_bytes)
357                .ok_or(ProtocolError::ElementLengthMismatch)?;
358            if expected != self.payload_bytes {
359                return Err(ProtocolError::ElementLengthMismatch);
360            }
361        }
362        if let Some(payload) = payload {
363            if payload.len() != payload_bytes {
364                return Err(ProtocolError::PayloadLengthMismatch {
365                    declared: payload_bytes,
366                    actual: payload.len(),
367                });
368            }
369            let actual = checksum(payload);
370            if actual != self.payload_checksum {
371                return Err(ProtocolError::ChecksumMismatch {
372                    declared: self.payload_checksum,
373                    actual,
374                });
375            }
376        }
377        Ok(())
378    }
379}
380
381#[derive(Debug, Clone, PartialEq, Eq)]
382pub struct Frame {
383    pub header: FrameHeader,
384    pub payload: Vec<u8>,
385}
386
387impl Frame {
388    pub fn new(header: FrameHeader, payload: Vec<u8>) -> Result<Self, ProtocolError> {
389        let header = if payload.is_empty() {
390            header.validate(Some(&payload))?;
391            header
392        } else {
393            header.with_payload(&payload)?
394        };
395        Ok(Self { header, payload })
396    }
397
398    pub fn encode(&self) -> Result<Vec<u8>, ProtocolError> {
399        self.header.validate(Some(&self.payload))?;
400        let mut bytes = Vec::with_capacity(FRAME_HEADER_BYTES + self.payload.len());
401        bytes.extend_from_slice(&self.header.encode()?);
402        bytes.extend_from_slice(&self.payload);
403        Ok(bytes)
404    }
405
406    pub(crate) fn encode_transport_header(
407        &self,
408    ) -> Result<[u8; FRAME_HEADER_BYTES], ProtocolError> {
409        let declared = usize::try_from(self.header.payload_bytes)
410            .map_err(|_| ProtocolError::PayloadTooLarge(usize::MAX))?;
411        if self.payload.len() != declared {
412            return Err(ProtocolError::PayloadLengthMismatch {
413                declared,
414                actual: self.payload.len(),
415            });
416        }
417        self.header.encode()
418    }
419
420    pub fn decode(bytes: &[u8]) -> Result<Self, ProtocolError> {
421        let header = FrameHeader::decode(bytes)?;
422        let payload_bytes = usize::try_from(header.payload_bytes)
423            .map_err(|_| ProtocolError::PayloadTooLarge(usize::MAX))?;
424        let expected = FRAME_HEADER_BYTES
425            .checked_add(payload_bytes)
426            .ok_or(ProtocolError::PayloadTooLarge(payload_bytes))?;
427        if bytes.len() != expected {
428            return Err(ProtocolError::FrameLengthMismatch {
429                declared: expected,
430                actual: bytes.len(),
431            });
432        }
433        let payload = bytes[FRAME_HEADER_BYTES..].to_vec();
434        header.validate(Some(&payload))?;
435        Ok(Self { header, payload })
436    }
437}
438
439#[derive(Debug, Clone, PartialEq, Eq)]
440pub struct OperationDescriptor {
441    pub opcode: Opcode,
442    pub element_type: ElementType,
443    pub root_rank: u32,
444    pub element_count: u64,
445    pub layout_hash: u64,
446}
447
448impl OperationDescriptor {
449    pub const fn new(
450        opcode: Opcode,
451        element_type: ElementType,
452        root_rank: u32,
453        element_count: u64,
454    ) -> Self {
455        Self {
456            opcode,
457            element_type,
458            root_rank,
459            element_count,
460            layout_hash: 0,
461        }
462    }
463
464    pub const fn with_layout_hash(mut self, layout_hash: u64) -> Self {
465        self.layout_hash = layout_hash;
466        self
467    }
468}
469
470/// Coordinator-side sequence agreement. An operation becomes ready only when
471/// every rank submitted exactly the same descriptor for the same sequence.
472#[derive(Debug)]
473pub struct CollectiveAgreement {
474    world_size: usize,
475    next_sequence: u64,
476    pending: BTreeMap<u64, Vec<Option<OperationDescriptor>>>,
477}
478
479impl CollectiveAgreement {
480    pub fn new(world_size: usize) -> Result<Self, ProtocolError> {
481        if world_size == 0 {
482            return Err(ProtocolError::EmptyWorld);
483        }
484        Ok(Self {
485            world_size,
486            next_sequence: 0,
487            pending: BTreeMap::new(),
488        })
489    }
490
491    pub const fn next_sequence(&self) -> u64 {
492        self.next_sequence
493    }
494
495    pub fn submit(
496        &mut self,
497        rank: usize,
498        sequence: u64,
499        descriptor: OperationDescriptor,
500    ) -> Result<bool, ProtocolError> {
501        if rank >= self.world_size {
502            return Err(ProtocolError::RankOutOfRange {
503                name: "source",
504                rank: rank as u32,
505                world_size: self.world_size as u32,
506            });
507        }
508        if sequence != self.next_sequence {
509            return Err(ProtocolError::SequenceMismatch {
510                expected: self.next_sequence,
511                actual: sequence,
512            });
513        }
514        let ranks = self
515            .pending
516            .entry(sequence)
517            .or_insert_with(|| vec![None; self.world_size]);
518        if ranks[rank].is_some() {
519            return Err(ProtocolError::DuplicateSubmission { rank, sequence });
520        }
521        if let Some(expected) = ranks.iter().flatten().next()
522            && expected != &descriptor
523        {
524            return Err(ProtocolError::CollectiveMismatch {
525                sequence,
526                expected: expected.clone(),
527                actual: descriptor,
528            });
529        }
530        ranks[rank] = Some(descriptor);
531        let ready = ranks.iter().all(Option::is_some);
532        if ready {
533            self.pending.remove(&sequence);
534            self.next_sequence = self
535                .next_sequence
536                .checked_add(1)
537                .ok_or(ProtocolError::SequenceOverflow)?;
538        }
539        Ok(ready)
540    }
541}
542
543#[derive(Debug, Clone, PartialEq, Eq)]
544pub enum ProtocolError {
545    TruncatedHeader(usize),
546    BadMagic([u8; 4]),
547    UnsupportedVersion(u16),
548    UnsupportedHeaderLength(usize),
549    UnknownOpcode(u8),
550    UnknownElementType(u8),
551    UnknownFlags(u16),
552    UnexpectedCountsPrefix(Opcode),
553    UnexpectedP2pChannelFlag(Opcode),
554    EmptyWorld,
555    RankOutOfRange {
556        name: &'static str,
557        rank: u32,
558        world_size: u32,
559    },
560    PayloadTooLarge(usize),
561    PayloadFlagMismatch,
562    ElementLengthMismatch,
563    PayloadLengthMismatch {
564        declared: usize,
565        actual: usize,
566    },
567    FrameLengthMismatch {
568        declared: usize,
569        actual: usize,
570    },
571    ChecksumMismatch {
572        declared: u32,
573        actual: u32,
574    },
575    SequenceMismatch {
576        expected: u64,
577        actual: u64,
578    },
579    DuplicateSubmission {
580        rank: usize,
581        sequence: u64,
582    },
583    CollectiveMismatch {
584        sequence: u64,
585        expected: OperationDescriptor,
586        actual: OperationDescriptor,
587    },
588    SequenceOverflow,
589}
590
591impl Display for ProtocolError {
592    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
593        match self {
594            Self::TruncatedHeader(actual) => write!(
595                formatter,
596                "GX collective header needs {FRAME_HEADER_BYTES} bytes, got {actual}"
597            ),
598            Self::BadMagic(actual) => write!(formatter, "bad GX collective magic {actual:?}"),
599            Self::UnsupportedVersion(version) => {
600                write!(
601                    formatter,
602                    "unsupported GX collective protocol version {version}"
603                )
604            }
605            Self::UnsupportedHeaderLength(length) => {
606                write!(
607                    formatter,
608                    "unsupported GX collective header length {length}"
609                )
610            }
611            Self::UnknownOpcode(opcode) => write!(formatter, "unknown collective opcode {opcode}"),
612            Self::UnknownElementType(element_type) => {
613                write!(formatter, "unknown collective element type {element_type}")
614            }
615            Self::UnknownFlags(flags) => {
616                write!(formatter, "unknown collective frame flags {flags:#06x}")
617            }
618            Self::UnexpectedCountsPrefix(opcode) => {
619                write!(formatter, "counts prefix is invalid for {opcode:?}")
620            }
621            Self::UnexpectedP2pChannelFlag(opcode) => {
622                write!(
623                    formatter,
624                    "point-to-point channel flag is invalid for {opcode:?}"
625                )
626            }
627            Self::EmptyWorld => write!(formatter, "collective protocol world cannot be empty"),
628            Self::RankOutOfRange {
629                name,
630                rank,
631                world_size,
632            } => write!(
633                formatter,
634                "{name} rank {rank} is outside protocol world size {world_size}"
635            ),
636            Self::PayloadTooLarge(bytes) => {
637                write!(
638                    formatter,
639                    "collective payload of {bytes} bytes is too large"
640                )
641            }
642            Self::PayloadFlagMismatch => {
643                write!(formatter, "collective payload flag and length disagree")
644            }
645            Self::ElementLengthMismatch => {
646                write!(
647                    formatter,
648                    "collective element count and payload length disagree"
649                )
650            }
651            Self::PayloadLengthMismatch { declared, actual } => write!(
652                formatter,
653                "collective payload declares {declared} bytes, got {actual}"
654            ),
655            Self::FrameLengthMismatch { declared, actual } => write!(
656                formatter,
657                "collective frame declares {declared} bytes, got {actual}"
658            ),
659            Self::ChecksumMismatch { declared, actual } => write!(
660                formatter,
661                "collective payload checksum {actual:#010x} does not match {declared:#010x}"
662            ),
663            Self::SequenceMismatch { expected, actual } => write!(
664                formatter,
665                "collective sequence {actual} does not match expected {expected}"
666            ),
667            Self::DuplicateSubmission { rank, sequence } => write!(
668                formatter,
669                "rank {rank} submitted collective sequence {sequence} twice"
670            ),
671            Self::CollectiveMismatch {
672                sequence,
673                expected,
674                actual,
675            } => write!(
676                formatter,
677                "collective sequence {sequence} mismatch: expected {expected:?}, got {actual:?}"
678            ),
679            Self::SequenceOverflow => write!(formatter, "collective sequence overflow"),
680        }
681    }
682}
683
684impl Error for ProtocolError {}
685
686fn validate_rank(
687    name: &'static str,
688    rank: u32,
689    world_size: u32,
690    allow_any: bool,
691) -> Result<(), ProtocolError> {
692    if (allow_any && rank == ANY_RANK) || rank < world_size {
693        Ok(())
694    } else {
695        Err(ProtocolError::RankOutOfRange {
696            name,
697            rank,
698            world_size,
699        })
700    }
701}
702
703fn checksum(bytes: &[u8]) -> u32 {
704    crc32fast::hash(bytes)
705}
706
707fn put_u16(bytes: &mut [u8], offset: usize, value: u16) {
708    bytes[offset..offset + 2].copy_from_slice(&value.to_le_bytes());
709}
710
711fn put_u32(bytes: &mut [u8], offset: usize, value: u32) {
712    bytes[offset..offset + 4].copy_from_slice(&value.to_le_bytes());
713}
714
715fn put_u64(bytes: &mut [u8], offset: usize, value: u64) {
716    bytes[offset..offset + 8].copy_from_slice(&value.to_le_bytes());
717}
718
719fn get_u16(bytes: &[u8], offset: usize) -> u16 {
720    u16::from_le_bytes(bytes[offset..offset + 2].try_into().unwrap())
721}
722
723fn get_u32(bytes: &[u8], offset: usize) -> u32 {
724    u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap())
725}
726
727fn get_u64(bytes: &[u8], offset: usize) -> u64 {
728    u64::from_le_bytes(bytes[offset..offset + 8].try_into().unwrap())
729}
730
731#[cfg(test)]
732mod tests {
733    use super::*;
734
735    #[test]
736    fn frame_round_trip_is_little_endian_and_checksum_protected() {
737        let unique_id = UniqueId::from_bytes([7; 16]);
738        let mut header = FrameHeader::collective(
739            unique_id,
740            Opcode::Send,
741            ElementType::U32,
742            1,
743            ANY_RANK,
744            4,
745            9,
746            3,
747        );
748        header.destination_rank = 2;
749        header.tag = 44;
750        let frame = Frame::new(header, vec![1, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0]).unwrap();
751        let encoded = frame.encode().unwrap();
752        assert_eq!(&encoded[0..4], b"GXCL");
753        assert_eq!(Frame::decode(&encoded).unwrap(), frame);
754
755        let mut corrupt = encoded;
756        *corrupt.last_mut().unwrap() ^= 0x80;
757        assert!(matches!(
758            Frame::decode(&corrupt),
759            Err(ProtocolError::ChecksumMismatch { .. })
760        ));
761    }
762
763    #[test]
764    fn protocol_v3_uses_crc32_and_transport_reuses_the_precomputed_checksum() {
765        assert_eq!(PROTOCOL_VERSION, 3);
766        assert_eq!(checksum(b"123456789"), 0xcbf4_3926);
767
768        let header = FrameHeader::collective(
769            UniqueId::from_bytes([3; 16]),
770            Opcode::Send,
771            ElementType::U8,
772            0,
773            ANY_RANK,
774            2,
775            1,
776            4,
777        );
778        let mut frame = Frame::new(header, vec![1, 2, 3, 4]).unwrap();
779        frame.payload[0] ^= 0xff;
780
781        let encoded_header = frame.encode_transport_header().unwrap();
782        let mut wire = encoded_header.to_vec();
783        wire.extend_from_slice(&frame.payload);
784        assert!(matches!(
785            Frame::decode(&wire),
786            Err(ProtocolError::ChecksumMismatch { .. })
787        ));
788
789        frame.payload.push(5);
790        assert!(matches!(
791            frame.encode_transport_header(),
792            Err(ProtocolError::PayloadLengthMismatch {
793                declared: 4,
794                actual: 5
795            })
796        ));
797    }
798
799    #[test]
800    fn extended_pytorch_element_types_round_trip_with_exact_widths() {
801        for (code, element_type, width) in [
802            (7, ElementType::Bool, 1),
803            (8, ElementType::I8, 1),
804            (9, ElementType::I16, 2),
805            (10, ElementType::I64, 8),
806            (11, ElementType::F64, 8),
807            (12, ElementType::U16, 2),
808            (13, ElementType::U64, 8),
809            (14, ElementType::Complex64, 8),
810            (15, ElementType::Complex128, 16),
811            (16, ElementType::F8E4M3Fn, 1),
812            (17, ElementType::F8E5M2, 1),
813            (18, ElementType::F8E4M3Fnuz, 1),
814            (19, ElementType::F8E5M2Fnuz, 1),
815            (20, ElementType::F8E8M0Fnu, 1),
816            (21, ElementType::F4E2M1FnX2, 1),
817        ] {
818            assert_eq!(ElementType::try_from(code).unwrap(), element_type);
819            assert_eq!(element_type.byte_width(), width);
820            assert_eq!(element_type.is_low_precision_storage(), code >= 16);
821        }
822    }
823
824    #[test]
825    fn decoder_rejects_version_rank_and_length_mismatches() {
826        let header = FrameHeader::collective(
827            UniqueId::from_bytes([1; 16]),
828            Opcode::AllReduce,
829            ElementType::F32,
830            0,
831            ANY_RANK,
832            2,
833            0,
834            8,
835        );
836        let mut encoded = Frame::new(header, vec![0; 32]).unwrap().encode().unwrap();
837        encoded[4..6].copy_from_slice(&99_u16.to_le_bytes());
838        assert!(matches!(
839            Frame::decode(&encoded),
840            Err(ProtocolError::UnsupportedVersion(99))
841        ));
842
843        let mut bad_rank = FrameHeader::collective(
844            UniqueId::from_bytes([1; 16]),
845            Opcode::Barrier,
846            ElementType::None,
847            2,
848            ANY_RANK,
849            2,
850            0,
851            0,
852        );
853        assert!(matches!(
854            bad_rank.encode(),
855            Err(ProtocolError::RankOutOfRange { .. })
856        ));
857        bad_rank.source_rank = 0;
858        bad_rank.flags = FLAG_PAYLOAD;
859        assert!(matches!(
860            bad_rank.encode(),
861            Err(ProtocolError::PayloadFlagMismatch)
862        ));
863    }
864
865    #[test]
866    fn agreement_detects_order_and_contract_mismatch() {
867        let mut agreement = CollectiveAgreement::new(3).unwrap();
868        let operation =
869            OperationDescriptor::new(Opcode::AllReduce, ElementType::BF16, ANY_RANK, 1024);
870        assert!(!agreement.submit(2, 0, operation.clone()).unwrap());
871        assert!(!agreement.submit(0, 0, operation.clone()).unwrap());
872        assert!(matches!(
873            agreement.submit(
874                1,
875                0,
876                OperationDescriptor::new(Opcode::AllGather, ElementType::BF16, ANY_RANK, 1024)
877            ),
878            Err(ProtocolError::CollectiveMismatch { .. })
879        ));
880        assert!(agreement.submit(1, 0, operation.clone()).unwrap());
881        assert_eq!(agreement.next_sequence(), 1);
882        assert!(matches!(
883            agreement.submit(0, 0, operation),
884            Err(ProtocolError::SequenceMismatch { .. })
885        ));
886    }
887}
ruccl/rank/protocol.rs

ruccl/rank/
protocol.rs
