diff --git a/docs/source/cli.md b/docs/source/cli.md index a6b746b..65a014b 100644 --- a/docs/source/cli.md +++ b/docs/source/cli.md @@ -48,7 +48,7 @@ $ kernels to-wheel drbh/img2grey 1.1.2 ### kernels upload Use `kernels upload --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**: diff --git a/src/kernels/check.py b/src/kernels/check.py index 1e091f5..581c6f6 100644 --- a/src/kernels/check.py +++ b/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 diff --git a/src/kernels/cli.py b/src/kernels/cli.py index b1f1be4..4ea55a5 100644 --- a/src/kernels/cli.py +++ b/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`.", diff --git a/tests/test_kernel_upload.py b/tests/test_kernel_upload.py index 4010a11..d962c92 100644 --- a/tests/test_kernel_upload.py +++ b/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)