From e0c179c7ccd685a74cd91cd51b5fc6b36cbdd8ff Mon Sep 17 00:00:00 2001
From: Philip Hyunsu Cho <chohyu01@cs.washington.edu>
Date: Mon, 13 Jul 2020 21:51:11 -0700
Subject: [PATCH] [CI] Enforce daily budget in Jenkins CI (#5884)

* [CI] Throttle Jenkins CI

* Don't use Jenkins master instance
---
 Jenkinsfile                   |  5 +++--
 Jenkinsfile-win64             |  5 +++--
 tests/jenkins_get_approval.py | 26 ++++++++++++++++++++++++++
 3 files changed, 32 insertions(+), 4 deletions(-)
 create mode 100644 tests/jenkins_get_approval.py

diff --git a/Jenkinsfile b/Jenkinsfile
index 40db6f9c0..9aacc6412 100644
--- a/Jenkinsfile
+++ b/Jenkinsfile
@@ -31,13 +31,14 @@ pipeline {
 
   // Build stages
   stages {
-    stage('Jenkins Linux: Get sources') {
-      agent { label 'linux && cpu' }
+    stage('Jenkins Linux: Initialize') {
+      agent { label 'job_initializer' }
       steps {
         script {
           checkoutSrcs()
           commit_id = "${GIT_COMMIT}"
         }
+        sh 'python3 tests/jenkins_get_approval.py'
         stash name: 'srcs'
         milestone ordinal: 1
       }
diff --git a/Jenkinsfile-win64 b/Jenkinsfile-win64
index eba5464ff..9cd12cf4f 100644
--- a/Jenkinsfile-win64
+++ b/Jenkinsfile-win64
@@ -12,13 +12,14 @@ pipeline {
   agent none
   // Build stages
   stages {
-    stage('Jenkins Win64: Get sources') {
-      agent { label 'win64 && build' }
+    stage('Jenkins Win64: Initialize') {
+      agent { label 'job_initializer' }
       steps {
         script {
           checkoutSrcs()
           commit_id = "${GIT_COMMIT}"
         }
+        sh 'python3 tests/jenkins_get_approval.py'
         stash name: 'srcs'
         milestone ordinal: 1
       }
diff --git a/tests/jenkins_get_approval.py b/tests/jenkins_get_approval.py
new file mode 100644
index 000000000..4a68722d9
--- /dev/null
+++ b/tests/jenkins_get_approval.py
@@ -0,0 +1,26 @@
+import boto3
+import json
+
+lambda_client = boto3.client('lambda', region_name='us-west-2')
+
+# Source code for the Lambda function is available at https://github.com/hcho3/xgboost-devops
+r = lambda_client.invoke(
+    FunctionName='XGBoostCICostWatcher',
+    InvocationType='RequestResponse',
+    Payload='{}'.encode('utf-8')
+)
+
+payload = r['Payload'].read().decode('utf-8')
+if 'FunctionError' in r:
+    msg = 'Error when invoking the Lambda function. Stack trace:\n'
+    error = json.loads(payload)
+    msg += f"    {error['errorType']}: {error['errorMessage']}\n"
+    for trace in error['stackTrace']:
+        for line in trace.split('\n'):
+            msg += f'    {line}\n'
+    raise RuntimeError(msg)
+response = json.loads(payload)
+if response['approved']:
+    print(f"Testing approved. Reason: {response['reason']}")
+else:
+    raise RuntimeError(f"Testing rejected. Reason: {response['reason']}")