From 95208b00a23dd82375e8bd18bc6547db88225b0e Mon Sep 17 00:00:00 2001
From: caetano melone <caetanomelone@gmail.com>
Date: Tue, 8 Oct 2024 15:51:59 -0300
Subject: [PATCH] return SPACK_BUILD_JOBS with resource allocation to control
 the build system

`SPACK_BUILD_JOBS=max(1, cpu_request)`
---
 gantry/routes/prediction.py     | 5 +++++
 gantry/tests/defs/prediction.py | 4 ++++
 2 files changed, 9 insertions(+)

diff --git a/gantry/routes/prediction.py b/gantry/routes/prediction.py
index afa941b..710a809 100644
--- a/gantry/routes/prediction.py
+++ b/gantry/routes/prediction.py
@@ -44,6 +44,7 @@ async def predict(db: aiosqlite.Connection, spec: dict) -> dict:
             "mem_request": DEFAULT_MEM_REQUEST,
             "cpu_limit": DEFAULT_CPU_LIMIT,
             "mem_limit": DEFAULT_MEM_LIMIT,
+            "build_jobs": DEFAULT_CPU_REQUEST,
         }
     else:
         # mapping of sample: [0] cpu_mean, [1] cpu_max, [2] mem_mean, [3] mem_max
@@ -54,6 +55,8 @@ async def predict(db: aiosqlite.Connection, spec: dict) -> dict:
             "cpu_limit": sum([build[1] for build in sample]) / n,
             "mem_limit": max([build[3] for build in sample]) * MEM_LIMIT_BUMP,
         }
+        # build jobs cannot be less than 1
+        predictions["build_jobs"] = max(1, round(predictions["cpu_request"]))
 
     if strategy == "ensure_higher":
         ensure_higher_pred(predictions, spec["pkg_name"])
@@ -71,6 +74,8 @@ async def predict(db: aiosqlite.Connection, spec: dict) -> dict:
             "KUBERNETES_MEMORY_REQUEST": predictions["mem_request"],
             "KUBERNETES_CPU_LIMIT": predictions["cpu_limit"],
             "KUBERNETES_MEMORY_LIMIT": predictions["mem_limit"],
+            "SPACK_BUILD_JOBS": predictions["build_jobs"],
+            "CI_JOB_SIZE": "custom",
         },
     }
 
diff --git a/gantry/tests/defs/prediction.py b/gantry/tests/defs/prediction.py
index db9899d..7ba2624 100644
--- a/gantry/tests/defs/prediction.py
+++ b/gantry/tests/defs/prediction.py
@@ -21,10 +21,12 @@
 # calculated by running the baseline prediction algorithm on the sample data in gantry/tests/sql/insert_prediction.sql
 NORMAL_PREDICTION = {
     "variables": {
+        "CI_JOB_SIZE": "custom",
         "KUBERNETES_CPU_LIMIT": "12001m",
         "KUBERNETES_CPU_REQUEST": "11779m",
         "KUBERNETES_MEMORY_LIMIT": "49424M",
         "KUBERNETES_MEMORY_REQUEST": "9577M",
+        "SPACK_BUILD_JOBS": 12,
     },
 }
 
@@ -32,9 +34,11 @@
 # that match what the client wants
 DEFAULT_PREDICTION = {
     "variables": {
+        "CI_JOB_SIZE": "custom",
         "KUBERNETES_CPU_LIMIT": "5000m",
         "KUBERNETES_CPU_REQUEST": "1000m",
         "KUBERNETES_MEMORY_LIMIT": "5000M",
         "KUBERNETES_MEMORY_REQUEST": "2000M",
+        "SPACK_BUILD_JOBS": 1,
     },
 }