[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/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 sys
-
-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)
+import pathlib
+from argparse import ArgumentParser
 
 
-whl_path = sys.argv[1]
-commit_id = sys.argv[2]
-platform_tag = sys.argv[3]
+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
 
-dirname, basename = os.path.dirname(whl_path), os.path.basename(whl_path)
-
-with DirectoryExcursion(dirname):
-    tokens = basename.split("-")
+    tokens = wheel_name.split("-")
     assert len(tokens) == 5
     version = tokens[1].split("+")[0]
     keywords = {
         "pkg_name": tokens[0],
         "version": version,
-        "commit_id": commit_id,
-        "platform_tag": platform_tag,
+        "commit_id": args.commit_hash,
+        "platform_tag": args.platform_tag,
     }
-    new_name = "{pkg_name}-{version}+{commit_id}-py3-none-{platform_tag}.whl".format(
-        **keywords
+    new_wheel_name = (
+        "{pkg_name}-{version}+{commit_id}-py3-none-{platform_tag}.whl".format(
+            **keywords
+        )
     )
-    print("Renaming {} to {}...".format(basename, new_name))
-    if os.path.isfile(new_name):
-        os.remove(new_name)
-    os.rename(basename, new_name)
+    new_wheel_path = wheel_dir / new_wheel_name
+    print(f"Renaming {wheel_name} to {new_wheel_name}...")
+    if new_wheel_path.is_file():
+        new_wheel_path.unlink()
+    wheel_path.rename(new_wheel_path)
 
-    filesize = os.path.getsize(new_name) / 1024 / 1024  # MB
-    print(f"Wheel size: {filesize}")
+    filesize = new_wheel_path.stat().st_size / 1024 / 1024  # MiB
+    print(f"Wheel size: {filesize:.2f} MiB")
 
-    msg = f"Limit of wheel size set by PyPI is exceeded. {new_name}: {filesize}"
-    assert filesize <= 300, msg
+    if filesize > 300:
+        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)