fandhe_ai/distributed/
config.rs

1//! 分散学習の設定（world_size・rank・集合点アドレス・opt-in・タイムアウト）。
2//!
3//! 役割: autodiff 内部の `RendezvousConfig`／`CollectiveOptions` を facade の builder で包み、
4//! `build()` の時点でソケットを開く前にすべて検証する（OWASP A03。決定記録 §16）。
5//! 既定は loopback。非 loopback は `allow_non_loopback(true)` の明示 opt-in だけで許す。
6//! タイムアウトの既定値は内部の既定（接続 10 秒・確立 60 秒・io 30 秒・all-reduce 1 回 120 秒）を
7//! そのまま使い、facade では独自の既定値を持たない。
8
9use std::net::SocketAddr;
10use std::time::Duration;
11
12use fandhe_ai_autodiff::distributed::{CollectiveOptions, RendezvousConfig};
13
14use super::error::DistributedError;
15
16/// 検証済みの分散学習設定。[`DistributedConfig::builder`] から作る。
17#[derive(Debug, Clone)]
18pub struct DistributedConfig {
19    pub(crate) rendezvous: RendezvousConfig,
20    pub(crate) collective: CollectiveOptions,
21}
22
23impl DistributedConfig {
24    /// builder を作る（world_size・rank・集合点アドレスは `build` で検証する）。
25    ///
26    /// 集合点アドレスはランク 0 が bind し、他ランクが接続する先。`SocketAddr` だけを受け、
27    /// ホスト名解決（DNS）は行わない。
28    pub fn builder(
29        world_size: usize,
30        rank: usize,
31        rendezvous_addr: SocketAddr,
32    ) -> DistributedConfigBuilder {
33        DistributedConfigBuilder {
34            world_size,
35            rank,
36            rendezvous_addr,
37            allow_non_loopback: false,
38            session_id: None,
39            connect_timeout: None,
40            establish_timeout: None,
41            io_timeout: None,
42            op_timeout: None,
43        }
44    }
45
46    /// world_size（全ランク数）。
47    pub fn world_size(&self) -> usize {
48        self.rendezvous.world_size()
49    }
50
51    /// 自ランク。
52    pub fn rank(&self) -> usize {
53        self.rendezvous.rank()
54    }
55
56    /// 集合点アドレス（IPv4-mapped IPv6 は IPv4 へ正規化済み）。
57    pub fn rendezvous_addr(&self) -> SocketAddr {
58        self.rendezvous.rendezvous_addr()
59    }
60
61    /// 非 loopback が明示 opt-in されているか。
62    pub fn is_non_loopback_allowed(&self) -> bool {
63        self.rendezvous.is_non_loopback_allowed()
64    }
65}
66
67/// [`DistributedConfig`] の builder。
68#[derive(Debug, Clone)]
69#[must_use]
70pub struct DistributedConfigBuilder {
71    world_size: usize,
72    rank: usize,
73    rendezvous_addr: SocketAddr,
74    allow_non_loopback: bool,
75    session_id: Option<u64>,
76    connect_timeout: Option<Duration>,
77    establish_timeout: Option<Duration>,
78    io_timeout: Option<Duration>,
79    op_timeout: Option<Duration>,
80}
81
82impl DistributedConfigBuilder {
83    /// 非 loopback の bind・接続を許可する。**認証も暗号化も無い**ため、なりすまし・盗聴の
84    /// リスクを受け入れられる信頼できるネットワークでだけ使うこと。
85    pub fn allow_non_loopback(mut self, allow: bool) -> Self {
86        self.allow_non_loopback = allow;
87        self
88    }
89
90    /// セッション識別子。誤接続の検出用で、**認証ではない**。
91    pub fn session_id(mut self, id: u64) -> Self {
92        self.session_id = Some(id);
93        self
94    }
95
96    /// 接続 1 回の上限。
97    pub fn connect_timeout(mut self, d: Duration) -> Self {
98        self.connect_timeout = Some(d);
99        self
100    }
101
102    /// 集合の確立全体の期限。
103    pub fn establish_timeout(mut self, d: Duration) -> Self {
104        self.establish_timeout = Some(d);
105        self
106    }
107
108    /// 1 回の読み書きの上限（確立時と all-reduce の両方に適用する）。
109    pub fn io_timeout(mut self, d: Duration) -> Self {
110        self.io_timeout = Some(d);
111        self
112    }
113
114    /// all-reduce 1 回の全体の絶対期限。
115    pub fn op_timeout(mut self, d: Duration) -> Self {
116        self.op_timeout = Some(d);
117        self
118    }
119
120    /// すべてをソケットを開く前に検証して確定する。
121    ///
122    /// # Errors
123    ///
124    /// world_size・rank の範囲外、タイムアウトの 0／巨大値、unspecified・マルチキャスト・
125    /// ブロードキャスト・リンクローカルのアドレス、ランク 1 以降のポート 0、opt-in 無しの
126    /// 非 loopback を [`DistributedError`] で拒否する。
127    pub fn build(self) -> Result<DistributedConfig, DistributedError> {
128        let mut rendezvous =
129            RendezvousConfig::new(self.world_size, self.rank, self.rendezvous_addr)
130                .map_err(DistributedError::from_inner)?
131                .allow_non_loopback(self.allow_non_loopback);
132        let mut collective = CollectiveOptions::default();
133        if let Some(id) = self.session_id {
134            rendezvous = rendezvous.session_id(id);
135        }
136        if let Some(d) = self.connect_timeout {
137            rendezvous = rendezvous.connect_timeout(d);
138        }
139        if let Some(d) = self.establish_timeout {
140            rendezvous = rendezvous.establish_timeout(d);
141        }
142        if let Some(d) = self.io_timeout {
143            rendezvous = rendezvous.io_timeout(d);
144            collective = collective.io_timeout(d);
145        }
146        if let Some(d) = self.op_timeout {
147            collective = collective.op_timeout(d);
148        }
149        rendezvous
150            .validate()
151            .map_err(DistributedError::from_inner)?;
152        collective
153            .validate()
154            .map_err(DistributedError::from_inner)?;
155        Ok(DistributedConfig {
156            rendezvous,
157            collective,
158        })
159    }
160}
fandhe_ai/distributed/config.rs

fandhe_ai/distributed/
config.rs
