Push model data files to a worker that doesn't have them cached.
(
stream: &mut TcpStream,
model_path: &Path,
layers: &[String],
worker_name: &str,
model_name: &str,
)
| 507 | |
| 508 | /// Push model data files to a worker that doesn't have them cached. |
| 509 | async fn push_model_data( |
| 510 | stream: &mut TcpStream, |
| 511 | model_path: &Path, |
| 512 | layers: &[String], |
| 513 | worker_name: &str, |
| 514 | model_name: &str, |
| 515 | ) -> Result<()> { |
| 516 | let overall_start = Instant::now(); |
| 517 | let mut overall_bytes: u64 = 0; |
| 518 | |
| 519 | let layer_range = if layers.is_empty() { |
| 520 | "(none)".to_string() |
| 521 | } else { |
| 522 | format!( |
| 523 | "{} — {} ({} layers)", |
| 524 | layers.first().unwrap(), |
| 525 | layers.last().unwrap(), |
| 526 | layers.len() |
| 527 | ) |
| 528 | }; |
| 529 | |
| 530 | log::info!( |
| 531 | "[{}] pushing {} [{}]", |
| 532 | worker_name, |
| 533 | model_name, |
| 534 | layer_range |
| 535 | ); |
| 536 | |
| 537 | // Always send config.json and tokenizer.json |
| 538 | let mut files_to_send: Vec<PathBuf> = vec![ |
| 539 | model_path.join("config.json"), |
| 540 | model_path.join("tokenizer.json"), |
| 541 | ]; |
| 542 | |
| 543 | // Determine which safetensors shard files contain the assigned layers |
| 544 | let index_path = model_path.join("model.safetensors.index.json"); |
| 545 | let mut filtered_index: Option<Vec<u8>> = None; |
| 546 | if index_path.exists() { |
| 547 | files_to_send.push(index_path.clone()); |
| 548 | let index_data = std::fs::read(&index_path)?; |
| 549 | let mut index_json: serde_json::Value = serde_json::from_slice(&index_data)?; |
| 550 | let weight_map = index_json |
| 551 | .get("weight_map") |
| 552 | .and_then(|v| v.as_object()) |
| 553 | .ok_or_else(|| anyhow!("no weight_map in model.safetensors.index.json"))? |
| 554 | .clone(); |
| 555 | |
| 556 | // Find shard files that contain tensors for the assigned layers |
| 557 | let mut needed_shards: HashSet<String> = HashSet::new(); |
| 558 | let mut needed_weights: serde_json::Map<String, serde_json::Value> = |
| 559 | serde_json::Map::new(); |
| 560 | for (tensor_name, shard_file) in &weight_map { |
| 561 | for layer in layers { |
| 562 | if tensor_name.starts_with(&format!("{}.", layer)) { |
| 563 | if let Some(filename) = shard_file.as_str() { |
| 564 | needed_shards.insert(filename.to_string()); |
| 565 | } |
| 566 | needed_weights.insert(tensor_name.clone(), shard_file.clone()); |
no test coverage detected