dynamo-mocker 1.5.0

Mock LLM scheduler and KV manager for testing
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Conversion from neutral replay metadata to Dynamo Router wire identities.

use aisimulate_core::replay::loadgen::Trace;
use anyhow::{Context, Result};
use dynamo_kv_router::LocalBlockHash;
use dynamo_kv_router::protocols::{
    BlockHashOptions, ExternalSequenceBlockHash, WorkerId, compute_block_hash_for_seq,
    compute_seq_hash_for_block,
};

#[derive(Debug, Clone, Copy)]
pub enum SequenceHashMode {
    Raw,
    Cumulative,
}

#[derive(Debug, Clone)]
pub struct RouterSequence {
    pub worker_id: WorkerId,
    pub local_hashes: Vec<LocalBlockHash>,
    pub external_hashes: Vec<ExternalSequenceBlockHash>,
}

/// Restore the legacy method-style API while keeping Router types out of the
/// neutral replay crate.
pub trait DynamoTraceRouterExt {
    fn to_router_sequences(
        &self,
        worker_id: WorkerId,
        hash_mode: SequenceHashMode,
    ) -> Result<Vec<RouterSequence>>;
}

impl DynamoTraceRouterExt for Trace {
    fn to_router_sequences(
        &self,
        worker_id: WorkerId,
        hash_mode: SequenceHashMode,
    ) -> Result<Vec<RouterSequence>> {
        to_router_sequences(self, worker_id, hash_mode)
    }
}

pub fn to_router_sequences(
    trace: &Trace,
    worker_id: WorkerId,
    hash_mode: SequenceHashMode,
) -> Result<Vec<RouterSequence>> {
    let mut sequences = Vec::new();
    for session in &trace.sessions {
        for turn in &session.turns {
            let local_hashes = turn
                .hash_ids
                .iter()
                .map(|&hash_id| local_block_hash_from_id(hash_id, trace.block_size))
                .collect::<Result<Vec<_>>>()?;
            let external_hashes = match hash_mode {
                SequenceHashMode::Raw => local_hashes
                    .iter()
                    .map(|hash| ExternalSequenceBlockHash(hash.0))
                    .collect(),
                SequenceHashMode::Cumulative => compute_seq_hash_for_block(&local_hashes)
                    .into_iter()
                    .map(ExternalSequenceBlockHash)
                    .collect(),
            };
            sequences.push(RouterSequence {
                worker_id,
                local_hashes,
                external_hashes,
            });
        }
    }
    Ok(sequences)
}

/// Convert neutral replay block identities once, at the Dynamo placement
/// boundary.
pub(crate) fn local_block_hashes(hashes: Vec<u64>) -> Vec<LocalBlockHash> {
    hashes.into_iter().map(LocalBlockHash).collect()
}

fn local_block_hash_from_id(hash_id: u32, block_size: usize) -> Result<LocalBlockHash> {
    let block_size = u32::try_from(block_size).context("trace block size does not fit in u32")?;
    let tokens = vec![hash_id; block_size as usize];
    compute_block_hash_for_seq(&tokens, block_size, BlockHashOptions::default())
        .into_iter()
        .next()
        .context("trace block size must be greater than zero")
}