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    /// `GET {base}/api/models/by-name/{name}` response - only the fields needed
45    /// to build the download URL are parsed.
46    #[derive(Debug, Deserialize)]
47    struct ModelLookupResponse {
48        id: String,
49    }
50
51    impl ModelService {
52        /// Resolve `model_name`'s registry `Model` id via `mecha10-api`, then
53        /// download `version`'s artifact and cache it at
54        /// `registry_model_dir(model_name, version)/model.onnx`. Returns the
55        /// cached path immediately (without a network round-trip) if it's
56        /// already present.
57        pub async fn pull_registry_version(
58            &self,
59            base_url: &str,
60            api_key: &str,
61            model_name: &str,
62            version: i32,
63            progress: Option<&ProgressBar>,
64        ) -> Result<PathBuf> {
65            let model_dir = self.registry_model_dir(model_name, version);
66            fs::create_dir_all(&model_dir).await?;
67            let output_path = model_dir.join("model.onnx");
68
69            if output_path.exists() {
70                if let Some(pb) = progress {
71                    pb.set_message(format!("Model '{model_name}' v{version} already cached"));
72                }
73                return Ok(output_path);
74            }
75
76            let base = base_url.trim_end_matches('/');
77            let client = reqwest::Client::builder()
78                .user_agent("mecha10-cli")
79                .build()
80                .context("Failed to build HTTP client")?;
81
82            let lookup: ModelLookupResponse = client
83                .get(format!("{base}/api/models/by-name/{model_name}"))
84                .bearer_auth(api_key)
85                .send()
86                .await
87                .context(format!("failed to reach the model registry for '{model_name}'"))?
88                .error_for_status()
89                .context(format!("model '{model_name}' was not found in the registry"))?
90                .json()
91                .await
92                .context("model registry returned an unexpected response")?;
93
94            if let Some(pb) = progress {
95                pb.set_style(
96                    ProgressStyle::default_spinner()
97                        .template("{spinner:.green} {msg}")
98                        .unwrap(),
99                );
100                pb.set_message(format!("Downloading {model_name} v{version} from registry"));
101            }
102
103            let download_url = format!("{base}/api/models/{}/versions/number/{version}/download", lookup.id);
104
105            tracing::info!("Downloading registry model from mecha10-api: {}", download_url);
106
107            let response = client
108                .get(&download_url)
109                .bearer_auth(api_key)
110                .send()
111                .await
112                .context(format!("failed to download registry model '{model_name}' v{version}"))?
113                .error_for_status()
114                .context(format!("registry rejected the download for '{model_name}' v{version}"))?;
115
116            let mut file = tokio::fs::File::create(&output_path).await?;
117            let mut stream = response.bytes_stream();
118
119            use futures_util::StreamExt;
120            use tokio::io::AsyncWriteExt;
121
122            while let Some(chunk) = stream.next().await {
123                let chunk = chunk.context("Error downloading registry model chunk")?;
124                file.write_all(&chunk).await?;
125            }
126            file.flush().await?;
127
128            if let Some(pb) = progress {
129                pb.set_message(format!("Downloaded {model_name} v{version} successfully"));
130            }
131
132            Ok(output_path)
133        }
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140
141    fn service() -> ModelService {
142        ModelService::with_models_dir(PathBuf::from("models")).unwrap()
143    }
144
145    #[test]
146    fn registry_model_dir_is_namespaced_separately_from_catalog() {
147        let service = service();
148        let registry_dir = service.registry_model_dir("rover-detector", 3);
149        let catalog_dir = service.catalog_model_dir("rover-detector");
150
151        assert_eq!(registry_dir, PathBuf::from("models/registry/rover-detector/3"));
152        assert_eq!(catalog_dir, PathBuf::from("models/catalog/rover-detector"));
153        assert_ne!(registry_dir, catalog_dir);
154    }
155
156    #[test]
157    fn registry_model_dir_is_namespaced_by_version() {
158        let service = service();
159        let v1 = service.registry_model_dir("rover-detector", 1);
160        let v2 = service.registry_model_dir("rover-detector", 2);
161
162        assert_ne!(v1, v2);
163    }
164}
mecha10_models/services/model_service/registry.rs

mecha10_models/services/model_service/
registry.rs
