From b2bbf49015e7c9b34524450755249ae23b519a89 Mon Sep 17 00:00:00 2001
From: Philip Hyunsu Cho <chohyu01@cs.washington.edu>
Date: Tue, 4 Oct 2022 04:03:38 -0700
Subject: [PATCH] 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            |  2 ++
 tests/buildkite/conftest.sh             | 16 +++++++++-------
 tests/buildkite/enforce_daily_budget.sh |  4 ++--
 tests/buildkite/pipeline-mgpu.yml       |  1 +
 tests/buildkite/pipeline-win64.yml      |  1 +
 tests/buildkite/pipeline.yml            |  1 +
 6 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/tests/buildkite/conftest.ps1 b/tests/buildkite/conftest.ps1
index c8dd5f1d1..bd623caf0 100644
--- a/tests/buildkite/conftest.ps1
+++ b/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
 }
diff --git a/tests/buildkite/conftest.sh b/tests/buildkite/conftest.sh
index db86b5bc6..dea3f75ac 100755
--- a/tests/buildkite/conftest.sh
+++ b/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_"* ]]
-  then
-    is_release_branch=1
-  else
-    is_release_branch=0
-  fi
+  is_release_branch=1
+  enforce_daily_budget=0
 else
   is_release_branch=0
+  enforce_daily_budget=1
+fi
+
+if [[ -n ${DISABLE_RELEASE:-} ]]
+then
+  is_release_branch=0
 fi
 
 set +x
diff --git a/tests/buildkite/enforce_daily_budget.sh b/tests/buildkite/enforce_daily_budget.sh
index 9be920d5a..8212f07c1 100755
--- a/tests/buildkite/enforce_daily_budget.sh
+++ b/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
diff --git a/tests/buildkite/pipeline-mgpu.yml b/tests/buildkite/pipeline-mgpu.yml
index be71033ad..fa8f4988b 100644
--- a/tests/buildkite/pipeline-mgpu.yml
+++ b/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 --------
diff --git a/tests/buildkite/pipeline-win64.yml b/tests/buildkite/pipeline-win64.yml
index 9b73baeb3..f8127285d 100644
--- a/tests/buildkite/pipeline-win64.yml
+++ b/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 --------
diff --git a/tests/buildkite/pipeline.yml b/tests/buildkite/pipeline.yml
index 1cd44db3d..206910147 100644
--- a/tests/buildkite/pipeline.yml
+++ b/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 --------