Skip to content

GitLab

Unverified Commit 3bde1e45 authored by Graham King's avatar Graham King Committed by GitHub
Browse files

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.
parent 5d5080ba
Show whitespace changes
Inline Side-by-side
Showing with 437 additions and 35 deletions
.gitignore
View file @ 3bde1e45
......@@ -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*
......
lib/llm/src/kv_router/indexer.rs
View file @ 3bde1e45
......@@ -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"
);
let scores = kv_indexer.find_matches(block_hashes.clone()).await.unwrap();
assert_eq!(scores.frequencies, vec![1, 1, 1, 1]);
// 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;
time::sleep(Duration::from_millis(100)).await;
// 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"
);
let scores = kv_indexer.find_matches(block_hashes.clone()).await.unwrap();
assert_eq!(scores.frequencies.len(), 0);
// 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"
);
// Let those two accesses expire
time::sleep(expiration + Duration::from_millis(10)).await;
// 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
View file @ 3bde1e45
......@@ -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 0 → 100644
View file @ 3bde1e45
---
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 0 → 100644
View file @ 3bde1e45
---
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 0 → 100644
View file @ 3bde1e45
---
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 0 → 100644
View file @ 3bde1e45
---
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 0 → 100644
View file @ 3bde1e45
---
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 0 → 100644
View file @ 3bde1e45
---
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|>
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment