allow fractional cpu predictions

also fixes a bug which rounded predictions < 0.5 to 0
spack · Aug 21, 2024 · c0e2977 · c0e2977
1 parent 32e451c
commit c0e2977
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 7 deletions.
diff --git a/gantry/routes/prediction/prediction.py b/gantry/routes/prediction/prediction.py
@@ -1,4 +1,5 @@
 import logging
+import math
 
 import aiosqlite
 
@@ -49,17 +50,20 @@ async def predict(db: aiosqlite.Connection, spec: dict, strategy: str = None) ->
         # mapping of sample: [0] cpu_mean, [1] cpu_max, [2] mem_mean, [3] mem_max
         predictions = {
             # averages the respective metric in the sample
-            # cpu should always be whole number
-            "cpu_request": round(sum([build[0] for build in sample]) / len(sample)),
+            # cpu rounded up to nearest 0.1
+            "cpu_request": math.ceil(
+                (sum([build[0] for build in sample]) / len(sample)) * 10
+            )
+            / 10,
             "mem_request": sum([build[2] for build in sample]) / len(sample),
         }
 
     if strategy == "ensure_higher":
         ensure_higher_pred(predictions, spec["pkg_name"])
 
     # warn if the prediction is below some thresholds
-    if predictions["cpu_request"] < 0.25:
-        logger.warning(f"Warning: CPU request for {spec} is below 0.25 cores")
+    if predictions["cpu_request"] < 0.2:
+        logger.warning(f"Warning: CPU request for {spec} is below 0.2 cores")
         predictions["cpu_request"] = DEFAULT_CPU_REQUEST
     if predictions["mem_request"] < 10_000_000:
         logger.warning(f"Warning: Memory request for {spec} is below 10MB")
@@ -68,7 +72,7 @@ async def predict(db: aiosqlite.Connection, spec: dict, strategy: str = None) ->
     # convert predictions to k8s friendly format
     for k, v in predictions.items():
         if k.startswith("cpu"):
-            predictions[k] = str(int(v))
+            predictions[k] = str(v)
         elif k.startswith("mem"):
             predictions[k] = k8s.convert_bytes(v)
 

diff --git a/gantry/tests/defs/prediction.py b/gantry/tests/defs/prediction.py
@@ -21,7 +21,7 @@
 # calculated by running the baseline prediction algorithm on the sample data in gantry/tests/sql/insert_prediction.sql
 NORMAL_PREDICTION = {
     "variables": {
-        "KUBERNETES_CPU_REQUEST": "12",
+        "KUBERNETES_CPU_REQUEST": "11.8",
         "KUBERNETES_MEMORY_REQUEST": "9576M",
     },
 }
@@ -30,7 +30,7 @@
 # that match what the client wants
 DEFAULT_PREDICTION = {
     "variables": {
-        "KUBERNETES_CPU_REQUEST": "1",
+        "KUBERNETES_CPU_REQUEST": "1.0",
         "KUBERNETES_MEMORY_REQUEST": "2000M",
     },
 }