[CI] add script to generate meta info and upload to s3 (#10295)

* [CI] add script to generate meta info and upload to s3

* Write Python script to generate meta.json

* Update other pipelines

* Add wheel_name field

* Add description

---------

Co-authored-by: Hyunsu Cho <phcho@nvidia.com>

tests/buildkite/build-cpu-arm64.sh

@@ -19,13 +19,17 @@ $command_wrapper bash -c "cd build && ctest --extra-verbose"
 echo "--- Build binary wheel"
 $command_wrapper bash -c \
   "cd python-package && rm -rf dist/* && pip wheel --no-deps -v . --wheel-dir dist/"
-$command_wrapper python tests/ci_build/rename_whl.py python-package/dist/*.whl \
+$command_wrapper python tests/ci_build/rename_whl.py  \
-  ${BUILDKITE_COMMIT} ${WHEEL_TAG}
+  --wheel-path python-package/dist/*.whl  \
+  --commit-hash ${BUILDKITE_COMMIT}  \
+  --platform-tag ${WHEEL_TAG}
 
 echo "--- Audit binary wheel to ensure it's compliant with manylinux2014 standard"
 $command_wrapper auditwheel repair --plat ${WHEEL_TAG} python-package/dist/*.whl
-$command_wrapper python tests/ci_build/rename_whl.py wheelhouse/*.whl \
+$command_wrapper python tests/ci_build/rename_whl.py  \
-  ${BUILDKITE_COMMIT} ${WHEEL_TAG}
+  --wheel-path wheelhouse/*.whl  \
+  --commit-hash ${BUILDKITE_COMMIT}  \
+  --platform-tag ${WHEEL_TAG}
 mv -v wheelhouse/*.whl python-package/dist/
 # Make sure that libgomp.so is vendored in the wheel
 $command_wrapper bash -c \

tests/buildkite/build-cuda-with-rmm.sh

@@ -36,14 +36,18 @@ $command_wrapper tests/ci_build/build_via_cmake.sh \
 echo "--- Build binary wheel"
 $command_wrapper bash -c \
   "cd python-package && rm -rf dist/* && pip wheel --no-deps -v . --wheel-dir dist/"
-$command_wrapper python tests/ci_build/rename_whl.py python-package/dist/*.whl \
+$command_wrapper python tests/ci_build/rename_whl.py  \
-  ${BUILDKITE_COMMIT} ${WHEEL_TAG}
+  --wheel-path python-package/dist/*.whl  \
+  --commit-hash ${BUILDKITE_COMMIT}  \
+  --platform-tag ${WHEEL_TAG}
 
 echo "--- Audit binary wheel to ensure it's compliant with manylinux2014 standard"
 tests/ci_build/ci_build.sh auditwheel_x86_64 auditwheel repair \
   --plat ${WHEEL_TAG} python-package/dist/*.whl
-$command_wrapper python tests/ci_build/rename_whl.py wheelhouse/*.whl \
+$command_wrapper python tests/ci_build/rename_whl.py  \
-  ${BUILDKITE_COMMIT} ${WHEEL_TAG}
+  --wheel-path wheelhouse/*.whl  \
+  --commit-hash ${BUILDKITE_COMMIT}  \
+  --platform-tag ${WHEEL_TAG}
 mv -v wheelhouse/*.whl python-package/dist/
 # Make sure that libgomp.so is vendored in the wheel
 tests/ci_build/ci_build.sh auditwheel_x86_64 bash -c \

tests/buildkite/build-cuda.sh

@@ -35,14 +35,18 @@ $command_wrapper tests/ci_build/build_via_cmake.sh \
 echo "--- Build binary wheel"
 $command_wrapper bash -c \
   "cd python-package && rm -rf dist/* && pip wheel --no-deps -v . --wheel-dir dist/"
-$command_wrapper python tests/ci_build/rename_whl.py python-package/dist/*.whl \
+$command_wrapper python tests/ci_build/rename_whl.py  \
-  ${BUILDKITE_COMMIT} ${WHEEL_TAG}
+  --wheel-path python-package/dist/*.whl  \
+  --commit-hash ${BUILDKITE_COMMIT}  \
+  --platform-tag ${WHEEL_TAG}
 
 echo "--- Audit binary wheel to ensure it's compliant with manylinux2014 standard"
 tests/ci_build/ci_build.sh auditwheel_x86_64 auditwheel repair \
   --plat ${WHEEL_TAG} python-package/dist/*.whl
-$command_wrapper python tests/ci_build/rename_whl.py wheelhouse/*.whl \
+$command_wrapper python tests/ci_build/rename_whl.py  \
-  ${BUILDKITE_COMMIT} ${WHEEL_TAG}
+  --wheel-path wheelhouse/*.whl  \
+  --commit-hash ${BUILDKITE_COMMIT}  \
+  --platform-tag ${WHEEL_TAG}
 mv -v wheelhouse/*.whl python-package/dist/
 # Make sure that libgomp.so is vendored in the wheel
 tests/ci_build/ci_build.sh auditwheel_x86_64 bash -c \
@@ -54,6 +58,15 @@ if [[ ($is_pull_request == 0) && ($is_release_branch == 1) ]]
 then
   aws s3 cp python-package/dist/*.whl s3://xgboost-nightly-builds/${BRANCH_NAME}/ \
     --acl public-read --no-progress
+
+  # Generate the meta info which includes xgboost version and the commit info
+  $command_wrapper python tests/ci_build/format_wheel_meta.py \
+    --wheel-path python-package/dist/*.whl  \
+    --commit-hash ${BUILDKITE_COMMIT}  \
+    --platform-tag ${WHEEL_TAG}  \
+    --meta-path python-package/dist/
+  aws s3 cp python-package/dist/meta.json s3://xgboost-nightly-builds/${BRANCH_NAME}/ \
+    --acl public-read --no-progress
 fi
 echo "-- Stash C++ test executable (testxgboost)"
 buildkite-agent artifact upload build/testxgboost

tests/buildkite/build-win64-gpu.ps1

@@ -29,7 +29,10 @@ conda activate
 & pip wheel --no-deps -v . --wheel-dir dist/
 Get-ChildItem . -Filter dist/*.whl |
 Foreach-Object {
-  & python ../tests/ci_build/rename_whl.py $_.FullName $Env:BUILDKITE_COMMIT win_amd64
+  & python ../tests/ci_build/rename_whl.py `
+    --wheel-path $_.FullName `
+    --commit-hash $Env:BUILDKITE_COMMIT `
+    --platform-tag win_amd64
   if ($LASTEXITCODE -ne 0) { throw "Last command failed" }
 }
 
@@ -44,7 +47,7 @@ if ( $is_release_branch -eq 1 ) {
   Get-ChildItem . -Filter python-package/dist/*.whl |
   Foreach-Object {
     & aws s3 cp python-package/dist/$_ s3://xgboost-nightly-builds/$Env:BUILDKITE_BRANCH/ `
-    --acl public-read --no-progress
+      --acl public-read --no-progress
     if ($LASTEXITCODE -ne 0) { throw "Last command failed" }
   }
 }

tests/ci_build/build_python_wheels.sh

@@ -67,4 +67,7 @@ fi
 
 python -m pip install cibuildwheel
 python -m cibuildwheel python-package --output-dir wheelhouse
-python tests/ci_build/rename_whl.py wheelhouse/*.whl ${commit_id} ${wheel_tag}
+python tests/ci_build/rename_whl.py  \
+  --wheel-path wheelhouse/*.whl  \
+  --commit-hash ${commit_id}  \
+  --platform-tag ${wheel_tag}

tests/ci_build/format_wheel_meta.py

@@ -0,0 +1,58 @@
+"""
+Script to generate meta.json to store metadata for a nightly build of
+XGBoost Python package.
+"""
+import json
+import pathlib
+from argparse import ArgumentParser
+
+
+def main(args):
+    wheel_path = pathlib.Path(args.wheel_path).expanduser().resolve()
+    if not wheel_path.exists():
+        raise ValueError(f"Wheel cannot be found at path {wheel_path}")
+    if not wheel_path.is_file():
+        raise ValueError(f"Path {wheel_path} is not a valid file")
+    wheel_dir, wheel_name = wheel_path.parent, wheel_path.name
+
+    meta_path = pathlib.Path(args.meta_path)
+    if not meta_path.exists():
+        raise ValueError(f"Path {meta_path} does not exist")
+    if not meta_path.is_dir():
+        raise ValueError(f"Path {meta_path} is not a valid directory")
+
+    tokens = wheel_name.split("-")
+    assert len(tokens) == 5
+    version = tokens[1].split("+")[0]
+
+    meta_info = {
+        "wheel_name": wheel_name,
+        "platform_tag": args.platform_tag,
+        "version": version,
+        "commit_id": args.commit_hash,
+    }
+    with open(meta_path / "meta.json", "w") as f:
+        json.dump(meta_info, f, indent=4)
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser(
+        description="Format meta.json encoding the latest nightly version of the Python wheel"
+    )
+    parser.add_argument(
+        "--wheel-path", type=str, required=True, help="Path to the wheel"
+    )
+    parser.add_argument(
+        "--commit-hash", type=str, required=True, help="Git commit hash"
+    )
+    parser.add_argument(
+        "--platform-tag",
+        type=str,
+        required=True,
+        help="Platform tag (e.g. manylinux2014_x86_64)",
+    )
+    parser.add_argument(
+        "--meta-path", type=str, required=True, help="Directory to place meta.json"
+    )
+    parsed_args = parser.parse_args()
+    main(parsed_args)

tests/ci_build/rename_whl.py

@@ -1,39 +1,59 @@
-import os
+import pathlib
-import sys
+from argparse import ArgumentParser
-
-from test_utils import DirectoryExcursion
-
-if len(sys.argv) != 4:
-    print("Usage: {} [wheel to rename] [commit id] [platform tag]".format(sys.argv[0]))
-    sys.exit(1)
 
 
-whl_path = sys.argv[1]
+def main(args):
-commit_id = sys.argv[2]
+    wheel_path = pathlib.Path(args.wheel_path).expanduser().resolve()
-platform_tag = sys.argv[3]
+    if not wheel_path.exists():
+        raise ValueError(f"Wheel cannot be found at path {wheel_path}")
+    if not wheel_path.is_file():
+        raise ValueError(f"Path {wheel_path} is not a valid file")
+    wheel_dir, wheel_name = wheel_path.parent, wheel_path.name
 
-dirname, basename = os.path.dirname(whl_path), os.path.basename(whl_path)
+    tokens = wheel_name.split("-")
-
-with DirectoryExcursion(dirname):
-    tokens = basename.split("-")
     assert len(tokens) == 5
     version = tokens[1].split("+")[0]
     keywords = {
         "pkg_name": tokens[0],
         "version": version,
-        "commit_id": commit_id,
+        "commit_id": args.commit_hash,
-        "platform_tag": platform_tag,
+        "platform_tag": args.platform_tag,
     }
-    new_name = "{pkg_name}-{version}+{commit_id}-py3-none-{platform_tag}.whl".format(
+    new_wheel_name = (
-        **keywords
+        "{pkg_name}-{version}+{commit_id}-py3-none-{platform_tag}.whl".format(
+            **keywords
+        )
     )
-    print("Renaming {} to {}...".format(basename, new_name))
+    new_wheel_path = wheel_dir / new_wheel_name
-    if os.path.isfile(new_name):
+    print(f"Renaming {wheel_name} to {new_wheel_name}...")
-        os.remove(new_name)
+    if new_wheel_path.is_file():
-    os.rename(basename, new_name)
+        new_wheel_path.unlink()
+    wheel_path.rename(new_wheel_path)
 
-    filesize = os.path.getsize(new_name) / 1024 / 1024  # MB
+    filesize = new_wheel_path.stat().st_size / 1024 / 1024  # MiB
-    print(f"Wheel size: {filesize}")
+    print(f"Wheel size: {filesize:.2f} MiB")
 
-    msg = f"Limit of wheel size set by PyPI is exceeded. {new_name}: {filesize}"
+    if filesize > 300:
-    assert filesize <= 300, msg
+        raise RuntimeError(
+            f"Limit of wheel size set by PyPI is exceeded. {new_wheel_name}: {filesize:.2f} MiB"
+        )
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser(
+        description="Format a Python wheel's name using the git commit hash and platform tag"
+    )
+    parser.add_argument(
+        "--wheel-path", type=str, required=True, help="Path to the wheel"
+    )
+    parser.add_argument(
+        "--commit-hash", type=str, required=True, help="Git commit hash"
+    )
+    parser.add_argument(
+        "--platform-tag",
+        type=str,
+        required=True,
+        help="Platform tag (e.g. manylinux2014_x86_64)",
+    )
+    parsed_args = parser.parse_args()
+    main(parsed_args)