[CI] Automate artifact fetch step in JVM release process

dev/prepare_jvm_release.py

@@ -1,9 +1,15 @@
+import re
 import os
 import sys
+import platform
 import errno
+import argparse
 import subprocess
 import glob
 import shutil
+import tempfile
+import zipfile
+from urllib.request import urlretrieve
 from contextlib import contextmanager
 
 def normpath(path):
@@ -44,7 +50,48 @@ def run(command, **kwargs):
     print(command)
     subprocess.check_call(command, shell=True, **kwargs)
 
+def get_current_git_tag():
+    out = subprocess.check_output(["git", "tag", "--points-at", "HEAD"])
+    return out.decode().split("\n")[0]
+
+def get_current_commit_hash():
+    out = subprocess.check_output(["git", "rev-parse", "HEAD"])
+    return out.decode().split("\n")[0]
+
+def get_current_git_branch():
+    out = subprocess.check_output(["git", "log", "-n", "1", "--pretty=%d", "HEAD"])
+    m = re.search(r"release_[0-9\.]+", out.decode())
+    if not m:
+        raise ValueError("Expected branch name of form release_xxx")
+    return m.group(0)
+
+def retrieve(url, filename=None):
+    print(f"{url} -> {filename}")
+    return urlretrieve(url, filename)
+
 def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--release-version", type=str, required=True,
+                        help="Version of the release being prepared")
+    args = parser.parse_args()
+
+    if sys.platform != "darwin" or platform.machine() != "x86_64":
+        raise NotImplementedError("Please run this script using an Intel Mac")
+
+    version = args.release_version
+    expected_git_tag = "v" + version
+    current_git_tag = get_current_git_tag()
+    if current_git_tag != expected_git_tag:
+        if not current_git_tag:
+            raise ValueError(f"Expected git tag {expected_git_tag} but current HEAD has no tag. "
+                             f"Run: git checkout {expected_git_tag}")
+        raise ValueError(f"Expected git tag {expected_git_tag} but current HEAD is at tag "
+                         f"{current_git_tag}. Run: git checkout {expected_git_tag}")
+
+    commit_hash = get_current_commit_hash()
+    git_branch = get_current_git_branch()
+    print(f"Using commit {commit_hash} of branch {git_branch}, git tag {current_git_tag}")
+
     with cd("jvm-packages/"):
         print("====copying pure-Python tracker====")
         for use_cuda in [True, False]:
@@ -67,12 +114,46 @@ def main():
                 cp(file, f"{xgboost4j_spark}/src/test/resources")
 
         print("====Creating directories to hold native binaries====")
-        for os, arch in [("linux", "x86_64"), ("windows", "x86_64"), ("macos", "x86_64")]:
-            output_dir = f"xgboost4j/src/main/resources/lib/{os}/{arch}"
+        for os_ident, arch in [("linux", "x86_64"), ("windows", "x86_64"), ("macos", "x86_64")]:
+            output_dir = f"xgboost4j/src/main/resources/lib/{os_ident}/{arch}"
             maybe_makedirs(output_dir)
-        for os, arch in [("linux", "x86_64")]:
-            output_dir = f"xgboost4j-gpu/src/main/resources/lib/{os}/{arch}"
+        for os_ident, arch in [("linux", "x86_64")]:
+            output_dir = f"xgboost4j-gpu/src/main/resources/lib/{os_ident}/{arch}"
             maybe_makedirs(output_dir)
+
+        print("====Downloading native binaries from CI====")
+        nightly_bucket_prefix = "https://s3-us-west-2.amazonaws.com/xgboost-nightly-builds"
+        maven_repo_prefix = "https://s3-us-west-2.amazonaws.com/xgboost-maven-repo/release/ml/dmlc"
+
+        retrieve(url=f"{nightly_bucket_prefix}/{git_branch}/xgboost4j_{commit_hash}.dll",
+                 filename="xgboost4j/src/main/resources/lib/windows/x86_64/xgboost4j.dll")
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            # libxgboost4j.so for Linux x86_64, CPU only
+            zip_path = os.path.join(tempdir, "xgboost4j_2.12.jar")
+            extract_dir = os.path.join(tempdir, "xgboost4j")
+            retrieve(url=f"{maven_repo_prefix}/xgboost4j_2.12/{version}/"
+                         f"xgboost4j_2.12-{version}.jar",
+                     filename=zip_path)
+            os.mkdir(extract_dir)
+            with zipfile.ZipFile(zip_path, "r") as t:
+                t.extractall(extract_dir)
+            cp(os.path.join(extract_dir, "lib", "linux", "x86_64", "libxgboost4j.so"),
+               "xgboost4j/src/main/resources/lib/linux/x86_64/libxgboost4j.so")
+
+            # libxgboost4j.so for Linux x86_64, GPU support
+            zip_path = os.path.join(tempdir, "xgboost4j-gpu_2.12.jar")
+            extract_dir = os.path.join(tempdir, "xgboost4j-gpu")
+            retrieve(url=f"{maven_repo_prefix}/xgboost4j-gpu_2.12/{version}/"
+                         f"xgboost4j-gpu_2.12-{version}.jar",
+                     filename=zip_path)
+            os.mkdir(extract_dir)
+            with zipfile.ZipFile(zip_path, "r") as t:
+                t.extractall(extract_dir)
+            cp(os.path.join(extract_dir, "lib", "linux", "x86_64", "libxgboost4j.so"),
+               "xgboost4j-gpu/src/main/resources/lib/linux/x86_64/libxgboost4j.so")
+
+
     print("====Next Steps====")
     print("1. Gain upload right to Maven Central repo.")
     print("1-1. Sign up for a JIRA account at Sonatype: ")
@@ -81,19 +162,9 @@ def main():
           "https://issues.sonatype.org/browse/OSSRH-67724")
     print("2. Store the Sonatype credentials in .m2/settings.xml. See insturctions in "
           "https://central.sonatype.org/publish/publish-maven/")
-    print("3. Obtain Linux and Windows binaries from the CI server")
-    print("3-1. Get xgboost4j_[commit].dll from "
-          "https://s3-us-west-2.amazonaws.com/xgboost-nightly-builds/list.html. Rename it to"
-          "xgboost4j.dll.")
-    print("3-2. For Linux binaries, go to "
-          "https://s3-us-west-2.amazonaws.com/xgboost-maven-repo/list.html and navigate to the "
-          "release/ directory. Find and download two JAR files: xgboost4j_2.12-[version].jar and "
-          "xgboost4j-gpu_2.12-[version].jar. Use unzip command to extract libxgboost4j.so (one "
-          "version compiled with GPU support and another compiled without).")
-    print("4. Put the binaries in xgboost4j(-gpu)/src/main/resources/lib/[os]/[arch]")
-    print("5. Now on a Mac machine, run:")
+    print("3. Now on a Mac machine, run:")
     print("   GPG_TTY=$(tty) mvn deploy -Prelease -DskipTests")
-    print("6. Log into https://oss.sonatype.org/. On the left menu panel, click Staging "
+    print("4. Log into https://oss.sonatype.org/. On the left menu panel, click Staging "
           "Repositories. Visit the URL https://oss.sonatype.org/content/repositories/mldmlc-1085 "
           "to inspect the staged JAR files. Finally, press Release button to publish the "
           "artifacts to the Maven Central repository.")