[bug] limit bootstrap room to to [0, 2^63 - 1] (#8684)
This commit is contained in:
@@ -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 {
|
|||||||
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 {
|
|||||||
#[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 {
|
|||||||
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");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user