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

aria_engine/
setup.rs
