Skip to main content

kernel/install/
plan.rs

1//! The plan for installing a model: which files will be fetched, and where.
2
3use crate::install::event::InstallProgress;
4use crate::install::file_selection::is_weight_path;
5use crate::install::provider::InstallProviderId;
6
7/// One file an install will fetch, with its size when the provider reported it.
8#[derive(Debug, Clone, PartialEq, Eq, Hash)]
9pub struct InstallPlanFile {
10    /// The file's path within the repository.
11    pub path: String,
12    /// The file's size in bytes, if known.
13    pub bytes: Option<i64>,
14}
15
16impl InstallPlanFile {
17    /// A plan file for `path`.
18    pub fn new(path: impl Into<String>, bytes: Option<i64>) -> Self {
19        Self {
20            path: path.into(),
21            bytes,
22        }
23    }
24
25    /// Whether this file is a model weight (by extension).
26    pub fn is_weight(&self) -> bool {
27        is_weight_path(&self.path)
28    }
29}
30
31/// A resolved install: the provider and reference, the files to fetch (and their
32/// total/remaining bytes), where they land, and whether authentication is needed.
33#[derive(Debug, Clone, PartialEq, Eq, Hash)]
34pub struct InstallPlan {
35    /// The provider that will fetch the model.
36    pub provider: InstallProviderId,
37    /// The reference being installed (repo or tag).
38    pub reference: String,
39    /// The name to show the user.
40    pub display_name: String,
41    /// The resolved revision/commit, if pinned.
42    pub revision: Option<String>,
43    /// The files to fetch.
44    pub files: Vec<InstallPlanFile>,
45    /// The total bytes across all files, if known.
46    pub total_bytes: Option<i64>,
47    /// The bytes still to fetch (total minus what's already on disk), if known.
48    pub remaining_bytes: Option<i64>,
49    /// Where the model will be written.
50    pub destination: String,
51    /// Whether the model is gated and needs a token.
52    pub requires_auth: bool,
53}
54
55/// A search result: a model the user could install, with its popularity signals.
56#[derive(Debug, Clone, PartialEq, Eq, Hash)]
57pub struct InstallSearchHit {
58    /// The provider that would install it.
59    pub provider: InstallProviderId,
60    /// The reference (repo or tag).
61    pub reference: String,
62    /// The name to show (the last path segment, usually).
63    pub name: String,
64    /// The download count, if reported.
65    pub downloads: Option<i64>,
66    /// The like count, if reported.
67    pub likes: Option<i64>,
68    /// When it was last updated, epoch milliseconds.
69    pub updated_at: Option<i64>,
70}
71
72impl InstallSearchHit {
73    /// A stable id: `provider|reference`.
74    pub fn id(&self) -> String {
75        format!("{}|{}", self.provider.as_str(), self.reference)
76    }
77}
78
79impl InstallPlan {
80    /// A plan with the required fields; the optional fields default to empty.
81    pub fn new(
82        provider: InstallProviderId,
83        reference: impl Into<String>,
84        display_name: impl Into<String>,
85        destination: impl Into<String>,
86    ) -> Self {
87        Self {
88            provider,
89            reference: reference.into(),
90            display_name: display_name.into(),
91            revision: None,
92            files: Vec::new(),
93            total_bytes: None,
94            remaining_bytes: None,
95            destination: destination.into(),
96            requires_auth: false,
97        }
98    }
99}
100
101/// The outcome of browsing for a model: the hits found, plus a hint when the
102/// lookup failed (kept separate so a partial/failed browse can still show hits).
103#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
104pub struct InstallBrowseResult {
105    /// The models found.
106    pub hits: Vec<InstallSearchHit>,
107    /// Why the browse failed, when it did.
108    pub failure_hint: Option<String>,
109}
110
111impl InstallBrowseResult {
112    /// A result carrying `hits` and no failure.
113    pub fn with_hits(hits: Vec<InstallSearchHit>) -> Self {
114        Self {
115            hits,
116            failure_hint: None,
117        }
118    }
119
120    /// A failed browse carrying only a hint.
121    pub fn failure(hint: impl Into<String>) -> Self {
122        Self {
123            hits: Vec::new(),
124            failure_hint: Some(hint.into()),
125        }
126    }
127}
128
129/// A running install: identity, what's being fetched, live progress, and when it
130/// started (epoch milliseconds).
131#[derive(Debug, Clone, PartialEq, Eq, Hash)]
132pub struct ActiveInstall {
133    /// The install's opaque id.
134    pub id: String,
135    /// The provider fetching it.
136    pub provider: InstallProviderId,
137    /// The reference being installed.
138    pub reference: String,
139    /// The name to show.
140    pub display_name: String,
141    /// The total to download, if known.
142    pub total_bytes: Option<i64>,
143    /// The latest progress.
144    pub progress: InstallProgress,
145    /// When it started, epoch milliseconds.
146    pub started_at: i64,
147}
148
149impl ActiveInstall {
150    /// A fresh install at zero progress, started at `started_at` (epoch millis).
151    pub fn new(
152        id: impl Into<String>,
153        provider: InstallProviderId,
154        reference: impl Into<String>,
155        display_name: impl Into<String>,
156        total_bytes: Option<i64>,
157        started_at: i64,
158    ) -> Self {
159        Self {
160            id: id.into(),
161            provider,
162            reference: reference.into(),
163            display_name: display_name.into(),
164            total_bytes,
165            progress: InstallProgress::default(),
166            started_at,
167        }
168    }
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174
175    #[test]
176    fn an_active_install_starts_at_zero_progress() {
177        let active = ActiveInstall::new(
178            "in-1",
179            InstallProviderId::huggingface(),
180            "org/Model",
181            "Model",
182            Some(1000),
183            42,
184        );
185        assert_eq!(active.progress, InstallProgress::default());
186        assert_eq!(active.progress.bytes_downloaded, 0);
187        assert_eq!(active.total_bytes, Some(1000));
188        assert_eq!(active.started_at, 42);
189    }
190
191    #[test]
192    fn a_browse_result_separates_hits_from_a_failure() {
193        assert!(InstallBrowseResult::default().hits.is_empty());
194        assert_eq!(
195            InstallBrowseResult::failure("down").failure_hint.as_deref(),
196            Some("down")
197        );
198        assert!(InstallBrowseResult::failure("down").hits.is_empty());
199    }
200}