dynamo_mocker/loadgen/
router.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Conversion from neutral replay metadata to Dynamo Router wire identities.
5
6use aisimulate_core::replay::loadgen::Trace;
7use anyhow::{Context, Result};
8use dynamo_kv_router::LocalBlockHash;
9use dynamo_kv_router::protocols::{
10    BlockHashOptions, ExternalSequenceBlockHash, WorkerId, compute_block_hash_for_seq,
11    compute_seq_hash_for_block,
12};
13
14#[derive(Debug, Clone, Copy)]
15pub enum SequenceHashMode {
16    Raw,
17    Cumulative,
18}
19
20#[derive(Debug, Clone)]
21pub struct RouterSequence {
22    pub worker_id: WorkerId,
23    pub local_hashes: Vec<LocalBlockHash>,
24    pub external_hashes: Vec<ExternalSequenceBlockHash>,
25}
26
27/// Restore the legacy method-style API while keeping Router types out of the
28/// neutral replay crate.
29pub trait DynamoTraceRouterExt {
30    fn to_router_sequences(
31        &self,
32        worker_id: WorkerId,
33        hash_mode: SequenceHashMode,
34    ) -> Result<Vec<RouterSequence>>;
35}
36
37impl DynamoTraceRouterExt for Trace {
38    fn to_router_sequences(
39        &self,
40        worker_id: WorkerId,
41        hash_mode: SequenceHashMode,
42    ) -> Result<Vec<RouterSequence>> {
43        to_router_sequences(self, worker_id, hash_mode)
44    }
45}
46
47pub fn to_router_sequences(
48    trace: &Trace,
49    worker_id: WorkerId,
50    hash_mode: SequenceHashMode,
51) -> Result<Vec<RouterSequence>> {
52    let mut sequences = Vec::new();
53    for session in &trace.sessions {
54        for turn in &session.turns {
55            let local_hashes = turn
56                .hash_ids
57                .iter()
58                .map(|&hash_id| local_block_hash_from_id(hash_id, trace.block_size))
59                .collect::<Result<Vec<_>>>()?;
60            let external_hashes = match hash_mode {
61                SequenceHashMode::Raw => local_hashes
62                    .iter()
63                    .map(|hash| ExternalSequenceBlockHash(hash.0))
64                    .collect(),
65                SequenceHashMode::Cumulative => compute_seq_hash_for_block(&local_hashes)
66                    .into_iter()
67                    .map(ExternalSequenceBlockHash)
68                    .collect(),
69            };
70            sequences.push(RouterSequence {
71                worker_id,
72                local_hashes,
73                external_hashes,
74            });
75        }
76    }
77    Ok(sequences)
78}
79
80/// Convert neutral replay block identities once, at the Dynamo placement
81/// boundary.
82pub(crate) fn local_block_hashes(hashes: Vec<u64>) -> Vec<LocalBlockHash> {
83    hashes.into_iter().map(LocalBlockHash).collect()
84}
85
86fn local_block_hash_from_id(hash_id: u32, block_size: usize) -> Result<LocalBlockHash> {
87    let block_size = u32::try_from(block_size).context("trace block size does not fit in u32")?;
88    let tokens = vec![hash_id; block_size as usize];
89    compute_block_hash_for_seq(&tokens, block_size, BlockHashOptions::default())
90        .into_iter()
91        .next()
92        .context("trace block size must be greater than zero")
93}
dynamo_mocker/loadgen/router.rs

dynamo_mocker/loadgen/
router.rs
