(
worker_name: &str,
cluster_key: &str,
bind_address: &str,
model_cache_dir: &Path,
on_progress: Option<&SetupProgressFn>,
)
| 826 | } |
| 827 | |
| 828 | pub async fn worker_setup_with_progress( |
| 829 | worker_name: &str, |
| 830 | cluster_key: &str, |
| 831 | bind_address: &str, |
| 832 | model_cache_dir: &Path, |
| 833 | on_progress: Option<&SetupProgressFn>, |
| 834 | ) -> Result<(Vec<String>, PathBuf, TcpListener)> { |
| 835 | // Detect GPUs |
| 836 | let gpus = discovery::detect_gpus(); |
| 837 | log::info!("detected {} GPU(s):", gpus.len()); |
| 838 | for gpu in &gpus { |
| 839 | log::info!( |
| 840 | " {} — {} (~{:.1} TFLOPS)", |
| 841 | &gpu.name, |
| 842 | human_bytes::human_bytes(gpu.vram_bytes as f64), |
| 843 | gpu.tflops |
| 844 | ); |
| 845 | } |
| 846 | |
| 847 | // Bind listener |
| 848 | let listener = TcpListener::bind(bind_address).await?; |
| 849 | let port = listener.local_addr()?.port(); |
| 850 | log::info!("listening on {} (setup mode)", bind_address); |
| 851 | |
| 852 | // Advertise via UDP broadcast |
| 853 | let _discovery = discovery::advertise_worker(worker_name, port, cluster_key, &gpus)?; |
| 854 | |
| 855 | log::info!("waiting for master to connect and assign layers..."); |
| 856 | if let Some(cb) = &on_progress { |
| 857 | cb("discovery", "Waiting for master...", 0.0); |
| 858 | } |
| 859 | |
| 860 | // Accept one setup connection from master |
| 861 | let (mut stream, client_addr) = listener.accept().await?; |
| 862 | let _ = stream.set_nodelay(true); |
| 863 | log::info!("[{}] master connected", client_addr); |
| 864 | if let Some(cb) = &on_progress { |
| 865 | cb("connected", &format!("Master connected ({})", client_addr), 0.0); |
| 866 | } |
| 867 | |
| 868 | // Authenticate |
| 869 | auth::authenticate_as_worker(&mut stream, cluster_key).await?; |
| 870 | log::info!("[{}] authenticated", client_addr); |
| 871 | if let Some(cb) = &on_progress { |
| 872 | cb("authenticated", "Authenticated with master", 0.0); |
| 873 | } |
| 874 | |
| 875 | // Receive layer assignment |
| 876 | let (_, msg) = Message::from_reader(&mut stream).await?; |
| 877 | let (layers, model_hash) = match msg { |
| 878 | Message::LayerAssignment { |
| 879 | layers, |
| 880 | model_hash, |
| 881 | } => (layers, model_hash), |
| 882 | other => { |
| 883 | return Err(anyhow!( |
| 884 | "expected LayerAssignment, got {:?}", |
| 885 | other |
no test coverage detected