|
|
@@ -57,6 +57,76 @@ fn marker_path(dir: &str) -> String {
|
|
|
format!("{}/.sm121-flare-package.json", dir.trim_end_matches('/'))
|
|
|
}
|
|
|
|
|
|
+/// A prepared package has two on-disk layouts and both must be handled:
|
|
|
+///
|
|
|
+/// * **blob store** (current `flare download`, format 2):
|
|
|
+/// `<root>/blobs/<sha>.safetensors` holds the data once, and
|
|
|
+/// `<root>/models/<version>/model.safetensors` + `weights/` + `deployments/`
|
|
|
+/// are **hardlinks** into `blobs/`. Copying only `models/<version>` with a
|
|
|
+/// hardlink-blind tool expands the links and duplicates the whole package
|
|
|
+/// (measured: a 488 GB package becomes 488 GB again, not a subset).
|
|
|
+/// * **flat** (legacy self-contained package, and the agents' own cache):
|
|
|
+/// `<version>/model.safetensors` sits beside its `weights/` and
|
|
|
+/// `deployments/` directories.
|
|
|
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
|
+pub enum Layout {
|
|
|
+ BlobStore,
|
|
|
+ Flat,
|
|
|
+}
|
|
|
+
|
|
|
+/// Segment that separates a blob store's model versions from its root.
|
|
|
+const MODELS: &str = "models";
|
|
|
+
|
|
|
+/// Which layout a package path lives in, judged by its `/models/` segment.
|
|
|
+pub fn layout_of(package: &str) -> Layout {
|
|
|
+ let mut parts = package.trim_end_matches('/').rsplit('/');
|
|
|
+ // …/<root>/models/<version>/model.safetensors — skip the file and version
|
|
|
+ let _file = parts.next();
|
|
|
+ let _version = parts.next();
|
|
|
+ match parts.next() {
|
|
|
+ Some(MODELS) => Layout::BlobStore,
|
|
|
+ _ => Layout::Flat,
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+/// Directory to copy when mirroring a package: the whole blob store (so the
|
|
|
+/// hardlinks stay links), or the package directory itself for a flat package.
|
|
|
+pub fn mirror_source(package: &str) -> String {
|
|
|
+ let trimmed = package.trim_end_matches('/');
|
|
|
+ match layout_of(trimmed) {
|
|
|
+ Layout::BlobStore => {
|
|
|
+ // strip `models/<version>/model.safetensors`
|
|
|
+ let mut parts = trimmed.rsplitn(4, '/');
|
|
|
+ let _ = (parts.next(), parts.next(), parts.next());
|
|
|
+ parts
|
|
|
+ .next()
|
|
|
+ .map(str::to_string)
|
|
|
+ .unwrap_or_else(|| "/".to_string())
|
|
|
+ }
|
|
|
+ Layout::Flat => trimmed
|
|
|
+ .rsplit_once('/')
|
|
|
+ .map(|(dir, _)| dir.to_string())
|
|
|
+ .unwrap_or_else(|| ".".to_string()),
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+/// The package path inside a mirror rooted at `dir`: the path relative to
|
|
|
+/// [`mirror_source`] keeps a blob store's `models/<version>/…` prefix.
|
|
|
+pub fn mirror_destination(dir: &str, package: &str) -> String {
|
|
|
+ let src = mirror_source(package);
|
|
|
+ let rel = package
|
|
|
+ .trim_end_matches('/')
|
|
|
+ .strip_prefix(&format!("{src}/"))
|
|
|
+ .unwrap_or_else(|| {
|
|
|
+ package
|
|
|
+ .trim_end_matches('/')
|
|
|
+ .rsplit_once('/')
|
|
|
+ .map(|(_, name)| name)
|
|
|
+ .unwrap_or(package)
|
|
|
+ });
|
|
|
+ format!("{}/{rel}", dir.trim_end_matches('/'))
|
|
|
+}
|
|
|
+
|
|
|
/// Current prepared versions in Flare's cache (`flare models list --json`).
|
|
|
/// The package version is a directory name; comparing it catches a re-cut
|
|
|
/// package under the same revision (the stale-NFS class, but for a package).
|
|
|
@@ -84,25 +154,108 @@ async fn current_versions(config: &Config, controller: &Node) -> Result<Vec<Stri
|
|
|
}
|
|
|
|
|
|
/// Test that a local mirror is complete enough to launch: the primary
|
|
|
-/// `model.safetensors`, the `weights/` companion and every rank deployment.
|
|
|
-async fn local_files_present(
|
|
|
- config: &Config,
|
|
|
- controller: &Node,
|
|
|
- dir: &str,
|
|
|
- package: &str,
|
|
|
-) -> Result<bool> {
|
|
|
- let cmd = format!(
|
|
|
+/// `model.safetensors`, the `weights/` companion and every rank deployment,
|
|
|
+/// plus the `blobs/` store when the package uses the blob layout.
|
|
|
+///
|
|
|
+/// The companions live beside the primary in **both** layouts, so the
|
|
|
+/// package's own directory is the one to probe — the mirror root is not.
|
|
|
+async fn local_files_present(config: &Config, controller: &Node, package: &str) -> Result<bool> {
|
|
|
+ let pkg = package.trim_end_matches('/');
|
|
|
+ let dir = pkg
|
|
|
+ .rsplit_once('/')
|
|
|
+ .map(|(d, _)| d.to_string())
|
|
|
+ .unwrap_or_default();
|
|
|
+ let mut cmd = format!(
|
|
|
"test -f {} && test -d {}/weights && test -d {}/deployments",
|
|
|
- shq(package),
|
|
|
- shq(dir),
|
|
|
- shq(dir)
|
|
|
+ shq(pkg),
|
|
|
+ shq(&dir),
|
|
|
+ shq(&dir)
|
|
|
);
|
|
|
+ if layout_of(pkg) == Layout::BlobStore {
|
|
|
+ // The blobs hold the bytes the hardlinked companions point at.
|
|
|
+ cmd.push_str(&format!(" && test -d {}/blobs", shq(&mirror_source(pkg))));
|
|
|
+ }
|
|
|
match crate::ssh::run_ssh(config, controller, cmd, false).await {
|
|
|
Ok(_) => Ok(true),
|
|
|
Err(_) => Ok(false),
|
|
|
}
|
|
|
}
|
|
|
|
|
|
+/// Resolve the prepared package an existing `flare download` left in the
|
|
|
+/// node's store: `<cache>/models/<model_key>/<version>/model.safetensors`.
|
|
|
+///
|
|
|
+/// `flare models list --json` names the cache root, the model key and the
|
|
|
+/// current version, so a store populated by an earlier serve (or by the
|
|
|
+/// operator by hand) is adopted without a second 500 GB verification pass —
|
|
|
+/// the same short-circuit the marker gives a mirror.
|
|
|
+///
|
|
|
+/// The model key is the checkpoint's own identity (`checkpoint-expanded-scales`
|
|
|
+/// for the gb10 package), not the recipe's repo name, so the match is made on
|
|
|
+/// the version: prepared versions are named `<hardware>-tp<N>-<stamp>`, and
|
|
|
+/// only a version carrying **our** revision and TP size may be adopted. Any
|
|
|
+/// other — or more than one — candidate means "download it".
|
|
|
+pub async fn cached_package(
|
|
|
+ config: &Config,
|
|
|
+ recipe: &Recipe,
|
|
|
+ node: &Node,
|
|
|
+) -> Result<Option<String>> {
|
|
|
+ let id = package_id(recipe)?;
|
|
|
+ let cmd = format!(
|
|
|
+ "{} models list --json 2>/dev/null",
|
|
|
+ crate::native::flare::binary(config)
|
|
|
+ );
|
|
|
+ let out = match crate::ssh::run_ssh(config, node, cmd, false).await {
|
|
|
+ Ok(o) => o,
|
|
|
+ Err(_) => return Ok(None),
|
|
|
+ };
|
|
|
+ let Ok(v) = serde_json::from_str::<serde_json::Value>(&out) else {
|
|
|
+ return Ok(None);
|
|
|
+ };
|
|
|
+ let cache = v.get("cache").and_then(|c| c.as_str()).unwrap_or_default();
|
|
|
+ if cache.is_empty() {
|
|
|
+ return Ok(None);
|
|
|
+ }
|
|
|
+ let prefix = format!("{}-tp{}", id.revision, id.tensor_parallel);
|
|
|
+ Ok(select_cached(cache, &v, &prefix))
|
|
|
+}
|
|
|
+
|
|
|
+/// Pick the single prepared version whose name starts with `prefix` (our
|
|
|
+/// hardware revision and TP size) out of a `flare models list --json` value.
|
|
|
+///
|
|
|
+/// One unambiguous candidate only: adopting a different package would serve
|
|
|
+/// the wrong model with no error at all, so two matches mean "download it".
|
|
|
+pub fn select_cached(cache: &str, listing: &serde_json::Value, prefix: &str) -> Option<String> {
|
|
|
+ let mut found: Vec<String> = Vec::new();
|
|
|
+ for model in listing
|
|
|
+ .get("models")
|
|
|
+ .and_then(|m| m.as_array())
|
|
|
+ .map(Vec::as_slice)
|
|
|
+ .unwrap_or_default()
|
|
|
+ {
|
|
|
+ let key = model
|
|
|
+ .get("model_key")
|
|
|
+ .and_then(|k| k.as_str())
|
|
|
+ .unwrap_or_default();
|
|
|
+ for ver in model
|
|
|
+ .get("versions")
|
|
|
+ .and_then(|x| x.as_array())
|
|
|
+ .map(Vec::as_slice)
|
|
|
+ .unwrap_or_default()
|
|
|
+ {
|
|
|
+ let Some(name) = ver.get("version").and_then(|x| x.as_str()) else {
|
|
|
+ continue;
|
|
|
+ };
|
|
|
+ if name.starts_with(prefix) {
|
|
|
+ found.push(format!("{cache}/models/{key}/{name}/model.safetensors"));
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ match found.len() {
|
|
|
+ 1 => found.pop(),
|
|
|
+ _ => None,
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
/// The local package path if the marker matches the recipe identity and the
|
|
|
/// current prepared version. `None` means a download/mirror pass is needed.
|
|
|
pub async fn local_package(
|
|
|
@@ -111,12 +264,37 @@ pub async fn local_package(
|
|
|
cluster: &str,
|
|
|
controller: &Node,
|
|
|
) -> Result<Option<String>> {
|
|
|
- let Some(dir) = crate::native::flare::local_dir_for(config, cluster) else {
|
|
|
- return Ok(None);
|
|
|
- };
|
|
|
- let dir = dir.replace('~', "$HOME");
|
|
|
+ if let Some(dir) = crate::native::flare::local_dir_for(config, cluster) {
|
|
|
+ let dir = dir.replace('~', "$HOME");
|
|
|
+ if let Some(pkg) = marked_package(config, recipe, controller, &dir).await? {
|
|
|
+ return Ok(Some(pkg));
|
|
|
+ }
|
|
|
+ }
|
|
|
+ // No usable mirror marker — but the store may already hold the current
|
|
|
+ // version (a plain `flare download` into `paths.package_cache` leaves no
|
|
|
+ // marker). Adopt it instead of re-verifying 500 GB.
|
|
|
+ if let Some(pkg) = cached_package(config, recipe, controller).await?
|
|
|
+ && local_files_present(config, controller, &pkg).await?
|
|
|
+ {
|
|
|
+ tracing::info!(
|
|
|
+ "FLARE package already prepared on {}: {pkg}",
|
|
|
+ controller.name
|
|
|
+ );
|
|
|
+ return Ok(Some(pkg));
|
|
|
+ }
|
|
|
+ Ok(None)
|
|
|
+}
|
|
|
+
|
|
|
+/// The package path recorded by a previous mirror, when the marker still
|
|
|
+/// matches the recipe and the files are intact.
|
|
|
+async fn marked_package(
|
|
|
+ config: &Config,
|
|
|
+ recipe: &Recipe,
|
|
|
+ controller: &Node,
|
|
|
+ dir: &str,
|
|
|
+) -> Result<Option<String>> {
|
|
|
let id = package_id(recipe)?;
|
|
|
- let mp = marker_path(&dir);
|
|
|
+ let mp = marker_path(dir);
|
|
|
let read = format!("test -f {} && cat {}", shq(&mp), shq(&mp));
|
|
|
let out = match crate::ssh::run_ssh(config, controller, read, false).await {
|
|
|
Ok(o) => o,
|
|
|
@@ -135,7 +313,7 @@ pub async fn local_package(
|
|
|
// Same revision, but the lab re-cut the package: the mirror is stale.
|
|
|
return Ok(None);
|
|
|
}
|
|
|
- if !local_files_present(config, controller, &dir, &marker.package).await? {
|
|
|
+ if !local_files_present(config, controller, &marker.package).await? {
|
|
|
return Ok(None);
|
|
|
}
|
|
|
Ok(Some(marker.package))
|
