Update stash-artifacts.sh script

.github/workflows/jvm_tests.yml

@@ -99,10 +99,8 @@ jobs:
           CONTAINER_ID: xgb-ci.jvm_gpu_build
       - run: bash ops/pipeline/build-jvm-gpu.sh
       - name: Stash files
-        run: bash ops/pipeline/stash-artifacts.sh lib/libxgboost4j.so
-        env:
-          COMMAND: upload
-          KEY: build-jvm-gpu
+        run: |
+          bash ops/pipeline/stash-artifacts.sh stash build-jvm-gpu lib/libxgboost4j.so
 
   build-jvm-mac:
     name: "Build libxgboost4j.dylib for ${{ matrix.description }}"
@@ -151,14 +149,13 @@ jobs:
         env:
           CONTAINER_ID: xgb-ci.jvm_gpu_build
       - name: Unstash files
-        run: bash ops/pipeline/stash-artifacts.sh lib/libxgboost4j.so
-        env:
-          COMMAND: download
-          KEY: build-jvm-gpu
+        run: |
+          bash ops/pipeline/stash-artifacts.sh unstash build-jvm-gpu lib/libxgboost4j.so
       - run: bash ops/pipeline/build-jvm-doc.sh
       - name: Upload JVM doc
         run: |
-          bash ops/pipeline/publish-artifact.sh jvm-packages/${{ env.BRANCH_NAME }}.tar.bz2 \
+          bash ops/pipeline/publish-artifact.sh \
+            jvm-packages/${{ env.BRANCH_NAME }}.tar.bz2 \
             s3://xgboost-docs/
 
   build-test-jvm-packages:
@@ -187,10 +184,9 @@ jobs:
         env:
           SCALA_VERSION: 2.13
       - name: Stash files
-        run: bash ops/pipeline/stash-artifacts.sh lib/libxgboost4j.so
-        env:
-          COMMAND: upload
-          KEY: build-test-jvm-packages
+        run: |
+          bash ops/pipeline/stash-artifacts.sh stash \
+            build-test-jvm-packages lib/libxgboost4j.so
 
   build-test-jvm-packages-other-os:
     name: Build and test JVM packages (${{ matrix.os }})
@@ -263,10 +259,8 @@ jobs:
         env:
           CONTAINER_ID: xgb-ci.jvm_gpu_build
       - name: Unstash files
-        run: bash ops/pipeline/stash-artifacts.sh lib/libxgboost4j.so
-        env:
-          COMMAND: download
-          KEY: build-jvm-gpu
+        run: |
+          bash ops/pipeline/stash-artifacts.sh unstash build-jvm-gpu lib/libxgboost4j.so
       - run: bash ops/pipeline/test-jvm-gpu.sh
 
   deploy-jvm-packages:
@@ -299,11 +293,10 @@ jobs:
           CONTAINER_ID: ${{ matrix.container_id }}
       - name: Unstash files
         run: |
-          bash ops/pipeline/stash-artifacts.sh lib/libxgboost4j.so
+          bash ops/pipeline/stash-artifacts.sh \
+            unstash ${{ matrix.artifact_from }} \
+            lib/libxgboost4j.so
           ls -lh lib/libxgboost4j.so
-        env:
-          COMMAND: download
-          KEY: ${{ matrix.artifact_from }}
       - name: Deploy JVM packages to S3
         run: |
           bash ops/pipeline/deploy-jvm-packages.sh ${{ matrix.variant }} \

.github/workflows/main.yml

@@ -69,10 +69,7 @@ jobs:
           CONTAINER_ID: xgb-ci.cpu
       - run: bash ops/pipeline/build-cpu.sh
       - name: Stash CLI executable
-        run: bash ops/pipeline/stash-artifacts.sh ./xgboost
-        env:
-          COMMAND: upload
-          KEY: build-cpu
+        run: bash ops/pipeline/stash-artifacts.sh stash build-cpu ./xgboost
 
   build-cpu-arm64:
     name: Build CPU ARM64 + manylinux_2_28_aarch64 wheel
@@ -93,10 +90,9 @@ jobs:
           CONTAINER_ID: xgb-ci.aarch64
       - run: bash ops/pipeline/build-cpu-arm64.sh
       - name: Stash files
-        run: bash ops/pipeline/stash-artifacts.sh ./xgboost python-package/dist/*.whl
-        env:
-          COMMAND: upload
-          KEY: build-cpu-arm64
+        run: |
+          bash ops/pipeline/stash-artifacts.sh stash build-cpu-arm64 \
+          ./xgboost python-package/dist/*.whl
       - name: Upload Python wheel
         run: |
           bash ops/pipeline/publish-artifact.sh python-package/dist/*.whl \
@@ -126,11 +122,8 @@ jobs:
       - run: bash ops/pipeline/build-cuda.sh
       - name: Stash files
         run: |
-          bash ops/pipeline/stash-artifacts.sh \
+          bash ops/pipeline/stash-artifacts.sh stash build-cuda \
             build/testxgboost ./xgboost python-package/dist/*.whl
-        env:
-          COMMAND: upload
-          KEY: build-cuda
       - name: Upload Python wheel
         run: |
           for file in python-package/dist/*.whl python-package/dist/meta.json
@@ -162,10 +155,8 @@ jobs:
           CONTAINER_ID: xgb-ci.manylinux_2_28_x86_64
       - run: bash ops/pipeline/build-cuda-with-rmm.sh
       - name: Stash files
-        run: bash ops/pipeline/stash-artifacts.sh build/testxgboost
-        env:
-          COMMAND: upload
-          KEY: build-cuda-with-rmm
+        run: bash ops/pipeline/stash-artifacts.sh  \
+          stash build-cuda-with-rmm build/testxgboost
       - name: Upload Python wheel
         run: |
           bash ops/pipeline/publish-artifact.sh python-package/dist/*.whl \
@@ -266,11 +257,9 @@ jobs:
           CONTAINER_ID: xgb-ci.gpu
       - name: Unstash gtest
         run: |
-          bash ops/pipeline/stash-artifacts.sh build/testxgboost
+          bash ops/pipeline/stash-artifacts.sh unstash ${{ matrix.artifact_from }}  \
+            build/testxgboost
           chmod +x build/testxgboost
-        env:
-          COMMAND: download
-          KEY: ${{ matrix.artifact_from }}
       - run: bash ops/pipeline/test-cpp-gpu.sh ${{ matrix.suite }}
 
   test-python-wheel:
@@ -328,10 +317,8 @@ jobs:
           CONTAINER_ID: ${{ matrix.container }}
       - name: Unstash Python wheel
         run: |
-          bash ops/pipeline/stash-artifacts.sh python-package/dist/*.whl ./xgboost
+          bash ops/pipeline/stash-artifacts.sh unstash ${{ matrix.artifact_from }} \
+            python-package/dist/*.whl ./xgboost
           chmod +x ./xgboost
-        env:
-          COMMAND: download
-          KEY: ${{ matrix.artifact_from }}
       - name: Run Python tests, ${{ matrix.description }}
         run: bash ops/pipeline/test-python-wheel.sh ${{ matrix.suite }} ${{ matrix.container }}

.github/workflows/windows.yml

@@ -31,12 +31,9 @@ jobs:
       - run: powershell ops/pipeline/build-win64-gpu.ps1
       - name: Stash files
         run: |
-          powershell ops/pipeline/stash-artifacts.ps1 `
+          powershell ops/pipeline/stash-artifacts.ps1 stash build-win64-gpu `
             build/testxgboost.exe xgboost.exe `
             (Get-ChildItem python-package/dist/*.whl | Select-Object -Expand FullName)
-        env:
-          COMMAND: upload
-          KEY: build-win64-gpu
 
   test-win64-gpu:
     name: Test XGBoost on Windows
@@ -51,9 +48,6 @@ jobs:
           submodules: "true"
       - name: Unstash files
         run: |
-          powershell ops/pipeline/stash-artifacts.ps1 `
+          powershell ops/pipeline/stash-artifacts.ps1 unstash build-wind64-gpu `
             build/testxgboost.exe xgboost.exe python-package/dist/*.whl
-        env:
-          COMMAND: download
-          KEY: build-win64-gpu
       - run: powershell ops/pipeline/test-win64-gpu.ps1

doc/contrib/ci.rst

@@ -264,35 +264,38 @@ We use `Amazon S3 <https://aws.amazon.com/s3/>`_ to store the stashed files.
 
 .. code-block:: bash
 
-  export COMMAND="upload"
-  export KEY="unique key to identify a group of files"
-  bash ops/pipeline/stash-artifacts.sh path/to/file
+  REMOTE_PREFIX="remote directory to place the artifact(s)"
+  bash ops/pipeline/stash-artifacts.sh stash "${REMOTE_PREFIX}" path/to/file
+
+The ``REMOTE_PREFIX`` argument, which is the second command-line argument
+for ``stash-artifacts.sh``, specifies the remote directory in which the artifact(s)
+should be placed. More precisely, the artifact(s) will be placed in
+``s3://{RUNS_ON_S3_BUCKET_CACHE}/cache/{GITHUB_REPOSITORY}/stash/{GITHUB_RUN_ID}/{REMOTE_PREFIX}/``
+where ``RUNS_ON_S3_BUCKET_CACHE``, ``GITHUB_REPOSITORY``, and ``GITHUB_RUN_ID`` are set by
+the CI. (RunsOn provisions an S3 bucket to stage cache, and its name is stored in the environment
+variable ``RUNS_ON_S3_BUCKET_CACHE``.)
 
 You can upload multiple files, possibly with wildcard globbing:
 
 .. code-block:: bash
 
-  export COMMAND="upload"
-  export KEY="build-cuda"
-  bash ops/pipeline/stash-artifacts.sh \
+  REMOTE_PREFIX="build-cuda"
+  bash ops/pipeline/stash-artifacts.sh stash "${REMOTE_PREFIX}" \
     build/testxgboost python-package/dist/*.whl
 
 **To unstash a file**:
 
 .. code-block:: bash
 
-  export COMMAND="download"
-  export KEY="unique key to identify a group of files"
-  bash ops/pipeline/stash-artifacts.sh path/to/file
+  REMOTE_PREFIX="unique key to identify a group of files"
+  bash ops/pipeline/stash-artifacts.sh unstash "${REMOTE_PREFIX}" path/to/file
 
-You can also use the wildcard globbing. The script will search for files in
-the S3 bucket whose path matches the pattern.
+You can also use the wildcard globbing. The script will download the matching artifacts
+from the remote directory.
 
 .. code-block:: bash
 
-  export COMMAND="download"
-  export KEY="build-cuda"
-  # Download all files whose path matches pattern
-  #   python-package/dist/*.whl
-  bash ops/pipeline/stash-artifacts.sh \
+  REMOTE_PREFIX="build-cuda"
+  # Download all files whose path matches the wildcard pattern python-package/dist/*.whl
+  bash ops/pipeline/stash-artifacts.sh unstash "${REMOTE_PREFIX}" \
     python-package/dist/*.whl

ops/pipeline/stash-artifacts.ps1

@@ -2,29 +2,30 @@
 Param(
     [Parameter(
       Mandatory=$true,
-      Position=0,
+      Position=0
+    )][string]$command,
+    [Parameter(
+      Mandatory=$true,
+      Position=1
+    )][string]$remote_prefix,
+    [Parameter(
+      Mandatory=$true,
+      Position=2,
       ValueFromRemainingArguments=$true
     )][string[]]$artifacts
 )
 
 ## Convenience wrapper for ops/pipeline/stash-artifacts.py
 ## Meant to be used inside GitHub Actions
 
-$ENV_VAR_DOC = @'
-Inputs
-  - COMMAND: Either "upload" or "download"
-  - KEY:     Unique string to identify a group of artifacts
-'@
-
 $ErrorActionPreference = "Stop"
 
 . ops/pipeline/enforce-ci.ps1
 
-foreach ($env in "COMMAND", "KEY", "GITHUB_REPOSITORY", "GITHUB_RUN_ID",
-                 "RUNS_ON_S3_BUCKET_CACHE") {
+foreach ($env in "GITHUB_REPOSITORY", "GITHUB_RUN_ID", "RUNS_ON_S3_BUCKET_CACHE") {
   $val = [Environment]::GetEnvironmentVariable($env)
   if ($val -eq $null) {
-    Write-Host "Error: $env must be set.`n${ENV_VAR_DOC}"
+    Write-Host "Error: $env must be set."
     exit 1
   }
 }
@@ -35,13 +36,13 @@ conda activate
 
 Write-Host @"
 python ops/pipeline/stash-artifacts.py `
-  --command "${Env:COMMAND}"  `
+  --command "${command}"  `
   --s3-bucket "${Env:RUNS_ON_S3_BUCKET_CACHE}" `
-  --prefix "${artifact_stash_prefix}/${Env:KEY}" `
+  --prefix "${artifact_stash_prefix}/${remote_prefix}" `
   -- $artifacts
 "@
 python ops/pipeline/stash-artifacts.py `
-  --command "${Env:COMMAND}"  `
+  --command "${command}"  `
   --s3-bucket "${Env:RUNS_ON_S3_BUCKET_CACHE}" `
-  --prefix "${artifact_stash_prefix}/${Env:KEY}" `
+  --prefix "${artifact_stash_prefix}/${remote_prefix}" `
   -- $artifacts

ops/pipeline/stash-artifacts.py

@@ -85,15 +85,15 @@ def aws_s3_download_with_wildcard(src: str, dest: Path) -> None:
 
 
 def upload(args: argparse.Namespace) -> None:
-    print(f"Uploading artifacts with prefix {args.prefix}...")
+    print(f"Stashing artifacts to prefix {args.prefix}...")
     for artifact in args.artifacts:
         artifact_path = Path(artifact)
         s3_url = compute_s3_url(args.s3_bucket, args.prefix, artifact_path)
         aws_s3_upload(artifact_path, s3_url)
 
 
 def download(args: argparse.Namespace) -> None:
-    print(f"Downloading artifacts with prefix {args.prefix}...")
+    print(f"Unstashing artifacts from prefix {args.prefix}...")
     for artifact in args.artifacts:
         artifact_path = Path(artifact)
         print(f"mkdir -p {str(artifact_path.parent)}")
@@ -117,9 +117,9 @@ def download(args: argparse.Namespace) -> None:
     parser.add_argument(
         "--command",
         type=str,
-        choices=["upload", "download"],
+        choices=["stash", "unstash"],
         required=True,
-        help="Whether to upload or download the artifact (upload/download)",
+        help="Whether to stash or unstash the artifact",
     )
     parser.add_argument(
         "--s3-bucket",
@@ -138,7 +138,7 @@ def download(args: argparse.Namespace) -> None:
     )
     parser.add_argument("artifacts", type=str, nargs="+", metavar="artifact")
     parsed_args = parser.parse_args()
-    if parsed_args.command == "upload":
+    if parsed_args.command == "stash":
         upload(parsed_args)
-    elif parsed_args.command == "download":
+    elif parsed_args.command == "unstash":
         download(parsed_args)

ops/pipeline/stash-artifacts.sh

@@ -3,37 +3,33 @@
 ## Convenience wrapper for ops/pipeline/stash-artifacts.py
 ## Meant to be used inside GitHub Actions
 
-ENV_VAR_DOC=$(
-cat <<-EOF
-Inputs
-  - COMMAND: Either "upload" or "download"
-  - KEY:     Unique string to identify a group of artifacts
-EOF
-)
-
 set -euo pipefail
 
-source ops/pipeline/enforce-ci.sh
+# source ops/pipeline/enforce-ci.sh
 
-if [ "$#" -lt 1 ]; then
-  echo "Usage: $0 [artifact] [artifact ...]"
+if [ "$#" -lt 3 ]; then
+  echo "Usage: $0 {stash,unstash} [remote_prefix] [artifact] [artifact ...]"
   exit 1
 fi
 
-for arg in "COMMAND" "KEY" "GITHUB_REPOSITORY" "GITHUB_RUN_ID" "RUNS_ON_S3_BUCKET_CACHE"
+command="$1"
+remote_prefix="$2"
+shift 2
+
+for arg in "GITHUB_REPOSITORY" "GITHUB_RUN_ID" "RUNS_ON_S3_BUCKET_CACHE"
 do
   if [[ -z "${!arg:-}" ]]
   then
-    echo -e "Error: $arg must be set.\n${ENV_VAR_DOC}"
-    exit 1
+    echo "Error: $arg must be set."
+    exit 2
   fi
 done
 
 artifact_stash_prefix="cache/${GITHUB_REPOSITORY}/stash/${GITHUB_RUN_ID}"
 
 set -x
 python3 ops/pipeline/stash-artifacts.py \
-  --command "${COMMAND}"  \
+  --command "${command}"  \
   --s3-bucket "${RUNS_ON_S3_BUCKET_CACHE}" \
-  --prefix "${artifact_stash_prefix}/${KEY}" \
+  --prefix "${artifact_stash_prefix}/${remote_prefix}" \
   -- "$@"