aria_engine/
setup.rs

1//! Instance-level Engine setup (in-memory; does not write engine.yml).
2
3use thiserror::Error;
4
5pub const INTL_SITE: &str = "https://ariacompute.com";
6pub const INTL_UPGRADE: &str = "https://github.com/ariacompute";
7pub const CN_SITE: &str = "https://ariacompute.cn";
8pub const CN_UPGRADE: &str = "https://gitee.com/ariacompute";
9
10#[derive(Debug, Error, Clone, PartialEq, Eq)]
11pub enum SetupError {
12    #[error("invalid compute: {0}")]
13    InvalidCompute(String),
14}
15
16#[derive(Debug, Clone, PartialEq, Eq)]
17pub struct SetupConfig {
18    pub router: String,
19    pub router_api_key: String,
20    pub site_url: String,
21    pub upgrade_url: String,
22    pub compute: String,
23    pub hf_token: String,
24    pub modelscope_api_token: String,
25}
26
27impl Default for SetupConfig {
28    fn default() -> Self {
29        Self {
30            router: String::new(),
31            router_api_key: String::new(),
32            site_url: String::new(),
33            upgrade_url: String::new(),
34            compute: "auto".into(),
35            hf_token: String::new(),
36            modelscope_api_token: String::new(),
37        }
38    }
39}
40
41/// Partial merge. `None` fields are omitted.
42#[derive(Debug, Clone, Default)]
43pub struct SetupUpdates {
44    pub router: Option<String>,
45    pub router_api_key: Option<String>,
46    pub site_url: Option<String>,
47    pub upgrade_url: Option<String>,
48    pub compute: Option<String>,
49    pub hf_token: Option<String>,
50    pub modelscope_api_token: Option<String>,
51}
52
53fn gateway_region(url: &str) -> Option<&'static str> {
54    let lower = url.to_ascii_lowercase();
55    if lower.contains("ariacompute.cn") || lower.contains("gitee.com/ariacompute") {
56        Some("cn")
57    } else if lower.contains("ariacompute.com") || lower.contains("github.com/ariacompute") {
58        Some("intl")
59    } else {
60        None
61    }
62}
63
64fn pair_urls(region: &str) -> (&'static str, &'static str) {
65    if region == "cn" {
66        (CN_SITE, CN_UPGRADE)
67    } else {
68        (INTL_SITE, INTL_UPGRADE)
69    }
70}
71
72/// Fill missing site/upgrade URLs from a provided TLD.
73pub fn fill_setup_urls(mut cfg: SetupConfig) -> SetupConfig {
74    let region = gateway_region(&cfg.site_url).or_else(|| gateway_region(&cfg.upgrade_url));
75    let Some(region) = region else {
76        return cfg;
77    };
78    let (site, upgrade) = pair_urls(region);
79    if cfg.site_url.is_empty() {
80        cfg.site_url = site.into();
81    }
82    if cfg.upgrade_url.is_empty() {
83        cfg.upgrade_url = upgrade.into();
84    }
85    cfg
86}
87
88/// Merge `updates` into `existing`. Validates; does not mutate `existing`.
89pub fn apply_setup(existing: &SetupConfig, updates: &SetupUpdates) -> Result<SetupConfig, SetupError> {
90    let mut out = existing.clone();
91    if let Some(v) = &updates.router {
92        out.router = v.clone();
93    }
94    if let Some(v) = &updates.router_api_key {
95        out.router_api_key = v.clone();
96    }
97    if let Some(v) = &updates.site_url {
98        out.site_url = v.clone();
99    }
100    if let Some(v) = &updates.upgrade_url {
101        out.upgrade_url = v.clone();
102    }
103    if let Some(v) = &updates.compute {
104        out.compute = v.clone();
105    }
106    if let Some(v) = &updates.hf_token {
107        out.hf_token = v.clone();
108    }
109    if let Some(v) = &updates.modelscope_api_token {
110        out.modelscope_api_token = v.clone();
111    }
112    match out.compute.as_str() {
113        "auto" | "cpu" | "cuda" => {}
114        other => return Err(SetupError::InvalidCompute(other.into())),
115    }
116    Ok(fill_setup_urls(out))
117}
aria_engine/setup.rs

aria_engine/
setup.rs
