torrents!
This commit is contained in:
Generated
+1271
-21
File diff suppressed because it is too large
Load Diff
@@ -2,6 +2,7 @@
|
||||
resolver = "3"
|
||||
members = [
|
||||
"rust/networking",
|
||||
"rust/downloads",
|
||||
"rust/exo_pyo3_bindings",
|
||||
"rust/system_custodian",
|
||||
"rust/util",
|
||||
@@ -25,6 +26,7 @@ opt-level = 3
|
||||
[workspace.dependencies]
|
||||
## Crate members as common dependencies
|
||||
networking = { path = "rust/networking" }
|
||||
downloads = { path = "rust/downloads" }
|
||||
system_custodian = { path = "rust/system_custodian" }
|
||||
util = { path = "rust/util" }
|
||||
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
[package]
|
||||
name = "downloads"
|
||||
version = { workspace = true }
|
||||
edition = { workspace = true }
|
||||
publish = false
|
||||
|
||||
[lib]
|
||||
doctest = false
|
||||
name = "downloads"
|
||||
path = "src/lib.rs"
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
[dependencies]
|
||||
# macro dependencies
|
||||
derive_more = { workspace = true }
|
||||
|
||||
# async
|
||||
tokio = { workspace = true, features = ["full"] }
|
||||
futures = { workspace = true }
|
||||
futures-util = { workspace = true }
|
||||
|
||||
# utility dependencies
|
||||
util = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
anyhow = { workspace = true }
|
||||
itertools = { workspace = true }
|
||||
|
||||
# tracing/logging
|
||||
log = { workspace = true }
|
||||
|
||||
# BitTorrent library
|
||||
librqbit = "8.1.1"
|
||||
|
||||
# Embed torrent files
|
||||
include_dir = "0.7"
|
||||
|
||||
# Serialization
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
@@ -0,0 +1 @@
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
//! Embedded torrent file access
|
||||
//!
|
||||
//! Provides access to .torrent files embedded in the binary at compile time
|
||||
|
||||
use include_dir::{Dir, include_dir};
|
||||
|
||||
/// Embedded torrent files directory
|
||||
static TORRENTS: Dir<'_> = include_dir!("$CARGO_MANIFEST_DIR/torrents");
|
||||
|
||||
/// Get an embedded torrent file by model_id and revision
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `model_id` - Model identifier (e.g., "mlx-community/Qwen3-30B-A3B-4bit")
|
||||
/// * `revision` - Git commit hash
|
||||
///
|
||||
/// # Returns
|
||||
/// The torrent file contents, or None if not found
|
||||
#[inline]
|
||||
pub fn get_embedded_torrent(model_id: &str, revision: &str) -> Option<Vec<u8>> {
|
||||
let path = format!("{}/{}.torrent", model_id, revision);
|
||||
TORRENTS
|
||||
.get_file(&path)
|
||||
.map(|file| file.contents().to_vec())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_get_embedded_torrent() {
|
||||
// Test with the Qwen3 torrent we have
|
||||
let result = get_embedded_torrent(
|
||||
"mlx-community/Qwen3-30B-A3B-4bit",
|
||||
"d388dead1515f5e085ef7a0431dd8fadf0886c57",
|
||||
);
|
||||
|
||||
assert!(result.is_some(), "Expected to find embedded torrent");
|
||||
let torrent_data = result.unwrap();
|
||||
assert!(!torrent_data.is_empty(), "Torrent data should not be empty");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_missing_torrent() {
|
||||
let result = get_embedded_torrent("nonexistent/model", "abc123");
|
||||
assert!(result.is_none(), "Expected None for missing torrent");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
//! BitTorrent-based download system for model shards using rqbit
|
||||
//!
|
||||
//! This crate provides:
|
||||
//! - Torrent session management via rqbit
|
||||
//! - Embedded torrent file access
|
||||
//! - Private tracker announce handling
|
||||
//! - Selective file download based on shard layer ranges
|
||||
|
||||
#![allow(clippy::missing_inline_in_public_items)]
|
||||
|
||||
pub mod bencode;
|
||||
pub mod embedded;
|
||||
pub mod progress;
|
||||
pub mod session;
|
||||
pub mod tracker;
|
||||
|
||||
pub use embedded::get_embedded_torrent;
|
||||
pub use session::{DownloadProgress, TorrentSession};
|
||||
pub use tracker::handle_announce;
|
||||
@@ -0,0 +1,77 @@
|
||||
//! Download progress tracking
|
||||
//!
|
||||
//! Types for tracking and reporting download progress to Python
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Progress update for a torrent download
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct DownloadProgress {
|
||||
/// Total bytes to download
|
||||
pub total_bytes: u64,
|
||||
|
||||
/// Bytes downloaded so far
|
||||
pub downloaded_bytes: u64,
|
||||
|
||||
/// Number of pieces completed
|
||||
pub pieces_completed: usize,
|
||||
|
||||
/// Total number of pieces
|
||||
pub total_pieces: usize,
|
||||
|
||||
/// Number of peers connected
|
||||
pub peers_connected: usize,
|
||||
|
||||
/// Download speed in bytes/second
|
||||
pub speed_bytes_per_sec: f64,
|
||||
|
||||
/// Estimated time remaining in seconds
|
||||
pub eta_seconds: Option<f64>,
|
||||
|
||||
/// Per-file progress
|
||||
pub files: HashMap<String, FileProgress>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct FileProgress {
|
||||
/// Total file size
|
||||
pub total_bytes: u64,
|
||||
|
||||
/// Bytes downloaded for this file
|
||||
pub downloaded_bytes: u64,
|
||||
|
||||
/// Whether the file is complete
|
||||
pub complete: bool,
|
||||
}
|
||||
|
||||
impl DownloadProgress {
|
||||
#[inline]
|
||||
pub fn new(total_bytes: u64, total_pieces: usize) -> Self {
|
||||
Self {
|
||||
total_bytes,
|
||||
downloaded_bytes: 0,
|
||||
pieces_completed: 0,
|
||||
total_pieces,
|
||||
peers_connected: 0,
|
||||
speed_bytes_per_sec: 0.0,
|
||||
eta_seconds: None,
|
||||
files: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub fn progress_fraction(&self) -> f64 {
|
||||
if self.total_bytes == 0 {
|
||||
0.0
|
||||
} else {
|
||||
#[allow(clippy::cast_precision_loss)]
|
||||
let fraction = self.downloaded_bytes as f64 / self.total_bytes as f64;
|
||||
fraction
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub fn is_complete(&self) -> bool {
|
||||
self.pieces_completed >= self.total_pieces
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,162 @@
|
||||
//! Torrent session management using rqbit
|
||||
//!
|
||||
//! Provides a wrapper around rqbit's Session for managing torrent downloads
|
||||
//! with persistent seeding and selective file downloads.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use librqbit::{AddTorrent, AddTorrentOptions, Api, ManagedTorrentHandle, Session, SessionOptions};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
/// Download progress information
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct DownloadProgress {
|
||||
pub downloaded_bytes: u64,
|
||||
pub total_bytes: u64,
|
||||
pub download_speed: f64,
|
||||
pub upload_speed: f64,
|
||||
pub peers_connected: usize,
|
||||
pub is_finished: bool,
|
||||
}
|
||||
|
||||
/// Torrent session handle for managing multiple torrents
|
||||
pub struct TorrentSession {
|
||||
session: Arc<Session>,
|
||||
api: Arc<Api>,
|
||||
session_dir: PathBuf,
|
||||
torrents: Arc<RwLock<HashMap<String, ManagedTorrentHandle>>>,
|
||||
}
|
||||
|
||||
impl TorrentSession {
|
||||
/// Create a new torrent session
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `session_dir` - Directory to store session state and downloaded files
|
||||
pub async fn new(session_dir: PathBuf) -> Result<Self> {
|
||||
std::fs::create_dir_all(&session_dir).context("Failed to create session directory")?;
|
||||
|
||||
let opts = SessionOptions {
|
||||
disable_dht: false,
|
||||
disable_dht_persistence: false,
|
||||
dht_config: None,
|
||||
persistence: true,
|
||||
fastresume: true,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let session = Session::new_with_opts(session_dir.clone(), opts)
|
||||
.await
|
||||
.context("Failed to create rqbit session")?;
|
||||
|
||||
let api = Api::new(Arc::clone(&session), None);
|
||||
|
||||
Ok(Self {
|
||||
session: Arc::new(session),
|
||||
api: Arc::new(api),
|
||||
session_dir,
|
||||
torrents: Arc::new(RwLock::new(HashMap::new())),
|
||||
})
|
||||
}
|
||||
|
||||
/// Add a torrent from raw bytes
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `torrent_data` - Raw .torrent file contents
|
||||
/// * `save_path` - Where to save the downloaded files
|
||||
/// * `file_indices` - Optional list of file indices to download (None = all files)
|
||||
///
|
||||
/// # Returns
|
||||
/// Info hash as hex string
|
||||
pub async fn add_torrent(
|
||||
&self,
|
||||
torrent_data: Vec<u8>,
|
||||
save_path: PathBuf,
|
||||
file_indices: Option<Vec<usize>>,
|
||||
) -> Result<String> {
|
||||
let opts = AddTorrentOptions {
|
||||
overwrite: false,
|
||||
only_files_regex: None,
|
||||
only_files: file_indices
|
||||
.map(|indices| librqbit::AddTorrentOptions::only_files_from_vec(indices)),
|
||||
output_folder: Some(save_path.to_string_lossy().to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let add_torrent = AddTorrent::from_bytes(torrent_data);
|
||||
|
||||
let handle = self
|
||||
.session
|
||||
.add_torrent(add_torrent, Some(opts))
|
||||
.await
|
||||
.context("Failed to add torrent")?;
|
||||
|
||||
let info_hash = handle.info_hash().as_string();
|
||||
|
||||
self.torrents
|
||||
.write()
|
||||
.await
|
||||
.insert(info_hash.clone(), handle);
|
||||
|
||||
Ok(info_hash)
|
||||
}
|
||||
|
||||
/// Get download progress for a torrent
|
||||
pub async fn get_progress(&self, info_hash: &str) -> Result<DownloadProgress> {
|
||||
let torrents = self.torrents.read().await;
|
||||
let handle = torrents.get(info_hash).context("Torrent not found")?;
|
||||
|
||||
let stats = handle.stats();
|
||||
let state = handle.state();
|
||||
|
||||
Ok(DownloadProgress {
|
||||
downloaded_bytes: stats.downloaded_bytes,
|
||||
total_bytes: stats.total_bytes,
|
||||
download_speed: stats.download_speed,
|
||||
upload_speed: stats.upload_speed,
|
||||
peers_connected: stats.peers_connected,
|
||||
is_finished: state.is_finished(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Wait until torrent download is completed
|
||||
pub async fn wait_until_completed(&self, info_hash: &str) -> Result<()> {
|
||||
let torrents = self.torrents.read().await;
|
||||
let handle = torrents.get(info_hash).context("Torrent not found")?;
|
||||
|
||||
handle
|
||||
.wait_until_completed()
|
||||
.await
|
||||
.context("Failed to wait for completion")?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Enable seeding for a completed torrent
|
||||
///
|
||||
/// Note: rqbit seeds by default after completion, this is a no-op
|
||||
/// but kept for API compatibility
|
||||
pub async fn enable_seeding(&self, _info_hash: &str) -> Result<()> {
|
||||
// rqbit automatically seeds after download completion
|
||||
// This is kept for API compatibility
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Remove a torrent from the session
|
||||
pub async fn remove_torrent(&self, info_hash: &str) -> Result<()> {
|
||||
let mut torrents = self.torrents.write().await;
|
||||
|
||||
if let Some(handle) = torrents.remove(info_hash) {
|
||||
drop(handle);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get list of all torrent info hashes in the session
|
||||
pub async fn list_torrents(&self) -> Vec<String> {
|
||||
self.torrents.read().await.keys().cloned().collect()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
|
||||
BIN
Binary file not shown.
@@ -23,6 +23,7 @@ workspace = true
|
||||
|
||||
[dependencies]
|
||||
networking = { workspace = true }
|
||||
downloads = { workspace = true }
|
||||
|
||||
# interop
|
||||
pyo3 = { version = "0.27.1", features = [
|
||||
|
||||
@@ -0,0 +1,327 @@
|
||||
//! Downloads module - BitTorrent downloads PyO3 bindings
|
||||
|
||||
use crate::ext::*;
|
||||
use downloads::bencode::AnnounceParams;
|
||||
use downloads::tracker::{PeerInfo, TopologyData, handle_announce as rust_handle_announce};
|
||||
use downloads::{DownloadProgress, TorrentSession};
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::{PyBytes, PyDict};
|
||||
use std::net::Ipv4Addr;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
/// Handle a tracker announce request
|
||||
///
|
||||
/// Args:
|
||||
/// params: Dictionary with announce parameters (info_hash, peer_id, port, etc.)
|
||||
/// peers: List of peer dictionaries (node_id, ip, port, has_complete, priority)
|
||||
///
|
||||
/// Returns:
|
||||
/// Bencoded announce response as bytes
|
||||
#[pyfunction]
|
||||
fn handle_tracker_announce(
|
||||
py: Python<'_>,
|
||||
params: &Bound<'_, PyDict>,
|
||||
peers: &Bound<'_, pyo3::types::PyList>,
|
||||
) -> PyResult<Py<PyBytes>> {
|
||||
// Parse announce params
|
||||
let info_hash = {
|
||||
let info_hash_item = params
|
||||
.get_item("info_hash")?
|
||||
.ok_or_else(|| pyo3::exceptions::PyValueError::new_err("Missing info_hash"))?;
|
||||
let info_hash_bytes: &[u8] = info_hash_item.extract()?;
|
||||
|
||||
if info_hash_bytes.len() != 20 {
|
||||
return Err(pyo3::exceptions::PyValueError::new_err(
|
||||
"info_hash must be 20 bytes",
|
||||
));
|
||||
}
|
||||
|
||||
let mut info_hash = [0u8; 20];
|
||||
info_hash.copy_from_slice(info_hash_bytes);
|
||||
info_hash
|
||||
};
|
||||
|
||||
let peer_id = {
|
||||
let peer_id_item = params
|
||||
.get_item("peer_id")?
|
||||
.ok_or_else(|| pyo3::exceptions::PyValueError::new_err("Missing peer_id"))?;
|
||||
let peer_id_bytes: &[u8] = peer_id_item.extract()?;
|
||||
|
||||
if peer_id_bytes.len() != 20 {
|
||||
return Err(pyo3::exceptions::PyValueError::new_err(
|
||||
"peer_id must be 20 bytes",
|
||||
));
|
||||
}
|
||||
|
||||
let mut peer_id = [0u8; 20];
|
||||
peer_id.copy_from_slice(peer_id_bytes);
|
||||
peer_id
|
||||
};
|
||||
|
||||
let port: u16 = params
|
||||
.get_item("port")?
|
||||
.ok_or_else(|| pyo3::exceptions::PyValueError::new_err("Missing port"))?
|
||||
.extract()?;
|
||||
|
||||
let uploaded: u64 = params
|
||||
.get_item("uploaded")?
|
||||
.ok_or_else(|| pyo3::exceptions::PyValueError::new_err("Missing uploaded"))?
|
||||
.extract()?;
|
||||
|
||||
let downloaded: u64 = params
|
||||
.get_item("downloaded")?
|
||||
.ok_or_else(|| pyo3::exceptions::PyValueError::new_err("Missing downloaded"))?
|
||||
.extract()?;
|
||||
|
||||
let left: u64 = params
|
||||
.get_item("left")?
|
||||
.ok_or_else(|| pyo3::exceptions::PyValueError::new_err("Missing left"))?
|
||||
.extract()?;
|
||||
|
||||
let compact: bool = params
|
||||
.get_item("compact")?
|
||||
.map(|v| v.extract().unwrap_or(true))
|
||||
.unwrap_or(true);
|
||||
|
||||
let announce_params = AnnounceParams {
|
||||
info_hash,
|
||||
peer_id,
|
||||
port,
|
||||
uploaded,
|
||||
downloaded,
|
||||
left,
|
||||
compact,
|
||||
event: None, // TODO: parse event if needed
|
||||
};
|
||||
|
||||
// Parse peer list
|
||||
let peer_infos: Result<Vec<PeerInfo>, PyErr> = peers
|
||||
.iter()
|
||||
.map(|peer_item| {
|
||||
let peer_dict: &Bound<'_, PyDict> = peer_item.downcast()?;
|
||||
let node_id: String = peer_dict
|
||||
.get_item("node_id")?
|
||||
.ok_or_else(|| pyo3::exceptions::PyValueError::new_err("Missing node_id"))?
|
||||
.extract()?;
|
||||
|
||||
let ip_str: String = peer_dict
|
||||
.get_item("ip")?
|
||||
.ok_or_else(|| pyo3::exceptions::PyValueError::new_err("Missing ip"))?
|
||||
.extract()?;
|
||||
|
||||
let ip: Ipv4Addr = ip_str
|
||||
.parse()
|
||||
.map_err(|_| pyo3::exceptions::PyValueError::new_err("Invalid IP address"))?;
|
||||
|
||||
let port: u16 = peer_dict
|
||||
.get_item("port")?
|
||||
.ok_or_else(|| pyo3::exceptions::PyValueError::new_err("Missing port"))?
|
||||
.extract()?;
|
||||
|
||||
let has_complete: bool = peer_dict
|
||||
.get_item("has_complete")?
|
||||
.map(|v: Bound<'_, pyo3::PyAny>| v.extract().unwrap_or(false))
|
||||
.unwrap_or(false);
|
||||
|
||||
let priority: i32 = peer_dict
|
||||
.get_item("priority")?
|
||||
.map(|v: Bound<'_, pyo3::PyAny>| v.extract().unwrap_or(0))
|
||||
.unwrap_or(0);
|
||||
|
||||
Ok(PeerInfo {
|
||||
node_id,
|
||||
ip,
|
||||
port,
|
||||
has_complete,
|
||||
priority,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
let peer_infos = peer_infos?;
|
||||
|
||||
let topology = TopologyData { peers: peer_infos };
|
||||
|
||||
// Call Rust tracker handler
|
||||
let response_bytes = rust_handle_announce(&announce_params, &topology).pyerr()?;
|
||||
|
||||
// Return as Python bytes
|
||||
Ok(PyBytes::new(py, &response_bytes).unbind())
|
||||
}
|
||||
|
||||
/// Get an embedded torrent file
|
||||
///
|
||||
/// Args:
|
||||
/// model_id: Model identifier (e.g., "mlx-community/Qwen3-30B-A3B-4bit")
|
||||
/// revision: Git commit hash
|
||||
///
|
||||
/// Returns:
|
||||
/// Torrent file contents as bytes, or None if not found
|
||||
#[pyfunction]
|
||||
fn get_embedded_torrent(
|
||||
py: Python<'_>,
|
||||
model_id: String,
|
||||
revision: String,
|
||||
) -> PyResult<Option<Py<PyBytes>>> {
|
||||
match downloads::get_embedded_torrent(&model_id, &revision) {
|
||||
Some(data) => Ok(Some(PyBytes::new(py, &data).unbind())),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
/// Python wrapper for TorrentSession
|
||||
#[pyclass]
|
||||
struct TorrentSessionHandle {
|
||||
session: Arc<Mutex<TorrentSession>>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl TorrentSessionHandle {
|
||||
/// Create a new torrent session
|
||||
///
|
||||
/// Args:
|
||||
/// session_dir: Directory to store session state and downloads
|
||||
#[new]
|
||||
fn new(session_dir: String) -> PyResult<Self> {
|
||||
let session_path = PathBuf::from(session_dir);
|
||||
|
||||
let session = tokio::runtime::Runtime::new()
|
||||
.pyerr()?
|
||||
.block_on(async { TorrentSession::new(session_path).await })
|
||||
.pyerr()?;
|
||||
|
||||
Ok(Self {
|
||||
session: Arc::new(Mutex::new(session)),
|
||||
})
|
||||
}
|
||||
|
||||
/// Add a torrent from bytes
|
||||
///
|
||||
/// Args:
|
||||
/// torrent_data: Raw .torrent file contents
|
||||
/// save_path: Where to save downloaded files
|
||||
/// file_indices: Optional list of file indices to download
|
||||
///
|
||||
/// Returns:
|
||||
/// Info hash as hex string
|
||||
fn add_torrent(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
torrent_data: Vec<u8>,
|
||||
save_path: String,
|
||||
file_indices: Option<Vec<usize>>,
|
||||
) -> PyResult<String> {
|
||||
let session = Arc::clone(&self.session);
|
||||
let save_path = PathBuf::from(save_path);
|
||||
|
||||
py.allow_threads(|| {
|
||||
tokio::runtime::Runtime::new()
|
||||
.pyerr()?
|
||||
.block_on(async {
|
||||
session
|
||||
.lock()
|
||||
.await
|
||||
.add_torrent(torrent_data, save_path, file_indices)
|
||||
.await
|
||||
})
|
||||
.pyerr()
|
||||
})
|
||||
}
|
||||
|
||||
/// Get download progress for a torrent
|
||||
///
|
||||
/// Args:
|
||||
/// info_hash: Torrent info hash
|
||||
///
|
||||
/// Returns:
|
||||
/// Dictionary with progress information
|
||||
fn get_progress(&self, py: Python<'_>, info_hash: String) -> PyResult<Py<PyDict>> {
|
||||
let session = Arc::clone(&self.session);
|
||||
|
||||
let progress: DownloadProgress = py.allow_threads(|| {
|
||||
tokio::runtime::Runtime::new()
|
||||
.pyerr()?
|
||||
.block_on(async { session.lock().await.get_progress(&info_hash).await })
|
||||
.pyerr()
|
||||
})?;
|
||||
|
||||
let dict = PyDict::new(py);
|
||||
dict.set_item("downloaded_bytes", progress.downloaded_bytes)?;
|
||||
dict.set_item("total_bytes", progress.total_bytes)?;
|
||||
dict.set_item("download_speed", progress.download_speed)?;
|
||||
dict.set_item("upload_speed", progress.upload_speed)?;
|
||||
dict.set_item("peers_connected", progress.peers_connected)?;
|
||||
dict.set_item("is_finished", progress.is_finished)?;
|
||||
|
||||
Ok(dict.unbind())
|
||||
}
|
||||
|
||||
/// Wait until torrent download is completed
|
||||
///
|
||||
/// Args:
|
||||
/// info_hash: Torrent info hash
|
||||
fn wait_until_completed(&self, py: Python<'_>, info_hash: String) -> PyResult<()> {
|
||||
let session = Arc::clone(&self.session);
|
||||
|
||||
py.allow_threads(|| {
|
||||
tokio::runtime::Runtime::new()
|
||||
.pyerr()?
|
||||
.block_on(async { session.lock().await.wait_until_completed(&info_hash).await })
|
||||
.pyerr()
|
||||
})
|
||||
}
|
||||
|
||||
/// Enable seeding for a torrent
|
||||
///
|
||||
/// Args:
|
||||
/// info_hash: Torrent info hash
|
||||
fn enable_seeding(&self, py: Python<'_>, info_hash: String) -> PyResult<()> {
|
||||
let session = Arc::clone(&self.session);
|
||||
|
||||
py.allow_threads(|| {
|
||||
tokio::runtime::Runtime::new()
|
||||
.pyerr()?
|
||||
.block_on(async { session.lock().await.enable_seeding(&info_hash).await })
|
||||
.pyerr()
|
||||
})
|
||||
}
|
||||
|
||||
/// Remove a torrent from the session
|
||||
///
|
||||
/// Args:
|
||||
/// info_hash: Torrent info hash
|
||||
fn remove_torrent(&self, py: Python<'_>, info_hash: String) -> PyResult<()> {
|
||||
let session = Arc::clone(&self.session);
|
||||
|
||||
py.allow_threads(|| {
|
||||
tokio::runtime::Runtime::new()
|
||||
.pyerr()?
|
||||
.block_on(async { session.lock().await.remove_torrent(&info_hash).await })
|
||||
.pyerr()
|
||||
})
|
||||
}
|
||||
|
||||
/// List all torrents in the session
|
||||
///
|
||||
/// Returns:
|
||||
/// List of info hashes
|
||||
fn list_torrents(&self, py: Python<'_>) -> PyResult<Vec<String>> {
|
||||
let session = Arc::clone(&self.session);
|
||||
|
||||
py.allow_threads(|| {
|
||||
tokio::runtime::Runtime::new()
|
||||
.pyerr()?
|
||||
.block_on(async { Ok(session.lock().await.list_torrents().await) })
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Downloads submodule
|
||||
pub(crate) fn downloads_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_function(wrap_pyfunction!(handle_tracker_announce, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(get_embedded_torrent, m)?)?;
|
||||
m.add_class::<TorrentSessionHandle>()?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -17,10 +17,12 @@
|
||||
|
||||
extern crate core;
|
||||
mod allow_threading;
|
||||
pub(crate) mod downloads;
|
||||
mod examples;
|
||||
pub(crate) mod networking;
|
||||
pub(crate) mod pylibp2p;
|
||||
|
||||
use crate::downloads::downloads_submodule;
|
||||
use crate::networking::networking_submodule;
|
||||
use crate::pylibp2p::ident::ident_submodule;
|
||||
use crate::pylibp2p::multiaddr::multiaddr_submodule;
|
||||
@@ -207,6 +209,7 @@ fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
ident_submodule(m)?;
|
||||
multiaddr_submodule(m)?;
|
||||
networking_submodule(m)?;
|
||||
downloads_submodule(m)?;
|
||||
|
||||
// top-level constructs
|
||||
// TODO: ...
|
||||
|
||||
Executable
+58
@@ -0,0 +1,58 @@
|
||||
#!/usr/bin/env nix-shell
|
||||
#!nix-shell -i bash -p mktorrent -p python3Packages.huggingface-hub -p git -p git-lfs
|
||||
set -euo pipefail
|
||||
set -x
|
||||
|
||||
MODEL="$1"
|
||||
|
||||
mkdir -p "$MODEL"
|
||||
|
||||
# Step 1: Clone/fetch the repo and get the hash of head
|
||||
mkdir -p "$MODEL"
|
||||
if test -d "$MODEL/git"; then
|
||||
# Assert that the origin is correct
|
||||
git -C "$MODEL/git" fetch
|
||||
else
|
||||
git clone "https://huggingface.co/$MODEL" "$MODEL/git"
|
||||
fi
|
||||
|
||||
HASH=$(git -C "$MODEL/git" rev-parse origin/main)
|
||||
LARGE_FILES=$(git -C "$MODEL/git" lfs ls-files --all --name-only)
|
||||
|
||||
SMALL_DIR="$MODEL/$HASH-small"
|
||||
LARGE_DIR="$MODEL/$HASH-large"
|
||||
mkdir -p "$SMALL_DIR" "$LARGE_DIR"
|
||||
|
||||
# Step 2: Prepare files. Two torrents: one for large files and one for metadata.
|
||||
git -C "$MODEL/git" archive "$HASH" | tar -x -C "$SMALL_DIR"
|
||||
echo "$LARGE_FILES" | xargs -I{} rm "$SMALL_DIR/{}"
|
||||
|
||||
echo "$LARGE_FILES" | xargs hf download "$MODEL" --revision "$HASH" --local-dir "$LARGE_DIR" --cache-dir "$(realpath .cache)" --include
|
||||
if test -d "$LARGE_DIR/.cache"; then
|
||||
echo ".cache created against our wishes, deleting it..."
|
||||
rm -r "$LARGE_DIR/.cache"
|
||||
fi
|
||||
|
||||
# Step 3: Create both torrents
|
||||
mkdir -p "torrents/$MODEL/"
|
||||
SMALL_TORRENT_PATH="torrents/$MODEL/${HASH}.small.torrent"
|
||||
LARGE_TORRENT_PATH="torrents/$MODEL/${HASH}.large.torrent"
|
||||
|
||||
mktorrent "$SMALL_DIR/" --output="$SMALL_TORRENT_PATH" \
|
||||
-n "$HASH" \
|
||||
--web-seed="https://huggingface.co/$MODEL/raw/" \
|
||||
--no-date \
|
||||
--announce="udp://tracker.opentrackr.org:1337/announce"
|
||||
# --private
|
||||
|
||||
mktorrent "$LARGE_DIR/" --output="$LARGE_TORRENT_PATH" \
|
||||
-n "$HASH" \
|
||||
--web-seed="https://huggingface.co/$MODEL/resolve/" \
|
||||
--piece-length=24 \
|
||||
--no-date \
|
||||
--announce="udp://tracker.opentrackr.org:1337/announce"
|
||||
# --private
|
||||
|
||||
echo "Successfully created torrent files in:"
|
||||
echo "$SMALL_TORRENT_PATH"
|
||||
echo "$LARGE_TORRENT_PATH"
|
||||
+12
-4
@@ -35,6 +35,7 @@ class Node:
|
||||
api: API | None
|
||||
|
||||
node_id: NodeId
|
||||
enable_torrents: bool
|
||||
_tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
|
||||
|
||||
@classmethod
|
||||
@@ -66,7 +67,8 @@ class Node:
|
||||
worker = Worker(
|
||||
node_id,
|
||||
session_id,
|
||||
exo_shard_downloader(),
|
||||
exo_shard_downloader(enable_torrents=args.enable_torrents),
|
||||
initial_connection_messages=[],
|
||||
connection_message_receiver=router.receiver(topics.CONNECTION_MESSAGES),
|
||||
global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
|
||||
local_event_sender=router.sender(topics.LOCAL_EVENTS),
|
||||
@@ -74,7 +76,6 @@ class Node:
|
||||
)
|
||||
else:
|
||||
worker = None
|
||||
|
||||
# We start every node with a master
|
||||
master = Master(
|
||||
node_id,
|
||||
@@ -98,7 +99,7 @@ class Node:
|
||||
election_result_sender=er_send,
|
||||
)
|
||||
|
||||
return cls(router, worker, election, er_recv, master, api, node_id)
|
||||
return cls(router, worker, election, er_recv, master, api, node_id, args.enable_torrents)
|
||||
|
||||
async def run(self):
|
||||
async with self._tg as tg:
|
||||
@@ -175,7 +176,7 @@ class Node:
|
||||
self.worker = Worker(
|
||||
self.node_id,
|
||||
result.session_id,
|
||||
exo_shard_downloader(),
|
||||
exo_shard_downloader(enable_torrents=self.enable_torrents),
|
||||
connection_message_receiver=self.router.receiver(
|
||||
topics.CONNECTION_MESSAGES
|
||||
),
|
||||
@@ -215,6 +216,7 @@ class Args(CamelCaseModel):
|
||||
api_port: PositiveInt = 52415
|
||||
tb_only: bool = False
|
||||
no_worker: bool = False
|
||||
enable_torrents: bool = False
|
||||
|
||||
@classmethod
|
||||
def parse(cls) -> Self:
|
||||
@@ -256,6 +258,12 @@ class Args(CamelCaseModel):
|
||||
"--no-worker",
|
||||
action="store_true",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-torrents",
|
||||
action="store_true",
|
||||
dest="enable_torrents",
|
||||
help="Enable BitTorrent-based downloads (experimental)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
return cls(**vars(args)) # pyright: ignore[reportAny] - We are intentionally validating here, we can't do it statically
|
||||
|
||||
+62
-2
@@ -5,9 +5,9 @@ from typing import cast
|
||||
import anyio
|
||||
from anyio import create_task_group
|
||||
from anyio.abc import TaskGroup
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi import FastAPI, HTTPException, Request
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import StreamingResponse
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from hypercorn.asyncio import serve # pyright: ignore[reportUnknownVariableType]
|
||||
from hypercorn.config import Config
|
||||
@@ -178,6 +178,7 @@ class API:
|
||||
self.app.post("/bench/chat/completions")(self.bench_chat_completions)
|
||||
self.app.get("/state")(lambda: self.state)
|
||||
self.app.get("/events")(lambda: self._event_log)
|
||||
self.app.get("/_internal/announce")(self.tracker_announce)
|
||||
|
||||
async def place_instance(self, payload: PlaceInstanceParams):
|
||||
command = PlaceInstance(
|
||||
@@ -622,6 +623,65 @@ class API:
|
||||
]
|
||||
)
|
||||
|
||||
async def tracker_announce(self, request: Request) -> Response:
|
||||
"""BitTorrent tracker announce endpoint for private tracker."""
|
||||
try:
|
||||
from exo_pyo3_bindings import handle_tracker_announce # type: ignore
|
||||
except ImportError as e:
|
||||
raise HTTPException(
|
||||
status_code=501,
|
||||
detail="Torrent support not available (exo_pyo3_bindings not installed)",
|
||||
) from e
|
||||
|
||||
# Parse announce parameters from query string
|
||||
query_params = dict(request.query_params)
|
||||
|
||||
# Extract required parameters
|
||||
try:
|
||||
info_hash_hex = query_params.get("info_hash", "")
|
||||
peer_id_hex = query_params.get("peer_id", "")
|
||||
|
||||
# URL decode and convert to bytes
|
||||
info_hash = bytes.fromhex(info_hash_hex) if info_hash_hex else b""
|
||||
peer_id = bytes.fromhex(peer_id_hex) if peer_id_hex else b""
|
||||
|
||||
if len(info_hash) != 20 or len(peer_id) != 20:
|
||||
raise ValueError("info_hash and peer_id must be 20 bytes")
|
||||
|
||||
params = {
|
||||
"info_hash": info_hash,
|
||||
"peer_id": peer_id,
|
||||
"port": int(query_params.get("port", "6881")),
|
||||
"uploaded": int(query_params.get("uploaded", "0")),
|
||||
"downloaded": int(query_params.get("downloaded", "0")),
|
||||
"left": int(query_params.get("left", "0")),
|
||||
"compact": query_params.get("compact", "1") == "1",
|
||||
}
|
||||
except (ValueError, KeyError) as e:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Invalid announce parameters: {e}",
|
||||
) from e
|
||||
|
||||
# Build peer list from topology
|
||||
# TODO: Implement _build_peer_list_from_topology() to extract peers from self.state.topology
|
||||
peers = [] # For now, return empty peer list
|
||||
|
||||
# Call Rust tracker handler
|
||||
try:
|
||||
response_bytes: bytes = handle_tracker_announce(params, peers) # type: ignore
|
||||
return Response(
|
||||
content=response_bytes,
|
||||
media_type="text/plain",
|
||||
headers={"Content-Type": "text/plain"},
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Tracker announce error: {e}")
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Tracker announce failed: {e}",
|
||||
) from e
|
||||
|
||||
async def run(self):
|
||||
cfg = Config()
|
||||
cfg.bind = f"0.0.0.0:{self.port}"
|
||||
|
||||
@@ -16,3 +16,4 @@ class ModelMetadata(CamelCaseModel):
|
||||
n_layers: PositiveInt
|
||||
hidden_size: PositiveInt
|
||||
supports_tensor: bool
|
||||
revision: str | None = None # Git commit hash for torrent lookup
|
||||
|
||||
@@ -12,10 +12,17 @@ from exo.worker.download.download_utils import RepoDownloadProgress, download_sh
|
||||
from exo.worker.download.shard_downloader import ShardDownloader
|
||||
|
||||
|
||||
def exo_shard_downloader(max_parallel_downloads: int = 8) -> ShardDownloader:
|
||||
return SingletonShardDownloader(
|
||||
CachedShardDownloader(ResumableShardDownloader(max_parallel_downloads))
|
||||
)
|
||||
def exo_shard_downloader(
|
||||
max_parallel_downloads: int = 8, enable_torrents: bool = False
|
||||
) -> ShardDownloader:
|
||||
if enable_torrents:
|
||||
from exo.worker.download.torrent_downloader import TorrentShardDownloader
|
||||
|
||||
base = TorrentShardDownloader(max_parallel_downloads)
|
||||
else:
|
||||
base = ResumableShardDownloader(max_parallel_downloads)
|
||||
|
||||
return SingletonShardDownloader(CachedShardDownloader(base))
|
||||
|
||||
|
||||
async def build_base_shard(model_id: str) -> ShardMetadata:
|
||||
|
||||
@@ -0,0 +1,212 @@
|
||||
"""Torrent-based shard downloader using BitTorrent protocol with private tracker.
|
||||
|
||||
This module implements downloading model shards via BitTorrent with:
|
||||
- Private tracker integration (/_internal/announce endpoint)
|
||||
- WebSeed (BEP 19) fallback to HTTPS
|
||||
- Selective file download based on shard layer ranges
|
||||
- Persistent seeding of downloaded content
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncIterator
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from exo.shared.constants import EXO_MODELS_DIR
|
||||
from exo.shared.types.worker.shards import ShardMetadata
|
||||
from exo.worker.download.download_utils import RepoDownloadProgress
|
||||
from exo.worker.download.shard_downloader import ShardDownloader
|
||||
|
||||
|
||||
class TorrentShardDownloader(ShardDownloader):
|
||||
"""Download model shards using BitTorrent with private tracker."""
|
||||
|
||||
def __init__(self, max_parallel_downloads: int = 8):
|
||||
self.max_parallel_downloads = max_parallel_downloads
|
||||
self._progress_callbacks: list[
|
||||
Callable[[ShardMetadata, RepoDownloadProgress], None]
|
||||
] = []
|
||||
|
||||
# Initialize TorrentSessionHandle
|
||||
try:
|
||||
from exo_pyo3_bindings import TorrentSessionHandle # type: ignore
|
||||
|
||||
session_dir = EXO_MODELS_DIR / "v2" / ".torrent_session"
|
||||
session_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.session = TorrentSessionHandle(str(session_dir))
|
||||
except ImportError:
|
||||
logger.error("exo_pyo3_bindings not available, torrent downloads disabled")
|
||||
self.session = None
|
||||
|
||||
def on_progress(
|
||||
self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
|
||||
) -> None:
|
||||
"""Register a progress callback."""
|
||||
self._progress_callbacks.append(callback)
|
||||
|
||||
async def ensure_shard(
|
||||
self, shard: ShardMetadata, config_only: bool = False
|
||||
) -> Path:
|
||||
"""Download a model shard using BitTorrent.
|
||||
|
||||
Args:
|
||||
shard: Shard metadata including model ID and layer range
|
||||
config_only: If True, only download config files (not implemented for torrents yet)
|
||||
|
||||
Returns:
|
||||
Path to the downloaded shard directory
|
||||
|
||||
Raises:
|
||||
RuntimeError: If torrent file is not found or session not initialized
|
||||
"""
|
||||
if self.session is None:
|
||||
raise RuntimeError(
|
||||
"TorrentSessionHandle not initialized. "
|
||||
"exo_pyo3_bindings module not available."
|
||||
)
|
||||
|
||||
model_id = str(shard.model_meta.model_id)
|
||||
|
||||
# Resolve "main" branch to commit hash
|
||||
revision = await self._resolve_revision(model_id, "main")
|
||||
|
||||
# Load embedded torrent file
|
||||
torrent_data = self._load_embedded_torrent(model_id, revision)
|
||||
if torrent_data is None:
|
||||
raise RuntimeError(
|
||||
f"Torrent not found for {model_id}@{revision}. "
|
||||
f"Expected at: rust/downloads/torrents/{model_id}/{revision}.torrent. "
|
||||
f"Please add the torrent file or use HTTP downloads (--enable-torrents=False)."
|
||||
)
|
||||
|
||||
# Build v2 path: ~/.exo/models/v2/{model_id}/{revision}
|
||||
save_path = EXO_MODELS_DIR / "v2" / model_id / revision
|
||||
save_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Calculate which files to download based on shard layers
|
||||
# For now, download all files (selective download will be implemented later)
|
||||
file_indices = None # None means download all files
|
||||
|
||||
# Add torrent and download
|
||||
logger.info(
|
||||
f"Adding torrent for {model_id}@{revision} to session, saving to {save_path}"
|
||||
)
|
||||
info_hash = self.session.add_torrent(torrent_data, str(save_path), file_indices)
|
||||
|
||||
# Wait for download with progress reporting
|
||||
logger.info(f"Starting download for {info_hash}")
|
||||
while True:
|
||||
progress_dict = self.session.get_progress(info_hash)
|
||||
|
||||
# Convert dict to RepoDownloadProgress for callbacks
|
||||
progress = RepoDownloadProgress(
|
||||
downloaded=progress_dict["downloaded_bytes"],
|
||||
total=progress_dict["total_bytes"],
|
||||
)
|
||||
self._report_progress(shard, progress)
|
||||
|
||||
if progress_dict["is_finished"]:
|
||||
logger.info(f"Download completed for {info_hash}")
|
||||
break
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Enable persistent seeding
|
||||
logger.info(f"Enabling seeding for {info_hash}")
|
||||
self.session.enable_seeding(info_hash)
|
||||
|
||||
return save_path
|
||||
|
||||
async def get_shard_download_status(
|
||||
self,
|
||||
) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]:
|
||||
"""Get download status for all shards."""
|
||||
if self.session is None:
|
||||
return
|
||||
|
||||
# List all torrents in the session
|
||||
info_hashes = self.session.list_torrents()
|
||||
|
||||
for info_hash in info_hashes:
|
||||
try:
|
||||
progress_dict = self.session.get_progress(info_hash)
|
||||
progress = RepoDownloadProgress(
|
||||
downloaded=progress_dict["downloaded_bytes"],
|
||||
total=progress_dict["total_bytes"],
|
||||
)
|
||||
|
||||
# We don't have a straightforward way to map info_hash back to path
|
||||
# This would require tracking in the session or metadata
|
||||
# For now, yield empty path
|
||||
yield (Path(), progress)
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting status for {info_hash}: {e}")
|
||||
|
||||
async def get_shard_download_status_for_shard(
|
||||
self, shard: ShardMetadata
|
||||
) -> RepoDownloadProgress:
|
||||
"""Get download status for a specific shard."""
|
||||
if self.session is None:
|
||||
return RepoDownloadProgress(downloaded=0, total=0)
|
||||
|
||||
model_id = str(shard.model_meta.model_id)
|
||||
revision = await self._resolve_revision(model_id, "main")
|
||||
|
||||
# We would need to track info_hash -> shard mapping
|
||||
# For now, return empty progress
|
||||
return RepoDownloadProgress(downloaded=0, total=0)
|
||||
|
||||
async def _resolve_revision(self, model_id: str, branch: str = "main") -> str:
|
||||
"""Resolve branch name to commit hash using HuggingFace API.
|
||||
|
||||
Args:
|
||||
model_id: Model identifier (e.g., "mlx-community/Qwen3-30B-A3B-4bit")
|
||||
branch: Branch name (default: "main")
|
||||
|
||||
Returns:
|
||||
Git commit hash (SHA)
|
||||
"""
|
||||
try:
|
||||
from huggingface_hub import model_info
|
||||
|
||||
info = model_info(model_id, revision=branch)
|
||||
return str(info.sha)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to resolve revision for {model_id}@{branch}: {e}")
|
||||
# Fallback to "main" as revision if API call fails
|
||||
logger.warning(f"Using branch name '{branch}' as revision")
|
||||
return branch
|
||||
|
||||
def _load_embedded_torrent(self, model_id: str, revision: str) -> bytes | None:
|
||||
"""Load embedded torrent file from Rust binary.
|
||||
|
||||
Args:
|
||||
model_id: Model identifier
|
||||
revision: Git commit hash
|
||||
|
||||
Returns:
|
||||
Torrent file contents, or None if not found
|
||||
"""
|
||||
try:
|
||||
from exo_pyo3_bindings import get_embedded_torrent # type: ignore
|
||||
|
||||
result: bytes | None = get_embedded_torrent(model_id, revision) # type: ignore
|
||||
return result
|
||||
except ImportError:
|
||||
logger.warning("exo_pyo3_bindings not available, cannot load torrents")
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading torrent for {model_id}@{revision}: {e}")
|
||||
return None
|
||||
|
||||
def _report_progress(
|
||||
self, shard: ShardMetadata, progress: RepoDownloadProgress
|
||||
) -> None:
|
||||
"""Report progress to all registered callbacks."""
|
||||
for callback in self._progress_callbacks:
|
||||
try:
|
||||
callback(shard, progress)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in progress callback: {e}")
|
||||
Reference in New Issue
Block a user