Fix that always the memoryVersion was used

Signed-off-by: Lehmann_Fabian <[email protected]>
CommonWorkflowScheduler · Mar 14, 2024 · 2ec2fec · 2ec2fec
1 parent 6ff81a3
commit 2ec2fec
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/src/main/java/cws/k8s/scheduler/prediction/MemoryScaler.java b/src/main/java/cws/k8s/scheduler/prediction/MemoryScaler.java
@@ -115,4 +115,9 @@ protected Predictor createPredictor( String taskName ) {
         return predictorBuilder.build();
     }
 
+    @Override
+    protected long getTaskVersionForPredictor( Task task ) {
+        return task.getMemoryPredictionVersion();
+    }
+
 }
diff --git a/src/main/java/cws/k8s/scheduler/prediction/TaskScaler.java b/src/main/java/cws/k8s/scheduler/prediction/TaskScaler.java
@@ -81,7 +81,7 @@ private void scaleTaskIntern( Task task ) {
 
         final long predictorVersion;
         //Do not predict if the predictor is not set or if the task was already predicted with the same version
-        if ( predictor == null || (predictorVersion = predictor.getVersion()) == task.getMemoryPredictionVersion() ) {
+        if ( predictor == null || (predictorVersion = predictor.getVersion()) == getTaskVersionForPredictor(task) ) {
             return;
         }
 
@@ -106,6 +106,8 @@ public synchronized void beforeTasksScheduled( final List<Task> unscheduledTasks
         log.debug("--- unscheduledTasks END ---");
     }
 
+    protected abstract long getTaskVersionForPredictor( Task task );
+
     public void deactivate(){
         active = false;
     }