make code dry, extensible

ci/compute-latest-versions.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# Computes versions used for "latest" tag based on "matrices/ci-matrix.yaml"
+# Computes versions used for "latest" tag based on "matrix.yaml"
 # values. Will also check to ensure that the "latest" values are
 # included in the matrix values.
 # Example Usage:
@@ -10,15 +10,17 @@ export LINUX_KEY="LINUX_VER"
 export CUDA_KEY="CUDA_VER"
 export PYTHON_KEY="PYTHON_VER"
 
+PREFIX="ci"
+
 # Get latest values
-LATEST_LINUX_VER=$(yq '.LATEST_VERSIONS.[strenv(LINUX_KEY)]' matrices/ci-matrix.yaml)
-LATEST_CUDA_VER=$(yq '.LATEST_VERSIONS.[strenv(CUDA_KEY)]' matrices/ci-matrix.yaml)
-LATEST_PYTHON_VER=$(yq '.LATEST_VERSIONS.[strenv(PYTHON_KEY)]' matrices/ci-matrix.yaml)
+LATEST_LINUX_VER=$(yq ".${PREFIX}.LATEST_VERSIONS.[strenv(LINUX_KEY)]" matrix.yaml)
+LATEST_CUDA_VER=$(yq ".${PREFIX}.LATEST_VERSIONS.[strenv(CUDA_KEY)]" matrix.yaml)
+LATEST_PYTHON_VER=$(yq ".${PREFIX}.LATEST_VERSIONS.[strenv(PYTHON_KEY)]" matrix.yaml)
 
 # Get matrix array values
-LINUX_VERS=$(yq '.[strenv(LINUX_KEY)]' matrices/ci-matrix.yaml)
-CUDA_VERS=$(yq '.[strenv(CUDA_KEY)]' matrices/ci-matrix.yaml)
-PYTHON_VERS=$(yq '.[strenv(PYTHON_KEY)]' matrices/ci-matrix.yaml)
+LINUX_VERS=$(yq ".${PREFIX}.[strenv(LINUX_KEY)]" matrix.yaml)
+CUDA_VERS=$(yq ".${PREFIX}.[strenv(CUDA_KEY)]" matrix.yaml)
+PYTHON_VERS=$(yq ".${PREFIX}.[strenv(PYTHON_KEY)]" matrix.yaml)
 
 # Ensure matrix array values contain latest values
 for KEY in "${LINUX_KEY}" "${CUDA_KEY}" "${PYTHON_KEY}"; do

ci/compute-matrix.jq

@@ -0,0 +1,89 @@
+include "ci/shared";
+
+# CI specific computations
+def compute_ci($x):
+  def compute_ci_arch($x):
+    ["amd64"] |
+    if
+        $x.CUDA_VER > "11.2.2" and
+        $x.LINUX_VER != "centos7"
+    then
+        . + ["arm64"]
+    else
+        .
+    end |
+    $x + {ARCHES: .};
+
+  def compute_image_name($x):
+    compute_repo($x) as $repo |
+    compute_tag_prefix($x) as $tag_prefix |
+    "rapidsai/" + $repo + ":" + $tag_prefix + "cuda" + $x.CUDA_VER + "-" + $x.LINUX_VER + "-" + "py" + $x.PYTHON_VER |
+    $x + {IMAGE_NAME: .};
+
+  ($x.exclude // []) as $excludes |
+  $x | del(.exclude) |
+  keys_unsorted as $matrix_keys |
+  to_entries |
+  map(.value) |
+  [
+    combinations |
+    lists2dict($matrix_keys; .) |
+    filter_excludes(.; $excludes) |
+    compute_ci_arch(.) |
+    compute_image_name(.)
+  ] |
+  {include: .};
+
+# Wheels specific computations
+def compute_wheels($x):
+  def compute_wheels_arch($x):
+    ["amd64"] |
+    if
+        ["ubuntu18.04", "centos7"] | index($x.LINUX_VER) != null
+    then
+        .
+    else
+        . + ["arm64"]
+    end |
+    $x + {ARCHES: .};
+
+  def compute_manylinux_version($x):
+    if
+      ["ubuntu18.04", "ubuntu20.04"] | index($x.LINUX_VER) != null
+    then
+      "manylinux_2_31"
+    else
+      "manylinux_2_17"
+    end |
+    $x + {MANYLINUX_VER: .};
+
+  def compute_image_name($x):
+    compute_repo($x) as $repo |
+    compute_tag_prefix($x) as $tag_prefix |
+    "rapidsai/" + $repo + ":" + $tag_prefix + "cuda" + $x.CUDA_VER + "-" + $x.LINUX_VER + "-" + "py" + $x.PYTHON_VER |
+    $x + {IMAGE_NAME: .};
+
+  ($x.exclude // []) as $excludes |
+  $x | del(.exclude) |
+  keys_unsorted as $matrix_keys |
+  to_entries |
+  map(.value) |
+  [
+    combinations |
+    lists2dict($matrix_keys; .) |
+    filter_excludes(.; $excludes) |
+    compute_wheels_arch(.) |
+    compute_manylinux_version(.) |
+    compute_image_name(.)
+  ] |
+  {include: .};
+
+# Main function to compute matrix
+def compute_matrix($type; $input):
+  if $type == "ci" then
+    compute_ci($input)
+  elif $type == "wheels" then
+    compute_wheels($input)
+  else
+    error("Unknown matrix type: " + $type)
+  end;

ci/compute-matrix.sh

@@ -13,12 +13,19 @@ case "${BUILD_TYPE}" in
     ;;
 esac
 
-CI_MATRIX=$(yq -o json '. | del(.LATEST_VERSIONS)' matrices/ci-matrix.yaml | jq -c 'include "ci/compute-ci-matrix"; compute_matrix(.)')
-WHEELS_MATRIX=$(yq -o json matrices/wheels-matrix.yaml | jq -c 'include "ci/compute-wheels-matrix"; compute_matrix(.)')
+COMBINED_MATRIX_YAML=$(yq -o json '. | del(.LATEST_VERSIONS)' 'matrix.yaml')
 
-COMBINED_MATRIX=$(jq -c -n \
-  --argjson ci_matrix "$CI_MATRIX" \
-  --argjson wheels_matrix "$WHEELS_MATRIX" \
+# Separate CI and Wheels axes
+CI_MATRIX=$(echo "$COMBINED_MATRIX_YAML" | jq -c '{ci: .ci} | .ci | del(.LATEST_VERSIONS)')
+WHEELS_MATRIX=$(echo "$COMBINED_MATRIX_YAML" | jq -c '{wheels: .wheels} | .wheels')
+
+CI_COMPUTED=$(echo "$CI_MATRIX" | jq -c --arg type "ci" 'include "ci/compute-matrix"; compute_matrix($type; .)')
+WHEELS_COMPUTED=$(echo "$WHEELS_MATRIX" | jq -c --arg type "wheels" 'include "ci/compute-matrix"; compute_matrix($type; .)')
+
+# Combine CI and Wheels matrices
+COMBINED_COMPUTED=$(jq -c -n \
+  --argjson ci_matrix "$CI_COMPUTED" \
+  --argjson wheels_matrix "$WHEELS_COMPUTED" \
   '{"include": ($ci_matrix.include + $wheels_matrix.include)}')
 
-echo "$COMBINED_MATRIX"
+echo "$COMBINED_COMPUTED"

ci/create-multiarch-manifest.sh

@@ -2,9 +2,9 @@
 set -euo pipefail
 
 if [[ "${IMAGE_REPO}" != "ci" ]]; then
-  LATEST_CUDA_VER=$(yq '.CUDA_VER | sort | .[-1]' matrices/wheels-matrix.yaml)
-  LATEST_PYTHON_VER=$(yq -o json '.PYTHON_VER' matrices/wheels-matrix.yaml | jq -r 'max_by(split(".") | map(tonumber))')
-  LATEST_UBUNTU_VER=$(yq '.LINUX_VER | map(select(. == "*ubuntu*")) | sort | .[-1]' matrices/wheels-matrix.yaml)
+  LATEST_CUDA_VER=$(yq '.wheels.CUDA_VER | sort | .[-1]' matrices/matrix.yaml)
+  LATEST_PYTHON_VER=$(yq -o json '.wheels.PYTHON_VER' matrices/matrix.yaml | jq -r 'max_by(split(".") | map(tonumber))')
+  LATEST_UBUNTU_VER=$(yq '.wheels.LINUX_VER | map(select(. == "*ubuntu*")) | sort | .[-1]' matrices/matrix.yaml)
 fi
 
 source_tags=()

ci/shared.jq

@@ -0,0 +1,30 @@
+def compute_repo($x):
+    if
+      env.BUILD_TYPE == "pull-request"
+    then
+      "staging"
+    else
+      $x.IMAGE_REPO
+    end;
+
+# Compute tag prefix
+def compute_tag_prefix($x):
+if
+    env.BUILD_TYPE == "branch"
+then
+    ""
+else
+    $x.IMAGE_REPO + "-" + env.PR_NUM + "-"
+end;
+
+# Checks the current entry to see if it matches the given exclude
+def matches($entry; $exclude):
+all($exclude | to_entries | .[]; $entry[.key] == .value);
+
+# Checks the current entry to see if it matches any of the excludes.
+def filter_excludes($entry; $excludes):
+select(any($excludes[]; matches($entry; .)) | not);
+
+# Convert lists to dictionary
+def lists2dict($keys; $values):
+reduce range($keys | length) as $ind ({}; . + {($keys[$ind]): $values[$ind]});