Sfoglia il codice sorgente

flare: handle cached packages and nested package companions

Alexander Korolev 19 ore fa
parent
commit
51f45ea9d2
1 ha cambiato i file con 39 aggiunte e 8 eliminazioni
  1. 39 8
      src/native/flare.rs

+ 39 - 8
src/native/flare.rs

@@ -110,6 +110,14 @@ pub fn download_cmd(config: &Config, recipe: &Recipe) -> Result<String> {
     if let Some(rev) = &f.revision {
         cmd.push_str(&format!(" --revision {}", shq(rev)));
     }
+    // A prepared package is built for one TP size; Spark has one visible GPU,
+    // so `flare download` defaults to TP1 and rejects a TP4 package
+    // ("package exports no deployment with tensor-parallel size 1").
+    // Use the launch-node count as the package's TP (a recipe with
+    // `required_nodes` set is a TP-N package; base recipes fall back to
+    // tensor_parallel_size).
+    let tp = recipe.required_nodes.unwrap_or(recipe.tensor_parallel_size).max(1);
+    cmd.push_str(&format!(" --tensor-parallel {}", tp));
     // NOTE: `flare download` does NOT accept the approval flags (it fails with
     // "unexpected argument") — the signature gate lives on `orbit run`, which
     // is where `run_cmd` passes them. A download has always been accepted
@@ -222,14 +230,16 @@ pub fn preflight_cmd(config: &Config, orbit_cluster: &str, controller: bool) ->
 pub fn mirror_cmd(local_dir: &str, remote_pkg: &str) -> String {
     let dir = shq(local_dir);
     let pkg = shq(remote_pkg);
+    // New packages keep companion data in nested `weights/` and
+    // `deployments/` directories, so a flat file loop is wrong: copy the
+    // package directory recursively, and skip the copy when the cache root
+    // is already the package directory itself.
     format!(
-        "set -e; src=$(dirname {pkg}); dst={dir}; base=$(basename {pkg}); \
+        "set -e; src=$(dirname {pkg}); dst={dir}; \
          mkdir -p \"$dst\"; \
-         cp -f {pkg} \"$dst/$base\"; \
-         for f in \"$src\"/*; do b=$(basename \"$f\"); [ \"$b\" = \"$base\" ] && continue; \
-           if [ ! -f \"$dst/$b\" ] || [ \"$(stat -c%s \"$f\")\" != \"$(stat -c%s \"$dst/$b\")\" ]; then \
-             cp -f \"$f\" \"$dst/$b\" || exit 1; fi; done; \
-         echo \"$dst/$base\""
+         if [ \"$src\" = \"$dst\" ]; then echo {pkg}; exit 0; fi; \
+         cp -a \"$src\"/. \"$dst\"/; \
+         echo \"$dst/$(basename {pkg})\""
     )
 }
 
@@ -238,11 +248,13 @@ pub fn mirror_cmd(local_dir: &str, remote_pkg: &str) -> String {
 /// ending in `.safetensors`. Taking the last line blindly handed `orbit run`
 /// a progress line and the group died silently — so scan backwards for a path.
 pub fn parse_package_path(stdout: &str) -> Option<String> {
+    // Cached downloads print only a "Resolved package: … / model.safetensors"
+    // line, so require an absolute path and ignore the metadata lines.
     stdout
         .lines()
         .rev()
         .map(str::trim)
-        .find(|l| l.ends_with(".safetensors") && l.contains('/'))
+        .find(|l| l.starts_with('/') && l.ends_with(".safetensors"))
         .map(str::to_string)
 }
 
@@ -285,7 +297,26 @@ pub async fn serve(
     }
     let dl = download_cmd(config, recipe)?;
     let package = crate::ssh::run_ssh(config, controller, dl, dry_run).await?;
-    let package = parse_package_path(&package).unwrap_or_default();
+    let mut package = parse_package_path(&package).unwrap_or_default();
+    // A cached download prints only "Resolved package: … / model.safetensors"
+    // and no final path; resolve the default prepared model from Flare's
+    // model cache instead of failing after a successful download.
+    if !dry_run && package.is_empty() {
+        let cache_cmd = format!(
+            "{} models list --json 2>/dev/null | python3 -c 'import json,sys; d=json.load(sys.stdin); m=d.get(\"models\") or []; vers=(m[0].get(\"versions\") or []) if m else []; v=vers[0] if vers else dict(); print(d.get(\"cache\",\"\") + \"/models/\" + m[0].get(\"model_key\",\"\") + \"/\" + v.get(\"version\",\"\"))'",
+            binary(config)
+        );
+        let dir = crate::ssh::run_ssh(config, controller, cache_cmd, false).await?;
+        let dir = dir.trim();
+        if !dir.is_empty() {
+            let find = format!(
+                "find {} -maxdepth 1 -name '*.safetensors' -print -quit",
+                shq(dir)
+            );
+            let primary = crate::ssh::run_ssh(config, controller, find, false).await?;
+            package = primary.trim().to_string();
+        }
+    }
     if !dry_run && package.is_empty() {
         bail!(
             "flare download printed no package path on {}",