rlmesh-runtime 0.1.0-rc.14

Internal RLMesh crate (unstable Rust API): runtime driver for evaluation sessions.
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
94
95
96
97
98

//! Episode bookkeeping: a stable, session-global record id per episode id.

use std::collections::BTreeMap;

/// One episode's session-global identity (`ep-NNNNNN` + index) and origin lane.
#[derive(Debug, Clone)]
pub(crate) struct EpisodeRecord {
    pub(crate) record_id: String,
    pub(crate) index: i64,
    pub(crate) env_index: i32,
    pub(crate) started_from_auto_reset: bool,
}

/// Assigns each distinct episode id a stable record id, session-global and
/// independent of which lane the episode ran in. A repeated id resolves to its
/// existing record (idempotent), so re-observing a slot never re-counts it.
#[derive(Debug, Default)]
pub(crate) struct EpisodeRecordRegistry {
    next_index: i64,
    by_episode_id: BTreeMap<String, EpisodeRecord>,
}

impl EpisodeRecordRegistry {
    /// `indices` aligns to `episode_ids`: a `Some` pins that episode's index
    /// (and record id) to its route-global slot, so ids stay unique and in
    /// start order across every lane; `None` falls back to the registry's
    /// running count.
    pub(crate) fn ensure_for_slots(
        &mut self,
        episode_ids: &[String],
        started_from_auto_reset: bool,
        indices: &[Option<i64>],
    ) -> (Vec<String>, Vec<(String, EpisodeRecord)>) {
        let mut record_ids = Vec::with_capacity(episode_ids.len());
        let mut started = Vec::new();
        for (env_index, episode_id) in episode_ids.iter().enumerate() {
            if episode_id.is_empty() {
                record_ids.push(String::new());
                continue;
            }
            if let Some(existing) = self.by_episode_id.get(episode_id) {
                record_ids.push(existing.record_id.clone());
                continue;
            }

            self.next_index += 1;
            let index = indices
                .get(env_index)
                .copied()
                .flatten()
                .unwrap_or(self.next_index);
            let record = EpisodeRecord {
                record_id: format!("ep-{index:06}"),
                index,
                env_index: env_index as i32,
                started_from_auto_reset,
            };
            record_ids.push(record.record_id.clone());
            self.by_episode_id
                .insert(episode_id.clone(), record.clone());
            started.push((episode_id.clone(), record));
        }
        (record_ids, started)
    }

    pub(crate) fn record_for(&self, episode_id: &str) -> Option<&EpisodeRecord> {
        self.by_episode_id.get(episode_id)
    }
}

#[cfg(test)]
mod tests {
    use super::EpisodeRecordRegistry;

    #[test]
    fn episode_record_ids_are_global_across_vector_slots() {
        let mut registry = EpisodeRecordRegistry::default();
        let (initial_ids, started) = registry.ensure_for_slots(
            &["runtime-a".to_string(), "runtime-b".to_string()],
            false,
            &[],
        );
        assert_eq!(initial_ids, ["ep-000001", "ep-000002"]);
        assert_eq!(started.len(), 2);
        assert!(!started[0].1.started_from_auto_reset);

        let (next_ids, started) = registry.ensure_for_slots(
            &["runtime-a".to_string(), "runtime-c".to_string()],
            true,
            &[],
        );
        assert_eq!(next_ids, ["ep-000001", "ep-000003"]);
        assert_eq!(started.len(), 1);
        assert_eq!(started[0].0, "runtime-c");
        assert_eq!(started[0].1.env_index, 1);
        assert!(started[0].1.started_from_auto_reset);
    }
}