fix: Fix race condition in kv_router unit test (#1174)

Removed the hard coded sleeps, explained what we're testing.

Closes https://github.com/ai-dynamo/dynamo/issues/1132

The race condition is that `apply_event` sends a message on a channel, it does not directly apply the event. At some later point the tokio runtime schedules the task running the channel receiver, which applies the event. If that had not happened yet the test would fail.

.gitignore

@@ -55,6 +55,7 @@ api_server_models/
 server/
 # Replay/Snapshot test artifacts
 *.new
+lib/llm/tests/data/sample-models/models--meta-llama--Llama-3.1-70B-Instruct/
 
 **/*backups*
 
@@ -90,4 +91,4 @@ TensorRT-LLM
 # Local build artifacts for devcontainer
 .build/
 # Copied binaries to ignore
-deploy/sdk/src/dynamo/sdk/cli/bin
+deploy/sdk/src/dynamo/sdk/cli/bin
\ No newline at end of file

lib/llm/src/kv_router/indexer.rs

@@ -752,7 +752,7 @@ impl KvIndexerSharded {
                                 }
 
                                 _ = cancel.cancelled() => {
-                                    tracing::debug!("KvCacheIndexer progress loop shutting down");
+                                    tracing::trace!("KvCacheIndexer progress loop shutting down");
                                     return;
                                 }
 
@@ -896,6 +896,10 @@ mod tests {
     use tokio::time;
     use tokio_util::sync::CancellationToken;
 
+    fn setup() {
+        dynamo_runtime::logging::init();
+    }
+
     fn make_blocks(hashes: Vec<u64>) -> Vec<KvCacheStoredBlockData> {
         hashes
             .iter()
@@ -948,6 +952,8 @@ mod tests {
 
     #[test]
     fn test_radix_tree() {
+        setup();
+
         let mut trie = RadixTree::new();
 
         let worker_1 = 0;
@@ -1151,6 +1157,7 @@ mod tests {
 
     #[test]
     fn test_remove_worker() {
+        setup();
         let mut trie = RadixTree::new();
 
         let worker_0 = 0;
@@ -1175,6 +1182,7 @@ mod tests {
 
     #[test]
     fn test_early_stopping() {
+        setup();
         let mut trie = RadixTree::new();
 
         let worker_0 = 0;
@@ -1203,6 +1211,7 @@ mod tests {
     #[case(32)]
     #[case(64)]
     fn test_compute_block_hash_for_seq(#[case] kv_block_size: usize) {
+        setup();
         // create a sequence of 64 elements
         let sequence = (0..kv_block_size).map(|i| i as u32).collect::<Vec<u32>>();
         let hashes = compute_block_hash_for_seq(&sequence, kv_block_size);
@@ -1250,6 +1259,7 @@ mod tests {
     #[tokio::test]
     #[apply(indexer_template)]
     async fn test_kv_indexer_new(num_shards: usize, kv_block_size: usize) {
+        setup();
         let token: CancellationToken = CancellationToken::new();
         let _ = make_indexer(&token, num_shards, kv_block_size);
     }
@@ -1257,6 +1267,7 @@ mod tests {
     #[tokio::test]
     #[apply(indexer_template)]
     async fn test_find_matches(num_shards: usize, kv_block_size: usize) {
+        setup();
         let token = CancellationToken::new();
         let kv_indexer = make_indexer(&token, num_shards, kv_block_size);
 
@@ -1269,6 +1280,7 @@ mod tests {
     #[tokio::test]
     #[apply(indexer_template)]
     async fn test_find_matches_for_request(num_shards: usize, kv_block_size: usize) {
+        setup();
         let token = CancellationToken::new();
         let kv_indexer = make_indexer(&token, num_shards, kv_block_size);
 
@@ -1281,6 +1293,7 @@ mod tests {
     #[tokio::test]
     #[apply(indexer_template)]
     async fn test_apply_event(num_shards: usize, kv_block_size: usize) {
+        setup();
         let worker_id = 0;
 
         let token = CancellationToken::new();
@@ -1295,6 +1308,7 @@ mod tests {
     #[tokio::test]
     #[apply(indexer_template)]
     async fn test_shutdown(num_shards: usize, kv_block_size: usize) {
+        setup();
         let token = CancellationToken::new();
         let mut kv_indexer = make_indexer(&token, num_shards, kv_block_size);
 
@@ -1304,65 +1318,107 @@ mod tests {
     #[tokio::test]
     #[apply(indexer_template)]
     async fn test_frequency(num_shards: usize, kv_block_size: usize) {
+        const ONE_MILLIS: Duration = Duration::from_millis(1);
+
+        setup();
         let mut kv_indexer: Box<dyn KvIndexerInterface>;
         let token = CancellationToken::new();
-        let duration = Some(Duration::from_millis(50));
+        let expiration = Duration::from_millis(50);
 
         if num_shards == 1 {
             kv_indexer = Box::new(KvIndexer::new_with_frequency(
                 token,
-                duration,
+                Some(expiration),
                 kv_block_size,
             ));
         } else {
             kv_indexer = Box::new(KvIndexerSharded::new_with_frequency(
                 token,
                 num_shards,
-                duration,
+                Some(expiration),
                 kv_block_size,
             ));
         }
 
-        let worker_id = 0;
-
-        let event = create_store_event(worker_id, 0, vec![1, 2, 3, 4], None);
-        kv_indexer.apply_event(event).await;
-
-        time::sleep(Duration::from_millis(5)).await;
-
+        // The blocks
         let block_hashes = vec![
             LocalBlockHash(1),
             LocalBlockHash(2),
             LocalBlockHash(3),
             LocalBlockHash(4),
         ];
-        let scores = kv_indexer.find_matches(block_hashes.clone()).await.unwrap();
 
-        assert_eq!(scores.frequencies.len(), 0);
+        let overlap = kv_indexer.find_matches(block_hashes.clone()).await.unwrap();
+        assert_eq!(
+            overlap.frequencies.len(),
+            0,
+            "Should be no cached blocks yet"
+        );
+
+        // Blocks go in cache
+        let worker_id = 0;
+        let event = create_store_event(worker_id, 0, vec![1, 2, 3, 4], None);
+        kv_indexer.apply_event(event).await;
 
-        let scores = kv_indexer.find_matches(block_hashes.clone()).await.unwrap();
-        assert_eq!(scores.frequencies, vec![1, 1, 1, 1]);
+        // First access
+        // The store event is applied async so poll briefly
+        let mut overlap = OverlapScores::default();
+        let timeout = Duration::from_millis(10);
+        let start = Instant::now();
+        while overlap.scores.is_empty() && Instant::now().duration_since(start) < timeout {
+            time::sleep(ONE_MILLIS).await;
+            overlap = kv_indexer.find_matches(block_hashes.clone()).await.unwrap();
+        }
+        assert_eq!(
+            overlap.scores.len(),
+            1,
+            "One worker has these blocks cached"
+        );
+        assert_eq!(
+            overlap.frequencies.len(),
+            0,
+            "Blocks have not previously been accessed"
+        );
+
+        // Second access
+        let overlap = kv_indexer.find_matches(block_hashes.clone()).await.unwrap();
+        assert_eq!(overlap.scores.len(), 1, "Still one worker matches");
+        assert_eq!(
+            overlap.frequencies,
+            vec![1, 1, 1, 1],
+            "We should see the first access now"
+        );
 
-        time::sleep(Duration::from_millis(100)).await;
+        // Let those two accesses expire
+        time::sleep(expiration + Duration::from_millis(10)).await;
 
-        let scores = kv_indexer.find_matches(block_hashes.clone()).await.unwrap();
-        assert_eq!(scores.frequencies.len(), 0);
+        // New first access
+        let overlap = kv_indexer.find_matches(block_hashes.clone()).await.unwrap();
+        assert_eq!(
+            overlap.frequencies.len(),
+            0,
+            "Blocks were accessed too long ago"
+        );
 
-        let scores = kv_indexer.find_matches(block_hashes.clone()).await.unwrap();
-        assert_eq!(scores.frequencies, vec![1, 1, 1, 1]);
+        // New second access
+        let _ = kv_indexer.find_matches(block_hashes.clone()).await.unwrap();
 
-        let scores = kv_indexer
+        // Access only the first three blocks
+        let overlap = kv_indexer
             .find_matches(block_hashes[0..3].to_vec())
             .await
             .unwrap();
-        assert_eq!(scores.frequencies, vec![2, 2, 2]);
+        // We see the previous two new accesses
+        assert_eq!(overlap.frequencies, vec![2, 2, 2]);
 
-        let scores = kv_indexer.find_matches(block_hashes.clone()).await.unwrap();
-        assert_eq!(scores.frequencies, vec![3, 3, 3, 2]);
+        // The third access did not touch the last block
+        let overlap = kv_indexer.find_matches(block_hashes.clone()).await.unwrap();
+        assert_eq!(overlap.frequencies, vec![3, 3, 3, 2]);
     }
 
     #[test]
     fn test_router_event_new() {
+        setup();
         let worker_id = 0;
         let kv_cache_event = KvCacheEvent {
             event_id: 1,
@@ -1392,6 +1448,7 @@ mod tests {
 
     #[test]
     fn test_radix_tree_default() {
+        setup();
         let radix_tree: RadixTree = Default::default();
         assert!(radix_tree.root.borrow().children.is_empty());
         assert!(radix_tree.root.borrow().workers.is_empty());
@@ -1400,6 +1457,7 @@ mod tests {
 
     #[test]
     fn test_overlap_scores_default() {
+        setup();
         let overlap_scores: OverlapScores = Default::default();
         assert!(overlap_scores.scores.is_empty());
     }

lib/llm/tests/preprocessor.rs

@@ -236,16 +236,19 @@ impl Request {
             serde_json::from_str(messages).unwrap();
         let tools: Option<Vec<async_openai::types::ChatCompletionTool>> =
             tools.map(|x| serde_json::from_str(x).unwrap());
-        let tools = tools.unwrap();
-        let tool_choice = tool_choice.unwrap();
-
-        let inner = async_openai::types::CreateChatCompletionRequestArgs::default()
-            .model(model)
-            .messages(messages)
-            .tools(tools)
-            .tool_choice(tool_choice)
-            .build()
-            .unwrap();
+        //let tools = tools.unwrap();
+        //let tool_choice = tool_choice.unwrap();
+
+        let mut inner = async_openai::types::CreateChatCompletionRequestArgs::default();
+        inner.model(model);
+        inner.messages(messages);
+        if let Some(tools) = tools {
+            inner.tools(tools);
+        }
+        if let Some(tool_choice) = tool_choice {
+            inner.tool_choice(tool_choice);
+        }
+        let inner = inner.build().unwrap();
 
         NvCreateChatCompletionRequest { inner, nvext: None }
     }

lib/llm/tests/snapshots/preprocessor__mulit_turn_with_system@meta-llama_llama-3_1-70b-instruct--1605565.snap

+---
+source: lib/llm/tests/preprocessor.rs
+expression: formatted_prompt
+info:
+  messages:
+    - role: system
+      content: You are a very helpful assistant!
+    - role: user
+      content: How do I reverse a string in Python?
+    - role: assistant
+      content: "You can reverse a string in Python using slicing:\n\n```python\nreversed_string = your_string[::-1]\n```\n\nAlternatively, you can use `reversed()` with `join()`:\n\n```python\nreversed_string = ''.join(reversed(your_string))\n```\n"
+    - role: user
+      content: What if I want to reverse each word in a sentence but keep their order?
+  model: meta-llama_llama-3_1-70b-instruct--1605565
+---
+<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+
+Cutting Knowledge Date: December 2023
+Today Date: <redacted>
+
+You are a very helpful assistant!<|eot_id|><|start_header_id|>user<|end_header_id|>
+
+How do I reverse a string in Python?<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+
+You can reverse a string in Python using slicing:
+
+```python
+reversed_string = your_string[::-1]
+```
+
+Alternatively, you can use `reversed()` with `join()`:
+
+```python
+reversed_string = ''.join(reversed(your_string))
+```<|eot_id|><|start_header_id|>user<|end_header_id|>
+
+What if I want to reverse each word in a sentence but keep their order?<|eot_id|><|start_header_id|>assistant<|end_header_id|>

lib/llm/tests/snapshots/preprocessor__mulit_turn_without_system@meta-llama_llama-3_1-70b-instruct--1605565.snap

+---
+source: lib/llm/tests/preprocessor.rs
+expression: formatted_prompt
+info:
+  messages:
+    - role: user
+      content: How do I reverse a string in Python?
+    - role: assistant
+      content: "You can reverse a string in Python using slicing:\n\n```python\nreversed_string = your_string[::-1]\n```\n\nAlternatively, you can use `reversed()` with `join()`:\n\n```python\nreversed_string = ''.join(reversed(your_string))\n```\n"
+    - role: user
+      content: What if I want to reverse each word in a sentence but keep their order?
+  model: meta-llama_llama-3_1-70b-instruct--1605565
+---
+<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+
+Cutting Knowledge Date: December 2023
+Today Date: <redacted>
+
+<|eot_id|><|start_header_id|>user<|end_header_id|>
+
+How do I reverse a string in Python?<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+
+You can reverse a string in Python using slicing:
+
+```python
+reversed_string = your_string[::-1]
+```
+
+Alternatively, you can use `reversed()` with `join()`:
+
+```python
+reversed_string = ''.join(reversed(your_string))
+```<|eot_id|><|start_header_id|>user<|end_header_id|>
+
+What if I want to reverse each word in a sentence but keep their order?<|eot_id|><|start_header_id|>assistant<|end_header_id|>

lib/llm/tests/snapshots/preprocessor__multi_turn_with_continuation@meta-llama_llama-3_1-70b-instruct--1605565.snap

+---
+source: lib/llm/tests/preprocessor.rs
+expression: formatted_prompt
+info:
+  messages:
+    - role: system
+      content: You are a very helpful assistant!
+    - role: user
+      content: How do I reverse a string in Python?
+    - role: assistant
+      content: "You can reverse a "
+  model: meta-llama_llama-3_1-70b-instruct--1605565
+---
+<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+
+Cutting Knowledge Date: December 2023
+Today Date: <redacted>
+
+You are a very helpful assistant!<|eot_id|><|start_header_id|>user<|end_header_id|>
+
+How do I reverse a string in Python?<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+
+You can reverse a<|eot_id|>

lib/llm/tests/snapshots/preprocessor__multi_turn_with_system_with_tools@meta-llama_llama-3_1-70b-instruct--1605565.snap

+---
+source: lib/llm/tests/preprocessor.rs
+expression: formatted_prompt
+info:
+  messages:
+    - role: system
+      content: You are a very helpful assistant!
+    - role: user
+      content: How do I reverse a string in Python?
+    - role: assistant
+      content: "You can reverse a string in Python using slicing:\n\n```python\nreversed_string = your_string[::-1]\n```\n\nAlternatively, you can use `reversed()` with `join()`:\n\n```python\nreversed_string = ''.join(reversed(your_string))\n```\n"
+    - role: user
+      content: What if I want to reverse each word in a sentence but keep their order?
+  model: meta-llama_llama-3_1-70b-instruct--1605565
+  tools:
+    - type: function
+      function:
+        name: get_current_temperature
+        description: Get the current temperature for a specific location
+        parameters:
+          type: object
+          properties:
+            location:
+              type: string
+              description: "The city and state, e.g., San Francisco, CA"
+            unit:
+              type: string
+              enum:
+                - Celsius
+                - Fahrenheit
+              description: "The temperature unit to use. Infer this from the user's location."
+          required:
+            - location
+            - unit
+    - type: function
+      function:
+        name: get_rain_probability
+        description: Get the probability of rain for a specific location
+        parameters:
+          type: object
+          properties:
+            location:
+              type: string
+              description: "The city and state, e.g., San Francisco, CA"
+          required:
+            - location
+  tool_choice: auto
+---
+<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+
+Environment: ipython
+Cutting Knowledge Date: December 2023
+Today Date: <redacted>
+
+You are a very helpful assistant!<|eot_id|><|start_header_id|>user<|end_header_id|>
+
+Given the following functions, please respond with a JSON for a function call with its proper arguments that best answers the given prompt.
+
+Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}.Do not use variables.
+
+{
+    "function": {
+        "description": "Get the current temperature for a specific location",
+        "name": "get_current_temperature",
+        "parameters": {
+            "properties": {
+                "location": {
+                    "description": "The city and state, e.g., San Francisco, CA",
+                    "type": "string"
+                },
+                "unit": {
+                    "description": "The temperature unit to use. Infer this from the user\u0027s location.",
+                    "enum": [
+                        "Celsius",
+                        "Fahrenheit"
+                    ],
+                    "type": "string"
+                }
+            },
+            "required": [
+                "location",
+                "unit"
+            ],
+            "type": "object"
+        }
+    },
+    "type": "function"
+}
+
+{
+    "function": {
+        "description": "Get the probability of rain for a specific location",
+        "name": "get_rain_probability",
+        "parameters": {
+            "properties": {
+                "location": {
+                    "description": "The city and state, e.g., San Francisco, CA",
+                    "type": "string"
+                }
+            },
+            "required": [
+                "location"
+            ],
+            "type": "object"
+        }
+    },
+    "type": "function"
+}
+
+How do I reverse a string in Python?<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+
+You can reverse a string in Python using slicing:
+
+```python
+reversed_string = your_string[::-1]
+```
+
+Alternatively, you can use `reversed()` with `join()`:
+
+```python
+reversed_string = ''.join(reversed(your_string))
+```<|eot_id|><|start_header_id|>user<|end_header_id|>
+
+What if I want to reverse each word in a sentence but keep their order?<|eot_id|><|start_header_id|>assistant<|end_header_id|>

lib/llm/tests/snapshots/preprocessor__single_turn@meta-llama_llama-3_1-70b-instruct--1605565.snap

+---
+source: lib/llm/tests/preprocessor.rs
+expression: formatted_prompt
+info:
+  messages:
+    - role: user
+      content: What is deep learning?
+  model: meta-llama_llama-3_1-70b-instruct--1605565
+---
+<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+
+Cutting Knowledge Date: December 2023
+Today Date: <redacted>
+
+<|eot_id|><|start_header_id|>user<|end_header_id|>
+
+What is deep learning?<|eot_id|><|start_header_id|>assistant<|end_header_id|>

lib/llm/tests/snapshots/preprocessor__single_turn_with_tools@meta-llama_llama-3_1-70b-instruct--1605565.snap

+---
+source: lib/llm/tests/preprocessor.rs
+expression: formatted_prompt
+info:
+  messages:
+    - role: user
+      content: What is deep learning?
+  model: meta-llama_llama-3_1-70b-instruct--1605565
+  tools:
+    - type: function
+      function:
+        name: get_current_temperature
+        description: Get the current temperature for a specific location
+        parameters:
+          type: object
+          properties:
+            location:
+              type: string
+              description: "The city and state, e.g., San Francisco, CA"
+            unit:
+              type: string
+              enum:
+                - Celsius
+                - Fahrenheit
+              description: "The temperature unit to use. Infer this from the user's location."
+          required:
+            - location
+            - unit
+    - type: function
+      function:
+        name: get_rain_probability
+        description: Get the probability of rain for a specific location
+        parameters:
+          type: object
+          properties:
+            location:
+              type: string
+              description: "The city and state, e.g., San Francisco, CA"
+          required:
+            - location
+  tool_choice: auto
+---
+<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+
+Environment: ipython
+Cutting Knowledge Date: December 2023
+Today Date: <redacted>
+
+<|eot_id|><|start_header_id|>user<|end_header_id|>
+
+Given the following functions, please respond with a JSON for a function call with its proper arguments that best answers the given prompt.
+
+Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}.Do not use variables.
+
+{
+    "function": {
+        "description": "Get the current temperature for a specific location",
+        "name": "get_current_temperature",
+        "parameters": {
+            "properties": {
+                "location": {
+                    "description": "The city and state, e.g., San Francisco, CA",
+                    "type": "string"
+                },
+                "unit": {
+                    "description": "The temperature unit to use. Infer this from the user\u0027s location.",
+                    "enum": [
+                        "Celsius",
+                        "Fahrenheit"
+                    ],
+                    "type": "string"
+                }
+            },
+            "required": [
+                "location",
+                "unit"
+            ],
+            "type": "object"
+        }
+    },
+    "type": "function"
+}
+
+{
+    "function": {
+        "description": "Get the probability of rain for a specific location",
+        "name": "get_rain_probability",
+        "parameters": {
+            "properties": {
+                "location": {
+                    "description": "The city and state, e.g., San Francisco, CA",
+                    "type": "string"
+                }
+            },
+            "required": [
+                "location"
+            ],
+            "type": "object"
+        }
+    },
+    "type": "function"
+}
+
+What is deep learning?<|eot_id|><|start_header_id|>assistant<|end_header_id|>