feat: allow kernels to be uploaded to a revision (#161)

* feat: allow kernels to be uploaded to a revision

* revision -> branch

docs/source/cli.md

@@ -48,7 +48,7 @@ $ kernels to-wheel drbh/img2grey 1.1.2
 ### kernels upload
 
 Use `kernels upload <dir_containing_build> --repo_id="hub-username/kernel"` to upload
-your kernel builds to the Hub.
+your kernel builds to the Hub. To know the supported arguments run: `kernels upload -h`.
 
 **Notes**:
 

src/kernels/check.py

@@ -1,18 +1,19 @@
-from pathlib import Path
 import sys
+from pathlib import Path
 
 from huggingface_hub import snapshot_download
-from kernels.utils import CACHE_DIR
 from kernel_abi_check import (
     BinaryFormat,
-    IncompatibleMacOSVersion,
-    ObjectFile,
     IncompatibleAbi3Symbol,
-    NonAbi3Symbol,
+    IncompatibleMacOSVersion,
     IncompatibleManylinuxSymbol,
     MissingMacOSVersion,
+    NonAbi3Symbol,
+    ObjectFile,
 )
 
+from kernels.utils import CACHE_DIR
+
 
 def check_kernel(
     *, macos: str, manylinux: str, python_abi: str, repo_id: str, revision: str

src/kernels/cli.py

@@ -69,6 +69,11 @@ def main():
         type=str,
         help="Repository ID to use to upload to the Hugging Face Hub",
     )
+    upload_parser.add_argument(
+        "--branch",
+        type=None,
+        help="If set, the upload will be made to a particular branch of the provided `repo_id`.",
+    )
     upload_parser.add_argument(
         "--private",
         action="store_true",
@@ -218,6 +223,7 @@ def upload_kernels(args):
     upload_folder(
         repo_id=repo_id,
         folder_path=build_dir,
+        revision=args.branch,
         path_in_repo="build",
         delete_patterns=list(delete_patterns),
         commit_message="Build uploaded using `kernels`.",

tests/test_kernel_upload.py

@@ -30,6 +30,7 @@ class UploadArgs:
     kernel_dir: None
     repo_id: None
     private: False
+    branch: None
 
 
 def next_filename(path: Path) -> Path:
@@ -70,14 +71,15 @@ def get_filenames_from_a_repo(repo_id: str) -> List[str]:
 
 @pytest.mark.token
 @pytest.mark.is_staging_test
-def test_kernel_upload_works_as_expected():
+@pytest.mark.parametrize("branch", (None, "foo"))
+def test_kernel_upload_works_as_expected(branch):
     with tempfile.TemporaryDirectory() as tmpdir:
         path = f"{tmpdir}/build/torch-universal/upload_test"
         build_dir = Path(path)
         build_dir.mkdir(parents=True, exist_ok=True)
         script_path = build_dir / "foo.py"
         script_path.write_text(PY_CONTENT)
-        upload_kernels(UploadArgs(tmpdir, REPO_ID, False))
+        upload_kernels(UploadArgs(tmpdir, REPO_ID, False, branch))
 
     repo_filenames = get_filenames_from_a_repo(REPO_ID)
     assert any(str(script_path.name) for f in repo_filenames)