mecha10_models/services/model_service/
registry.rs

1//! Registry model artifact download (`LAB-1802`)
2//!
3//! Distinct from `download.rs`'s public-catalog HuggingFace pulls: this hits the
4//! private, versioned, per-project model registry
5//! (`packages/mecha10-api/src/models/registry`) via the same by-name -> version ->
6//! download resolution the `mecha10 model set` CLI command
7//! (`packages/cli/src/handlers/model.rs`) already used ad hoc. Pulling that
8//! resolution into `ModelService` lets any caller that knows a robot/environment
9//! is source-aware (`node-runner`'s preflight, the CLI's bundled-node runner) fetch
10//! a *specific* registry version instead of only being able to assume every model
11//! name is a catalog entry.
12//!
13//! Artifacts are cached at `models_dir/registry/<name>/<version>/model.onnx` -
14//! namespaced separately from the public catalog's `models_dir/catalog/<name>/`
15//! tree so a name collision between the two systems can never silently overwrite
16//! the wrong artifact. See the module docs on [`super::ModelService`].
17
18use super::ModelService;
19use std::path::PathBuf;
20
21impl ModelService {
22    /// `models_dir/registry/<model_name>/<version>/` - a registry model's
23    /// namespaced, per-version subdirectory (`LAB-1802`). Pure path computation
24    /// (no I/O), so it's available regardless of the `model-download` feature -
25    /// callers that only need to check/derive paths (e.g. `node-runner`'s
26    /// preflight existence check) don't need the download machinery linked in.
27    pub fn registry_model_dir(&self, model_name: &str, version: i32) -> PathBuf {
28        self.models_dir
29            .join("registry")
30            .join(model_name)
31            .join(version.to_string())
32    }
33}
34
35#[cfg(feature = "model-download")]
36mod download {
37    use super::ModelService;
38    use anyhow::{Context, Result};
39    use indicatif::{ProgressBar, ProgressStyle};
40    use serde::Deserialize;
41    use std::path::PathBuf;
42    use tokio::fs;
43
44    #[derive(Debug, Deserialize)]
45    struct ModelLookupResponse {
46        id: String,
47    }
48
49    impl ModelService {
50        /// Resolve `model_name`'s registry `Model` id via `mecha10-api`, then
51        /// download `version`'s artifact and cache it at
52        /// `registry_model_dir(model_name, version)/model.onnx`. Returns the
53        /// cached path immediately (without a network round-trip) if it's
54        /// already present.
55        pub async fn pull_registry_version(
56            &self,
57            base_url: &str,
58            api_key: &str,
59            model_name: &str,
60            version: i32,
61            progress: Option<&ProgressBar>,
62        ) -> Result<PathBuf> {
63            let model_dir = self.registry_model_dir(model_name, version);
64            fs::create_dir_all(&model_dir).await?;
65            let output_path = model_dir.join("model.onnx");
66
67            if output_path.exists() {
68                if let Some(pb) = progress {
69                    pb.set_message(format!("Model '{model_name}' v{version} already cached"));
70                }
71                return Ok(output_path);
72            }
73
74            let base = base_url.trim_end_matches('/');
75            let client = reqwest::Client::builder()
76                .user_agent("mecha10-cli")
77                .build()
78                .context("Failed to build HTTP client")?;
79
80            let lookup: ModelLookupResponse = client
81                .get(format!("{base}/api/models/by-name/{model_name}"))
82                .bearer_auth(api_key)
83                .send()
84                .await
85                .context(format!("failed to reach the model registry for '{model_name}'"))?
86                .error_for_status()
87                .context(format!("model '{model_name}' was not found in the registry"))?
88                .json()
89                .await
90                .context("model registry returned an unexpected response")?;
91
92            if let Some(pb) = progress {
93                pb.set_style(
94                    ProgressStyle::default_spinner()
95                        .template("{spinner:.green} {msg}")
96                        .unwrap(),
97                );
98                pb.set_message(format!("Downloading {model_name} v{version} from registry"));
99            }
100
101            let download_url = format!("{base}/api/models/{}/versions/number/{version}/download", lookup.id);
102
103            tracing::info!("Downloading registry model from mecha10-api: {}", download_url);
104
105            let response = client
106                .get(&download_url)
107                .bearer_auth(api_key)
108                .send()
109                .await
110                .context(format!("failed to download registry model '{model_name}' v{version}"))?
111                .error_for_status()
112                .context(format!("registry rejected the download for '{model_name}' v{version}"))?;
113
114            let mut file = tokio::fs::File::create(&output_path).await?;
115            let mut stream = response.bytes_stream();
116
117            use futures_util::StreamExt;
118            use tokio::io::AsyncWriteExt;
119
120            while let Some(chunk) = stream.next().await {
121                let chunk = chunk.context("Error downloading registry model chunk")?;
122                file.write_all(&chunk).await?;
123            }
124            file.flush().await?;
125
126            if let Some(pb) = progress {
127                pb.set_message(format!("Downloaded {model_name} v{version} successfully"));
128            }
129
130            Ok(output_path)
131        }
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138
139    fn service() -> ModelService {
140        ModelService::with_models_dir(PathBuf::from("models")).unwrap()
141    }
142
143    #[test]
144    fn registry_model_dir_is_namespaced_separately_from_catalog() {
145        let service = service();
146        let registry_dir = service.registry_model_dir("rover-detector", 3);
147        let catalog_dir = service.catalog_model_dir("rover-detector");
148
149        assert_eq!(registry_dir, PathBuf::from("models/registry/rover-detector/3"));
150        assert_eq!(catalog_dir, PathBuf::from("models/catalog/rover-detector"));
151        assert_ne!(registry_dir, catalog_dir);
152    }
153
154    #[test]
155    fn registry_model_dir_is_namespaced_by_version() {
156        let service = service();
157        let v1 = service.registry_model_dir("rover-detector", 1);
158        let v2 = service.registry_model_dir("rover-detector", 2);
159
160        assert_ne!(v1, v2);
161    }
162}
mecha10_models/services/model_service/registry.rs

mecha10_models/services/model_service/
registry.rs
