torrents!

This commit is contained in:
Jake Hillion
2026-01-12 15:12:56 +00:00
parent c9d95859e8
commit ba708c2ccd
19 changed files with 2308 additions and 31 deletions
Generated
+1271 -21
View File
File diff suppressed because it is too large Load Diff
+2
View File
@@ -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" }
+40
View File
@@ -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"] }
+1
View File
@@ -0,0 +1 @@
+48
View File
@@ -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");
}
}
+19
View File
@@ -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;
+77
View File
@@ -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
}
}
+162
View File
@@ -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()
}
}
+1
View File
@@ -0,0 +1 @@
+1
View File
@@ -23,6 +23,7 @@ workspace = true
[dependencies]
networking = { workspace = true }
downloads = { workspace = true }
# interop
pyo3 = { version = "0.27.1", features = [
+327
View File
@@ -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(())
}
+3
View File
@@ -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: ...
+58
View File
@@ -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
View File
@@ -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
View File
@@ -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}"
+1
View File
@@ -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}")