Unverified Commit ab9b893e authored by Simo Lin's avatar Simo Lin Committed by GitHub
Browse files

[bug] limit bootstrap room to to [0, 2^63 - 1] (#8684)

parent 6a7528e6
...@@ -102,10 +102,8 @@ pub trait Bootstrap: Send + Sync { ...@@ -102,10 +102,8 @@ pub trait Bootstrap: Send + Sync {
BootstrapRoom::Batch( BootstrapRoom::Batch(
(0..batch_size) (0..batch_size)
.map(|_| { .map(|_| {
// Combine multiple sources of randomness for better distribution // Generate a value in the range [0, 2^63 - 1] to match Python's random.randint(0, 2**63 - 1)
let r1 = rand::random::<u64>(); rand::random::<u64>() & (i64::MAX as u64)
let r2 = rand::random::<u64>();
r1.wrapping_add(r2.rotate_left(32))
}) })
.collect(), .collect(),
), ),
...@@ -114,12 +112,10 @@ pub trait Bootstrap: Send + Sync { ...@@ -114,12 +112,10 @@ pub trait Bootstrap: Send + Sync {
self.set_bootstrap_info( self.set_bootstrap_info(
BootstrapHost::Single(hostname), BootstrapHost::Single(hostname),
BootstrapPort::Single(bootstrap_port), BootstrapPort::Single(bootstrap_port),
BootstrapRoom::Single({ BootstrapRoom::Single(
// Use high-quality random number for single requests too // Generate a value in the range [0, 2^63 - 1] to match Python's random.randint(0, 2**63 - 1)
let r1 = rand::random::<u64>(); rand::random::<u64>() & (i64::MAX as u64),
let r2 = rand::random::<u64>(); ),
r1.wrapping_add(r2.rotate_left(32))
}),
); );
} }
Ok(()) Ok(())
...@@ -279,6 +275,7 @@ impl Bootstrap for CompletionRequest { ...@@ -279,6 +275,7 @@ impl Bootstrap for CompletionRequest {
#[cfg(test)] #[cfg(test)]
mod bootstrap_tests { mod bootstrap_tests {
use super::*; use super::*;
use crate::core::BasicWorker;
use crate::openai_api_types::StringOrArray; use crate::openai_api_types::StringOrArray;
#[test] #[test]
...@@ -465,4 +462,62 @@ mod bootstrap_tests { ...@@ -465,4 +462,62 @@ mod bootstrap_tests {
assert_eq!(rooms[0].as_u64().unwrap(), 12345); assert_eq!(rooms[0].as_u64().unwrap(), 12345);
assert_eq!(rooms[1].as_u64().unwrap(), 67890); assert_eq!(rooms[1].as_u64().unwrap(), 67890);
} }
#[test]
fn test_bootstrap_room_range() {
// Test that bootstrap_room values are within the expected range [0, 2^63 - 1]
let worker = BasicWorker::new(
"http://test:8000".to_string(),
WorkerType::Prefill {
bootstrap_port: Some(8080),
},
);
// Test single request
let mut single_req = GenerateReqInput {
text: Some(InputText::Single("test".to_string())),
input_ids: None,
stream: false,
bootstrap_host: None,
bootstrap_port: None,
bootstrap_room: None,
other: Value::Object(serde_json::Map::new()),
};
for _ in 0..200000 {
single_req.add_bootstrap_info(&worker).unwrap();
if let Some(BootstrapRoom::Single(room)) = single_req.bootstrap_room {
// Verify the room value is within signed 64-bit range
assert!(room <= i64::MAX as u64, "Room {} exceeds i64::MAX", room);
} else {
panic!("Expected single bootstrap room");
}
}
// Test batch request
let mut batch_req = GenerateReqInput {
text: Some(InputText::Batch(vec![
"test1".to_string(),
"test2".to_string(),
])),
input_ids: None,
stream: false,
bootstrap_host: None,
bootstrap_port: None,
bootstrap_room: None,
other: Value::Object(serde_json::Map::new()),
};
for _ in 0..200000 {
batch_req.add_bootstrap_info(&worker).unwrap();
if let Some(BootstrapRoom::Batch(rooms)) = &batch_req.bootstrap_room {
for room in rooms {
// Verify each room value is within signed 64-bit range
assert!(*room <= i64::MAX as u64, "Room {} exceeds i64::MAX", room);
}
} else {
panic!("Expected batch bootstrap rooms");
}
}
}
} }
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