kernel-builder: use branch from build.toml when specified (#485)

Fixes #477.

Author: danieldk

kernel-builder/src/upload.rs

@@ -61,45 +61,57 @@ pub struct UploadArgs {
     pub repo_type: RepoTypeArg,
 }
 
+/// Get repository and branch from the given arguments, or fallback to
+/// reading `build.toml` and/or metadata otherwise.
+fn get_repo_and_branch(
+    kernel_dir: &Path,
+    repo_id: Option<String>,
+    branch: Option<String>,
+    variants: &[PathBuf],
+) -> Result<(String, Option<String>)> {
+    let build = parse_build(kernel_dir);
+
+    let build_branch = build
+        .as_ref()
+        .ok()
+        .and_then(|b| b.branch().map(ToOwned::to_owned));
+    let arg_branch = branch.or(build_branch);
+
+    let resolved_repo_id = match repo_id {
+        Some(id) => id,
+        None => build
+            .context("--repo-id is not provided and cannot parse build.toml.")?
+            .repo_id()
+            .ok_or_else(|| {
+                eyre::eyre!("No `general.hub.repo-id` in build.toml. Use --repo-id to specify it.")
+            })?
+            .to_owned(),
+    };
+
+    let version_branch =
+        arg_branch.map_or_else(|| detect_branch_from_metadata(variants), |b| Ok(Some(b)))?;
+
+    Ok((resolved_repo_id, version_branch))
+}
+
 pub fn run_upload(args: UploadArgs) -> Result<()> {
     let api = hf::api()?;
     let repo_type: RepoType = args.repo_type.into();
     let kernel_dir = check_or_infer_kernel_dir(args.kernel_dir)?;
     let kernel_dir = fs::canonicalize(&kernel_dir)
         .wrap_err_with(|| format!("Cannot resolve kernel directory `{}`", kernel_dir.display()))?;
 
-    let arg_repo_id = match args.repo_id {
-        Some(id) => id,
-        None =>
-        // WARN: parsing must not be moved out of this branch, we want users
-        //       to be able to upload without `build.toml` as long as they
-        //       provide a repo id.
-        {
-            parse_build(&kernel_dir)
-                .context("--repo-id is not provided and cannot parse build.toml.")?
-                .repo_id()
-                .ok_or_else(|| {
-                    eyre::eyre!(
-                        "No `general.hub.repo-id` in build.toml. Use --repo-id to specify it."
-                    )
-                })?
-                .to_owned()
-        }
-    };
-
     let (build_dir, variants) = discover_variants(&kernel_dir)?;
     eprintln!(
         "Found {} build variant(s) in {}",
         variants.len(),
         build_dir.display()
     );
 
-    let version_branch = args
-        .branch
-        .map_or_else(|| detect_branch_from_metadata(&variants), |b| Ok(Some(b)))?;
+    let (repo_id, branch) = get_repo_and_branch(&kernel_dir, args.repo_id, args.branch, &variants)?;
 
     let params = CreateRepoParams::builder()
-        .repo_id(&arg_repo_id)
+        .repo_id(&repo_id)
         .repo_type(repo_type)
         .private(args.private)
         .exist_ok(true)
@@ -113,12 +125,12 @@ pub fn run_upload(args: UploadArgs) -> Result<()> {
         .trim_end_matches('/')
         .strip_prefix("https://huggingface.co/")
         .map(|s| s.strip_prefix("kernels/").unwrap_or(s))
-        .unwrap_or(&arg_repo_id)
+        .unwrap_or(&repo_id)
         .to_owned();
 
     let repo = repo_handle(&api, repo_type, &repo_id);
 
-    let is_new_version_branch = if let Some(ref branch) = version_branch {
+    let is_new_version_branch = if let Some(ref branch) = branch {
         let refs_params = RepoListRefsParams::builder().build();
         let refs = repo
             .list_refs(&refs_params)
@@ -149,7 +161,7 @@ pub fn run_upload(args: UploadArgs) -> Result<()> {
             .or_default(),
     );
 
-    if let Some(ref branch) = version_branch {
+    if let Some(ref branch) = branch {
         let params = RepoListFilesParams {
             revision: Some(branch.clone()),
         };
@@ -227,7 +239,7 @@ pub fn run_upload(args: UploadArgs) -> Result<()> {
             RepoType::Kernel => "kernels/",
             _ => "",
         };
-        let tree_path = version_branch
+        let tree_path = branch
             .as_ref()
             .map_or(String::new(), |b| format!("/tree/{b}"));
         println!("Kernel uploaded: https://hf.co/{type_prefix}{repo_id}{tree_path}");
@@ -781,4 +793,76 @@ mod tests {
         let result = discover_build_file(temp_dir.path(), "CARD.md");
         assert!(result.is_err());
     }
+
+    #[test]
+    fn test_branch_from_build_toml() {
+        let temp_dir = tempfile::tempdir().unwrap();
+        let kernel_dir = temp_dir.path();
+
+        fs::write(
+            kernel_dir.join("build.toml"),
+            r#"[general]
+name = "test-kernel"
+backends = ["cuda"]
+
+[general.hub]
+repo-id = "test/kernel"
+branch = "custom-branch"
+"#,
+        )
+        .unwrap();
+
+        let build_dir = kernel_dir.join("build");
+        let variant = build_dir.join("torch-cuda");
+        fs::create_dir_all(&variant).unwrap();
+        fs::write(variant.join("metadata.json"), METADATA_V3).unwrap();
+        fs::write(variant.join("kernel.so"), "binary").unwrap();
+
+        let variants = vec![variant.clone()];
+        let (repo_id, branch) = get_repo_and_branch(kernel_dir, None, None, &variants).unwrap();
+
+        assert_eq!(repo_id, "test/kernel");
+        assert_eq!(branch, Some("custom-branch".to_owned()));
+
+        // Verify commit ops are generated - these would be uploaded to the branch above.
+        let mut operations = vec![];
+        collect_build_commit_ops(&build_dir, &variants, &[], false, &mut operations).unwrap();
+        assert!(!operations.is_empty());
+    }
+
+    #[test]
+    fn test_args_take_priority_over_files() {
+        let temp_dir = tempfile::tempdir().unwrap();
+        let kernel_dir = temp_dir.path();
+
+        fs::write(
+            kernel_dir.join("build.toml"),
+            r#"[general]
+name = "test-kernel"
+backends = ["cuda"]
+
+[general.hub]
+repo-id = "build-toml/kernel"
+branch = "build-toml-branch"
+"#,
+        )
+        .unwrap();
+
+        let build_dir = kernel_dir.join("build");
+        let variant = build_dir.join("torch-cuda");
+        fs::create_dir_all(&variant).unwrap();
+        fs::write(variant.join("metadata.json"), METADATA_V3).unwrap();
+
+        let variants = vec![variant];
+        let (repo_id, branch) = get_repo_and_branch(
+            kernel_dir,
+            Some("args/kernel".to_owned()),
+            Some("args-branch".to_owned()),
+            &variants,
+        )
+        .unwrap();
+
+        assert_eq!(repo_id, "args/kernel");
+        assert_eq!(branch, Some("args-branch".to_owned()));
+    }
 }

kernels-data/src/config/mod.rs

@@ -53,6 +53,10 @@ impl Build {
         self.kernels.is_empty()
     }
 
+    pub fn branch(&self) -> Option<&str> {
+        self.general.hub.as_ref().and_then(|h| h.branch.as_deref())
+    }
+
     pub fn repo_id(&self) -> Option<&str> {
         self.general.hub.as_ref().and_then(|h| h.repo_id.as_deref())
     }