Additional improvements to CI (#8303)

* Wait until budget check is complete

* Ensure that multi-GPU tests run for the master branch

* Fix

tests/buildkite/conftest.ps1

@@ -6,6 +6,8 @@ if ( $Env:BUILDKITE_PULL_REQUEST -and ($Env:BUILDKITE_PULL_REQUEST -ne "false")
 
 if ( ($Env:BUILDKITE_BRANCH -eq "master") -or ($Env:BUILDKITE_BRANCH -match "release_.+") ) {
   $is_release_branch = 1
+  $enforce_daily_budget = 0
 } else {
   $is_release_branch = 0
+  $enforce_daily_budget = 1
 }

tests/buildkite/conftest.sh

@@ -24,16 +24,18 @@ else
   export BRANCH_NAME=$BUILDKITE_BRANCH
 fi
 
-if [[ -z ${DISABLE_RELEASE:-} ]]
+if [[ $BUILDKITE_BRANCH == "master" || $BUILDKITE_BRANCH == "release_"* ]]
 then
-  if [[ $BUILDKITE_BRANCH == "master" || $BUILDKITE_BRANCH == "release_"* ]]
+  is_release_branch=1
-  then
+  enforce_daily_budget=0
-    is_release_branch=1
-  else
-    is_release_branch=0
-  fi
 else
   is_release_branch=0
+  enforce_daily_budget=1
+fi
+
+if [[ -n ${DISABLE_RELEASE:-} ]]
+then
+  is_release_branch=0
 fi
 
 set +x

tests/buildkite/enforce_daily_budget.sh

@@ -6,9 +6,9 @@ echo "--- Enforce daily budget"
 
 source tests/buildkite/conftest.sh
 
-if [[ $is_release_branch == 1 ]]
+if [[ $enforce_daily_budget == 0 ]]
 then
-  echo "Automatically approving all test jobs for release branches"
+  echo "Automatically approving all test jobs for trunk branches"
 else
   aws lambda invoke --function-name XGBoostCICostWatcher --invocation-type RequestResponse --region us-west-2 response.json
   python3 tests/buildkite/enforce_daily_budget.py --response response.json

tests/buildkite/pipeline-mgpu.yml

@@ -10,6 +10,7 @@ steps:
     key: enforce-daily-budget
     agents:
       queue: pipeline-loader
+  - wait
   - block: ":rocket: Run this test job"
     if: build.pull_request.repository.fork == true
   #### -------- BUILD --------

tests/buildkite/pipeline-win64.yml

@@ -4,6 +4,7 @@ steps:
     key: enforce-daily-budget
     agents:
       queue: pipeline-loader
+  - wait
   - block: ":rocket: Run this test job"
     if: build.pull_request.repository.fork == true
   #### -------- BUILD --------

tests/buildkite/pipeline.yml

@@ -7,6 +7,7 @@ steps:
     key: enforce-daily-budget
     agents:
       queue: pipeline-loader
+  - wait
   - block: ":rocket: Run this test job"
     if: build.pull_request.repository.fork == true
   #### -------- BUILD --------